Update app.py

app.py

@@ -18,19 +18,29 @@ from safetensors.torch import load_file
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16
 
-step = 4
-model_id = "black-forest-labs/FLUX.1-dev"
+step = 2
+repo = "ByteDance/AnimateDiff-Lightning"
+ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
+model_id = "black-forest-labs/FLUX.1-schnell"
 
-adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2")
+adapter = MotionAdapter().to(device, dtype)
+adapter.load_state_dict(load_file(hf_hub_download(repo ,ckpt), device=device))
 
 pipe = AnimateDiffPipeline.from_pretrained(
-    model_id, motion_adapter=adapter
-    ,token=os.getenv("hf_token")
+    model_id,
+    motion_adapter=adapter,
+    token=os.getenv("hf_token")
 )
+
 scheduler = DDIMScheduler.from_pretrained(
-    model_id, subfolder="scheduler", clip_sample=False, timestep_spacing="linspace", steps_offset=1
-    ,token=os.getenv("hf_token")
+    model_id,
+    subfolder="scheduler", 
+    clip_sample=False,
+    timestep_spacing="linspace", 
+    steps_offset=1,
+    token=os.getenv("hf_token")
 )
+
 pipe.scheduler = scheduler
 pipe.enable_vae_slicing()
 pipe.enable_vae_tiling()