Update app.py

app.py

@@ -78,7 +78,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(), selected_state
+        return "### No LoRA selected", gr.update(), None
     
     lora_repo = flux_loras[selected_state.index]["repo"]
     trigger_word = flux_loras[selected_state.index]["trigger_word"]
@@ -86,7 +86,7 @@ def update_selection(selected_state: gr.SelectData, flux_loras):
     updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
     new_placeholder = f"Enter your editing prompt{f' (use {trigger_word} for best results)' if trigger_word else ''}"
     
-    return updated_text, gr.update(placeholder=new_placeholder), selected_state
+    return updated_text, gr.update(placeholder=new_placeholder), selected_state.index
 
 def get_huggingface_lora(link):
     """Download LoRA from HuggingFace link"""
@@ -117,7 +117,7 @@ def get_huggingface_lora(link):
 def load_custom_lora(link):
     """Load custom LoRA from user input"""
     if not link:
-        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### Click on a LoRA in the gallery to select it"
+        return gr.update(visible=False), "", gr.update(visible=False), None, gr.Gallery(selected_index=None), "### Click on a LoRA in the gallery to select it", None
     
     try:
         repo_name, weights_file, trigger_word = get_huggingface_lora(link)
@@ -138,22 +138,26 @@ def load_custom_lora(link):
             "trigger_word": trigger_word
         }
         
-        return gr.update(visible=True), card, gr.update(visible=True), custom_lora_data, gr.Gallery(selected_index=None), f"Custom: {repo_name}"
+        return gr.update(visible=True), card, gr.update(visible=True), custom_lora_data, gr.Gallery(selected_index=None), f"Custom: {repo_name}", None
     
     except Exception as e:
-        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### Click on a LoRA in the gallery to select it"
+        return gr.update(visible=True), f"Error: {str(e)}", gr.update(visible=False), None, gr.update(), "### Click on a LoRA in the gallery to select it", None
 
 def remove_custom_lora():
     """Remove custom LoRA"""
-    return "", gr.update(visible=False), gr.update(visible=False), None
+    return "", gr.update(visible=False), gr.update(visible=False), None, None
 
 def classify_gallery(flux_loras):
     """Sort gallery by likes"""
     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.0, 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, flux_loras, progress)
+
 @spaces.GPU
-def infer_with_lora(input_image, prompt, selected_state, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
+def infer_with_lora(input_image, prompt, selected_index, custom_lora, seed=42, randomize_seed=False, guidance_scale=2.5, lora_scale=1.0, flux_loras=None, progress=gr.Progress(track_tqdm=True)):
     """Generate image with selected LoRA"""
     global current_lora, pipe
     
@@ -164,10 +168,8 @@ def infer_with_lora(input_image, prompt, selected_state, custom_lora, seed=42, r
     lora_to_use = None
     if custom_lora:
         lora_to_use = custom_lora
-    elif selected_state and flux_loras:
-        selected_index = selected_state.index if hasattr(selected_state, 'index') else None
-        if selected_index is not None and selected_index < len(flux_loras):
-            lora_to_use = flux_loras[selected_index]
+    elif selected_index is not None and flux_loras and selected_index < len(flux_loras):
+        lora_to_use = flux_loras[selected_index]
     
     # Load LoRA if needed
     if lora_to_use and lora_to_use != current_lora:
@@ -265,8 +267,8 @@ with gr.Blocks(css=css) as demo:
         <br><small style="font-size: 13px; opacity: 0.75;"></small></h1>""",
     )
     
-    selected_state = gr.State()
-    custom_loaded_lora = gr.State()
+    selected_state = gr.State(value=None)
+    custom_loaded_lora = gr.State(value=None)
     
     with gr.Row(elem_id="main_app"):
         with gr.Column(scale=4, elem_id="box_column"):
@@ -340,12 +342,12 @@ with gr.Blocks(css=css) as demo:
     custom_model.input(
         fn=load_custom_lora,
         inputs=[custom_model],
-        outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery, prompt_title],
+        outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery, prompt_title, selected_state],
     )
     
     custom_model_button.click(
         fn=remove_custom_lora,
-        outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora]
+        outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora, selected_state]
     )
     
     gallery.select(
@@ -357,7 +359,7 @@ with gr.Blocks(css=css) as demo:
     
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn=infer_with_lora,
+        fn=infer_with_lora_wrapper,
         inputs=[input_image, prompt, selected_state, custom_loaded_lora, seed, randomize_seed, guidance_scale, lora_scale, gr_flux_loras],
         outputs=[result, seed, reuse_button]
     )