Update app_g.py

app_g.py

@@ -10,6 +10,42 @@ import numpy as np
 from diffsynth import ModelManager, PusaMultiFramesPipeline, PusaV2VPipeline, WanVideoPusaPipeline, save_video
 import tempfile
 
+
+
+# Constants
+import os
+from huggingface_hub import snapshot_download
+
+# Constants
+MODEL_ZOO_DIR = "./model_zoo"
+PUSA_DIR = os.path.join(MODEL_ZOO_DIR, "PusaV1")
+WAN_SUBFOLDER = "Wan2.1-T2V-14B"
+WAN_MODEL_PATH = os.path.join(PUSA_DIR, WAN_SUBFOLDER)
+LORA_PATH = os.path.join(PUSA_DIR, "pusa_v1.pt")
+
+# Ensure model and weights are downloaded
+def ensure_model_downloaded():
+    if not os.path.exists(PUSA_DIR):
+        print("Downloading RaphaelLiu/PusaV1 to ./model_zoo/PusaV1 ...")
+        snapshot_download(
+            repo_id="RaphaelLiu/PusaV1",
+            local_dir=PUSA_DIR,
+            repo_type="model",
+            local_dir_use_symlinks=False,
+        )
+        print("✅ PusaV1 downloaded.")
+
+    if not os.path.exists(WAN_MODEL_PATH):
+        print("Downloading Wan-AI/Wan2.1-T2V-14B to ./model_zoo/PusaV1/Wan2.1-T2V-14B ...")
+        snapshot_download(
+            repo_id="Wan-AI/Wan2.1-T2V-14B",
+            local_dir=WAN_MODEL_PATH,  # Changed from WAN_DIR to WAN_MODEL_PATH
+            repo_type="model",
+            local_dir_use_symlinks=False,
+        )
+        print("✅ Wan2.1-T2V-14B downloaded.")
+
+
 class PusaVideoDemo:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -38,6 +74,8 @@ class PusaVideoDemo:
                 torch_dtype=torch.bfloat16,
             )
             print("Models loaded successfully!")
+
+            
         
     def load_lora_and_get_pipe(self, pipe_type, lora_path, lora_alpha):
         """Load LoRA and return appropriate pipeline"""
@@ -244,7 +282,7 @@ class PusaVideoDemo:
             
         except Exception as e:
             return None, f"Error: {str(e)}"
-
+    
     def generate_t2v_video(self, prompt, lora_alpha, num_inference_steps, 
                           negative_prompt, progress=gr.Progress()):
         """Generate video from text prompt"""
@@ -1191,6 +1229,7 @@ def create_demo():
     return demo
 
 if __name__ == "__main__":
+    ensure_model_downloaded()
     demo = create_demo()
     demo.launch(
         share=False,