Update app.py

app.py

@@ -206,8 +206,7 @@ def main():
 
     adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2")
     vae = AutoencoderKL.from_single_file("https://huggingface.co/stabilityai/sd-vae-ft-mse-original/vae-ft-mse-840000-ema-pruned.safetensors")
-    unet = UNet3DConditionModel.to(device, dtype)
-    unet.load_state_dict(load_file(hf_hub_download("ByteDance/SDXL-Lightning", f"sdxl_lightning_{step}step_unet.safetensors"), device=device), strict=False)
+    unet = UNet3DConditionModel().to(device, dtype).load_state_dict(load_file(hf_hub_download("ByteDance/SDXL-Lightning", f"sdxl_lightning_{step}step_unet.safetensors"), device=device), strict=False)
     
     repo = "ByteDance/AnimateDiff-Lightning"
     ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"