Update app.py

app.py

@@ -6,7 +6,7 @@ import spaces
 import torch
 import random
 from PIL import Image
-from typing import Iterable
+from typing import Iterable, List, Dict
 
 from diffusers import Flux2KleinPipeline
 from diffusers.utils import load_image
@@ -61,21 +61,18 @@ class OrangeRedTheme(Soft):
 orange_red_theme = OrangeRedTheme()
 MAX_SEED = np.iinfo(np.int32).max
 
-ADAPTER_SPECS = {
-    "Klein-Delight-Style": {
+# Defined LoRA List for Gallery
+loras = [
+    {
+        "image": "examples/2.jpg", 
+        "title": "Klein Delight",
         "repo": "linoyts/Flux2-Klein-Delight-LoRA",
         "weights": "pytorch_lora_weights.safetensors",
         "adapter_name": "klein-delight"
     },
-}
-
-GALLERY_ITEMS = [
-    ("examples/1.jpg", "None"),
-    ("examples/2.jpg", "Klein-Delight-Style")
+    # Add more LoRAs here in the future
 ]
 
-GALLERY_KEYS = ["None", "Klein-Delight-Style"]
-
 LOADED_ADAPTERS = set()
 
 print("Loading FLUX.2 Klein 9B model base...")
@@ -101,11 +98,20 @@ def update_dimensions_on_upload(image):
     
     return new_width, new_height
 
+def update_selection(evt: gr.SelectData):
+    if evt.index is not None and evt.index < len(loras):
+        selected_lora = loras[evt.index]
+        return evt.index, f"Selected Style: {selected_lora['title']}"
+    return None, "Selected Style: None (Base Model)"
+
+def clear_selection():
+    return None, "Selected Style: None (Base Model)", gr.Gallery(selected_index=None)
+
 @spaces.GPU
 def infer(
     input_image, 
     prompt, 
-    lora_adapter, 
+    selected_index, 
     seed=42, 
     randomize_seed=True, 
     guidance_scale=1.0, 
@@ -118,16 +124,18 @@ def infer(
     if not input_image:
         raise gr.Error("Please upload an image to apply a style to.")
 
-    if lora_adapter == "None" or lora_adapter is None:
-        print("Selection is None. Disabling LoRA adapters.")
+    # LoRA Loading Logic
+    if selected_index is None:
+        print("No selection or cleared. Disabling LoRA adapters.")
         pipe.disable_lora()
     else:
-        spec = ADAPTER_SPECS.get(lora_adapter)
-        if spec:
+        # Get spec from list using index
+        if 0 <= selected_index < len(loras):
+            spec = loras[selected_index]
             adapter_name = spec["adapter_name"]
             
             if adapter_name not in LOADED_ADAPTERS:
-                print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
+                print(f"--- Downloading and Loading Adapter: {spec['title']} ---")
                 try:
                     pipe.load_lora_weights(
                         spec["repo"], 
@@ -136,14 +144,14 @@ def infer(
                     )
                     LOADED_ADAPTERS.add(adapter_name)
                 except Exception as e:
-                    raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
+                    raise gr.Error(f"Failed to load adapter {spec['title']}: {e}")
             else:
-                print(f"--- Adapter {lora_adapter} is already loaded. ---")
+                print(f"--- Adapter {spec['title']} is already loaded. ---")
                 
             print(f"Activating LoRA: {adapter_name}")
             pipe.set_adapters([adapter_name], adapter_weights=[1.0])
         else:
-            print("Adapter not found in specs. Disabling LoRA.")
+            print("Invalid selection index. Disabling LoRA.")
             pipe.disable_lora()
 
     if randomize_seed:
@@ -172,14 +180,30 @@ def infer(
         torch.cuda.empty_cache()
 
 @spaces.GPU
-def infer_example(input_image, prompt, lora_adapter):
+def infer_example(input_image, prompt, selected_index_dummy):
+    # This wrapper allows examples to run. 
+    # Note: Examples usually pass static values. Logic adapted to assume specific style for examples if needed.
+    # For now, we run base model (None) for examples unless specific index provided.
     if input_image is None: 
         return None, 0
-        
+    
+    # Map string name from example to index if necessary, or just use 0 (Delight)
+    # For the example 1 (Delight), index is 0. For example 2 (Polaroid), index is None.
+    
+    target_index = None
+    if isinstance(selected_index_dummy, str):
+        # Reverse lookup for examples compatibility
+        for i, item in enumerate(loras):
+            if item["title"] == selected_index_dummy or item["adapter_name"] == selected_index_dummy:
+                target_index = i
+                break
+    else:
+        target_index = selected_index_dummy
+
     image, seed = infer(
         input_image=input_image, 
         prompt=prompt, 
-        lora_adapter=lora_adapter, 
+        selected_index=target_index, 
         seed=0, 
         randomize_seed=True,
         guidance_scale=1.0,
@@ -187,14 +211,10 @@ def infer_example(input_image, prompt, lora_adapter):
     )
     return image, seed
 
-def update_gallery_selection(evt: gr.SelectData):
-    if evt.index < len(GALLERY_KEYS):
-        return GALLERY_KEYS[evt.index]
-    return "None"
-
 css="""
 #col-container { margin: 0 auto; max-width: 960px; }
 #main-title h1 { font-size: 2.2em !important; }
+#gallery-label { font-size: 0.9em; font-weight: bold; margin-bottom: 5px; }
 """
 
 with gr.Blocks() as demo:
@@ -202,6 +222,9 @@ with gr.Blocks() as demo:
         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.")
         
+        # State to store selected LoRA index (Default None)
+        selected_index = gr.State(None)
+
         with gr.Row(equal_height=True):
             with gr.Column():
                 input_image = gr.Image(
@@ -230,31 +253,39 @@ with gr.Blocks() as demo:
                 output_image = gr.Image(label="Output Image", interactive=False, format="png", height=358)
                 used_seed = gr.Textbox(label="Used Seed", interactive=False, visible=False)
 
-                gr.Markdown("### Edit Style Gallery")
-                lora_gallery = gr.Gallery(
-                    label="Choose Editing Style",
-                    value=GALLERY_ITEMS,
-                    selected_index=0,
+                # Replaced Dropdown with Gallery
+                gr.Markdown("Edit Style Gallery", elem_id="gallery-label")
+                gallery = gr.Gallery(
+                    [(item["image"], item["title"]) for item in loras],
+                    label="Edit Style Gallery",
+                    columns=3,
                     allow_preview=False,
-                    columns=2,
-                    height=220
+                    show_label=False,
+                    elem_id="gallery"
                 )
-                # Hidden textbox to store selection for the inference function
-                lora_adapter_state = gr.Textbox(value="None", visible=False)
-
-        # Event: Update hidden state when gallery item is selected
-        lora_gallery.select(
-            fn=update_gallery_selection,
-            inputs=None,
-            outputs=lora_adapter_state
+                selected_info = gr.Markdown("Selected Style: None (Base Model)")
+                clear_btn = gr.Button("Clear Style Selection", size="sm", variant="secondary")
+                    
+        # Update selection when gallery is clicked
+        gallery.select(
+            fn=update_selection,
+            inputs=[],
+            outputs=[selected_index, selected_info]
         )
         
+        # Clear selection logic
+        clear_btn.click(
+            fn=clear_selection,
+            inputs=[],
+            outputs=[selected_index, selected_info, gallery]
+        )
+
         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.", "Klein-Delight-Style"],
-                ["examples/1.jpg", "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed by prithivMLmods preserving realistic texture and details", "None"], 
+                ["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.", "Klein Delight"],
+                ["examples/1.jpg", "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed by prithivMLmods preserving realistic texture and details", None], 
             ],
-            inputs=[input_image, prompt, lora_adapter_state],
+            inputs=[input_image, prompt, selected_index], # Maps string to index in infer_example wrapper
             outputs=[output_image, used_seed],
             fn=infer_example,
             cache_examples=False,
@@ -264,7 +295,7 @@ with gr.Blocks() as demo:
         
     run_button.click(
         fn=infer,
-        inputs=[input_image, prompt, lora_adapter_state, seed, randomize_seed, guidance_scale, steps],
+        inputs=[input_image, prompt, selected_index, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_image, used_seed]
     )