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, Any
 
 from diffusers import Flux2KleinPipeline
 from diffusers.utils import load_image
@@ -61,13 +61,24 @@ class OrangeRedTheme(Soft):
 orange_red_theme = OrangeRedTheme()
 MAX_SEED = np.iinfo(np.int32).max
 
-ADAPTER_SPECS = {
-    "Klein-Delight-Style": {
+loras_list = [
+    {
+        "image": "https://huggingface.co/spaces/stabilityai/stable-diffusion/resolve/main/assets/0.png", # Placeholder for None
+        "title": "None (Base Model)",
+        "repo": None,
+        "weights": None,
+        "adapter_name": None,
+        "description": "Original FLUX.2 Klein model without modifications."
+    },
+    {
+        "image": "https://huggingface.co/linoyts/Flux2-Klein-Delight-LoRA/resolve/main/images/example_image.jpg", # Placeholder/Real if available
+        "title": "Klein Delight Style",
         "repo": "linoyts/Flux2-Klein-Delight-LoRA",
         "weights": "pytorch_lora_weights.safetensors",
-        "adapter_name": "klein-delight"
+        "adapter_name": "klein-delight",
+        "description": "Removes existing lighting and applies neutral, studio-like illumination."
     },
-}
+]
 
 LOADED_ADAPTERS = set()
 
@@ -94,11 +105,17 @@ def update_dimensions_on_upload(image):
     
     return new_width, new_height
 
+def update_selection(evt: gr.SelectData):
+    """Updates the UI when a gallery item is selected."""
+    selected_lora = loras_list[evt.index]
+    info_text = f"### Selected: {selected_lora['title']}\n{selected_lora['description']}"
+    return evt.index, info_text
+
 @spaces.GPU
 def infer(
     input_image, 
     prompt, 
-    lora_adapter, 
+    selected_index, 
     seed=42, 
     randomize_seed=True, 
     guidance_scale=1.0, 
@@ -110,35 +127,41 @@ def infer(
 
     if not input_image:
         raise gr.Error("Please upload an image to apply a style to.")
+    
+    # Ensure index is valid
+    if selected_index is None:
+        selected_index = 0
+    
+    selected_lora = loras_list[selected_index]
+    repo = selected_lora["repo"]
 
-    if lora_adapter == "None":
-        print("Selection is None. Disabling LoRA adapters.")
+    # --- LoRA Loading Logic ---
+    if repo is None:
+        print("Selection is None (Base). Disabling LoRA adapters.")
         pipe.disable_lora()
     else:
-        spec = ADAPTER_SPECS.get(lora_adapter)
-        if spec:
-            adapter_name = spec["adapter_name"]
-            
-            if adapter_name not in LOADED_ADAPTERS:
-                print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
-                try:
-                    pipe.load_lora_weights(
-                        spec["repo"], 
-                        weight_name=spec["weights"], 
-                        adapter_name=adapter_name
-                    )
-                    LOADED_ADAPTERS.add(adapter_name)
-                except Exception as e:
-                    raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
-            else:
-                print(f"--- Adapter {lora_adapter} is already loaded. ---")
-                
-            print(f"Activating LoRA: {adapter_name}")
-            pipe.set_adapters([adapter_name], adapter_weights=[1.0])
+        adapter_name = selected_lora["adapter_name"]
+        weights = selected_lora["weights"]
+        
+        if adapter_name not in LOADED_ADAPTERS:
+            print(f"--- Downloading and Loading Adapter: {selected_lora['title']} ---")
+            try:
+                pipe.load_lora_weights(
+                    repo, 
+                    weight_name=weights, 
+                    adapter_name=adapter_name
+                )
+                LOADED_ADAPTERS.add(adapter_name)
+            except Exception as e:
+                raise gr.Error(f"Failed to load adapter {selected_lora['title']}: {e}")
         else:
-            print("Adapter not found in specs. Disabling LoRA.")
-            pipe.disable_lora()
+            print(f"--- Adapter {adapter_name} is already loaded. ---")
+            
+        print(f"Activating LoRA: {adapter_name}")
+        # Enable specific adapter
+        pipe.set_adapters([adapter_name], adapter_weights=[1.0])
 
+    # --- Generation Logic ---
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
@@ -165,14 +188,17 @@ def infer(
         torch.cuda.empty_cache()
 
 @spaces.GPU
-def infer_example(input_image, prompt, lora_adapter):
-    if input_image is None: 
-        return None, 0
-        
+def infer_example(input_image, prompt, selected_title_str):
+    idx = 0
+    for i, item in enumerate(loras_list):
+        if item["title"] == selected_title_str:
+            idx = i
+            break
+            
     image, seed = infer(
         input_image=input_image, 
         prompt=prompt, 
-        lora_adapter=lora_adapter, 
+        selected_index=idx, 
         seed=0, 
         randomize_seed=True,
         guidance_scale=1.0,
@@ -181,87 +207,97 @@ def infer_example(input_image, prompt, lora_adapter):
     return image, seed
 
 css="""
-#col-container { margin: 0 auto; max-width: 960px; }
-#main-title h1 { font-size: 2.2em !important; }
+#col-container { margin: 0 auto; max-width: 1024px; }
+#main-title h1 { font-size: 2.5em !important; text-align: center; }
+#gallery { min-height: 200px; }
+.selected-info { margin-top: 10px; padding: 10px; background: rgba(0,0,0,0.05); border-radius: 8px; }
 """
 
 with gr.Blocks() as demo:
+    
+    selected_index_state = gr.State(value=0)
+
     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.")
         
-        with gr.Row(equal_height=True):
-            with gr.Column():
-                input_image = gr.Image(
-                    label="Upload Image", 
-                    type="pil", 
-                    height=290, 
-                    sources=["upload", "webcam", "clipboard"]
-                )
-                with gr.Row():
+        with gr.Row():
+            with gr.Column(scale=3):
+                with gr.Group():
+                    input_image = gr.Image(
+                        label="Upload Image", 
+                        type="pil", 
+                        height=400, 
+                        sources=["upload", "webcam", "clipboard"]
+                    )
                     prompt = gr.Text(
                         label="Edit Prompt", 
-                        max_lines=1,
-                        show_label=True, 
-                        placeholder="e.g., a man with a red superhero mask"
+                        max_lines=2,
+                        placeholder="Describe the desired change (e.g., 'Relight the image with soft studio lighting')..."
                     )
-                    
-                run_button = gr.Button("Apply Style", variant="primary")
-
-                with gr.Accordion("Advanced Settings", open=False, visible=False):
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0)       
-                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
-                    
-            with gr.Column():
-                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=358)
+            
+            with gr.Column(scale=2):
+                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=400)
                 used_seed = gr.Textbox(label="Used Seed", interactive=False, visible=False)
 
-                with gr.Row():
-                    # Hidden textbox to store the selected style name for the infer function
-                    lora_adapter = gr.Textbox(label="Selected Style", value="None", visible=False)
-                    
-                    # Preparing gallery data
-                    gallery_items = [("examples/1.jpg", "None")]
-                    for key in ADAPTER_SPECS.keys():
-                        # Using example image 2 as thumbnail for Klein style
-                        thumb = "examples/2.jpg" if key == "Klein-Delight-Style" else "examples/1.jpg"
-                        gallery_items.append((thumb, key))
-
-                    style_gallery = gr.Gallery(
-                        label="Edit Style Gallery",
-                        value=gallery_items,
-                        columns=3,
-                        allow_preview=False,
-                        show_label=True,
-                        height="auto"
-                    )
-
-    # Event logic for gallery selection
-    def on_style_select(evt: gr.SelectData):
-        return gallery_items[evt.index][1]
+        # Gallery Selection Section
+        gr.Markdown("### 🎨 Edit Style Gallery", elem_id="gallery_header")
+        with gr.Row():
+            gallery = gr.Gallery(
+                value=[(item["image"], item["title"]) for item in loras_list],
+                label="Choose a Style",
+                allow_preview=False,
+                columns=4,
+                elem_id="gallery",
+                show_label=False,
+                height="auto"
+            )
+        
+        selection_info = gr.Markdown("### Selected: None (Base Model)", elem_classes="selected-info")
+        
+        run_button = gr.Button("✨ Apply Style", variant="primary", size="lg")
 
-    style_gallery.select(fn=on_style_select, inputs=None, outputs=lora_adapter)
+        with gr.Accordion("Advanced Settings", open=False):
+            with gr.Row():
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            with gr.Row():
+                guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0)       
+                steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
+        
+        # Examples
+        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.", 
+                    "Klein Delight Style"
+                ],
+                [
+                    "examples/1.jpg", 
+                    "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed by prithivMLmods", 
+                    "None (Base Model)"
+                ], 
+            ],
+            inputs=[input_image, prompt, gr.Textbox(visible=False)],
+            outputs=[output_image, used_seed],
+            fn=infer_example,
+            cache_examples=False,
+        )
+        
+        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.")
 
-    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"], 
-        ],
-        inputs=[input_image, prompt, lora_adapter],
-        outputs=[output_image, used_seed],
-        fn=infer_example,
-        cache_examples=False,
+    gallery.select(
+        fn=update_selection,
+        inputs=None,
+        outputs=[selected_index_state, selection_info]
     )
-    
-    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.")
-        
+
     run_button.click(
         fn=infer,
-        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
+        inputs=[input_image, prompt, selected_index_state, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_image, used_seed]
     )
 
 if __name__ == "__main__":
-    demo.queue().launch(css=css, theme=orange_red_theme, show_error=True)
+    demo.queue().launch(theme=orange_red_theme, css=css, show_error=True)