update app

app.py

@@ -151,19 +151,29 @@ def process_gallery_images(images):
     
     return pil_images
 
+def get_style_config(name):
+    """Retrieve style configuration by title."""
+    for style in LORA_STYLES:
+        if style["title"] == name:
+            return style
+    return LORA_STYLES[0]  # Default to None if not found
+
 def update_style_selection(evt: gr.SelectData):
-    """Update selected style based on gallery click and optionally update prompt"""
+    """Update selected style based on gallery click"""
     selected_style = LORA_STYLES[evt.index]
     
-    if selected_style["title"] == "None":
-        info_text = "### Selected: None (FLUX.2-klein-9B) ✅"
-    else:
-        info_text = f"### Selected: {selected_style['title']} ✅"
-    
     # Get default prompt if available
     default_prompt = selected_style.get("default_prompt", None)
     
-    return info_text, evt.index, default_prompt if default_prompt else gr.update()
+    # Return name (for State) and prompt
+    return selected_style["title"], default_prompt if default_prompt else gr.update()
+
+def update_info_text(style_name):
+    """Updates the info text when the style name changes (via Gallery or Example)"""
+    if style_name == "None" or not style_name:
+        return "### Selected: None (FLUX.2-klein-9B) ✅"
+    else:
+        return f"### Selected: {style_name} ✅"
 
 def get_image_count_info(images):
     """Return info about uploaded images"""
@@ -182,7 +192,7 @@ def get_image_count_info(images):
 def infer(
     input_images, 
     prompt, 
-    selected_style_index,
+    selected_style_name,
     seed=42, 
     randomize_seed=True, 
     guidance_scale=1.0, 
@@ -201,10 +211,11 @@ def infer(
     if not pil_images:
         raise gr.Error("Could not process uploaded images.")
 
-    if selected_style_index is None:
-        selected_style_index = 0
-    
-    selected_style = LORA_STYLES[selected_style_index]
+    # Resolve Style Config from Name
+    if selected_style_name is None:
+        selected_style_name = "None"
+        
+    selected_style = get_style_config(selected_style_name)
     
     # Check if Face Swap is selected and validate image count
     if selected_style["adapter_name"] == "face-swap":
@@ -272,15 +283,14 @@ def infer(
         torch.cuda.empty_cache()
 
 @spaces.GPU
-def infer_example(input_images, prompt, style_index):
+def infer_example(input_images, prompt, style_name):
+    """Wrapper for examples to handle paths and names"""
     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
@@ -288,7 +298,7 @@ def infer_example(input_images, prompt, style_index):
     image, seed = infer(
         input_images=input_images, 
         prompt=prompt, 
-        selected_style_index=style_index, 
+        selected_style_name=style_name, 
         seed=0, 
         randomize_seed=True,
         guidance_scale=1.0,
@@ -308,7 +318,7 @@ 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.")
         
-        selected_style_index = gr.State(0)
+        selected_style_name = gr.State("None")
         
         with gr.Row(equal_height=True):
             with gr.Column():
@@ -335,7 +345,7 @@ with gr.Blocks() as demo:
                 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)       
+                    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():
@@ -354,12 +364,19 @@ 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.", 1],
-                [["examples/1.jpg"], "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed by prithivMLmods preserving realistic texture and details", 1],
-                [["examples/1.jpg", "examples/2.jpg"], "Apply the style from image 2 to image 1, preserving the subject's identity.", 1],
-                [["examples/face1.jpg", "examples/face2.jpg"], FACE_SWAP_PROMPT, 2],
+                # Example 1: 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.", "None"],
+                
+                # Example 2: 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", "Klein-Delight-Style"],
+                
+                # Example 3: Klein-Delight-Style (Multiple images scenario)
+                [["examples/1.jpg", "examples/2.jpg"], "Apply the style from image 2 to image 1, preserving the subject's identity.", "Klein-Delight-Style"],
+                
+                # Example 4: Best-Face-Swap
+                [["examples/face1.jpg", "examples/face2.jpg"], FACE_SWAP_PROMPT, "Best-Face-Swap"],
             ],
-            inputs=[input_images, prompt, selected_style_index],
+            inputs=[input_images, prompt, selected_style_name],
             outputs=[output_image, used_seed],
             fn=infer_example,
             cache_examples=False,
@@ -368,7 +385,6 @@ with gr.Blocks() as demo:
         
         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.")
     
-    # Update image count info when gallery changes
     input_images.change(
         fn=get_image_count_info,
         inputs=[input_images],
@@ -376,12 +392,18 @@ with gr.Blocks() as demo:
     
     style_gallery.select(
         update_style_selection,
-        outputs=[selected_style_info, selected_style_index, prompt]
+        outputs=[selected_style_name, prompt]
+    )
+    
+    selected_style_name.change(
+        fn=update_info_text,
+        inputs=[selected_style_name],
+        outputs=[selected_style_info]
     )
     
     run_button.click(
         fn=infer,
-        inputs=[input_images, prompt, selected_style_index, seed, randomize_seed, guidance_scale, steps],
+        inputs=[input_images, prompt, selected_style_name, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_image, used_seed]
     )