enable image url as input

app.py

@@ -22,15 +22,63 @@ def base64_to_image(base64_str, output_path):  # Remove 'self'
     image.save(output_path)
     return image
 
-def run_viton(model_image_path, garment_image_path, 
+def download_image_from_url(url, output_path):
+    """Download image from URL and save to local path"""
+    try:
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()
+        
+        # Save the image
+        with open(output_path, 'wb') as f:
+            f.write(response.content)
+        
+        # Verify it's a valid image
+        image = Image.open(output_path)
+        return output_path
+    except Exception as e:
+        print(f"Error downloading image from {url}: {str(e)}")
+        return None
+
+def url_to_base64(url):
+    """Convert image URL to base64 string"""
+    try:
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()
+        return base64.b64encode(response.content).decode()
+    except Exception as e:
+        print(f"Error converting URL to base64: {str(e)}")
+        return None
+
+def run_viton(model_image_path, garment_image_path, model_url, garment_url,
                 n_steps=20, image_scale=2.0, seed=-1):
     try:
         api_url = os.environ.get("SERVER_URL")
         print(f"Using API URL: {api_url}")  # Add this to debug
         
-        # Convert images to base64 (remove 'self.')
-        model_b64 = image_to_base64(model_image_path)
-        garment_b64 = image_to_base64(garment_image_path)
+        # Determine which inputs to use (file upload or URL)
+        model_b64 = None
+        garment_b64 = None
+        
+        # Handle model image
+        if model_url and model_url.strip():
+            print(f"Using model URL: {model_url}")
+            model_b64 = url_to_base64(model_url.strip())
+        elif model_image_path:
+            print(f"Using model file: {model_image_path}")
+            model_b64 = image_to_base64(model_image_path)
+        
+        # Handle garment image
+        if garment_url and garment_url.strip():
+            print(f"Using garment URL: {garment_url}")
+            garment_b64 = url_to_base64(garment_url.strip())
+        elif garment_image_path:
+            print(f"Using garment file: {garment_image_path}")
+            garment_b64 = image_to_base64(garment_image_path)
+        
+        # Check if we have both images
+        if not model_b64 or not garment_b64:
+            print("Error: Missing model or garment image")
+            return []
 
         # Prepare request
         request_data = {
@@ -73,14 +121,17 @@ def run_viton(model_image_path, garment_image_path,
         return []  # Fix: should return list, not dict for gallery
 
 block = gr.Blocks().queue()
-default_model = os.path.join(example_path, 'model/model_8.png')
-default_garment = os.path.join(example_path, 'garment/00055_00.jpg')
 with block:
     with gr.Row():
         gr.Markdown("# Virtual Try-On")
+    with gr.Row():
+        gr.Markdown("**Instructions:** You can either upload images using the file upload interface or provide direct URLs to images. URL inputs will take priority over uploaded files.")
     with gr.Row():
         with gr.Column():
-            vton_img = gr.Image(label="Model", sources=['upload', 'webcam'], type="filepath", height=384, value=default_model)
+            model_url = gr.Textbox(
+                label="Enter Model Image URL", 
+            )
+            vton_img = gr.Image(label="Model", sources=['upload', 'webcam'], type="filepath", height=384)
             example = gr.Examples(
                 inputs=vton_img,
                 examples_per_page=5,
@@ -92,7 +143,10 @@ with block:
                     os.path.join(example_path, 'model/model_5.png'),
                 ])
         with gr.Column():
-            garm_img = gr.Image(label="Garment", sources=['upload', 'webcam'], type="filepath", height=384, value=default_garment)
+            garment_url = gr.Textbox(
+                label="Enter Garment Image URL", 
+            )
+            garm_img = gr.Image(label="Garment", sources=['upload', 'webcam'], type="filepath", height=384)
             example = gr.Examples(
                 inputs=garm_img,
                 examples_per_page=5,
@@ -111,7 +165,7 @@ with block:
         image_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=5.0, value=2.0, step=0.1)
         seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=-1)
         
-    ips = [vton_img, garm_img, n_steps, image_scale, seed]
+    ips = [vton_img, garm_img, model_url, garment_url, n_steps, image_scale, seed]
     run_button.click(fn=run_viton, inputs=ips, outputs=result_gallery)
 
 block.launch(mcp_server=True)