Update app.py

app.py

@@ -1,75 +1,66 @@
-import os
-import sys
 import gradio as gr
+import os
 import tempfile
 from huggingface_hub import snapshot_download
+from diffsynth import ModelManager, WanVideoPusaPipeline, save_video
 import spaces
 
+# Constants
+WAN_SUBFOLDER = "Wan2.1-T2V-14B"
+MODEL_REPO_ID = "RaphaelLiu/PusaV1"
+MODEL_ZOO_DIR = "./model_zoo"
+WAN_MODEL_PATH = os.path.join(MODEL_ZOO_DIR, WAN_SUBFOLDER)
+LORA_PATH = os.path.join(MODEL_ZOO_DIR, "PusaV1", "pusa_v1.pt")
 
-from diffsynth import ModelManager, WanVideoPusaPipeline, save_video, VideoData
-
-WAN_MODEL_DIR = "/tmp/model_zoo/Wan2.1-T2V-14B"
-LORA_DIR = "/tmp/model_zoo/PusaV1"
-LORA_PATH = os.path.join(LORA_DIR, "pusa_v1.pt")
-
-@spaces.GPU
-def generate_video(prompt, lora_upload):
-    # Download Wan2.1 model only if missing
-    if not os.path.exists(WAN_MODEL_DIR):
+# Ensure model is downloaded
+def ensure_model_downloaded():
+    if not os.path.exists(WAN_MODEL_PATH):
+        print("Downloading Wan2.1-T2V-14B from HuggingFace Hub...")
         snapshot_download(
-            repo_id="RaphaelLiu/PusaV1",
-            allow_patterns=["Wan2.1-T2V-14B/*"],
-            local_dir=WAN_MODEL_DIR,
+            repo_id=MODEL_REPO_ID,
+            local_dir=MODEL_ZOO_DIR,
+            repo_type="model",
+            allow_patterns=[f"{WAN_SUBFOLDER}/**"],
             local_dir_use_symlinks=False,
-            resume_download=True,
         )
+        print("Model downloaded.")
 
-    # Handle LoRA file (upload or default download + stitch)
-    if lora_upload is not None:
-        lora_path = lora_upload
-    else:
-        if not os.path.exists(LORA_PATH):
-            os.makedirs(LORA_DIR, exist_ok=True)
-            snapshot_download(
-                repo_id="RaphaelLiu/PusaV1",
-                allow_patterns=["PusaV1/pusa_v1.pt.part*"],
-                local_dir=LORA_DIR,
-                local_dir_use_symlinks=False,
-            )
-            # Stitch parts
-            part_files = sorted(
-                f for f in os.listdir(LORA_DIR) if f.startswith("pusa_v1.pt.part")
-            )
-            with open(LORA_PATH, "wb") as wfd:
-                for part in part_files:
-                    with open(os.path.join(LORA_DIR, part), "rb") as fd:
-                        wfd.write(fd.read())
+# Video generation logic
+@spaces.GPU
+def generate_video(prompt: str):
+    ensure_model_downloaded()
 
-        lora_path = LORA_PATH
+    # Load model
+    manager = ModelManager(pretrained_model_dir=WAN_MODEL_PATH)
+    model = manager.load_model()
 
-    # Load model and pipeline
-    manager = ModelManager(pretrained_model_dir=WAN_MODEL_DIR)
-    pipe = WanVideoPusaPipeline(model_manager=manager)
-    pipe.set_lora_adapters(lora_path)
+    # Set up pipeline
+    pipeline = WanVideoPusaPipeline(model=model)
+    pipeline.set_lora_adapters(LORA_PATH)
 
-    # Run generation
-    result: VideoData = pipe(prompt)
+    # Generate video
+    result = pipeline(prompt)
 
-    # Save video to temp file
+    # Save video
     tmp_dir = tempfile.mkdtemp()
-    video_path = os.path.join(tmp_dir, "output.mp4")
-    save_video(result.frames, video_path, fps=8)
-
-    return video_path
+    output_path = os.path.join(tmp_dir, "video.mp4")
+    save_video(result.frames, output_path, fps=8)
 
+    return output_path
 
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎥 Pusa Text-to-Video (Wan2.1-T2V-14B)")
-    prompt = gr.Textbox(label="Prompt", lines=4)
-    lora_file = gr.File(label="Upload LoRA .pt (optional)", file_types=[".pt"])
-    generate_btn = gr.Button("Generate")
-    output_video = gr.Video(label="Generated Video")
+    gr.Markdown("## 🎥 Wan2.1-T2V-14B with Pusa LoRA | Text-to-Video Generator")
+
+    prompt_input = gr.Textbox(
+        lines=4,
+        label="Prompt",
+        placeholder="Describe your video (e.g. A coral reef full of colorful fish...)"
+    )
+
+    generate_btn = gr.Button("Generate Video")
+    video_output = gr.Video(label="Output")
 
-    generate_btn.click(fn=generate_video, inputs=[prompt, lora_file], outputs=output_video)
+    generate_btn.click(fn=generate_video, inputs=prompt_input, outputs=video_output)
 
 demo.launch()