Fix user scale override

app.py

@@ -55,7 +55,7 @@ def load_lora_weights(repo_id, weights_filename):
 def update_selection(selected_state: gr.SelectData, flux_loras):
     """Update UI when a LoRA is selected"""
     if selected_state.index >= len(flux_loras):
-        return "### No LoRA selected", gr.update(), None
+        return "### No LoRA selected", gr.update(), None, gr.update()
     
     lora_repo = flux_loras[selected_state.index]["repo"]
     trigger_word = flux_loras[selected_state.index]["trigger_word"]
@@ -67,7 +67,10 @@ def update_selection(selected_state: gr.SelectData, flux_loras):
     else:
         new_placeholder = f"opt - describe the person/subject, e.g. 'a man with glasses and a beard'"
     
-    return updated_text, gr.update(placeholder=new_placeholder), selected_state.index
+    optimal_scale = flux_loras[selected_state.index].get("lora_scale_config", 1.0)
+    
+    return updated_text, gr.update(placeholder=new_placeholder), selected_state.index, gr.update(value=optimal_scale)
+
 
 def get_huggingface_lora(link):
     """Download LoRA from HuggingFace link"""
@@ -133,12 +136,12 @@ def classify_gallery(flux_loras):
     sorted_gallery = sorted(flux_loras, key=lambda x: x.get("likes", 0), reverse=True)
     return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
 
-def infer_with_lora_wrapper(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75,portrait_mode=False, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
+def infer_with_lora_wrapper(input_image, prompt, selected_index, lora_state, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.75,portrait_mode=False, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
     """Wrapper function to handle state serialization"""
-    return infer_with_lora(input_image, prompt, selected_index, custom_lora, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, flux_loras, progress)
+    return infer_with_lora(input_image, prompt, selected_index, lora_state, custom_lora, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, flux_loras, progress)
 
 @spaces.GPU
-def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, portrait_mode=False, flux_loras=None,  progress=gr.Progress(track_tqdm=True)):
+def infer_with_lora(input_image, prompt, selected_index, lora_state, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, portrait_mode=False, flux_loras=None,  progress=gr.Progress(track_tqdm=True)):
     """Generate image with selected LoRA"""
     global current_lora, pipe
     
@@ -155,14 +158,9 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
     # Load LoRA if needed
     if lora_to_use and lora_to_use != current_lora:
         try:
-            # Unload current LoRA
             if current_lora:
                 pipe.unload_lora_weights()
             
-            # Load new LoRA
-            if lora_to_use["lora_scale_config"]:
-                lora_scale = lora_to_use["lora_scale_config"]
-                print("lora scale loaded from config", lora_scale)
             lora_path = load_lora_weights(lora_to_use["repo"], lora_to_use["weights"])
             if lora_path:
                 pipe.load_lora_weights(lora_path, adapter_name="selected_lora")
@@ -173,8 +171,9 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
         except Exception as e:
             print(f"Error loading LoRA: {e}")
             # Continue without LoRA
-    else:
-        print(f"using already loaded lora: {lora_to_use}")
+    elif lora_scale != lora_state:
+        pipe.set_adapters(["selected_lora"], adapter_weights=[lora_scale])
+        print(f"using already loaded lora: {lora_to_use}, udpated {lora_scale} based on user preference")
     
     input_image = input_image.convert("RGB")
     # Add trigger word to prompt
@@ -204,7 +203,7 @@ def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, r
             height=input_image.size[1],
             prompt=prompt,
             guidance_scale=guidance_scale,
-            generator=torch.Generator().manual_seed(seed),
+            generator=torch.Generator().manual_seed(seed)
         ).images[0]
         
         return image, seed, gr.update(visible=True), lora_scale
@@ -264,6 +263,7 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
     
     selected_state = gr.State(value=None)
     custom_loaded_lora = gr.State(value=None)
+    lora_state = gr.State(value=1.0)
     
     with gr.Row(elem_id="main_app"):
         with gr.Column(scale=4, elem_id="box_column"):
@@ -348,15 +348,15 @@ with gr.Blocks(css=css, theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Lexend
     gallery.select(
         fn=update_selection,
         inputs=[gr_flux_loras],
-        outputs=[prompt_title, prompt, selected_state],
+        outputs=[prompt_title, prompt, selected_state, lora_scale],
         show_progress=False
     )
     
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer_with_lora_wrapper,
-        inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, gr_flux_loras],
-        outputs=[result, seed, reuse_button, lora_scale]
+        inputs=[input_image, prompt, selected_state, lora_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, portrait_mode, gr_flux_loras],
+        outputs=[result, seed, reuse_button, lora_state]
     )
     
     reuse_button.click(