add .gitignore to exclude .DS_Store files

.gitignore

@@ -0,0 +1 @@
+.DS_Store

app.py

@@ -7,6 +7,7 @@ from leffa.inference import LeffaInference
 from utils.garment_agnostic_mask_predictor import AutoMasker
 from utils.densepose_predictor import DensePosePredictor
 from utils.utils import resize_and_center
+import spaces
 
 import gradio as gr
 
@@ -35,7 +36,7 @@ pt_model = LeffaModel(
 )
 pt_inference = LeffaInference(model=pt_model)
 
-
+@spaces.GPU
 def leffa_predict(src_image_path, ref_image_path, control_type):
     assert control_type in [
         "virtual_tryon", "pose_transfer"], "Invalid control type: {}".format(control_type)
@@ -83,145 +84,220 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
     # gen_image.save("gen_image.png")
     return np.array(gen_image)
 
-
+@spaces.GPU
 def leffa_predict_vt(src_image_path, ref_image_path):
     return leffa_predict(src_image_path, ref_image_path, "virtual_tryon")
 
+@spaces.GPU
+def leffa_predict_vt_image_url(src_image_path, ref_image_path):
+    src_image = fetch_image_from_url(src_image_path)
+    print("fetched person image")
+    ref_image = fetch_image_from_url(ref_image_path)
+    print("fetched garment image")
+
+    return leffa_predict(src_image, ref_image, "virtual_tryon")
 
+@spaces.GPU
 def leffa_predict_pt(src_image_path, ref_image_path):
     return leffa_predict(src_image_path, ref_image_path, "pose_transfer")
 
