|
|
|
|
|
"""
|
|
|
Download model weights listed in a config YAML (replicates Dockerfile download behavior without Docker).
|
|
|
|
|
|
Usage:
|
|
|
python download_models.py --config config.production.yaml
|
|
|
|
|
|
This script uses huggingface_hub.hf_hub_download to download specified .pth files to the
|
|
|
model's DOWNLOAD_MODEL_DIR (or ./models by default).
|
|
|
"""
|
|
|
import argparse
|
|
|
import os
|
|
|
import yaml
|
|
|
import time
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
|
|
|
|
|
def main():
|
|
|
parser = argparse.ArgumentParser()
|
|
|
parser.add_argument("--config", default="config.production.yaml")
|
|
|
parser.add_argument("--token", default=None, help="Hugging Face token (optional)")
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
with open(args.config, "r", encoding="utf-8") as f:
|
|
|
cfg = yaml.safe_load(f.read())
|
|
|
|
|
|
models = cfg.get("MODELS", [])
|
|
|
if len(models) == 0:
|
|
|
print("No models found in config. Nothing to download.")
|
|
|
return
|
|
|
|
|
|
for m in models:
|
|
|
repo_id = m.get("DOWNLOAD_MODEL_REPO_ID")
|
|
|
filename = m.get("DOWNLOAD_MODEL_FILE_NAME")
|
|
|
local_dir = m.get("DOWNLOAD_MODEL_DIR", "./models")
|
|
|
|
|
|
if repo_id is None or filename is None:
|
|
|
print(f"Skipping model with incomplete download info: {m}")
|
|
|
continue
|
|
|
|
|
|
os.makedirs(local_dir, exist_ok=True)
|
|
|
print(f"Downloading {filename} from repo {repo_id} into {local_dir} ...")
|
|
|
os.makedirs(local_dir, exist_ok=True)
|
|
|
|
|
|
|
|
|
max_attempts = 5
|
|
|
for attempt in range(1, max_attempts + 1):
|
|
|
try:
|
|
|
path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir, token=args.token)
|
|
|
print(f"Downloaded file to {path}")
|
|
|
break
|
|
|
except Exception as e:
|
|
|
print(f"Attempt {attempt} failed to download {filename} from {repo_id}: {e}")
|
|
|
if attempt < max_attempts:
|
|
|
print(f"Retrying in {attempt*5} seconds...")
|
|
|
time.sleep(attempt * 5)
|
|
|
else:
|
|
|
print(f"Failed after {max_attempts} attempts. Skipping {filename}.")
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main()
|
|
|
|
