Update app.py

app.py

@@ -1,33 +1,106 @@
 import gradio as gr
 import numpy as np
-
 import spaces
 import torch
 import random
 from PIL import Image
-
 from kontext_pipeline import FluxKontextPipeline
 from diffusers import FluxTransformer2DModel
 from diffusers.utils import load_image
-
 from huggingface_hub import hf_hub_download
 
-
 kontext_path = hf_hub_download(repo_id="diffusers/kontext", filename="kontext.safetensors")
-
 MAX_SEED = np.iinfo(np.int32).max
-
 transformer = FluxTransformer2DModel.from_single_file(kontext_path, torch_dtype=torch.bfloat16)
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")
 
+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
+
 @spaces.GPU
-def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
+def infer(input_images, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        
-    input_image = input_image.convert("RGB")
-    # original_width, original_height = input_image.size
+    
+    # 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]
+    
+    # Filter out None images
+    valid_images = [img for img in input_images if img is not None]
+    
+    if not valid_images:
+        raise gr.Error("Please upload at least one valid image.")
+    
+    # Concatenate images horizontally
+    concatenated_image = concatenate_images(valid_images, "horizontal")
+    
+    if concatenated_image is None:
+        raise gr.Error("Failed to process the input images.")
+    
+    # original_width, original_height = concatenated_image.size
     
     # if original_width >= original_height:
     #     new_width = 1024
@@ -38,15 +111,17 @@ def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5
     #     new_width = int(original_width * (new_height / original_height))
     #     new_width = round(new_width / 64) * 64
     
-    #input_image_resized = input_image.resize((new_width, new_height), Image.LANCZOS)
+    #concatenated_image_resized = concatenated_image.resize((new_width, new_height), Image.LANCZOS)
+    
     image = pipe(
-        image=input_image, 
+        image=concatenated_image, 
         prompt=prompt,
         guidance_scale=guidance_scale,
         # width=new_width,
         # height=new_height,
         generator=torch.Generator().manual_seed(seed),
     ).images[0]
+    
     return image, seed, gr.update(visible=True)
 
 css="""
@@ -59,12 +134,24 @@ css="""
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.1 Kontext [dev]
+        gr.Markdown(f"""# FLUX.1 Kontext [dev] - Multi-Image
+        Upload one or multiple images. 
         """)
-
         with gr.Row():
             with gr.Column():
-                input_image = gr.Image(label="Upload the image for editing", 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",
+                    type="pil"
+                )
+                
+
+                
                 with gr.Row():
                     prompt = gr.Text(
                         label="Prompt",
@@ -74,6 +161,7 @@ with gr.Blocks(css=css) as demo:
                         container=False,
                     )
                     run_button = gr.Button("Run", scale=0)
+                    
                 with gr.Accordion("Advanced Settings", open=False):
             
                     seed = gr.Slider(
@@ -99,17 +187,17 @@ with gr.Blocks(css=css) as demo:
                 reuse_button = gr.Button("Reuse this image", visible=False)
         
         
-
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [input_image, prompt, seed, randomize_seed, guidance_scale],
+        inputs = [input_images, prompt, seed, randomize_seed, guidance_scale],
         outputs = [result, seed, reuse_button]
     )
+    
     reuse_button.click(
-        fn = lambda image: image,
+        fn = lambda image: [image] if image is not None else [],  # Convert single image to list for gallery
         inputs = [result],
-        outputs = [input_image]
+        outputs = [input_images]
     )
 
 demo.launch()