Update app.py

app.py

@@ -1,10 +1,9 @@
-import spaces
 import os
 import subprocess
 import tempfile
 from huggingface_hub import snapshot_download
 import gradio as gr
-import spaces  # Hugging Face Spaces GPU decorator
+import spaces
 
 # ---------------- Step 1: Download Model ----------------
 repo_id = "Wan-AI/Wan2.2-TI2V-5B"
@@ -13,31 +12,44 @@ ckpt_dir = snapshot_download(repo_id, local_dir_use_symlinks=False)
 print(f"Using checkpoints from {ckpt_dir}")
 
 # ---------------- 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
-    """
+def get_duration(prompt, size, duration_seconds):
     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
+        h, w = 704, 1280
+    duration = int(duration_seconds) * 50 * 2.25 + 5  # 50 is fixed sample_steps
     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."
+# ---------------- Step 3: Helpers ----------------
+def find_generated_mp4():
+    mp4_files = [f for f in os.listdir(".") if f.lower().endswith(".mp4")]
+    if not mp4_files:
+        return None
+    mp4_files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
+    return mp4_files[0]
 
+def run_generate_command(cmd):
+    print(f"[RUN] {' '.join(cmd)}")
+    try:
+        subprocess.run(cmd, check=True)
+    except subprocess.CalledProcessError as e:
+        return None, None, f"Error: {e}"
+
+    mp4_file = find_generated_mp4()
+    if not mp4_file:
+        return None, None, "No output video found."
+    
     temp_dir = tempfile.mkdtemp()
-    output_path = os.path.join(temp_dir, "output.mp4")
+    output_path = os.path.join(temp_dir, mp4_file)
+    os.rename(mp4_file, output_path)
+    return output_path, output_path, "Generation successful!"
+
+# ---------------- Step 4: Generation Functions ----------------
+@spaces.GPU(duration=get_duration)
+def generate_t2v(prompt, size="1280*704", duration_seconds=5):
+    if not prompt.strip():
+        return None, None, "Please enter a prompt."
 
     cmd = [
         "python", "generate.py",
@@ -47,36 +59,21 @@ def generate_t2v(prompt, size="1280*704", duration_seconds=5, steps=25):
         "--offload_model", "True",
         "--convert_model_dtype",
         "--t5_cpu",
-        "--prompt", prompt
+        "--prompt", prompt,
+        "--steps", "50",
+        "--guidance_scale", "5.0"
     ]
-
-    print(f"[T2V] Running command: {' '.join(cmd)}")
-    try:
-        subprocess.run(cmd, check=True)
-    except subprocess.CalledProcessError as e:
-        return None, f"Error during T2V generation: {e}"
-
-    if os.path.exists("output.mp4"):
-        os.rename("output.mp4", output_path)
-    elif os.path.exists(output_path):
-        pass
-    else:
-        return None, "Generation finished but output file not found."
-
-    return output_path, "Text-to-Video generated successfully!"
+    return run_generate_command(cmd)
 
 @spaces.GPU(duration=get_duration)
-def generate_i2v(image, prompt, size="1280*704", duration_seconds=5, steps=25):
-    """Image-to-Video generation."""
+def generate_i2v(image, prompt, size="1280*704", duration_seconds=5):
     if image is None or not prompt.strip():
-        return None, "Please upload an image and enter a prompt."
+        return None, None, "Please upload an image and enter a prompt."
 
     temp_dir = tempfile.mkdtemp()
     image_path = os.path.join(temp_dir, "input.jpg")
     image.save(image_path)
 
-    output_path = os.path.join(temp_dir, "output.mp4")
-
     cmd = [
         "python", "generate.py",
         "--task", "ti2v-5B",
@@ -86,28 +83,16 @@ def generate_i2v(image, prompt, size="1280*704", duration_seconds=5, steps=25):
         "--convert_model_dtype",
         "--t5_cpu",
         "--image", image_path,
-        "--prompt", prompt
+        "--prompt", prompt,
+        "--steps", "50",
+        "--guidance_scale", "5.0"
     ]
+    return run_generate_command(cmd)
 
-    print(f"[I2V] Running command: {' '.join(cmd)}")
-    try:
-        subprocess.run(cmd, check=True)
-    except subprocess.CalledProcessError as e:
-        return None, f"Error during I2V generation: {e}"
-
-    if os.path.exists("output.mp4"):
-        os.rename("output.mp4", output_path)
-    elif os.path.exists(output_path):
-        pass
-    else:
-        return None, "Generation finished but output file not found."
-
-    return output_path, "Image-to-Video generated successfully!"
-
-# ---------------- Step 4: Gradio UI ----------------
+# ---------------- Step 5: 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.")
+    gr.Markdown("Generate AI videos from text or image prompts with download option.")
 
     with gr.Tab("Text-to-Video"):
         t2v_prompt = gr.Textbox(
@@ -116,14 +101,14 @@ with gr.Blocks() as demo:
         )
         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_video = gr.Video(label="Generated Video", autoplay=True)
+        t2v_download = gr.File(label="Download Video")
         t2v_status = gr.Textbox(label="Status")
         t2v_btn.click(
             generate_t2v,
-            [t2v_prompt, t2v_size, t2v_duration, t2v_steps],
-            [t2v_video, t2v_status]
+            [t2v_prompt, t2v_size, t2v_duration],
+            [t2v_video, t2v_download, t2v_status]
         )
 
     with gr.Tab("Image-to-Video"):
@@ -141,15 +126,15 @@ 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_video = gr.Video(label="Generated Video", autoplay=True)
+        i2v_download = gr.File(label="Download Video")
         i2v_status = gr.Textbox(label="Status")
         i2v_btn.click(
             generate_i2v,
-            [i2v_image, i2v_prompt, i2v_size, i2v_duration, i2v_steps],
-            [i2v_video, i2v_status]
+            [i2v_image, i2v_prompt, i2v_size, i2v_duration],
+            [i2v_video, i2v_download, i2v_status]
         )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()