add AI generated garment

app.py

@@ -49,8 +49,35 @@ def url_to_base64(url):
         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):
+def run_viton(model_image_path: str = None, 
+              garment_image_path: str = None, 
+              model_url: str = None, 
+              garment_url: str = None,
+              n_steps=20, 
+              image_scale=2.0, 
+              seed=-1
+            ):
+    """
+    Run the Virtual Try-On model with provided images path or URLs.
+
+    Args:
+        model_image_path (str): Path to the model image file.
+        garment_image_path (str): Path to the garment image file.
+        model_url (str): URL of the model image.
+        garment_url (str): URL of the garment image.
+        n_steps (int): Number of steps for the model.
+        image_scale (float): Scale for the generated images.
+        seed (int): Random seed for reproducibility.
+
+    Returns:
+        list: List of generated images in base64 format.
+    """
+    if not model_image_path and not model_url:
+        print("Error: No model image provided")
+        return []
+    if not garment_image_path and not garment_url:
+        print("Error: No garment image provided")
+        return []
     try:
         api_url = os.environ.get("SERVER_URL")
         print(f"Using API URL: {api_url}")  # Add this to debug
@@ -110,6 +137,94 @@ def run_viton(model_image_path, garment_image_path, model_url, garment_url,
                 img = base64_to_image(img_b64, output_path)  # Remove 'self.'
                 generated_images.append(img)
             
+            print(f"Successfully generated {len(generated_images)} images")
+            return generated_images
+        else:
+            print(f"Request failed with status code: {response.status_code}")
+            return []  
+            
+    except Exception as e:
+        print(f"Exception occurred: {str(e)}")  # Add this
+        return []  
+    
+def run_new_garment(model_image_path: str = None, 
+              garment_prompt: str = None, 
+              model_url: str = None, 
+              n_steps=20, 
+              image_scale=2.0, 
+              seed=-1
+            ):
+    """
+    Run the Virtual Try-On model with provided model image and garment image generated using FLUX.1-dev
+    based on description of the garment obtained from the user
+
+    Args:
+        model_image_path (str): Path to the model image file.
+        garment_prompt (str): Description of the garment.
+        model_url (str): URL of the model image.
+        n_steps (int): Number of steps for the model.
+        image_scale (float): Scale for the generated images.
+        seed (int): Random seed for reproducibility.
+
+    Returns:
+        list: List of generated images in base64 format.
+    """
+    if not model_image_path and not model_url:
+        print("Error: No model image provided")
+        return []
+    if not garment_prompt or not garment_prompt.strip():
+        print("Error: No garment description provided")
+        return []
+    try:
+        api_url = os.environ.get("SERVER_URL")
+        print(f"Using API URL: {api_url}")  # Add this to debug
+        
+        # Determine which inputs to use (file upload or URL)
+        model_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)
+        
+        # Check if we have both images
+        if not model_b64 or not garment_prompt:
+            print("Error: Missing model or garment description")
+            return []
+
+        # Prepare request
+        request_data = {
+            "model_image_base64": model_b64,
+            "garment_prompt": garment_prompt.strip(),
+            "n_samples": 1,
+            "n_steps": n_steps,
+            "image_scale": image_scale,
+            "seed": seed
+        }
+                
+        # Send request
+        response = requests.post(f"{api_url}/new-garment", 
+                                json=request_data, 
+                                timeout=300)
+        
+        print(f"Request sent to {api_url}/new-garment")
+        print(f"Response status code: {response.status_code}")
+        
+        if response.status_code == 200:
+            result = response.json()
+            if result.get("error"):
+                print(f"Error: {result['error']}")
+                return []
+            
+            generated_images = []
+            for i, img_b64 in enumerate(result.get("images_base64", [])):
+                output_path = f"flux_output_{i}.png"
+                img = base64_to_image(img_b64, output_path)  # Remove 'self.'
+                generated_images.append(img)
+            
             print(f"Successfully generated {len(generated_images)} images")
             return generated_images
         else:
@@ -124,48 +239,53 @@ block = gr.Blocks().queue()
 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():
+            gr.Markdown("### Provide image or URL of upper body photo")
             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,
+                examples_per_page=4,
                 examples=[
-                    os.path.join(example_path, 'model/model_8.png'),
                     os.path.join(example_path, 'model/model_2.png'),
                     os.path.join(example_path, 'model/model_7.png'),
                     os.path.join(example_path, 'model/model_4.png'),
                     os.path.join(example_path, 'model/model_5.png'),
                 ])
         with gr.Column():
+            gr.Markdown("### Provide image, URL or description of a garment")
             garment_url = gr.Textbox(
-                label="Enter Garment Image URL", 
+                label="Enter Garment Image URL",
+            )
+            garment_promt = gr.Textbox(
+                label="Describe Garment",
             )
             garm_img = gr.Image(label="Garment", sources=['upload', 'webcam'], type="filepath", height=384)
             example = gr.Examples(
                 inputs=garm_img,
-                examples_per_page=5,
+                examples_per_page=4,
                 examples=[
-                    os.path.join(example_path, 'garment/00055_00.jpg'),
                     os.path.join(example_path, 'garment/07764_00.jpg'),
                     os.path.join(example_path, 'garment/03032_00.jpg'),
                     os.path.join(example_path, 'garment/048554_1.jpg'),
                     os.path.join(example_path, 'garment/049805_1.jpg'),
                 ])
         with gr.Column():
-            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True, scale=1)   
+            gr.Markdown("### 2D Result")
+            result_gallery = gr.Gallery(label='Output 2D', show_label=False, elem_id="gallery", preview=True, scale=1)   
     with gr.Column():
-        run_button = gr.Button(value="Run")
+        run_button = gr.Button(value="Try On with your garment")
+        run_button2 = gr.Button(value="Try On with AI generated garment")
         n_steps = gr.Slider(label="Steps", minimum=20, maximum=40, value=20, step=1)
         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, model_url, garment_url, n_steps, image_scale, seed]
-    run_button.click(fn=run_viton, inputs=ips, outputs=result_gallery)
+    ips1 = [vton_img, garm_img, model_url, garment_url, n_steps, image_scale, seed]
+    run_button.click(fn=run_viton, inputs=ips1, outputs=result_gallery)
+    ips2 = [vton_img, garment_promt, model_url, n_steps, image_scale, seed]
+    run_button2.click(fn=run_new_garment, inputs=ips2, outputs=result_gallery)
 
 block.launch(mcp_server=True)