Update app.py

app.py

@@ -201,7 +201,7 @@ def main():
     ckpt = f"sdxl_lightning_{step}step_unet.safetensors"
 
     unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(device, dtype)
-    unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location=device), strict=False)
+    unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), map_location=device), strict=False)
 
     repo = "ByteDance/AnimateDiff-Lightning"
     ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"