Update app.py

app.py

@@ -4,6 +4,7 @@ import subprocess
 import tempfile
 from huggingface_hub import snapshot_download
 import gradio as gr
+import spaces  # Hugging Face Spaces GPU decorator
 
 # ---------------- Step 1: Download Model ----------------
 repo_id = "Wan-AI/Wan2.2-TI2V-5B"
@@ -11,9 +12,26 @@ print(f"Downloading/loading checkpoints for {repo_id}...")
 ckpt_dir = snapshot_download(repo_id, local_dir_use_symlinks=False)
 print(f"Using checkpoints from {ckpt_dir}")
 
-# ---------------- Step 2: Generation Functions ----------------
-@spaces.GPU
-def generate_t2v(prompt, size="1280*704"):
+# ---------------- Step 2: Duration Calculation ----------------
+def get_duration(prompt, size, duration_seconds, steps):
+    """
+    Calculate GPU duration dynamically.
+    - duration_seconds: estimated video length in seconds
+    - steps: number of inference steps
+    """
+    try:
+        h, w = size.lower().replace(" ", "").split("*")
+        h, w = int(h), int(w)
+    except Exception:
+        h, w = 704, 1280  # fallback
+
+    # Simple rule: time grows with steps and video duration
+    duration = int(duration_seconds) * int(steps) * 2.25 + 5
+    return duration
+
+# ---------------- Step 3: Generation Functions ----------------
+@spaces.GPU(duration=get_duration)
+def generate_t2v(prompt, size="1280*704", duration_seconds=5, steps=25):
     """Text-to-Video generation."""
     if not prompt.strip():
         return None, "Please enter a prompt."
@@ -47,9 +65,8 @@ def generate_t2v(prompt, size="1280*704"):
 
     return output_path, "Text-to-Video generated successfully!"
 
-
-@spaces.GPU(duration=100)    
-def generate_i2v(image, prompt, size="1280*704"):
+@spaces.GPU(duration=get_duration)
+def generate_i2v(image, prompt, size="1280*704", duration_seconds=5, steps=25):
     """Image-to-Video generation."""
     if image is None or not prompt.strip():
         return None, "Please upload an image and enter a prompt."
@@ -87,7 +104,7 @@ def generate_i2v(image, prompt, size="1280*704"):
 
     return output_path, "Image-to-Video generated successfully!"
 
-# ---------------- Step 3: Gradio UI ----------------
+# ---------------- Step 4: Gradio UI ----------------
 with gr.Blocks() as demo:
     gr.Markdown("## 🎥 Wan2.2-TI2V-5B Video Generator")
     gr.Markdown("Choose **Text-to-Video** or **Image-to-Video** mode below.")
@@ -98,10 +115,16 @@ with gr.Blocks() as demo:
             value="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage"
         )
         t2v_size = gr.Textbox(label="Video Size", value="1280*704")
+        t2v_duration = gr.Number(label="Video Length (seconds)", value=5)
+        t2v_steps = gr.Number(label="Inference Steps", value=25)
         t2v_btn = gr.Button("Generate from Text")
         t2v_video = gr.Video(label="Generated Video")
         t2v_status = gr.Textbox(label="Status")
-        t2v_btn.click(generate_t2v, [t2v_prompt, t2v_size], [t2v_video, t2v_status])
+        t2v_btn.click(
+            generate_t2v,
+            [t2v_prompt, t2v_size, t2v_duration, t2v_steps],
+            [t2v_video, t2v_status]
+        )
 
     with gr.Tab("Image-to-Video"):
         i2v_image = gr.Image(type="pil", label="Upload Image")
@@ -117,10 +140,16 @@ with gr.Blocks() as demo:
             )
         )
         i2v_size = gr.Textbox(label="Video Size", value="1280*704")
+        i2v_duration = gr.Number(label="Video Length (seconds)", value=5)
+        i2v_steps = gr.Number(label="Inference Steps", value=25)
         i2v_btn = gr.Button("Generate from Image")
         i2v_video = gr.Video(label="Generated Video")
         i2v_status = gr.Textbox(label="Status")
-        i2v_btn.click(generate_i2v, [i2v_image, i2v_prompt, i2v_size], [i2v_video, i2v_status])
+        i2v_btn.click(
+            generate_i2v,
+            [i2v_image, i2v_prompt, i2v_size, i2v_duration, i2v_steps],
+            [i2v_video, i2v_status]
+        )
 
 if __name__ == "__main__":
     demo.launch()