Update app.py

app.py

@@ -1,73 +1,75 @@
 import os
 import sys
 import gradio as gr
-import subprocess
+import tempfile
 from huggingface_hub import snapshot_download
-from gradio import Spaces
+import spaces
 
-# Use GPU
-@Spaces.GPU
-def dummy(): pass
-
-# Add PusaV1 to Python path
 sys.path.append(os.path.abspath("PusaV1"))
+from diffsynth import ModelManager, WanVideoPusaPipeline, save_video, VideoData
 
-# 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
+LORA_DIR = "/tmp/model_zoo/PusaV1"
+LORA_PATH = os.path.join(LORA_DIR, "pusa_v1.pt")
 
+@spaces.GPU
 def generate_video(prompt, lora_upload):
-    # Prepare LoRA path
+    # Download Wan2.1 model only if missing
+    if not os.path.exists(WAN_MODEL_DIR):
+        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,
+        )
+
+    # Handle LoRA file (upload or default download + stitch)
     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)
+        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())
 
-        # 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
+        lora_path = LORA_PATH
 
-    # Run pipeline
-    model_manager = ModelManager(pretrained_model_dir=WAN_MODEL_DIR)
-    pipe = WanVideoPusaPipeline(model_manager=model_manager, lora_path=lora_path)
+    # Load model and pipeline
+    manager = ModelManager(pretrained_model_dir=WAN_MODEL_DIR)
+    pipe = WanVideoPusaPipeline(model_manager=manager)
+    pipe.set_lora_adapters(lora_path)
 
+    # Run generation
     result: VideoData = pipe(prompt)
-    video_path = "/tmp/pusa_output.mp4"
+
+    # Save video to temp file
+    tmp_dir = tempfile.mkdtemp()
+    video_path = os.path.join(tmp_dir, "output.mp4")
     save_video(result.frames, video_path, fps=8)
+
     return video_path
 
-# Gradio UI
+
 with gr.Blocks() as demo:
-    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")
-    output_video = gr.Video(label="Output")
+    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")
 
-    generate_btn.click(fn=generate_video, inputs=[prompt, lora_upload], outputs=output_video)
+    generate_btn.click(fn=generate_video, inputs=[prompt, lora_file], outputs=output_video)
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()