scripts/download_artifacts.py

#!/usr/bin/env python3
"""Download and verify the released RAESR and RAESR-D checkpoints."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

DELIVERY_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DELIVERY_ROOT))

from raesr.artifacts import decoder_spec, frontend_spec, verify_decoder, verify_frontend


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=DELIVERY_ROOT / "artifacts")
    parser.add_argument("--revision", help="override the repository revision recorded in the manifest")
    parser.add_argument("--local-files-only", action="store_true")
    args = parser.parse_args()

    try:
        from huggingface_hub import hf_hub_download
    except ImportError as error:
        raise SystemExit("Install huggingface-hub from environment/requirements-runtime.txt") from error

    args.output_dir.mkdir(parents=True, exist_ok=True)
    records = {}
    for name, spec, verify in (
        ("decoder", decoder_spec(), verify_decoder),
        ("raesr_d_frontend", frontend_spec(), verify_frontend),
    ):
        path = Path(
            hf_hub_download(
                repo_id=spec["repo_id"],
                filename=spec["filename"],
                revision=args.revision or spec["revision"],
                local_dir=str(args.output_dir),
                local_files_only=args.local_files_only,
            )
        )
        records[name] = verify(path)
    print(json.dumps(records, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()