Update app.py

app.py

@@ -180,10 +180,71 @@ except Exception as e:
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max
 
+
+def concatenate_images(images, direction="horizontal"):
+    """
+    Concatenate multiple PIL images either horizontally or vertically.
+    
+    Args:
+        images: List of PIL Images
+        direction: "horizontal" or "vertical"
+    
+    Returns:
+        PIL Image: Concatenated image
+    """
+    if not images:
+        return None
+    
+    # Filter out None images
+    valid_images = [img for img in images if img is not None]
+    
+    if not valid_images:
+        return None
+    
+    if len(valid_images) == 1:
+        return valid_images[0].convert("RGB")
+    
+    # Convert all images to RGB
+    valid_images = [img.convert("RGB") for img in valid_images]
+    
+    if direction == "horizontal":
+        # Calculate total width and max height
+        total_width = sum(img.width for img in valid_images)
+        max_height = max(img.height for img in valid_images)
+        
+        # Create new image
+        concatenated = Image.new('RGB', (total_width, max_height), (255, 255, 255))
+        
+        # Paste images
+        x_offset = 0
+        for img in valid_images:
+            # Center image vertically if heights differ
+            y_offset = (max_height - img.height) // 2
+            concatenated.paste(img, (x_offset, y_offset))
+            x_offset += img.width
+            
+    else:  # vertical
+        # Calculate max width and total height
+        max_width = max(img.width for img in valid_images)
+        total_height = sum(img.height for img in valid_images)
+        
+        # Create new image
+        concatenated = Image.new('RGB', (max_width, total_height), (255, 255, 255))
+        
+        # Paste images
+        y_offset = 0
+        for img in valid_images:
+            # Center image horizontally if widths differ
+            x_offset = (max_width - img.width) // 2
+            concatenated.paste(img, (x_offset, y_offset))
+            y_offset += img.height
+    
+    return concatenated
+
 # --- Main Inference Function ---
 @spaces.GPU(duration=60)
 def infer(
-    image,
+    input_images,
     prompt,
     seed=42,
     randomize_seed=False,
@@ -203,6 +264,20 @@ def infer(
 
     # Set up the generator for reproducibility
     generator = torch.Generator(device=device).manual_seed(seed)
+
+    # Handle input_images - it could be a single image or a list of images
+    if input_images is None:
+        raise gr.Error("Please upload at least one image.")
+    
+    # If it's a single image (not a list), convert to list
+    if not isinstance(input_images, list):
+        input_images = [input_images]
+    
+    # Concatenate images horizontally
+    concatenated_image = concatenate_images(input_images, "horizontal")
+    
+    if concatenated_image is None:
+        raise gr.Error("Failed to process the input images.")
     
     print(f"Original prompt: '{prompt}'")
     print(f"Negative Prompt: '{negative_prompt}'")
@@ -215,7 +290,7 @@ def infer(
     # Generate the edited image - always generate just 1 image
     try:
         images = pipe(
-            image,
+            concatenated_image,
             prompt=prompt,
             negative_prompt=negative_prompt,
             num_inference_steps=num_inference_steps,
@@ -268,10 +343,16 @@ with gr.Blocks(css=css) as demo:
         
         with gr.Row():
             with gr.Column():
-                input_image = gr.Image(
-                    label="Input Image", 
-                    show_label=True, 
-                    type="pil"
+                input_images = gr.Gallery(
+                    label="Upload image(s) for editing", 
+                    show_label=True,
+                    elem_id="gallery_input",
+                    columns=3,
+                    rows=2,
+                    object_fit="contain",
+                    height="auto",
+                    file_types=['image'],
+                    type='pil'
                 )
             # Changed from Gallery to Image
             result = gr.Image(
@@ -329,7 +410,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
-            input_image,
+            input_images,
             prompt,
             seed,
             randomize_seed,