Update app.py

app.py

@@ -12,7 +12,7 @@ from preprocess.openpose.run_openpose import OpenPose
 
 import gradio as gr
 
-# Download checkpoints
+# Download checkpoints once at startup
 snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
 
 
@@ -34,64 +34,66 @@ class LeffaPredictor:
             body_model_path="./ckpts/openpose/body_pose_model.pth",
         )
 
-        vt_model_hd = LeffaModel(
+        # Virtual try‑on HD
+        vt_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)
+        self.vt_hd_inf = LeffaInference(model=vt_hd)
 
-        vt_model_dc = LeffaModel(
+        # Virtual try‑on DressCode
+        vt_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)
+        self.vt_dc_inf = LeffaInference(model=vt_dc)
 
-        pt_model = LeffaModel(
+        # Pose transfer
+        pt = 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"]
-        src = Image.open(src_image_path)
-        ref = Image.open(ref_image_path)
+        self.pt_inf = LeffaInference(model=pt)
+
+    def _prepare(self, src_path, ref_path):
+        src = Image.open(src_path)
+        ref = Image.open(ref_path)
         src = resize_and_center(src, 768, 1024)
         ref = resize_and_center(ref, 768, 1024)
-        arr = np.array(src)
-
-        # Mask
-        if control_type == "virtual_tryon":
-            src_rgb = src.convert("RGB")
-            parse, _ = self.parsing(src_rgb.resize((384, 512)))
-            kpt = self.openpose(src_rgb.resize((384, 512)))
-            if vt_model_type == "viton_hd":
-                mask = get_agnostic_mask_hd(parse, kpt, vt_garment_type)
-            else:
-                mask = get_agnostic_mask_dc(parse, kpt, vt_garment_type)
-            mask = mask.resize((768, 1024))
+        return src, ref
+
+    def predict_virtual_tryon(
+        self, src_path, ref_path,
+        accelerate_ref, steps, scale, seed,
+        model_type, garment_type, repaint
+    ):
+        src, ref = self._prepare(src_path, ref_path)
+        src_arr = np.array(src.convert("RGB"))
+
+        # 1) parsing + keypoints → agnostic mask
+        parse, _ = self.parsing(src.resize((384, 512)))
+        kpts = self.openpose(src.resize((384, 512)))
+        if model_type == "viton_hd":
+            mask = get_agnostic_mask_hd(parse, kpts, garment_type)
         else:
-            mask = Image.fromarray(np.ones_like(arr) * 255)
-
-        # DensePose
-        if control_type == "virtual_tryon":
-            if vt_model_type == "viton_hd":
-                seg = self.densepose_predictor.predict_seg(arr)[:, :, ::-1]
-                densepose = Image.fromarray(seg)
-            else:
-                iuv = self.densepose_predictor.predict_iuv(arr)
-                seg = np.repeat(iuv[:, :, 0:1], 3, axis=-1)
-                densepose = Image.fromarray(seg)
+            mask = get_agnostic_mask_dc(parse, kpts, garment_type)
+        mask = mask.resize((768, 1024))
+
+        # 2) DensePose → seg or IUV
+        if model_type == "viton_hd":
+            seg = self.densepose_predictor.predict_seg(src_arr)[:, :, ::-1]
+            densepose = Image.fromarray(seg)
+            inf = self.vt_hd_inf
         else:
-            iuv = self.densepose_predictor.predict_iuv(arr)[:, :, ::-1]
-            densepose = Image.fromarray(iuv)
+            iuv = self.densepose_predictor.predict_iuv(src_arr)
+            seg = np.concatenate([iuv[:, :, :1]] * 3, axis=-1)
+            densepose = Image.fromarray(seg)
+            inf = self.vt_dc_inf
 
-        # Inference
+        # 3) run Leffa
         data = {
             "src_image": [src],
             "ref_image": [ref],
@@ -99,100 +101,123 @@ class LeffaPredictor:
             "densepose": [densepose],
         }
         data = LeffaTransform()(data)
-
-        if control_type == "virtual_tryon":
-            inf = self.vt_inference_hd if vt_model_type == "viton_hd" else self.vt_inference_dc
-        else:
-            inf = self.pt_inference
-
         out = inf(
             data,
-            ref_acceleration=ref_acceleration,
-            num_inference_steps=step,
-            guidance_scale=scale,
-            seed=seed,
-            repaint=vt_repaint,
+            ref_acceleration=accelerate_ref,
+            num_inference_steps=int(steps),
+            guidance_scale=float(scale),
+            seed=int(seed),
+            repaint=repaint,
         )
-        img = out["generated_image"][0]
-        return np.array(img), np.array(mask), np.array(densepose)
-
-    def leffa_predict_vt(self, src, ref, accel, step, scale, seed, mtype, gtype, repaint):
-        return self.leffa_predict(src, ref, "virtual_tryon", accel, step, scale, seed, mtype, gtype, repaint)
+        gen = out["generated_image"][0]
+        return np.array(gen), np.array(mask), np.array(densepose)
+
+    def predict_pose_transfer(
+        self, src_path, ref_path,
+        accelerate_ref, steps, scale, seed
+    ):
+        src, ref = self._prepare(src_path, ref_path)
+        src_arr = np.array(src)
+        mask = Image.fromarray(np.ones_like(src_arr) * 255)
+        iuv = self.densepose_predictor.predict_iuv(src_arr)[:, :, ::-1]
+        densepose = Image.fromarray(iuv)
 
-    def leffa_predict_pt(self, src, ref, accel, step, scale, seed):
-        return self.leffa_predict(src, ref, "pose_transfer", accel, step, scale, seed)
+        data = {
+            "src_image": [src],
+            "ref_image": [ref],
+            "mask": [mask],
+            "densepose": [densepose],
+        }
+        data = LeffaTransform()(data)
+        out = self.pt_inf(
+            data,
+            ref_acceleration=accelerate_ref,
+            num_inference_steps=int(steps),
+            guidance_scale=float(scale),
+            seed=int(seed),
+        )
+        gen = out["generated_image"][0]
+        return np.array(gen), np.array(mask), np.array(densepose)
 
 
 if __name__ == "__main__":
     lp = LeffaPredictor()
-    ex = "./ckpts/examples"
-    p1 = list_dir(f"{ex}/person1")
-    p2 = list_dir(f"{ex}/person2")
-    g = list_dir(f"{ex}/garment")
-
-    with gr.Blocks(
-        theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, secondary_hue=gr.themes.colors.red)
-    ).queue() as demo:
-
-        gr.Markdown("## Leffa: Learning Flow Fields in Attention …")
-        gr.Markdown(
-            "[📚 Paper](https://arxiv.org/abs/2412.08486) • "
-            "[🤖 Code](https://github.com/franciszzj/Leffa) • "
-            "[🤗 Model](https://huggingface.co/franciszzj/Leffa)"
-        )
+    examples = "./ckpts/examples"
+    person1 = list_dir(f"{examples}/person1")
+    person2 = list_dir(f"{examples}/person2")
+    garments = list_dir(f"{examples}/garment")
+
+    title = "## Leffa: Controllable Person Image Generation"
+    note  = "Note: Virtual Try‑On uses VITON‑HD/DressCode; Pose Transfer uses DeepFashion."
+
+    with gr.Blocks(theme=gr.themes.Default(
+        primary_hue=gr.themes.colors.pink,
+        secondary_hue=gr.themes.colors.red
+    )).queue() as demo:
 