-
-if __name__ == "__main__":
-    # import sys
-
-    # src_image_path = sys.argv[1]
-    # ref_image_path = sys.argv[2]
-    # control_type = sys.argv[3]
-    # leffa_predict(src_image_path, ref_image_path, control_type)
-
-    title = "## Leffa: Learning Flow Fields in Attention for Controllable Person Image Generation"
-    link = "[📚 Paper](https://arxiv.org/abs/2412.08486) - [🔥 Demo](https://huggingface.co/spaces/franciszzj/Leffa) - [🤗 Model](https://huggingface.co/franciszzj/Leffa)"
-    description = "Leffa is a unified framework for controllable person image generation that enables precise manipulation of both appearance (i.e., virtual try-on) and pose (i.e., pose transfer)."
-    note = "Note: The models used in the demo are trained solely on academic datasets. Virtual try-on uses VITON-HD, and 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:
+def fetch_image_from_url(url):
+    try:
+        response = requests.get(url)
+        img = Image.open(BytesIO(response.content))
+        return img
+    except Exception as e:
+        print(e)
+        return None
+
+def handle_image_input(image_input):
+    if image_input.startswith('http'):
+        return fetch_image_from_url(image_input)
+    else:
+        return Image.open(image_input)
+
+# if __name__ == "__main__":
+#     # import sys
+
+#     # src_image_path = sys.argv[1]
+#     # ref_image_path = sys.argv[2]
+#     # control_type = sys.argv[3]
+#     # leffa_predict(src_image_path, ref_image_path, control_type)
+
+#     title = "## Leffa: Learning Flow Fields in Attention for Controllable Person Image Generation"
+#     link = "[📚 Paper](https://arxiv.org/abs/2412.08486) - [🔥 Demo](https://huggingface.co/spaces/franciszzj/Leffa) - [🤗 Model](https://huggingface.co/franciszzj/Leffa)"
+#     description = "Leffa is a unified framework for controllable person image generation that enables precise manipulation of both appearance (i.e., virtual try-on) and pose (i.e., pose transfer)."
+#     note = "Note: The models used in the demo are trained solely on academic datasets. Virtual try-on uses VITON-HD, and 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:
+#         gr.Markdown(title)
+#         gr.Markdown(link)
+#         gr.Markdown(description)
+
+#         with gr.Tab("Control Appearance (Virtual Try-on)"):
+#             with gr.Row():
+#                 with gr.Column():
+#                     gr.Markdown("#### Person Image")
+#                     vt_src_image = gr.Image(
+#                         sources=["upload", "url"],
+#                         type="filepath",
+#                         label="Person Image",
+#                         width=512,
+#                         height=512,
+#                     )
+
+#                     gr.Examples(
+#                         inputs=vt_src_image,
+#                         examples_per_page=5,
+#                         examples=["./ckpts/examples/person1/01350_00.jpg",
+#                                   "./ckpts/examples/person1/01376_00.jpg",
+#                                   "./ckpts/examples/person1/01416_00.jpg",
+#                                   "./ckpts/examples/person1/05976_00.jpg",
+#                                   "./ckpts/examples/person1/06094_00.jpg",],
+#                     )
+
+#                 with gr.Column():
+#                     gr.Markdown("#### Garment Image")
+#                     vt_ref_image = gr.Image(
+#                         sources=["upload", "url"],
+#                         type="filepath",
+#                         label="Garment Image",
+#                         width=512,
+#                         height=512,
+#                     )
+
+#                     gr.Examples(
+#                         inputs=vt_ref_image,
+#                         examples_per_page=5,
+#                         examples=["./ckpts/examples/garment/01449_00.jpg",
+#                                   "./ckpts/examples/garment/01486_00.jpg",
+#                                   "./ckpts/examples/garment/01853_00.jpg",
+#                                   "./ckpts/examples/garment/02070_00.jpg",
+#                                   "./ckpts/examples/garment/03553_00.jpg",],
+#                     )
+
+#                 with gr.Column():
+#                     gr.Markdown("#### Generated Image")
+#                     vt_gen_image = gr.Image(
+#                         label="Generated Image",
+#                         width=512,
+#                         height=512,
+#                     )
+
+#                     with gr.Row():
+#                         vt_gen_button = gr.Button("Generate")
+
+#                 vt_gen_button.click(fn=leffa_predict_vt, inputs=[
+#                     vt_src_image, vt_ref_image], outputs=[vt_gen_image])
+
+#         with gr.Tab("Control Pose (Pose Transfer)"):
+#             with gr.Row():
+#                 with gr.Column():
+#                     gr.Markdown("#### Person Image")
+#                     pt_ref_image = gr.Image(
+#                         sources=["upload"],
+#                         type="filepath",
+#                         label="Person Image",
+#                         width=512,
+#                         height=512,
+#                     )
+
+#                     gr.Examples(
+#                         inputs=pt_ref_image,
+#                         examples_per_page=5,
+#                         examples=["./ckpts/examples/person1/01350_00.jpg",
+#                                   "./ckpts/examples/person1/01376_00.jpg",
+#                                   "./ckpts/examples/person1/01416_00.jpg",
+#                                   "./ckpts/examples/person1/05976_00.jpg",
+#                                   "./ckpts/examples/person1/06094_00.jpg",],
+#                     )
+
+#                 with gr.Column():
+#                     gr.Markdown("#### Target Pose Person Image")
+#                     pt_src_image = gr.Image(
+#                         sources=["upload"],
+#                         type="filepath",
+#                         label="Target Pose Person Image",
+#                         width=512,
+#                         height=512,
+#                     )
+
+#                     gr.Examples(
+#                         inputs=pt_src_image,
+#                         examples_per_page=5,
+#                         examples=["./ckpts/examples/person2/01850_00.jpg",
+#                                   "./ckpts/examples/person2/01875_00.jpg",
+#                                   "./ckpts/examples/person2/02532_00.jpg",
+#                                   "./ckpts/examples/person2/02902_00.jpg",
+#                                   "./ckpts/examples/person2/05346_00.jpg",],
+#                     )
+
+#                 with gr.Column():
+#                     gr.Markdown("#### Generated Image")
+#                     pt_gen_image = gr.Image(
+#                         label="Generated Image",
+#                         width=512,
+#                         height=512,
+#                     )
+
+#                     with gr.Row():
+#                         pose_transfer_gen_button = gr.Button("Generate")
+
+#                 pose_transfer_gen_button.click(fn=leffa_predict_pt, inputs=[
+#                     pt_src_image, pt_ref_image], outputs=[pt_gen_image])
+
+#         gr.Markdown(note)
+
+#         demo.launch(share=True, server_port=7860)
+
+def create_demo():
+    title = "## Virtual Try-on with URLs"
+    description = "Enter URLs for both the person image and the garment image to generate a virtual try-on result."
+    
+    with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink)) as demo:
         gr.Markdown(title)
-        gr.Markdown(link)
         gr.Markdown(description)
+        
+        with gr.Row():
+            with gr.Column():
+                person_url = gr.Textbox(
+                    label="Person Image URL",
+                    placeholder="Enter URL of the person image..."
+                )
+                garment_url = gr.Textbox(
+                    label="Garment Image URL",
+                    placeholder="Enter URL of the garment image..."
+                )
+                
+                # Example URLs
+                gr.Examples(
+                    inputs=[person_url, garment_url],
+                    examples=[
+                        ["https://example.com/person1.jpg", "https://example.com/garment1.jpg"],
+                        ["https://example.com/person2.jpg", "https://example.com/garment2.jpg"],
+                    ],
+                    label="Example URLs"
+                )
+                
+                generate_btn = gr.Button("Generate Try-on")
+            
+            with gr.Column():
+                output_image = gr.Image(
+                    label="Generated Result",
+                    width=512,
+                    height=512
+                )
+        
+        generate_btn.click(
+            fn=virtual_tryon_from_urls,
+            inputs=[person_url, garment_url],
+            outputs=output_image
+        )
+        
+        gr.Markdown("Note: This model is trained solely on academic datasets (VITON-HD).")
+    
+    return demo
 
