Update app.py

app.py

@@ -257,8 +257,37 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     et = time.time()
     elapsed_time = et - st
     print('Zoe Depth calculations took: ', elapsed_time, 'seconds')
-    if last_lora != repo_name:
-        if(last_fused):
+    
+    # Only handle lora if we have weights
+    if loaded_state_dict is not None:
+        if last_lora != repo_name:
+            if(last_fused):
+                st = time.time()
+                pipe.unfuse_lora()
+                pipe.unload_lora_weights()
+                pipe.unload_textual_inversion()
+                et = time.time()
+                elapsed_time = et - st
+                print('Unfuse and unload LoRA took: ', elapsed_time, 'seconds')
+            st = time.time()
+            pipe.load_lora_weights(loaded_state_dict)
+            pipe.fuse_lora(lora_scale)
+            et = time.time()
+            elapsed_time = et - st
+            print('Fuse and load LoRA took: ', elapsed_time, 'seconds')
+            last_fused = True
+            is_pivotal = sdxl_loras[selected_state_index].get("is_pivotal", False)
+            if(is_pivotal):
+                #Add the textual inversion embeddings from pivotal tuning models
+                text_embedding_name = sdxl_loras[selected_state_index].get("text_embedding_weights")
+                if text_embedding_name:
+                    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)
+    else:
+        # No lora to load, unfuse any existing lora
+        if last_fused:
             st = time.time()
             pipe.unfuse_lora()
             pipe.unload_lora_weights()
@@ -266,21 +295,7 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
             et = time.time()
             elapsed_time = et - st
             print('Unfuse and unload LoRA took: ', elapsed_time, 'seconds')
-        st = time.time()
-        pipe.load_lora_weights(loaded_state_dict)
-        pipe.fuse_lora(lora_scale)
-        et = time.time()
-        elapsed_time = et - st
-        print('Fuse and load LoRA took: ', elapsed_time, 'seconds')
-        last_fused = True
-        is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
-        if(is_pivotal):
-            #Add the textual inversion embeddings from pivotal tuning models
-            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)
+            last_fused = False
             
     print("Processing prompt...")
     st = time.time()
@@ -359,8 +374,12 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
         full_path_lora = custom_lora_path
     else:
         repo_name = sdxl_loras[selected_state_index]["repo"]
-        weight_name = sdxl_loras[selected_state_index]["weights"]
-        full_path_lora = state_dicts[repo_name]["saved_name"]
+        weight_name = sdxl_loras[selected_state_index].get("weights", "")
+        if weight_name and repo_name in state_dicts:
+            full_path_lora = state_dicts[repo_name]["saved_name"]
+        else:
+            # No weights available, use base model without lora
+            full_path_lora = None
     print("Full path LoRA ", full_path_lora)
     #loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
     cross_attention_kwargs = None
@@ -376,11 +395,11 @@ run_lora.zerogpu = True
 
 def shuffle_gallery(sdxl_loras):
     random.shuffle(sdxl_loras)
-    return [(item["image"], item["title"]) for item in sdxl_loras], sdxl_loras
+    return [(item.get("image") or None, item["title"]) for item in sdxl_loras], sdxl_loras
 
 def classify_gallery(sdxl_loras):
     sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get("likes", 0), reverse=True)
-    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
+    return [(item.get("image") or None, item["title"]) for item in sorted_gallery], sorted_gallery
 
 def swap_gallery(order, sdxl_loras):
     if(order == "random"):
@@ -575,7 +594,7 @@ with gr.Blocks(css="custom.css") as demo:
         fn=update_selection,
         inputs=[gr_sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative],
         outputs=[prompt_title, prompt, face_strength, image_strength, weight, depth_control_scale, negative, selected_state],
-        show_progress=False
+        show_progress=True
     )
     #new_gallery.select(
     #    fn=update_selection,
@@ -587,7 +606,7 @@ with gr.Blocks(css="custom.css") as demo:
     prompt.submit(
         fn=check_selected,
         inputs=[selected_state, custom_loaded_lora],
-        show_progress=False
+        show_progress=True
     ).success(
         fn=run_lora,
         inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora],
@@ -596,7 +615,7 @@ with gr.Blocks(css="custom.css") as demo:
     button.click(
         fn=check_selected,
         inputs=[selected_state, custom_loaded_lora],
-        show_progress=False
+        show_progress=True
     ).success(
         fn=run_lora,
         inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora],