Update app.py

app.py

@@ -98,15 +98,19 @@ def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lor
   print("Loading state dicts...")
   start_time = time()  
   state_dict_1 = copy.deepcopy(state_dicts[repo_id_1]["state_dict"])
+  state_dict_1 = {k: v.to(device="cuda", dtype=torch.float16) for k,v in state_dict_1.items() if torch.is_tensor(v)}
   state_dict_2 = copy.deepcopy(state_dicts[repo_id_2]["state_dict"])
+  state_dict_2 = {k: v.to(device="cuda", dtype=torch.float16) for k,v in state_dict_2.items() if torch.is_tensor(v)}
   state_dict_time = time() - start_time
   print(f"State Dict time: {state_dict_time}")  
   #pipe = copy.deepcopy(original_pipe)
   start_time = time()
   unet = copy.deepcopy(original_pipe.unet)
+  text_encoder=copy.deepcopy(original_pipe.text_encoder)
+  text_encoder_2=copy.deepcopy(original_pipe.text_encoder_2)
   pipe = StableDiffusionXLPipeline(vae=original_pipe.vae,
-                                   text_encoder=original_pipe.text_encoder,
-                                   text_encoder_2=original_pipe.text_encoder_2,
+                                   text_encoder=text_encoder,
+                                   text_encoder_2=text_encoder_2,
                                    scheduler=original_pipe.scheduler,
                                    tokenizer=original_pipe.tokenizer,
                                    tokenizer_2=original_pipe.tokenizer_2,