Update app.py

app.py

@@ -4,11 +4,11 @@ import torch
 from diffusers import DiffusionPipeline
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
-print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
 if torch.cuda.is_available():
-    pipe_sd = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, revision="fp16", use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"))
-    pipe_vq = DiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq", torch_dtype=torch.float16, revision="fp16")
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+    pipe_sd = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, revision="fp16", use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN")).to("cuda")
+    pipe_vq = DiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq", torch_dtype=torch.float16, revision="fp16").to("cuda")
 else:
     pipe_sd = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"))
     pipe_vq = DiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq")