Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from time import sleep
 from diffusers import DiffusionPipeline
 from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
 import torch
 import json
@@ -28,12 +29,11 @@ with open(lora_list, "r") as file:
         for item in data
     ]
 
-saved_names = [
-    hf_hub_download(item["repo"], item["weights"]) for item in sdxl_loras
-]
-
-for item, saved_name in zip(sdxl_loras, saved_names):
+for item in sdxl_loras:
+    saved_name = hf_hub_download(item["repo"], item["weights"])
     item["saved_name"] = saved_name
+    state_dict = load_file(saved_name)
+    item["state_dict"] = {k: v.to(device="cuda", dtype=torch.float16) for k, v in state_dict.items() if torch.is_tensor(v)}
 
 css = '''
 #title{text-align:center;}
@@ -49,23 +49,21 @@ css = '''
 }
 '''
 
-#@spaces.GPU
 pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
 original_pipe = copy.deepcopy(pipe)
 
 def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, progress=gr.Progress(track_tqdm=True)):
   pipe = copy.deepcopy(original_pipe)
-  pipe.to("cuda")  
-  print("Loading LoRAs")
-  pipe.load_lora_weights(shuffled_items[0]['saved_name'])
+  pipe.to("cuda")
+  pipe.load_lora_weights(shuffled_items[0]['state_dict'])
   pipe.fuse_lora(lora_1_scale)
-  pipe.load_lora_weights(shuffled_items[1]['saved_name'])
+  pipe.load_lora_weights(shuffled_items[1]['state_dict'])
   pipe.fuse_lora(lora_2_scale)
   
   if negative_prompt == "":
     negative_prompt = False
       
-  image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=25).images[0]
+  image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=22, width=768, height=768).images[0]
   del pipe
   gc.collect()
   torch.cuda.empty_cache()