-        with gr.Tab("Control Appearance (Virtual Try-on)"):
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("#### Person Image")
-                    vt_src_image = gr.Image(
-                        sources=["upload"],
-                        type="filepath",
-                        label="Person Image",
-                        width=512,
-                        height=512,
-                    )
-
-                    gr.Examples(
-                        inputs=vt_src_image,
-                        examples_per_page=5,
-                        examples=["./ckpts/examples/person1/01350_00.jpg",
-                                  "./ckpts/examples/person1/01376_00.jpg",
-                                  "./ckpts/examples/person1/01416_00.jpg",
-                                  "./ckpts/examples/person1/05976_00.jpg",
-                                  "./ckpts/examples/person1/06094_00.jpg",],
-                    )
-
-                with gr.Column():
-                    gr.Markdown("#### Garment Image")
-                    vt_ref_image = gr.Image(
-                        sources=["upload"],
-                        type="filepath",
-                        label="Garment Image",
-                        width=512,
-                        height=512,
-                    )
-
-                    gr.Examples(
-                        inputs=vt_ref_image,
-                        examples_per_page=5,
-                        examples=["./ckpts/examples/garment/01449_00.jpg",
-                                  "./ckpts/examples/garment/01486_00.jpg",
-                                  "./ckpts/examples/garment/01853_00.jpg",
-                                  "./ckpts/examples/garment/02070_00.jpg",
-                                  "./ckpts/examples/garment/03553_00.jpg",],
-                    )
-
-                with gr.Column():
-                    gr.Markdown("#### Generated Image")
-                    vt_gen_image = gr.Image(
-                        label="Generated Image",
-                        width=512,
-                        height=512,
-                    )
-
-                    with gr.Row():
-                        vt_gen_button = gr.Button("Generate")
-
-                vt_gen_button.click(fn=leffa_predict_vt, inputs=[
-                    vt_src_image, vt_ref_image], outputs=[vt_gen_image])
-
-        with gr.Tab("Control Pose (Pose Transfer)"):
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("#### Person Image")
-                    pt_ref_image = gr.Image(
-                        sources=["upload"],
-                        type="filepath",
-                        label="Person Image",
-                        width=512,
-                        height=512,
-                    )
-
-                    gr.Examples(
-                        inputs=pt_ref_image,
-                        examples_per_page=5,
-                        examples=["./ckpts/examples/person1/01350_00.jpg",
-                                  "./ckpts/examples/person1/01376_00.jpg",
-                                  "./ckpts/examples/person1/01416_00.jpg",
-                                  "./ckpts/examples/person1/05976_00.jpg",
-                                  "./ckpts/examples/person1/06094_00.jpg",],
-                    )
-
-                with gr.Column():
-                    gr.Markdown("#### Target Pose Person Image")
-                    pt_src_image = gr.Image(
-                        sources=["upload"],
-                        type="filepath",
-                        label="Target Pose Person Image",
-                        width=512,
-                        height=512,
-                    )
-
-                    gr.Examples(
-                        inputs=pt_src_image,
-                        examples_per_page=5,
-                        examples=["./ckpts/examples/person2/01850_00.jpg",
-                                  "./ckpts/examples/person2/01875_00.jpg",
-                                  "./ckpts/examples/person2/02532_00.jpg",
-                                  "./ckpts/examples/person2/02902_00.jpg",
-                                  "./ckpts/examples/person2/05346_00.jpg",],
-                    )
-
-                with gr.Column():
-                    gr.Markdown("#### Generated Image")
-                    pt_gen_image = gr.Image(
-                        label="Generated Image",
-                        width=512,
-                        height=512,
-                    )
-
-                    with gr.Row():
-                        pose_transfer_gen_button = gr.Button("Generate")
-
-                pose_transfer_gen_button.click(fn=leffa_predict_pt, inputs=[
-                    pt_src_image, pt_ref_image], outputs=[pt_gen_image])
-
-        gr.Markdown(note)
-
-        demo.launch(share=True, server_port=7860)
+if __name__ == "__main__":
+    demo = create_demo()
+    demo.launch(share=True, server_port=7860)