Update app.py

app.py

@@ -15,7 +15,7 @@ repo = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_1step_unet_x0.safetensors" 
 
 unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(torch.float16)
-unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+unet.load_state_dict(load_file(hf_hub_download(repo, ckpt)))
 pipe_lightning = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16")#.to("cuda")
 del unet
 pipe_lightning.scheduler = EulerDiscreteScheduler.from_config(pipe_lightning.scheduler.config, timestep_spacing="trailing", prediction_type="sample")
@@ -26,7 +26,7 @@ repo_name = "ByteDance/Hyper-SD"
 ckpt_name = "Hyper-SDXL-1step-Unet.safetensors"
 
 unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(torch.float16)
-unet.load_state_dict(load_file(hf_hub_download(repo_name, ckpt_name))
+unet.load_state_dict(load_file(hf_hub_download(repo_name, ckpt_name)))
 pipe_hyper = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16")#.to("cuda")
 pipe_hyper.scheduler = LCMScheduler.from_config(pipe_hyper.scheduler.config)
 #pipe_hyper.to("cuda")