add gradio demo

.gitignore

@@ -49,4 +49,7 @@ coverage.xml
 .Spotlight-V100
 .Trashes
 ehthumbs.db
-Thumbs.db 
+Thumbs.db 
+
+# Gradio cache
+.gradio/

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

README.md

@@ -25,6 +25,7 @@ pip install -r requirements.txt
 
 ## Usage
 
+Run the following command to try on an image:
 ```bash
 python tryon_inference.py \
 --image ./example/person/00008_00.jpg \
@@ -33,9 +34,15 @@ python tryon_inference.py \
 --seed 42
 ```
 
+Run the following command to start a gradio demo:
+```bash
+python app.py
+```
+
+
 ## TODO:
 - [ ] Release the FID score
-- [ ] Add gradio demo
+- [x] Add gradio demo
 - [ ] Release updated weights with better performance
 
 ## Citation

app.py

@@ -0,0 +1,150 @@
+import gradio as gr
+from tryon_inference import run_inference
+import os
+import numpy as np
+from PIL import Image
+import tempfile
+
+def gradio_inference(
+    image_data, 
+    garment, 
+    num_steps=50, 
+    guidance_scale=30.0, 
+    seed=-1, 
+    size=(576,768)
+):
+    """Wrapper function for Gradio interface"""
+    # Use temporary directory
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        # Save inputs to temp directory
+        temp_image = os.path.join(tmp_dir, "image.png")
+        temp_mask = os.path.join(tmp_dir, "mask.png")
+        temp_garment = os.path.join(tmp_dir, "garment.png")
+        
+        # Extract image and mask from ImageEditor data
+        image = image_data["background"]
+        mask = image_data["layers"][0]  # First layer contains the mask
+        
+        # Convert to numpy array and process mask
+        mask_array = np.array(mask)
+        is_black = np.all(mask_array < 10, axis=2)
+        mask = Image.fromarray(((~is_black) * 255).astype(np.uint8))
+        
+        # Save files to temp directory
+        image.save(temp_image)
+        mask.save(temp_mask)
+        garment.save(temp_garment)
+        
+        try:
+            # Run inference
+            _, tryon_result = run_inference(
+                image_path=temp_image,
+                mask_path=temp_mask,
+                garment_path=temp_garment,
+                num_steps=num_steps,
+                guidance_scale=guidance_scale,
+                seed=seed,
+                size=size
+            )
+            return tryon_result
+        except Exception as e:
+            raise gr.Error(f"Error during inference: {str(e)}")
+
+def create_demo():
+    with gr.Blocks() as demo:
+        gr.Markdown("""
+        # CATVTON FLUX Virtual Try-On Demo
+        Upload a model image, an agnostic mask, and a garment image to generate virtual try-on results.
+        """)
+        
+        with gr.Column():
+            with gr.Row():
+                with gr.Column():
+                    image_input = gr.ImageMask(
+                        label="Model Image (Draw mask where garment should go)", 
+                        type="pil",
+                        height=576,
+                    )
+                    gr.Examples(
+                        examples=[
+                            ["./example/person/00008_00.jpg"],
+                            ["./example/person/00055_00.jpg"],
+                            ["./example/person/00057_00.jpg"],
+                            ["./example/person/00067_00.jpg"],
+                            ["./example/person/00069_00.jpg"],
+                        ],
+                        inputs=[image_input],
+                        label="Person Images",
+                    )
+                with gr.Column():
+                    garment_input = gr.Image(label="Garment Image", type="pil", height=576)
+                    gr.Examples(
+                        examples=[
+                            ["./example/garment/04564_00.jpg"],
+                            ["./example/garment/00055_00.jpg"],
+                            ["./example/garment/00057_00.jpg"],
+                            ["./example/garment/00067_00.jpg"],
+                            ["./example/garment/00069_00.jpg"],
+                        ],
+                        inputs=[garment_input],
+                        label="Garment Images",
+                    ) 
+                
+            with gr.Row():
+                num_steps = gr.Slider(
+                    minimum=1, 
+                    maximum=100, 
+                    value=50, 
+                    step=1, 
+                    label="Number of Steps"
+                )
+                guidance_scale = gr.Slider(
+                    minimum=1.0, 
+                    maximum=50.0, 
+                    value=30.0, 
+                    step=0.5, 
+                    label="Guidance Scale"
+                )
+                seed = gr.Slider(
+                    minimum=-1,
+                    maximum=2147483647,
+                    step=1,
+                    value=-1,
+                    label="Seed (-1 for random)"
+                )
+            
+            submit_btn = gr.Button("Generate Try-On", variant="primary")
+            
+            with gr.Column():
+                tryon_output = gr.Image(label="Try-On Result")
+                
+        with gr.Row():
+            gr.Markdown("""
+            ### Notes:
+            - The model image should be a full-body photo
+            - The mask should indicate the region where the garment will be placed
+            - The garment image should be on a clean background
+            """)
+        
+        submit_btn.click(
+            fn=gradio_inference,
+            inputs=[
+                image_input,
+                garment_input,
+                num_steps,
+                guidance_scale,
+                seed
+            ],
+            outputs=[tryon_output],
+            api_name="try-on"
+        )
+    
+    return demo
+
+if __name__ == "__main__":
+    demo = create_demo()
+    demo.queue()  # Enable queuing for multiple users
+    demo.launch(
+        share=True,
+        server_name="0.0.0.0"  # Makes the server accessible from other machines
+    ) 

requirements.txt

@@ -94,5 +94,7 @@ yarl==1.9.4
 zipp==3.20.0
 peft==0.13.2
 bitsandbytes==0.44.1
+gradio==5.6.0
+gradio_client==1.4.3
 prodigyopt
 git+https://github.com/huggingface/diffusers.git

tryon_inference.py

@@ -10,8 +10,6 @@ def run_inference(
     image_path,
     mask_path,
     garment_path,
-    output_garment_path=None,
-    output_tryon_path='flux_inpaint_tryon.png',
     size=(576, 768),
     num_steps=50,
     guidance_scale=30,
@@ -82,9 +80,7 @@ def run_inference(
     garment_result = result.crop((0, 0, width, size[1]))
     tryon_result = result.crop((width, 0, width * 2, size[1]))
 
-    if output_garment_path is not None:
-        garment_result.save(output_garment_path)
-    tryon_result.save(output_tryon_path)
+    
     return garment_result, tryon_result
 
 def main():
@@ -115,6 +111,13 @@ def main():
         seed=args.seed,
         size=(args.width, args.height)
     )
+    output_garment_path=args.output_garment,
+    output_tryon_path=args.output_tryon,
+    
+    if output_garment_path is not None:
+        garment_result.save(output_garment_path)
+    tryon_result.save(output_tryon_path)
+    
     print("Successfully saved garment and try-on images")
 
 if __name__ == "__main__":