-        with gr.Tab("Virtual Try-On"):
+        gr.Markdown(title)
+
+        with gr.Tab("Virtual Try‑On"):
             with gr.Row():
                 with gr.Column():
-                    vt_src = gr.Image(type="filepath", label="Person Image")
-                    gr.Examples(inputs=vt_src, examples_per_page=6, examples=p1)
+                    vt_src = gr.Image(source="upload", type="filepath", label="Person")
+                    gr.Examples(examples=person1, inputs=vt_src)
+
                 with gr.Column():
-                    vt_ref = gr.Image(type="filepath", label="Garment Image")
-                    gr.Examples(inputs=vt_ref, examples_per_page=6, examples=g)
+                    vt_ref = gr.Image(source="upload", type="filepath", label="Garment")
+                    gr.Examples(examples=garments, inputs=vt_ref)
+
                 with gr.Column():
-                    vt_out = gr.Image(label="Generated")
-                    btn = gr.Button("Generate")
-                    with gr.Accordion("Advanced"):
-                        vt_model_type = gr.Radio(
-                            ["viton_hd", "dress_code"], label="Model Type", value="viton_hd"
-                        )
-                        vt_garment_type = gr.Radio(
-                            ["upper_body", "lower_body", "dresses"],
-                            label="Garment Type",
-                            value="upper_body",
-                        )
-                        vt_accel = gr.Checkbox(label="Accelerate UNet", value=False)
-                        vt_repaint = gr.Checkbox(label="Repaint Mode", value=False)
-                        vt_steps = gr.Number(label="Steps", value=30)
-                        vt_scale = gr.Number(label="Scale", value=2.5)
-                        vt_seed = gr.Number(label="Seed", value=42)
-                    btn.click(
-                        lp.leffa_predict_vt,
-                        inputs=[vt_src, vt_ref, vt_accel, vt_steps, vt_scale, vt_seed, vt_model_type, vt_garment_type, vt_repaint],
-                        outputs=[vt_out],
-                    )
+                    vt_out  = gr.Image(label="Result")
+                    vt_mask = gr.Image(label="Mask")
+                    vt_dp   = gr.Image(label="DensePose")
+                    vt_btn  = gr.Button("Generate")
+
+            with gr.Accordion("Advanced Options", open=False):
+                vt_model     = gr.Radio(["viton_hd","dress_code"], value="viton_hd", label="Model")
+                vt_garment   = gr.Radio(["upper_body","lower_body","dresses"], value="upper_body", label="Garment Type")
+                vt_accel_ref = gr.Checkbox(label="Accelerate Reference UNet")
+                vt_repaint   = gr.Checkbox(label="Repaint Mode")
+                vt_steps     = gr.Slider(30,100,value=30,step=1,label="Steps")
+                vt_scale     = gr.Slider(0.1,5.0,value=2.5,step=0.1,label="Guidance Scale")
+                vt_seed      = gr.Number(value=42, label="Seed")
+
+            vt_btn.click(
+                fn=lp.predict_virtual_tryon,
+                inputs=[vt_src, vt_ref, vt_accel_ref, vt_steps, vt_scale, vt_seed, vt_model, vt_garment, vt_repaint],
+                outputs=[vt_out, vt_mask, vt_dp],
+            )
 
         with gr.Tab("Pose Transfer"):
             with gr.Row():
                 with gr.Column():
