Update app.py

app.py

@@ -45,12 +45,12 @@ with open("sdxl_loras.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
         {
-            "image": item.get("image", ""),
-            "title": item.get("nickname", item.get("title", "")),
-            "repo": item.get("model", item.get("repo", "")),
-            "trigger_word": item.get("prompt", item.get("trigger_word", "")),
-            "weights": item.get("weights", ""),
-            "is_compatible": item.get("is_compatible", True),
+            "image": item["image"],
+            "title": item["title"],
+            "repo": item["repo"],
+            "trigger_word": item["trigger_word"],
+            "weights": item["weights"],
+            "is_compatible": item["is_compatible"],
             "is_pivotal": item.get("is_pivotal", False),
             "text_embedding_weights": item.get("text_embedding_weights", None),
             "likes": item.get("likes", 0),
@@ -70,9 +70,6 @@ device = "cuda"
 state_dicts = {}
 
 for item in sdxl_loras_raw:
-    if not item["weights"]:
-        continue
-        
     saved_name = hf_hub_download(item["repo"], item["weights"])
     
     if not saved_name.endswith('.safetensors'):
@@ -134,8 +131,7 @@ elapsed_time = et - st
 print('Loading VAE took: ', elapsed_time, 'seconds')
 st = time.time()
 
-#pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("stablediffusionapi/albedobase-xl-v21",
-pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("frankjoshua/albedobaseXL_v21",
+pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("rubbrband/albedobaseXL_v21",
                                                                  vae=vae,
                                                                  controlnet=[identitynet, zoedepthnet],
                                                                  torch_dtype=torch.float16)
@@ -174,7 +170,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"]:
@@ -238,7 +234,7 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 del weights_sd
                 del lora_model
 
-@spaces.GPU(duration=100)
+@spaces.GPU(duration=80)
 def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index, st):
     print(loaded_state_dict)
     et = time.time()
@@ -257,37 +253,8 @@ 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')
-    
-    # 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:
+    if last_lora != repo_name:
+        if(last_fused):
             st = time.time()
             pipe.unfuse_lora()
             pipe.unload_lora_weights()
@@ -295,7 +262,21 @@ 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')
-            last_fused = False
+        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)
             
     print("Processing prompt...")
     st = time.time()
@@ -320,7 +301,7 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
         image=face_image,
         strength=1-image_strength,
         control_image=images,
-        num_inference_steps=36,
+        num_inference_steps=20,
         guidance_scale = guidance_scale,
         controlnet_conditioning_scale=[face_strength, depth_control_scale],
     ).images[0]
@@ -374,12 +355,8 @@ 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].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
+        weight_name = sdxl_loras[selected_state_index]["weights"]
+        full_path_lora = state_dicts[repo_name]["saved_name"]
     print("Full path LoRA ", full_path_lora)
     #loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
     cross_attention_kwargs = None
@@ -395,11 +372,11 @@ run_lora.zerogpu = True
 
 def shuffle_gallery(sdxl_loras):
     random.shuffle(sdxl_loras)
-    return [(item.get("image") or None, item["title"]) for item in sdxl_loras], sdxl_loras
+    return [(item["image"], 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.get("image") or None, item["title"]) for item in sorted_gallery], sorted_gallery
+    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
 
 def swap_gallery(order, sdxl_loras):
     if(order == "random"):
@@ -447,10 +424,10 @@ def get_civitai_safetensors(link):
   if(x.status_code != 200):
       raise Exception("Invalid CivitAI URL")
   model_data = x.json()
-  #if(model_data["nsfw"] == True or model_data["nsfwLevel"] > 20):
-   # gr.Warning("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
-   # raise Exception("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
-  if(model_data["type"] != "LORA"):
+  if(model_data["nsfw"] == True or model_data["nsfwLevel"] > 20):
+    gr.Warning("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
+    raise Exception("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
+  elif(model_data["type"] != "LORA"):
     gr.Warning("The model isn't tagged at CivitAI as a LoRA")
     raise Exception("The model isn't tagged at CivitAI as a LoRA")
   model_link_download = None
@@ -519,12 +496,12 @@ with gr.Blocks(css="custom.css") as demo:
     gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
     title = gr.HTML(
         """<h1><img src="https://i.imgur.com/DVoGw04.png">
-<span>Face to All<br><small style="
+<span>Face to All SDXL<br><small style="
     font-size: 13px;
     display: block;
     font-weight: normal;
     opacity: 0.75;
-">🧨 diffusers InstantID + ControlNet<br> inspired by fofr's <a href="https://github.com/fofr/cog-face-to-many" target="_blank">face-to-many</a></small></span></h1>""",
+">🧨 diffusers InstantID + ControlNet<br> inspired by fofr's <a href="https://github.com/fofr/cog-face-to-many" target="_blank">face-to-many</a></small></span></h1>""",
         elem_id="title",
     )
     selected_state = gr.State()
@@ -594,7 +571,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=True
+        show_progress=False
     )
     #new_gallery.select(
     #    fn=update_selection,
@@ -606,7 +583,7 @@ with gr.Blocks(css="custom.css") as demo:
     prompt.submit(
         fn=check_selected,
         inputs=[selected_state, custom_loaded_lora],
-        show_progress=True
+        show_progress=False
     ).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],
@@ -615,7 +592,7 @@ with gr.Blocks(css="custom.css") as demo:
     button.click(
         fn=check_selected,
         inputs=[selected_state, custom_loaded_lora],
-        show_progress=True
+        show_progress=False
     ).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],
@@ -625,4 +602,5 @@ with gr.Blocks(css="custom.css") as demo:
     demo.load(fn=classify_gallery, inputs=[gr_sdxl_loras], outputs=[gallery, gr_sdxl_loras], js=js)
 
 demo.queue(default_concurrency_limit=None, api_open=True)
-demo.launch(share=True)
+demo.launch(share=True)
+