Update app.py

app.py

@@ -1,9 +1,9 @@
+import spaces
 import os
 import subprocess
 import tempfile
 from huggingface_hub import snapshot_download
 import gradio as gr
-import spaces
 
 # ---------------- Step 1: Download Model ----------------
 repo_id = "Wan-AI/Wan2.2-TI2V-5B"
@@ -12,45 +12,23 @@ 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):
+def get_duration(prompt, size, duration_seconds, steps):
     try:
         h, w = size.lower().replace(" ", "").split("*")
         h, w = int(h), int(w)
     except Exception:
         h, w = 704, 1280
-    duration = int(duration_seconds) * 50 * 2.25 + 5  # 50 is fixed sample_steps
-    return duration
-
-# ---------------- 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}"
+    return int(duration_seconds) * int(steps) * 2.25 + 5
 
-    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, mp4_file)
-    os.rename(mp4_file, output_path)
-    return output_path, output_path, "Generation successful!"
-
-# ---------------- Step 4: Generation Functions ----------------
+# ---------------- Step 3: Generation Functions ----------------
 @spaces.GPU(duration=get_duration)
-def generate_t2v(prompt, size="1280*704", duration_seconds=5):
+def generate_t2v(prompt, size="1280*704", duration_seconds=5, steps=25):
     if not prompt.strip():
         return None, None, "Please enter a prompt."
 
+    temp_dir = tempfile.mkdtemp()
+    output_path = os.path.join(temp_dir, "output.mp4")
+
     cmd = [
         "python", "generate.py",
         "--task", "ti2v-5B",
@@ -59,14 +37,25 @@ def generate_t2v(prompt, size="1280*704", duration_seconds=5):
         "--offload_model", "True",
         "--convert_model_dtype",
         "--t5_cpu",
-        "--prompt", prompt,
-        "--steps", "50",
-        "--guidance_scale", "5.0"
+        "--prompt", prompt
     ]
-    return run_generate_command(cmd)
+    print(f"[T2V] Running command: {' '.join(cmd)}")
+
+    try:
+        subprocess.run(cmd, check=True)
+    except subprocess.CalledProcessError as e:
+        return None, None, f"Error during T2V generation: {e}"
+
+    if os.path.exists("output.mp4"):
+        os.rename("output.mp4", output_path)
+    if not os.path.exists(output_path):
+        return None, None, "Generation finished but output file not found."
+
+    download_link = f"<a href='file/{output_path}' download>📥 Download Video</a>"
+    return output_path, download_link, "Text-to-Video generated successfully!"
 
 @spaces.GPU(duration=get_duration)
-def generate_i2v(image, prompt, size="1280*704", duration_seconds=5):
+def generate_i2v(image, prompt, size="1280*704", duration_seconds=5, steps=25):
     if image is None or not prompt.strip():
         return None, None, "Please upload an image and enter a prompt."
 
@@ -74,6 +63,8 @@ def generate_i2v(image, prompt, size="1280*704", duration_seconds=5):
     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",
@@ -83,16 +74,27 @@ def generate_i2v(image, prompt, size="1280*704", duration_seconds=5):
         "--convert_model_dtype",
         "--t5_cpu",
         "--image", image_path,
-        "--prompt", prompt,
-        "--steps", "50",
-        "--guidance_scale", "5.0"
+        "--prompt", prompt
     ]
-    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, None, f"Error during I2V generation: {e}"
+
+    if os.path.exists("output.mp4"):
+        os.rename("output.mp4", output_path)
+    if not os.path.exists(output_path):
+        return None, None, "Generation finished but output file not found."
+
+    download_link = f"<a href='file/{output_path}' download>📥 Download Video</a>"
+    return output_path, download_link, "Image-to-Video generated successfully!"
 
-# ---------------- Step 5: Gradio UI ----------------
+# ---------------- Step 4: Gradio UI ----------------
 with gr.Blocks() as demo:
     gr.Markdown("## 🎥 Wan2.2-TI2V-5B Video Generator")
-    gr.Markdown("Generate AI videos from text or image prompts with download option.")
+    gr.Markdown("Choose **Text-to-Video** or **Image-to-Video** mode below.")
 
     with gr.Tab("Text-to-Video"):
         t2v_prompt = gr.Textbox(
@@ -101,13 +103,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", autoplay=True)
-        t2v_download = gr.File(label="Download Video")
+        t2v_video = gr.Video(label="Generated Video")
+        t2v_download = gr.HTML()
         t2v_status = gr.Textbox(label="Status")
         t2v_btn.click(
             generate_t2v,
-            [t2v_prompt, t2v_size, t2v_duration],
+            [t2v_prompt, t2v_size, t2v_duration, t2v_steps],
             [t2v_video, t2v_download, t2v_status]
         )
 
@@ -126,13 +129,14 @@ 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", autoplay=True)
-        i2v_download = gr.File(label="Download Video")
+        i2v_video = gr.Video(label="Generated Video")
+        i2v_download = gr.HTML()
         i2v_status = gr.Textbox(label="Status")
         i2v_btn.click(
             generate_i2v,
-            [i2v_image, i2v_prompt, i2v_size, i2v_duration],
+            [i2v_image, i2v_prompt, i2v_size, i2v_duration, i2v_steps],
             [i2v_video, i2v_download, i2v_status]
         )