Update app.py

app.py

@@ -36,8 +36,8 @@ base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_4step_unet.safetensors"
 
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
-unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+unet = UNet2DConditionModel.from_config(base, subfolder="unet")
+unet.load_state_dict(load_file(hf_hub_download(repo, ckpt)))
 image_pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16")
 image_pipe.scheduler = EulerDiscreteScheduler.from_config(image_pipe.scheduler.config, timestep_spacing="trailing")