Update download_models.py

download_models.py

@@ -1,30 +1,86 @@
-import os
-from huggingface_hub import snapshot_download
-
-# Target directory matches where MuseTalk expects to find models
-TARGET_DIR = "./Musetalk/models"
-
-def download_models():
-    print(f"🚀 Downloading MuseTalk models to {TARGET_DIR}...")
-    
-    # Ensure directory exists
-    os.makedirs(TARGET_DIR, exist_ok=True)
-
-    # Download specific subfolders from the official repo
-    # This gets: musetalk, dwpose, sd-vae-ft-mse, and face-parse-bisent
-    snapshot_download(
-        repo_id="TMElyralab/MuseTalk",
-        local_dir=TARGET_DIR,
-        allow_patterns=[
-            "musetalk/*", 
-            "sd-vae-ft-mse/*", 
-            "dwpose/*", 
-            "face-parse-bisent/*"
-        ],
-        local_dir_use_symlinks=False  # Crucial for Docker
-    )
-    
-    print("✅ All MuseTalk models downloaded successfully.")
-
-if __name__ == "__main__":
-    download_models()
+import os
+import requests
+from huggingface_hub import hf_hub_download
+
+# Target directory matches where MuseTalk expects to find models
+TARGET_DIR = "./Musetalk/models"
+
+def download_file(url, dest_path):
+    """Download a file from URL to destination path."""
+    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+    print(f"📥 Downloading {os.path.basename(dest_path)}...")
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+    with open(dest_path, 'wb') as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
+    print(f"✅ {os.path.basename(dest_path)} downloaded")
+
+def download_models():
+    print(f"🚀 Downloading MuseTalk models to {TARGET_DIR}...")
+    os.makedirs(TARGET_DIR, exist_ok=True)
+
+    # 1. DWPose models
+    print("\n📦 Downloading DWPose models...")
+    download_file(
+        "https://huggingface.co/yzd-v/DWPose/resolve/main/dw-ll_ucoco_384.pth",
+        f"{TARGET_DIR}/dwpose/dw-ll_ucoco_384.pth"
+    )
+    download_file(
+        "https://huggingface.co/yzd-v/DWPose/resolve/main/yolox_l.onnx",
+        f"{TARGET_DIR}/dwpose/yolox_l.onnx"
+    )
+
+    # 2. Face Parsing models
+    print("\n📦 Downloading Face Parsing models...")
+    download_file(
+        "https://huggingface.co/leonelhs/faceparser/resolve/main/79999_iter.pth",
+        f"{TARGET_DIR}/face-parse-bisent/79999_iter.pth"
+    )
+    download_file(
+        "https://download.pytorch.org/models/resnet18-5c106cde.pth",
+        f"{TARGET_DIR}/face-parse-bisent/resnet18-5c106cde.pth"
+    )
+
+    # 3. VAE models
+    print("\n📦 Downloading VAE models...")
+    download_file(
+        "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/config.json",
+        f"{TARGET_DIR}/sd-vae-ft-mse/config.json"
+    )
+    download_file(
+        "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.bin",
+        f"{TARGET_DIR}/sd-vae-ft-mse/diffusion_pytorch_model.bin"
+    )
+    download_file(
+        "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.safetensors",
+        f"{TARGET_DIR}/sd-vae-ft-mse/diffusion_pytorch_model.safetensors"
+    )
+
+    # 4. MuseTalk checkpoints
+    print("\n📦 Downloading MuseTalk checkpoints...")
+    hf_hub_download(
+        repo_id="TMElyralab/MuseTalk",
+        filename="musetalk/musetalk.json",
+        local_dir=TARGET_DIR,
+        local_dir_use_symlinks=False
+    )
+    hf_hub_download(
+        repo_id="TMElyralab/MuseTalk",
+        filename="musetalk/pytorch_model.bin",
+        local_dir=TARGET_DIR,
+        local_dir_use_symlinks=False
+    )
+
+    # 5. Whisper model (CRITICAL - was missing!)
+    print("\n📦 Downloading Whisper model...")
+    download_file(
+        "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt",
+        f"{TARGET_DIR}/whisper/tiny.pt"
+    )
+
+    print("\n✅ All MuseTalk models downloaded successfully!")
+    print(f"📂 Models location: {TARGET_DIR}")
+
+if __name__ == "__main__":
+    download_models()