aeayüz düzenleme 2

app.py

@@ -9,9 +9,144 @@ from leffa_utils.densepose_predictor import DensePosePredictor
 from leffa_utils.utils import resize_and_center, list_dir, get_agnostic_mask_hd, get_agnostic_mask_dc
 from preprocess.humanparsing.run_parsing import Parsing
 from preprocess.openpose.run_openpose import OpenPose
-
 import gradio as gr
 
+# Download checkpoints
+snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
+
+class LeffaPredictor(object):
+    def __init__(self):
+        self.mask_predictor = AutoMasker(
+            densepose_path="./ckpts/densepose",
+            schp_path="./ckpts/schp",
+        )
+
+        self.densepose_predictor = DensePosePredictor(
+            config_path="./ckpts/densepose/densepose_rcnn_R_50_FPN_s1x.yaml",
+            weights_path="./ckpts/densepose/model_final_162be9.pkl",
+        )
+
+        self.parsing = Parsing(
+            atr_path="./ckpts/humanparsing/parsing_atr.onnx",
+            lip_path="./ckpts/humanparsing/parsing_lip.onnx",
+        )
+
+        self.openpose = OpenPose(
+            body_model_path="./ckpts/openpose/body_pose_model.pth",
+        )
+
+        vt_model_hd = LeffaModel(
+            pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
+            pretrained_model="./ckpts/virtual_tryon.pth",
+            dtype="float16",
+        )
+        self.vt_inference_hd = LeffaInference(model=vt_model_hd)
+
+        vt_model_dc = LeffaModel(
+            pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
+            pretrained_model="./ckpts/virtual_tryon_dc.pth",
+            dtype="float16",
+        )
+        self.vt_inference_dc = LeffaInference(model=vt_model_dc)
+
+        pt_model = LeffaModel(
+            pretrained_model_name_or_path="./ckpts/stable-diffusion-xl-1.0-inpainting-0.1",
+            pretrained_model="./ckpts/pose_transfer.pth",
+            dtype="float16",
+        )
+        self.pt_inference = LeffaInference(model=pt_model)
+
+    def leffa_predict(
+        self,
+        src_image_path,
+        ref_image_path,
+        control_type,
+        ref_acceleration=False,
+        step=50,
+        scale=2.5,
+        seed=42,
+        vt_model_type="viton_hd",
+        vt_garment_type="upper_body",
+        vt_repaint=False
+    ):
+        assert control_type in [
+            "virtual_tryon", "pose_transfer"], "Invalid control type: {}".format(control_type)
+        src_image = Image.open(src_image_path)
+        ref_image = Image.open(ref_image_path)
+        src_image = resize_and_center(src_image, 768, 1024)
+        ref_image = resize_and_center(ref_image, 768, 1024)
+
+        src_image_array = np.array(src_image)
+
+        # Mask
+        if control_type == "virtual_tryon":
+            src_image = src_image.convert("RGB")
+            model_parse, _ = self.parsing(src_image.resize((384, 512)))
+            keypoints = self.openpose(src_image.resize((384, 512)))
+            if vt_model_type == "viton_hd":
+                mask = get_agnostic_mask_hd(
+                    model_parse, keypoints, vt_garment_type)
+            elif vt_model_type == "dress_code":
+                mask = get_agnostic_mask_dc(
+                    model_parse, keypoints, vt_garment_type)
+            mask = mask.resize((768, 1024))
+        elif control_type == "pose_transfer":
+            mask = Image.fromarray(np.ones_like(src_image_array) * 255)
+
+        # DensePose
+        if control_type == "virtual_tryon":
+            if vt_model_type == "viton_hd":
+                src_image_seg_array = self.densepose_predictor.predict_seg(
+                    src_image_array)[:, :, ::-1]
+                src_image_seg = Image.fromarray(src_image_seg_array)
+                densepose = src_image_seg
+            elif vt_model_type == "dress_code":
+                src_image_iuv_array = self.densepose_predictor.predict_iuv(
+                    src_image_array)
+                src_image_seg_array = src_image_iuv_array[:, :, 0:1]
+                src_image_seg_array = np.concatenate(
+                    [src_image_seg_array] * 3, axis=-1)
+                src_image_seg = Image.fromarray(src_image_seg_array)
+                densepose = src_image_seg
+        elif control_type == "pose_transfer":
+            src_image_iuv_array = self.densepose_predictor.predict_iuv(
+                src_image_array)[:, :, ::-1]
+            src_image_iuv = Image.fromarray(src_image_iuv_array)
+            densepose = src_image_iuv
+
+        # Leffa
+        transform = LeffaTransform()
+
+        data = {
+            "src_image": [src_image],
+            "ref_image": [ref_image],
+            "mask": [mask],
+            "densepose": [densepose],
+        }
+        data = transform(data)
+        if control_type == "virtual_tryon":
+            if vt_model_type == "viton_hd":
+                inference = self.vt_inference_hd
+            elif vt_model_type == "dress_code":
+                inference = self.vt_inference_dc
+        elif control_type == "pose_transfer":
+            inference = self.pt_inference
+        output = inference(
+            data,
+            ref_acceleration=ref_acceleration,
+            num_inference_steps=step,
+            guidance_scale=scale,
+            seed=seed,
+            repaint=vt_repaint,)
+        gen_image = output["generated_image"][0]
+        return np.array(gen_image), np.array(mask), np.array(densepose)
+
+    def dehasoft(self, src_image_path, ref_image_path, ref_acceleration, step, scale, seed, vt_model_type, vt_garment_type, vt_repaint):
+        return self.leffa_predict(src_image_path, ref_image_path, "virtual_tryon", ref_acceleration, step, scale, seed, vt_model_type, vt_garment_type, vt_repaint)
+
+    def leffa_predict_pt(self, src_image_path, ref_image_path, ref_acceleration, step, scale, seed):
+        return self.leffa_predict(src_image_path, ref_image_path, "pose_transfer", ref_acceleration, step, scale, seed)
+
 if __name__ == "__main__":
     leffa_predictor = LeffaPredictor()
     example_dir = "./ckpts/examples"
