Update app.py

app.py

@@ -53,7 +53,7 @@ for item in sdxl_loras:
     if not saved_name.endswith('.safetensors'):
         state_dict = torch.load(saved_name, map_location=torch.device('cpu'))
     else:
-        state_dict = load_file(saved_name)
+        state_dict = load_file(saved_name, device="cpu")
     
     state_dicts[item["repo"]] = {
         "saved_name": saved_name,
@@ -94,8 +94,8 @@ def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lor
   repo_id_1 = shuffled_items[0]['repo']
   repo_id_2 = shuffled_items[1]['repo']
   print("Loading state dicts...")
-  state_dict_1 = copy.deepcopy(state_dicts[repo_id_1]["state_dict"])
-  state_dict_2 = copy.deepcopy(state_dicts[repo_id_2]["state_dict"])
+  state_dict_1 = copy.deepcopy(state_dicts[repo_id_1]["state_dict"].to("cuda"))
+  state_dict_2 = copy.deepcopy(state_dicts[repo_id_2]["state_dict"].to("cuda"))
   print("Loaded state dicts.")
   pipe = copy.deepcopy(original_pipe)
   pipe.to("cuda")