fix gallery bug

app.py

@@ -244,7 +244,7 @@ class ClothingInpainter:
 
         return images_output
 
-def process_segmentation(image, dilation_iterations=2):
+def process_segmentation(image, clothing_parts=None, prompt=None, dilation_iterations=2):
     try:
         if image is None:
             raise gr.Error("Please upload an image")
@@ -260,13 +260,25 @@ def process_segmentation(image, dilation_iterations=2):
         if not all_masks:
             logger.error("No clothing detected in the image")
             raise gr.Error("No clothing detected in the image. Please try a different image.")
+            
         inpainter.last_mask = all_masks
-        # Only show main clothing parts for selection
+        
+        # If clothing parts are selected and prompt is provided, generate variations
+        if clothing_parts and prompt:
+            # Convert selected_parts to lowercase/dash format
+            selected_parts = [p.lower() for p in clothing_parts]
+            prompt_dict = {'pos': prompt}
+            
+            # Generate inpainted images
+            images = inpainter.inpaint(prompt_dict, image, selected_parts, dilation_iterations)
+            return images
+            
+        # Otherwise, just show the segmentation visualization
         main_parts = ['upper-clothes', 'dress', 'coat', 'pants', 'skirt']
         masks = {k: v for k, v in all_masks.items() if k in main_parts}
-        vis_image = inpainter.visualize_segmentation(image, masks, selected_parts=None)
-        detected_parts = [k for k in masks.keys()]
-        return vis_image, gr.update(choices=detected_parts, value=[])
+        vis_image = inpainter.visualize_segmentation(image, masks, selected_parts=clothing_parts)
+        return [vis_image]
+        
     except gr.Error as e:
         raise e
     except Exception as e:
@@ -360,9 +372,10 @@ def create_interface():
                 # Input Section
                 with gr.Group():
                     gr.Markdown("### Input Image")
+                    input_image = gr.Image(type="pil", label="Upload Image")
                     with gr.Row():
-                        input_image = gr.Image(type="pil", label="Upload Image")
-                        example_btn = gr.Button("Load Example Image", variant="secondary")
+                        example_btn = gr.Button("Load Example Image", variant="secondary", size="sm")
+                        gr.Markdown("or upload your own image above", size="sm")
                     
                     # Clothing Selection
                     gr.Markdown("### Select Clothing Parts")