Update app.py

app.py

@@ -12,6 +12,7 @@ import json
 import random
 import copy
 import gc
+import spaces
 
 lora_list = hf_hub_download(repo_id="multimodalart/LoraTheExplorer", filename="sdxl_loras.json", repo_type="space")
 
@@ -44,16 +45,20 @@ with open(lora_list, "r") as file:
         for item in data
     ]
 
+state_dicts = {}
+
 for item in sdxl_loras:
     saved_name = hf_hub_download(item["repo"], item["weights"])
     
-    if saved_name.endswith('.safetensors'):
-        state_dict = load_file(saved_name)
-    else:
+    if not saved_name.endswith('.safetensors'):
         state_dict = torch.load(saved_name)
-        
-    item["saved_name"] = saved_name
-    item["state_dict"] = state_dict #{k: v.to(device="cuda", dtype=torch.float16) for k, v in state_dict.items() if torch.is_tensor(v)}
+    else:
+        state_dict = load_file(saved_name)
+    
+    state_dicts[item["repo"]] = {
+        "saved_name": saved_name,
+        "state_dict": state_dict
+    }
 
 css = '''
 .gradio-container{max-width: 650px! important}
@@ -82,9 +87,12 @@ div#share-btn-container > div {flex-direction: row;background: black;align-items
 
 original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
 
+@spaces.GPU
 def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, seed=-1, progress=gr.Progress(track_tqdm=True)):
-  state_dict_1 = copy.deepcopy(shuffled_items[0]['state_dict'])
-  state_dict_2 = copy.deepcopy(shuffled_items[1]['state_dict'])
+  repo_id_1 = shuffled_items[0]['repo']
+  repo_id_2 = shuffled_items[1]['repo']
+  state_dict_1 = copy.deepcopy(state_dicts[repo_id_1]["state_dict"])
+  state_dict_2 = copy.deepcopy(state_dicts[repo_id_2]["state_dict"])
   pipe = copy.deepcopy(original_pipe)
   pipe.to("cuda") 
   
@@ -102,7 +110,6 @@ def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lor
   image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=20, width=768, height=768, generator=generator).images[0]
   del pipe
   gc.collect()
-  torch.cuda.empty_cache()
   return image, gr.update(visible=True), seed
 
 def get_description(item):