Update app.py

app.py

@@ -1,60 +1,82 @@
 import os
-import gradio as gr
-import subprocess
 import shutil
-import tempfile
+import subprocess
+import gradio as gr
 from huggingface_hub import snapshot_download
 
-# Paths
-PUSA_REPO = "./PusaV1"
-PUSA_SCRIPT = os.path.join(PUSA_REPO, "examples/pusavideo/wan_14b_text_to_video_pusa.py")
-MODEL_DIR = "./model_zoo/PusaV1"
-MODEL_PATH = os.path.join(MODEL_DIR, "pusa_v1.pt")
-OUTPUT_VIDEO_PATH = os.path.join(tempfile.gettempdir(), "output.mp4")
-
-def setup_dependencies():
-    subprocess.run(["pip", "install", "xfuser>=0.4.3", "absl-py", "peft", "lightning", "pandas", "deepspeed", "wandb", "av"])
-    subprocess.run(
-        'pip install flash-attn --no-build-isolation',
-        shell=True,
-        env={**os.environ, "FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}
-    )
 
-def download_and_stitch_model():
-    if os.path.exists(MODEL_PATH):
+MODEL_SUBFOLDER = "Wan2.1-T2V-14B"
+HF_REPO = "RaphaelLiu/PusaV1"
+MODEL_ZOO_DIR = "./model_zoo"
+MODEL_PARTS_DIR = os.path.join(MODEL_ZOO_DIR, MODEL_SUBFOLDER)
+FINAL_MODEL_PATH = os.path.join(MODEL_ZOO_DIR, "PusaV1", "pusa_v1.pt")
+PUSA_SCRIPT_PATH = "PusaV1/examples/pusavideo/wan_14b_text_to_video_pusa.py"
+
+
+def download_model_subset():
+    if os.path.exists(FINAL_MODEL_PATH):
+        print("✅ Model already exists. Skipping download.")
         return
-    os.makedirs(MODEL_DIR, exist_ok=True)
-    snapshot_download("RaphaelLiu/PusaV1", local_dir=MODEL_DIR, local_dir_use_symlinks=False)
-    part_files = sorted([f for f in os.listdir(MODEL_DIR) if f.startswith("pusa_v1.pt.part")])
-    with open(MODEL_PATH, 'wb') as outfile:
+
+    print("⏬ Downloading model parts...")
+    snapshot_download(
+        repo_id=HF_REPO,
+        repo_type="model",
+        local_dir=MODEL_ZOO_DIR,
+        local_dir_use_symlinks=False,
+        allow_patterns=[f"{MODEL_SUBFOLDER}/*"]
+    )
+    os.makedirs(os.path.dirname(FINAL_MODEL_PATH), exist_ok=True)
+
+    part_files = sorted([
+        os.path.join(MODEL_PARTS_DIR, f)
+        for f in os.listdir(MODEL_PARTS_DIR)
+        if f.startswith("pusa_v1.pt.part")
+    ])
+
+    print("🧩 Stitching model parts...")
+    with open(FINAL_MODEL_PATH, 'wb') as f_out:
         for part in part_files:
-            with open(os.path.join(MODEL_DIR, part), 'rb') as infile:
-                shutil.copyfileobj(infile, outfile)
+            with open(part, 'rb') as f_in:
+                shutil.copyfileobj(f_in, f_out)
+
+    print(f"✅ Final model saved at {FINAL_MODEL_PATH}")
 
-def generate_video_from_text(prompt: str):
-    download_and_stitch_model()
-    cmd = [
-        "python", PUSA_SCRIPT,
+
+def generate_video(prompt):
+    download_model_subset()
+
+    temp_output_dir = "/tmp/pusa_video_output"
+    os.makedirs(temp_output_dir, exist_ok=True)
+
+    command = [
+        "python", PUSA_SCRIPT_PATH,
         "--prompt", prompt,
-        "--lora_path", MODEL_PATH
+        "--lora_path", FINAL_MODEL_PATH,
+        "--output_dir", temp_output_dir
     ]
-    result = subprocess.run(cmd, capture_output=True, text=True)
-    if os.path.exists(OUTPUT_VIDEO_PATH):
-        return OUTPUT_VIDEO_PATH
-    return f"Video generation failed:\n{result.stderr}"
-
-def main():
-    setup_dependencies()
-    download_and_stitch_model()
-
-    with gr.Blocks() as demo:
-        gr.Markdown("## 🎥 PusaV1 - Text to Video Generator")
-        prompt = gr.Textbox(lines=6, label="Enter your scene prompt")
-        output = gr.Video(label="Generated Video")
-        run = gr.Button("Generate Video")
-        run.click(generate_video_from_text, inputs=prompt, outputs=output)
-
-    demo.launch()
-
-if __name__ == "__main__":
-    main()
+
+    try:
+        print("🚀 Running inference...")
+        subprocess.run(command, check=True)
+
+        # Return first mp4 video found
+        for file in os.listdir(temp_output_dir):
+            if file.endswith(".mp4"):
+                return os.path.join(temp_output_dir, file)
+
+        return "❌ No video generated."
+
+    except subprocess.CalledProcessError as e:
+        return f"❌ Inference failed: {str(e)}"
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧘‍♂️ PusaV1 Text-to-Video Generator (Wan2.1-T2V-14B)")
+    prompt_input = gr.Textbox(label="Enter your prompt", lines=4, placeholder="A coral reef full of colorful fish...")
+    generate_button = gr.Button("Generate Video")
+    video_output = gr.Video(label="Generated Video")
+
+    generate_button.click(fn=generate_video, inputs=prompt_input, outputs=video_output)
+
+demo.launch()