Update app.py

app.py

@@ -60,8 +60,8 @@ def load_pipeline():
     torch_dtype = torch.float16
 
     # torch_dtype = torch.float16
-    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype).to("cuda")
-    pipe.vae = AutoencoderKL.from_pretrained(vae_model_id, subfolder=vae_folder, torch_dtype=torch_dtype).to("cuda")
+    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+    pipe.vae = AutoencoderKL.from_pretrained(vae_model_id, subfolder=vae_folder, torch_dtype=torch_dtype)
     pipe.load_lora_weights(
         hf_hub_download(repo_id="jiaxiangc/res-adapter", subfolder=resadapter_model_name, filename="pytorch_lora_weights.safetensors"), 
         adapter_name="res_adapter",
@@ -118,6 +118,7 @@ def resize_image_with_aspect(image, res_range_min=128, res_range_max=1024):
 def reconstruct(input_img, caption):
 
     pipe, inverse_scheduler, scheduler = load_pipeline()
+    pipe.to("cuda")
 
     global weights
     weights = {}
@@ -346,7 +347,8 @@ def replace_attention_processor(unet, clear=False, blur_sigma=None):
 @spaces.GPU()
 def apply_prompt(meta_data, new_prompt):
 
-    pipe, inverse_scheduler, scheduler = load_pipeline()
+    pipe, _, scheduler = load_pipeline()
+    pipe.to("cuda")
 
     caption, real_image_initial_latents, inversed_latents, _ = meta_data
     negative_prompt = ""