Update app.py

app.py

@@ -5,14 +5,18 @@ from diffusers import DiffusionPipeline
 
 print("Loading pipeline...")
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+dtype = torch.float16 if device == "cuda" else torch.float32
+
 pipe = DiffusionPipeline.from_pretrained(
     "stable-diffusion-v1-5/stable-diffusion-v1-5",
     torch_dtype=torch.bfloat16,
 )
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe.to("cuda")
+pipe.to(device)
 
 pipe.unet.load_attn_procs(
 "./pytorch_custom_diffusion_weights.bin"