Update app.py

app.py

@@ -91,26 +91,29 @@ original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-
 
 @spaces.GPU
 def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, seed=-1):
-  print("Run this?")
+  
   repo_id_1 = shuffled_items[0]['repo']
   repo_id_2 = shuffled_items[1]['repo']
   print("Loading state dicts...")
-  state_dict_1 = state_dicts[repo_id_1]["state_dict"]
-  state_dict_2 = state_dicts[repo_id_2]["state_dict"]
-  print("Loaded state dicts.")
+  start_time = time()  
+  state_dict_1 = pickle.loads(pickle.dumps(state_dicts[repo_id_1]["state_dict"]))
+  state_dict_2 = pickle.loads(pickle.dumps(state_dicts[repo_id_2]["state_dict"]))
+  state_dict_time = time() - start_time
+  print(f"State Dict time: {state_dict_time}")  
   #pipe = copy.deepcopy(original_pipe)
-  # Time for pickle
-  start_time = time.time()
+  start_time = time()
   pipe = pickle.loads(pickle.dumps(original_pipe))
-  pickle_time = time.time() - start_time
+  pickle_time = time() - start_time
   print(f"Pickle time: {pickle_time}")  
   pipe.to("cuda") 
+  start_time = time()
   print("Loading LoRA weights...")
   pipe.load_lora_weights(state_dict_1)
   pipe.fuse_lora(lora_1_scale)
   pipe.load_lora_weights(state_dict_2)
   pipe.fuse_lora(lora_2_scale)
-  print("Loaded LoRA weights.")
+  lora_time = time() - start_time
+  print(f"Loaded LoRAs time: {lora_time}") 
   if negative_prompt == "":
     negative_prompt = None