Update app.py

app.py

@@ -269,7 +269,7 @@ def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, i
     lora_repo = sdxl_loras[selected_state.index]["repo"]
     new_placeholder = "Type a prompt to use your selected LoRA"
     weight_name = sdxl_loras[selected_state.index]["weights"]
-    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
+    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
 
     for lora_list in lora_defaults:
         if lora_list["model"] == sdxl_loras[selected_state.index]["repo"]:
@@ -355,7 +355,8 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     # Enhanced prompt processing
     if custom_lora_path and custom_lora[1]:
         prompt = f"{prompt} {custom_lora[1]}"
-    else:
+    elif selected_state_index >= 0 and selected_state_index < len(sdxl_loras):
+        # Only apply default prompt template if we have a valid selection
         for lora_list in lora_defaults:
             if lora_list["model"] == sdxl_loras[selected_state_index]["repo"]:
                 prompt_full = lora_list.get("prompt", None)
@@ -385,9 +386,11 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     elif custom_lora_path:
         repo_name = custom_lora_path
         full_path_lora = custom_lora_path
-    else:
+    elif selected_state_index >= 0 and selected_state_index < len(sdxl_loras):
         repo_name = sdxl_loras[selected_state_index]["repo"]
         full_path_lora = state_dicts[repo_name]["saved_name"]
+    else:
+        raise gr.Error("Invalid style selection. Please select a style again.")
 
     repo_name = repo_name.rstrip("/").lower()
 
@@ -469,29 +472,33 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
         # Load LoRA with better error handling
         try:
             pipe.load_lora_weights(full_path_lora)
-            pipe.fuse_lora(lora_scale)
+            pipe.fuse_lora(lora_scale=lora_scale)
             last_fused = True
             
             # Handle pivotal tuning embeddings (if needed for future LoRAs)
-            is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
-            if is_pivotal:
-                text_embedding_name = sdxl_loras[selected_state_index]["text_embedding_weights"]
-                embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
-                state_dict_embedding = load_file(embedding_path)
-                pipe.load_textual_inversion(
-                    state_dict_embedding["clip_l" if "clip_l" in state_dict_embedding else "text_encoders_0"],
-                    token=["<s0>", "<s1>"],
-                    text_encoder=pipe.text_encoder,
-                    tokenizer=pipe.tokenizer
-                )
-                pipe.load_textual_inversion(
-                    state_dict_embedding["clip_g" if "clip_g" in state_dict_embedding else "text_encoders_1"],
-                    token=["<s0>", "<s1>"],
-                    text_encoder=pipe.text_encoder_2,
-                    tokenizer=pipe.tokenizer_2
-                )
+            # Only check this if we're not using a custom LoRA and have a valid index
+            if selected_state_index >= 0 and selected_state_index < len(sdxl_loras):
+                is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
+                if is_pivotal:
+                    text_embedding_name = sdxl_loras[selected_state_index]["text_embedding_weights"]
+                    embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
+                    state_dict_embedding = load_file(embedding_path)
+                    pipe.load_textual_inversion(
+                        state_dict_embedding["clip_l" if "clip_l" in state_dict_embedding else "text_encoders_0"],
+                        token=["<s0>", "<s1>"],
+                        text_encoder=pipe.text_encoder,
+                        tokenizer=pipe.tokenizer
+                    )
+                    pipe.load_textual_inversion(
+                        state_dict_embedding["clip_g" if "clip_g" in state_dict_embedding else "text_encoders_1"],
+                        token=["<s0>", "<s1>"],
+                        text_encoder=pipe.text_encoder_2,
+                        tokenizer=pipe.tokenizer_2
+                    )
         except Exception as e:
             print(f"Error loading LoRA: {e}")
+            import traceback
+            traceback.print_exc()
             raise gr.Error(f"Failed to load LoRA: {str(e)}")
 
     print("Processing prompt...")
@@ -651,7 +658,7 @@ with gr.Blocks(css="custom.css") as demo:
     display: block;
     font-weight: normal;
     opacity: 0.75;
-">🔥 Improved: Better face identity preservation, Enhanced landscape mode, Multiple face support<br>AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
+">🔥 Improved: Better face identity preservation, Enhanced landscape mode, Multiple face support<br>AlbedoBase XL v2.1 + InstantID + ControlNet</small></span></h1>""",
         elem_id="title",
     )
     selected_state = gr.State()