Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 from diffusers import StableDiffusionXLPipeline, AutoencoderKL
 from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 from share_btn import community_icon_html, loading_icon_html, share_js
 from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
 import lora
@@ -26,12 +27,19 @@ with open("sdxl_loras.json", "r") as file:
         }
         for item in data
     ]
-print(sdxl_loras)
-saved_names = [
-    hf_hub_download(item["repo"], item["weights"]) for item in sdxl_loras
-]
 
-device = "cuda"  # replace this to `mps` if on a MacOS Silicon
+device = "cuda" 
+
+for item in sdxl_loras:
+    saved_name = hf_hub_download(item["repo"], item["weights"])
+    
+    if not saved_name.endswith('.safetensors'):
+        state_dict = torch.load(saved_name)
+    else:
+        state_dict = load_file(saved_name)
+        
+    item["saved_name"] = saved_name
+    item["state_dict"] = state_dict #{k: v.to(device=device, dtype=torch.float16) for k, v in state_dict.items() if torch.is_tensor(v)}
 
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
@@ -40,14 +48,13 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     vae=vae,
     torch_dtype=torch.float16,
-).to("cpu")
+)
 original_pipe = copy.deepcopy(pipe)
 pipe.to(device)
 
 last_lora = ""
 last_merged = False
 
-
 def update_selection(selected_state: gr.SelectData):
     lora_repo = sdxl_loras[selected_state.index]["repo"]
     instance_prompt = sdxl_loras[selected_state.index]["trigger_word"]
@@ -128,7 +135,7 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 del lora_model
                 gc.collect()
 
-def run_lora(prompt, negative, lora_scale, selected_state):
+def run_lora(prompt, negative, lora_scale, selected_state, progress=gr.Progress(track_tqdm=True)):
     global last_lora, last_merged, pipe
 
     if negative == "":
@@ -138,7 +145,8 @@ def run_lora(prompt, negative, lora_scale, selected_state):
         raise gr.Error("You must select a LoRA")
     repo_name = sdxl_loras[selected_state.index]["repo"]
     weight_name = sdxl_loras[selected_state.index]["weights"]
-    full_path_lora = saved_names[selected_state.index]
+    full_path_lora = sdxl_loras[selected_state.index]["saved_name"]
+    loaded_state_dict = sdxl_loras[selected_state.index]["state_dict"]
     cross_attention_kwargs = None
     if last_lora != repo_name:
         if last_merged:
@@ -148,17 +156,17 @@ def run_lora(prompt, negative, lora_scale, selected_state):
             pipe.to(device)
         else:
             pipe.unload_lora_weights()
+            pipe.unfuse_lora()
         is_compatible = sdxl_loras[selected_state.index]["is_compatible"]
         if is_compatible:
-            pipe.load_lora_weights(full_path_lora)
-            cross_attention_kwargs = {"scale": lora_scale}
+            pipe.load_lora_weights(loaded_state_dict)
+            pipe.fuse_lora(lora_scale)
         else:
             is_pivotal = sdxl_loras[selected_state.index]["is_pivotal"]
             if(is_pivotal):
+                pipe.load_lora_weights(loaded_state_dict)
+                pipe.fuse_lora(lora_scale)
                 
-                pipe.load_lora_weights(full_path_lora)
-                cross_attention_kwargs = {"scale": lora_scale}
-
                 #Add the textual inversion embeddings from pivotal tuning models
                 text_embedding_name = sdxl_loras[selected_state.index]["text_embedding_weights"]
                 text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
@@ -177,7 +185,6 @@ def run_lora(prompt, negative, lora_scale, selected_state):
         height=768,
         num_inference_steps=20,
         guidance_scale=7.5,
-        cross_attention_kwargs=cross_attention_kwargs,
     ).images[0]
     last_lora = repo_name
     gc.collect()