@@ -47,13 +182,10 @@ if __name__ == "__main__":
     """
 
     with gr.Blocks(theme=theme, title="Dehasoft AI Studio") as demo:
-        # Başlık ve Açıklama
         gr.Markdown(title, elem_classes=["title"])
         gr.Markdown(description, elem_classes=["description"])
 
-        # Sekmeler
         with gr.Tabs(elem_classes=["tabs"]):
-            # Virtual Try-on Sekmesi
             with gr.TabItem("Virtual Try-On", elem_id="vt_tab"):
                 with gr.Row(equal_height=True):
                     with gr.Column(scale=1):
@@ -163,7 +295,6 @@ if __name__ == "__main__":
                     _js="() => { document.querySelector('.generate-btn').classList.add('loading'); setTimeout(() => document.querySelector('.generate-btn').classList.remove('loading'), 5000); }"
                 )
 
-            # Pose Transfer Sekmesi
             with gr.TabItem("Pose Transfer", elem_id="pt_tab"):
                 with gr.Row(equal_height=True):
                     with gr.Column(scale=1):
@@ -256,10 +387,8 @@ if __name__ == "__main__":
                     _js="() => { document.querySelector('.generate-btn').classList.add('loading'); setTimeout(() => document.querySelector('.generate-btn').classList.remove('loading'), 5000); }"
                 )
 
-        # Altbilgi
         gr.Markdown(footer_note, elem_classes=["footer"])
 
-        # Özel CSS
         demo.css = """
         .title { text-align: center; font-size: 2.5em; margin-bottom: 10px; color: #4f46e5; }
         .description { text-align: center; font-size: 1.2em; margin-bottom: 20px; color: #374151; }