Update

app.py

@@ -12,13 +12,6 @@ from PIL import Image
 
 # Constants
 base = "frankjoshua/toonyou_beta6"
-repo = "ByteDance/AnimateDiff-Lightning"
-checkpoints = {
-    "1-Step" : ["animatediff_lightning_1step_diffusers.safetensors", 1],
-    "2-Step" : ["animatediff_lightning_2step_diffusers.safetensors", 2],
-    "4-Step" : ["animatediff_lightning_4step_diffusers.safetensors", 4],
-    "8-Step" : ["animatediff_lightning_8step_diffusers.safetensors", 8],
-}
 loaded = None
 
 # Ensure model and scheduler are initialized in GPU-enabled function
@@ -33,24 +26,21 @@ else:
 
 # Function 
 @spaces.GPU(enable_queue=True)
-def generate_image(prompt, ckpt):
+def generate_image(prompt, step):
     global loaded
-    print(prompt, ckpt)
+    print(prompt, step)
 
-    checkpoint = checkpoints[ckpt][0]
-    num_inference_steps = checkpoints[ckpt][1]
+    if loaded != step:
+        repo = "ByteDance/AnimateDiff-Lightning"
+        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
+        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
+        loaded = step
 
-    if loaded != num_inference_steps:
-        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device=device), strict=False)
-        loaded = num_inference_steps
-
-    output = pipe(prompt=prompt, guidance_scale=1.0, num_inference_steps=num_inference_steps)
+    output = pipe(prompt=prompt, guidance_scale=1.0, num_inference_steps=step)
     
     name = str(uuid.uuid4()).replace("-", "")
     path = f"/tmp/{name}.mp4"
-
     export_to_video(output.frames[0], path, fps=10)
-
     return path
 
 
@@ -62,10 +52,28 @@ with gr.Blocks(css="style.css") as demo:
     gr.HTML("<p><center>Lightning-fast text-to-video generation</center></p><p><center><a href='https://huggingface.co/ByteDance/AnimateDiff-Lightning'>https://huggingface.co/ByteDance/AnimateDiff-Lightning</a></center></p>")
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(label='Enter your prompt (English)', scale=8)
-            ckpt = gr.Dropdown(label='Select inference steps',choices=['1-Step', '2-Step', '4-Step', '8-Step'], value='4-Step', interactive=True)
-            submit = gr.Button(scale=1, variant='primary')
-    video = gr.Video(label='AnimateDiff-Lightning Generated Image')
+            prompt = gr.Textbox(
+                label='Enter your prompt (English)',
+                scale=8
+            )
+            ckpt = gr.Dropdown(
+                label='Select inference steps',
+                choices=[
+                    ('1-Step', 1), 
+                    ('2-Step', 2),
+                    ('4-Step', 4),
+                    ('8-Step', 8)],
+                value='4-Step',
+                interactive=True
+            )
+            submit = gr.Button(
+                scale=1,
+                variant='primary'
+            )
+    video = gr.Video(
+        label='AnimateDiff-Lightning',
+        autoplay=True,
+    )
 
     prompt.submit(
         fn=generate_image,