Update app.py

app.py

@@ -41,7 +41,7 @@ vae = AutoencoderKL.from_pretrained(
 
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
-    unet = UNet2DConditionModel.from_pretrained(model, subfolder="unet").to("cuda", torch.float16)
+    unet = UNet2DConditionModel.from_pretrained(model, subfolder="unet", torch_dtype=torch.float16).to("cuda")
     pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, unet=unet, torch_dtype=torch.float16).to("cuda")
 
 pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)