File size: 1,612 Bytes
45d3d29 | 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 | #!/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()
|