Update app.py

app.py

@@ -46,6 +46,10 @@ signal_value = 2.0
 blur_value = None
 allowed_res_max = 1.0
 
+# Device configuration
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+print(f"Using device: {device}")
+
 
 def weight_population(layer_type, resolution, depth, value):
     # Check if layer_type exists, if not, create it
@@ -100,9 +104,9 @@ def reconstruct(input_img, caption):
     ])
 
     if torch_dtype == torch.float16:
-        loaded_image = transform(img).half().to("cuda").unsqueeze(0)
+        loaded_image = transform(img).half().to(device).unsqueeze(0)
     else:
-        loaded_image = transform(img).to("cuda").unsqueeze(0)
+        loaded_image = transform(img).to(device).unsqueeze(0)
 
     if loaded_image.shape[1] == 4:
         loaded_image = loaded_image[:,:3,:,:]
@@ -114,7 +118,7 @@ def reconstruct(input_img, caption):
 
     # notice we disabled the CFG here by setting guidance scale as 1
     guidance_scale = 1.0
-    inverse_scheduler.set_timesteps(num_inference_steps, device="cuda")
+    inverse_scheduler.set_timesteps(num_inference_steps, device=device)
     timesteps = inverse_scheduler.timesteps
 
     latents = real_image_latents
@@ -148,7 +152,7 @@ def reconstruct(input_img, caption):
     real_image_initial_latents = latents
 
     guidance_scale = guidance_scale_value
-    scheduler.set_timesteps(num_inference_steps, device="cuda")
+    scheduler.set_timesteps(num_inference_steps, device=device)
     timesteps = scheduler.timesteps
 
     def adjust_latent(pipe, step, timestep, callback_kwargs):
@@ -319,7 +323,7 @@ def apply_prompt(meta_data, new_prompt):
     inference_steps = len(inversed_latents)
 
     guidance_scale = guidance_scale_value
-    scheduler.set_timesteps(inference_steps, device="cuda")
+    scheduler.set_timesteps(inference_steps, device=device)
     timesteps = scheduler.timesteps
 
     initial_latents = torch.cat([real_image_initial_latents] * 2)
@@ -470,8 +474,8 @@ if __name__ == "__main__":
     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).to(device)
+    pipe.vae = AutoencoderKL.from_pretrained(vae_model_id, subfolder=vae_folder, torch_dtype=torch_dtype).to(device)
     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",