Update app.py

app.py

@@ -91,16 +91,18 @@ original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-
 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)):
   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"])
+  print("Loaded state dicts.")
   pipe = copy.deepcopy(original_pipe)
   pipe.to("cuda") 
-  
-  pipe.load_lora_weights(state_dict_1)
+  print("Loading LoRA weights...")
+  pipe.load_lora_weights(state_dict_1.to("cuda"))
   pipe.fuse_lora(lora_1_scale)
-  pipe.load_lora_weights(state_dict_2)
+  pipe.load_lora_weights(state_dict_2.to("cuda"))
   pipe.fuse_lora(lora_2_scale)
-  
+  print("Loaded LoRA weights.")
   if negative_prompt == "":
     negative_prompt = None