Update app.py

app.py

@@ -39,24 +39,37 @@ class PatchedModelManager(ModelManager):
 
 
 # Constants
-WAN_SUBFOLDER = "Wan2.1-T2V-14B"
-MODEL_REPO_ID = "RaphaelLiu/PusaV1"
+import os
+from huggingface_hub import snapshot_download
+
+# Constants
 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")
+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 is downloaded
+# 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 Wan2.1-T2V-14B from HuggingFace Hub...")
+        print("Downloading Wan-AI/Wan2.1-T2V-14B to ./model_zoo/PusaV1 ...")
         snapshot_download(
-            repo_id=MODEL_REPO_ID,
-            local_dir=MODEL_ZOO_DIR,
+            repo_id="Wan-AI/Wan2.1-T2V-14B",
+            local_dir=os.path.join(PUSA_DIR, WAN_SUBFOLDER),
             repo_type="model",
-            allow_patterns=[f"{WAN_SUBFOLDER}/**"],
             local_dir_use_symlinks=False,
         )
-        print("Model downloaded.")
+        print("✅ Wan2.1-T2V-14B downloaded.")
 
 # Subclass ModelManager to force WanModelPusa