Update app.py

app.py

@@ -28,6 +28,8 @@ JS = """function () {
   }
 }"""
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 # Load VAE component
 vae = AutoencoderKL.from_pretrained(
     vae_model, 
@@ -36,7 +38,7 @@ vae = AutoencoderKL.from_pretrained(
 
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16).to(0)
+    pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16).to(device)
 
 pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)