Create app.py

app.py

@@ -0,0 +1,151 @@
+import gradio as gr
+from gradio_client import Client, handle_file
+import re
+import time
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Get Hugging Face token from environment variable
+hf_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
+
+# Initialize client with auth
+client = Client(
+    "levihsu/OOTDiffusion",
+    hf_token=hf_token
+)
+
+
+def generate_outfit(model_image, garment_image, n_samples=1, n_steps=20, image_scale=2, seed=-1):
+    if model_image is None or garment_image is None:
+        return None, "Please upload both model and garment images"
+        
+    max_retries = 3
+    for attempt in range(max_retries):
+        try:
+            # Use the client to predict
+            result = client.predict(
+                vton_img=handle_file(model_image),
+                garm_img=handle_file(garment_image),
+                n_samples=n_samples,
+                n_steps=n_steps,
+                image_scale=image_scale,
+                seed=seed,
+                api_name="/process_hd"
+            )
+            
+            # If result is a list, get the first item
+            if isinstance(result, list):
+                result = result[0]
+            
+            # If result is a dictionary, try to get the image path
+            if isinstance(result, dict):
+                if 'image' in result:
+                    return result['image'], None
+                else:
+                    return None, "API returned unexpected format"
+                
+            return result, None
+            
+        except Exception as e:
+            error_msg = str(e)
+            if "exceeded your GPU quota" in error_msg:
+                wait_time_match = re.search(r'retry in (\d+:\d+:\d+)', error_msg)
+                wait_time = wait_time_match.group(1) if wait_time_match else "60:00"  # Default to 1 hour
+                wait_seconds = sum(int(x) * 60 ** i for i, x in enumerate(reversed(wait_time.split(':'))))  # Convert wait time to seconds
+                if attempt < max_retries - 1:
+                    time.sleep(wait_seconds)  # Wait before retrying
+                return None, f"GPU quota exceeded. Please wait {wait_time} before trying again."
+            else:
+                return None, f"Error: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    ## Outfit Diffusion - Try On Virtual Outfits
+
+    ⚠️ **Note**: This demo uses free GPU quota which is limited. To avoid errors:
+    - Use lower values for Steps (10-15) and Scale (1-2)
+    - Wait between attempts if you get a quota error
+    - Sign up for a Hugging Face account for more quota
+
+  
+    """)
+
+    with gr.Row():
+        with gr.Column():
+            model_image = gr.Image(
+                label="Upload Model Image (person wearing clothes)", 
+                type="filepath",
+                height=300
+                
+            )
+            model_examples = [
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/ba5ba7978e7302e8ab5eb733cc7221394c4e6faf/model_5.png",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/40dade4a04a827c0fdf63c6c70b42ef26480f391/01861_00.jpg",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/3c4639c5fab3cdcd3239609dca5afee7b0677286/model_6.png",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/0089171df270f4532eec3d80a8f36cc8218c6840/01008_00.jpg"
+            ]
+            gr.Examples(examples=model_examples, inputs=model_image)
+
+            garment_image = gr.Image(
+                label="Upload Garment Image (clothing item)", 
+                type="filepath",
+                height=300
+            )
+            garment_examples = [
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/180d4e2a1139071a8685a5edee7ab24bcf1639f5/03244_00.jpg",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/584dda2c5ee1d8271a6cd06225c07db89c79ca03/04825_00.jpg",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/a51938ec99f13e548d365a9ca6d794b6fe7462af/049949_1.jpg",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/2d64241101189251ce415df84dc9205cda9a36ca/03032_00.jpg",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/44aee6b576cae51eeb979311306375b56b7e0d8b/02305_00.jpg",
+                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/578dfa869dedb649e91eccbe566fc76435bb6bbe/049920_1.jpg"
+            ]
+            gr.Examples(examples=garment_examples, inputs=garment_image)
+
+        
+        with gr.Column():
+            output_image = gr.Image(label="Generated Output")
+            error_text = gr.Markdown()  # Add error display
+    
+    with gr.Row():
+        with gr.Column():
+            n_samples = gr.Slider(
+                label="Number of Samples", 
+                minimum=1, 
+                maximum=5, 
+                step=1, 
+                value=1
+            )
+            n_steps = gr.Slider(
+                label="Steps (lower = faster, try 10-15)", 
+                minimum=1, 
+                maximum=50, 
+                step=1, 
+                value=10  # Reduced default
+            )
+            image_scale = gr.Slider(
+                label="Scale (lower = faster, try 1-2)", 
+                minimum=1, 
+                maximum=5, 
+                step=1, 
+                value=1  # Reduced default
+            )
+            seed = gr.Number(
+                label="Random Seed (-1 for random)", 
+                value=-1
+            )
+    
+    generate_button = gr.Button("Generate Outfit")
+
+    # Set up the action for the button
+    generate_button.click(
+        fn=generate_outfit,
+        inputs=[model_image, garment_image, n_samples, n_steps, image_scale, seed],
+        outputs=[output_image, error_text]
+    )
+
+# Launch the app
+demo.launch()