Update app.py

app.py

@@ -1,34 +1,122 @@
 import os
 import subprocess
+import tempfile
 from huggingface_hub import snapshot_download
+import gradio as gr
 
-# Step 1: Download model
+# ---------------- Step 1: Download Model ----------------
 repo_id = "Wan-AI/Wan2.2-TI2V-5B"
 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: Define command arguments (same as your CLI example)
-cmd = [
-    "python", "generate.py",
-    "--task", "ti2v-5B",
-    "--size", "1280*704",
-    "--ckpt_dir", ckpt_dir,
-    "--offload_model", "True",
-    "--convert_model_dtype",
-    "--t5_cpu",
-    "--image", "examples/i2v_input.JPG",
-    "--prompt", (
-        "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. "
-        "The fluffy-furred feline gazes directly at the camera with a relaxed expression. "
-        "Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, "
-        "and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, "
-        "as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's "
-        "intricate details and the refreshing atmosphere of the seaside."
-    )
-]
-
-# Step 3: Run the command
-print("Running generate.py...")
-subprocess.run(cmd, check=True)
-print("Generation complete!")
+# ---------------- Step 2: Generation Functions ----------------
+def generate_t2v(prompt, size="1280*704"):
+    """Text-to-Video generation."""
+    if not prompt.strip():
+        return 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",
+        "--size", size,
+        "--ckpt_dir", ckpt_dir,
+        "--offload_model", "True",
+        "--convert_model_dtype",
+        "--t5_cpu",
+        "--prompt", prompt
+    ]
+
+    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!"
+
+def generate_i2v(image, prompt, size="1280*704"):
+    """Image-to-Video generation."""
+    if image is None or not prompt.strip():
+        return 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",
+        "--size", size,
+        "--ckpt_dir", ckpt_dir,
+        "--offload_model", "True",
+        "--convert_model_dtype",
+        "--t5_cpu",
+        "--image", image_path,
+        "--prompt", prompt
+    ]
+
+    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 3: 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.")
+
+    with gr.Tab("Text-to-Video"):
+        t2v_prompt = gr.Textbox(
+            label="Prompt",
+            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_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])
+
+    with gr.Tab("Image-to-Video"):
+        i2v_image = gr.Image(type="pil", label="Upload Image")
+        i2v_prompt = gr.Textbox(
+            label="Prompt",
+            value=(
+                "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. "
+                "The fluffy-furred feline gazes directly at the camera with a relaxed expression. "
+                "Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, "
+                "and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, "
+                "as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's "
+                "intricate details and the refreshing atmosphere of the seaside."
+            )
+        )
+        i2v_size = gr.Textbox(label="Video Size", value="1280*704")
+        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])
+
+if __name__ == "__main__":
+    demo.launch()