update app

app.py

@@ -103,6 +103,31 @@ def update_dimensions_on_upload(image):
     
     return new_width, new_height
 
+def process_gallery_images(images):
+    """Process images from gallery input and return list of PIL images."""
+    if not images:
+        return []
+    
+    pil_images = []
+    for item in images:
+        try:
+            if isinstance(item, tuple) or isinstance(item, list):
+                path_or_img = item[0]
+            else:
+                path_or_img = item
+
+            if isinstance(path_or_img, str):
+                pil_images.append(Image.open(path_or_img).convert("RGB"))
+            elif isinstance(path_or_img, Image.Image):
+                pil_images.append(path_or_img.convert("RGB"))
+            else:
+                pil_images.append(Image.open(path_or_img.name).convert("RGB"))
+        except Exception as e:
+            print(f"Skipping invalid image item: {e}")
+            continue
+    
+    return pil_images
+
 def update_style_selection(evt: gr.SelectData):
     """Update selected style based on gallery click"""
     selected_style = LORA_STYLES[evt.index]
@@ -116,7 +141,7 @@ def update_style_selection(evt: gr.SelectData):
 
 @spaces.GPU
 def infer(
-    input_image, 
+    input_images, 
     prompt, 
     selected_style_index,
     seed=42, 
@@ -128,8 +153,14 @@ def infer(
     gc.collect()
     torch.cuda.empty_cache()
 
-    if not input_image:
-        raise gr.Error("Please upload an image to apply a style to.")
+    if not input_images:
+        raise gr.Error("Please upload at least one image to apply a style to.")
+
+    # Process gallery images
+    pil_images = process_gallery_images(input_images)
+    
+    if not pil_images:
+        raise gr.Error("Could not process uploaded images.")
 
     if selected_style_index is None:
         selected_style_index = 0
@@ -162,12 +193,21 @@ def infer(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    width, height = update_dimensions_on_upload(input_image)
-    processed_input = input_image.resize((width, height), Image.LANCZOS).convert("RGB")
+    # Get dimensions from first image
+    width, height = update_dimensions_on_upload(pil_images[0])
+    
+    # Process all images to the same dimensions
+    processed_images = [
+        img.resize((width, height), Image.LANCZOS).convert("RGB") 
+        for img in pil_images
+    ]
     
     try:
+        # Pass single image or list based on count
+        image_input = processed_images if len(processed_images) > 1 else processed_images[0]
+        
         image = pipe(
-            image=processed_input, 
+            image=image_input, 
             prompt=prompt,
             guidance_scale=guidance_scale,
             width=width,
@@ -185,12 +225,21 @@ def infer(
         torch.cuda.empty_cache()
 
 @spaces.GPU
-def infer_example(input_image, prompt, style_index):
-    if input_image is None: 
+def infer_example(input_images, prompt, style_index):
+    if not input_images: 
+        return None, 0
+    
+    # Handle both single image path and list of paths for examples
+    if isinstance(input_images, str):
+        input_images = [input_images]
+    elif isinstance(input_images, list) and len(input_images) > 0:
+        # If it's already a list, keep it as is
+        pass
+    else:
         return None, 0
         
     image, seed = infer(
-        input_image=input_image, 
+        input_images=input_images, 
         prompt=prompt, 
         selected_style_index=style_index, 
         seed=0, 
@@ -200,26 +249,30 @@ def infer_example(input_image, prompt, style_index):
     )
     return image, seed
 
-css="""
+css = """
 #col-container { margin: 0 auto; max-width: 960px; }
 #main-title h1 { font-size: 2.2em !important; }
-#style_gallery .grid-wrap{height: 10vh}
+#style_gallery .grid-wrap { height: 10vh }
+#input_gallery .grid-wrap { min-height: 200px }
 """
 
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# **FLUX.2-Klein-LoRA-Studio**", elem_id="main-title")
-        gr.Markdown("Perform diverse image edits using specialized [LoRAs](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.2-klein-9B) adapters for the [FLUX.2-Klein-Distilled](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) model.")
+        gr.Markdown("Perform diverse image edits using specialized [LoRAs](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.2-klein-9B) adapters for the [FLUX.2-Klein-Distilled](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) model. Supports multiple image inputs for style transfer and reference-based editing.")
         
         selected_style_index = gr.State(0)
         
         with gr.Row(equal_height=True):
             with gr.Column():
-                input_image = gr.Image(
-                    label="Upload Image", 
-                    type="pil", 
-                    height=290, 
-                    sources=["upload", "webcam", "clipboard"]
+                input_images = gr.Gallery(
+                    label="Upload Images", 
+                    type="filepath", 
+                    columns=2, 
+                    rows=1, 
+                    height=290,
+                    allow_preview=True,
+                    elem_id="input_gallery"
                 )
                 with gr.Row():
                     prompt = gr.Text(
@@ -253,13 +306,15 @@ with gr.Blocks() as demo:
                     
         gr.Examples(
             examples=[
-                ["examples/2.jpg", "Relight the image to remove all existing lighting conditions and replace them with neutral, uniform illumination. Apply soft, evenly distributed lighting with no directional shadows, no harsh highlights, and no dramatic contrast. Maintain the original identity of all subjects exactly—preserve facial structure, skin tone, proportions, expressions, hair, clothing, and textures. Do not alter pose, camera angle, background geometry, or image composition. Lighting should appear balanced, and studio-neutral, similar to diffuse overcast or a soft lightbox setup. Ensure consistent exposure across the entire image with realistic depth and subtle shading only where necessary for form."],
-                ["examples/1.jpg", "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed by prithivMLmods preserving realistic texture and details"], 
+                [["examples/2.jpg"], "Relight the image to remove all existing lighting conditions and replace them with neutral, uniform illumination. Apply soft, evenly distributed lighting with no directional shadows, no harsh highlights, and no dramatic contrast. Maintain the original identity of all subjects exactly—preserve facial structure, skin tone, proportions, expressions, hair, clothing, and textures. Do not alter pose, camera angle, background geometry, or image composition. Lighting should appear balanced, and studio-neutral, similar to diffuse overcast or a soft lightbox setup. Ensure consistent exposure across the entire image with realistic depth and subtle shading only where necessary for form."],
+                [["examples/1.jpg"], "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed by prithivMLmods preserving realistic texture and details"],
+                [["examples/1.jpg", "examples/2.jpg"], "Apply the style from image 2 to image 1, preserving the subject's identity."],
             ],
-            inputs=[input_image, prompt],
+            inputs=[input_images, prompt],
             outputs=[output_image, used_seed],
             fn=infer_example,
             cache_examples=False,
+            label="Examples"
         )
         
         gr.Markdown("[*](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B)This is still an experimental Space for FLUX.2-Klein-9B. More adapters will be added soon.")
@@ -271,7 +326,7 @@ with gr.Blocks() as demo:
     
     run_button.click(
         fn=infer,
-        inputs=[input_image, prompt, selected_style_index, seed, randomize_seed, guidance_scale, steps],
+        inputs=[input_images, prompt, selected_style_index, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_image, used_seed]
     )