-                    pt_ref = gr.Image(type="filepath", label="Person Image")
-                    gr.Examples(inputs=pt_ref, examples_per_page=6, examples=p1)
+                    pt_src = gr.Image(source="upload", type="filepath", label="Source Pose")
+                    gr.Examples(examples=person2, inputs=pt_src)
+
                 with gr.Column():
-                    pt_src = gr.Image(type="filepath", label="Target Pose Image")
-                    gr.Examples(inputs=pt_src, examples_per_page=6, examples=p2)
+                    pt_ref = gr.Image(source="upload", type="filepath", label="Target Person")
+                    gr.Examples(examples=person1, inputs=pt_ref)
+
                 with gr.Column():
-                    pt_out = gr.Image(label="Generated")
-                    btn2 = gr.Button("Generate")
-                    with gr.Accordion("Advanced"):
-                        pt_accel = gr.Checkbox(label="Accelerate UNet", value=False)
-                        pt_steps = gr.Number(label="Steps", value=30)
-                        pt_scale = gr.Number(label="Scale", value=2.5)
-                        pt_seed = gr.Number(label="Seed", value=42)
-                    btn2.click(
-                        lp.leffa_predict_pt,
-                        inputs=[pt_src, pt_ref, pt_accel, pt_steps, pt_scale, pt_seed],
-                        outputs=[pt_out],
-                    )
-
-        gr.Markdown("Note: Virtual try-on uses VITON-HD/DressCode; pose transfer uses DeepFashion.")
-        demo.launch(server_port=7860, allowed_paths=["./ckpts/examples"])
+                    pt_out  = gr.Image(label="Result")
+                    pt_mask = gr.Image(label="Mask")
+                    pt_dp   = gr.Image(label="DensePose")
+                    pt_btn  = gr.Button("Generate")
+
+            with gr.Accordion("Advanced Options", open=False):
+                pt_accel_ref = gr.Checkbox(label="Accelerate Reference UNet")
+                pt_steps     = gr.Slider(30,100,value=30,step=1,label="Steps")
+                pt_scale     = gr.Slider(0.1,5.0,value=2.5,step=0.1,label="Guidance Scale")
+                pt_seed      = gr.Number(value=42, label="Seed")
+
+            pt_btn.click(
+                fn=lp.predict_pose_transfer,
+                inputs=[pt_src, pt_ref, pt_accel_ref, pt_steps, pt_scale, pt_seed],
+                outputs=[pt_out, pt_mask, pt_dp],
+            )
+
+        gr.Markdown(note)
+
+        # expose publicly
+        demo.launch(share=True, server_port=7860,
+                    allowed_paths=["./ckpts/examples"])