62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
from huggingface_hub import snapshot_download
|
|
|
|
from plan_hf_downloads import remote_plan
|
|
from common import iter_model_card_paths, read_yaml_mapping
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Download missing Hugging Face model weights from model cards.")
|
|
parser.add_argument("--model", action="append", dest="models", help="Model card id to download. Can be repeated.")
|
|
parser.add_argument("--include-ready", action="store_true")
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
selected = set(args.models or [])
|
|
token = os.environ.get("HF_TOKEN") or None
|
|
plans = []
|
|
for card_path in iter_model_card_paths():
|
|
card = read_yaml_mapping(card_path)
|
|
card_id = str(card.get("id") or "")
|
|
if selected and card_id not in selected:
|
|
continue
|
|
plan = remote_plan(card_path, include_ready=args.include_ready)
|
|
if plan and plan["missing_bytes"]:
|
|
plans.append(plan)
|
|
|
|
if not plans:
|
|
print("No missing downloads.")
|
|
return 0
|
|
|
|
for plan in plans:
|
|
missing_files = [item["name"] for item in plan["files"] if not item["complete"]]
|
|
if not missing_files:
|
|
continue
|
|
print("")
|
|
print(f"== {plan['id']} ==")
|
|
print(f"Repository: {plan['repo_id']}")
|
|
print(f"Destination: {plan['storage_path']}")
|
|
print(f"Missing: {plan['missing_gb']} GB")
|
|
for name in missing_files:
|
|
print(f"- {name}")
|
|
if args.dry_run:
|
|
continue
|
|
snapshot_download(
|
|
repo_id=plan["repo_id"],
|
|
local_dir=plan["storage_path"],
|
|
token=token,
|
|
allow_patterns=missing_files,
|
|
ignore_patterns=["*.msgpack", "*.h5", "*.ot", "*.tflite", "*.onnx", "*.pb", "*.ckpt"],
|
|
)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|