File size: 2,115 Bytes
27caffe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | #!/usr/bin/env python3
"""Download ML weights from Hugging Face Hub into local paths expected by the app."""
from __future__ import annotations
import os
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
# Update these after you upload models to your Hugging Face account.
DEFAULT_REPOS = {
"f5tts": "IamSamkk/police-bot-f5tts",
"sadtalker": "IamSamkk/sadtalker-checkpoints",
"gfpgan": "IamSamkk/gfpgan-weights",
}
LOCAL_DIRS = {
"f5tts": ROOT / "my_finetuned_model",
"sadtalker": ROOT / "sadtalker+wav2lip" / "sadtalker" / "checkpoints",
"gfpgan": ROOT / "gfpgan" / "weights",
}
ENV_REPO_KEYS = {
"f5tts": "HF_REPO_F5TTS",
"sadtalker": "HF_REPO_SADTALKER",
"gfpgan": "HF_REPO_GFPGAN",
}
def _repo_id(key: str) -> str:
env_key = ENV_REPO_KEYS[key]
return os.environ.get(env_key, DEFAULT_REPOS[key])
def download_all(token: str | None = None, only: list[str] | None = None) -> None:
try:
from huggingface_hub import snapshot_download
except ImportError:
print("Install huggingface_hub: pip install huggingface_hub", file=sys.stderr)
sys.exit(1)
keys = only if only else list(DEFAULT_REPOS.keys())
for key in keys:
if key not in DEFAULT_REPOS:
print(f"Unknown model key: {key}", file=sys.stderr)
sys.exit(1)
repo_id = _repo_id(key)
target = LOCAL_DIRS[key]
target.mkdir(parents=True, exist_ok=True)
print(f"Downloading {repo_id} -> {target}")
snapshot_download(
repo_id=repo_id,
local_dir=str(target),
token=token or os.environ.get("HF_TOKEN"),
)
print("Done.")
def main() -> None:
import argparse
parser = argparse.ArgumentParser(description="Download model weights from Hugging Face Hub")
parser.add_argument(
"--only",
nargs="+",
choices=list(DEFAULT_REPOS.keys()),
help="Download only selected model groups",
)
args = parser.parse_args()
download_all(only=args.only)
if __name__ == "__main__":
main()
|