Update app.py

app.py

@@ -1,99 +1,73 @@
-import os
-import shutil
-import subprocess
-import gradio as gr
-from huggingface_hub import snapshot_download
-from huggingface_hub import snapshot_download
-import spaces
-
 import os
 import sys
 import gradio as gr
-import tempfile
-import torch
-import spaces
-# Add PusaV1 to path to resolve diffsynth imports
-sys.path.append(os.path.abspath("PusaV1"))
-
-# Import the actual model runner
-from diffsynth import ModelManager, WanVideoPusaPipeline, save_video
-
-# Define paths
-WAN_MODEL_DIR = "./model_zoo/Wan2.1-T2V-14B"
-LORA_PATH = "./model_zoo/PusaV1/pusa_v1.pt"
-
-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
-
-    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(part, 'rb') as f_in:
-                shutil.copyfileobj(f_in, f_out)
-
-    print(f"✅ Final model saved at {FINAL_MODEL_PATH}")
-
-
-@spaces.GPU
-def generate_video(prompt: str):
-    try:
-        # Load model manager
-        manager = ModelManager(base_model_dir=WAN_MODEL_DIR)
-        model = manager.load_model()
-
-        # Create video pipeline and apply LoRA
-        pipeline = WanVideoPusaPipeline(model=model)
-        pipeline.set_lora_adapters(LORA_PATH)
-
-        # Generate video
-        result = pipeline(prompt=prompt)
+import subprocess
+from huggingface_hub import snapshot_download
+from gradio import Spaces
 
-        # Save video to a temporary file
-        tmp_dir = tempfile.mkdtemp()
-        video_path = os.path.join(tmp_dir, "output.mp4")
-        save_video(result, video_path)
+# Use GPU
+@Spaces.GPU
+def dummy(): pass
 
-        return video_path
+# Add PusaV1 to Python path
+sys.path.append(os.path.abspath("PusaV1"))
 
-    except Exception as e:
-        print(f"[ERROR] {e}")
-        return None
+# Install flash-attn in a safe way
+subprocess.run('pip install flash-attn --no-build-isolation', shell=True, env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"})
+
+# Download Wan2.1-T2V-14B model only (not the full repo)
+WAN_MODEL_DIR = "/tmp/model_zoo/Wan2.1-T2V-14B"
+os.makedirs(WAN_MODEL_DIR, exist_ok=True)
+snapshot_download(
+    repo_id="RaphaelLiu/PusaV1",
+    allow_patterns=["Wan2.1-T2V-14B/*"],
+    local_dir=WAN_MODEL_DIR,
+    local_dir_use_symlinks=False,
+    resume_download=True
+)
+
+# Pusa Inference imports
+from diffsynth import ModelManager, WanVideoPusaPipeline, save_video, VideoData
+
+def generate_video(prompt, lora_upload):
+    # Prepare LoRA path
+    if lora_upload is not None:
+        lora_path = lora_upload
+    else:
+        # Default PusaV1 LoRA
+        default_lora_dir = "/tmp/model_zoo/PusaV1"
+        os.makedirs(default_lora_dir, exist_ok=True)
+
+        # Download all pusa_v1.pt.part* and merge them
+        snapshot_download(
+            repo_id="RaphaelLiu/PusaV1",
+            allow_patterns=["PusaV1/pusa_v1.pt.part*"],
+            local_dir=default_lora_dir,
+            local_dir_use_symlinks=False
+        )
+        merged_path = os.path.join(default_lora_dir, "pusa_v1.pt")
+        os.system(f"cat {default_lora_dir}/pusa_v1.pt.part* > {merged_path}")
+        lora_path = merged_path
+
+    # Run pipeline
+    model_manager = ModelManager(pretrained_model_dir=WAN_MODEL_DIR)
+    pipe = WanVideoPusaPipeline(model_manager=model_manager, lora_path=lora_path)
+
+    result: VideoData = pipe(prompt)
+    video_path = "/tmp/pusa_output.mp4"
+    save_video(result.frames, video_path, fps=8)
+    return video_path
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎥 PusaV1 Text-to-Video Generator")
-    gr.Markdown("Describe a scene and generate a short video using Wan2.1-T2V + Pusa LoRA!")
-
-    prompt_input = gr.Textbox(label="Enter Prompt", lines=4, placeholder="E.g. A coral reef full of colorful fish...")
+    gr.Markdown("# 🎬 Pusa Text-to-Video Generator")
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", value="A vibrant coral reef with sea turtles and sunlight.")
+        lora_upload = gr.File(label="Upload .pt LoRA (optional)", file_types=[".pt"])
     generate_btn = gr.Button("Generate Video")
-    video_output = gr.Video(label="Generated Video")
+    output_video = gr.Video(label="Output")
 
-    generate_btn.click(fn=generate_video, inputs=prompt_input, outputs=video_output)
+    generate_btn.click(fn=generate_video, inputs=[prompt, lora_upload], outputs=output_video)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()