download_models.py

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()