Update app.py

app.py

@@ -62,16 +62,17 @@ def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
     return scheduler_factory_map.get(name, lambda: None)()
 
 def load_pipeline(model_name):
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16,
-        custom_pipeline="lpw_stable_diffusion_xl",
-        safety_checker = None,
-        use_safetensors=True,
-        add_watermarker=False,
-        use_auth_token=HF_TOKEN,
-    )
-    pipe.to(device)
+    if torch.cuda.is_available():
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            custom_pipeline="lpw_stable_diffusion_xl",
+            safety_checker = None,
+            use_safetensors=True,
+            add_watermarker=False,
+            use_auth_token=HF_TOKEN
+        )
+    pipe.to("cuda:0")
     return pipe
 
 @spaces.GPU