scripts/verify_artifacts.py

#!/usr/bin/env python3
"""Verify released RAESR artifacts without downloading or constructing the model."""

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 (
    DINO_S_WEIGHT_NAME,
    DINO_WEIGHT_NAME,
    DELIVERY_ROOT as PACKAGE_ROOT,
)
from raesr.artifacts import (
    runtime_record,
    verify_decoder,
    verify_dinov3,
    verify_frontend,
    verify_runtime,
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--decoder", type=Path, default=PACKAGE_ROOT / "artifacts" / "ema_decoder.pt")
    parser.add_argument(
        "--frontend",
        type=Path,
        default=PACKAGE_ROOT / "weights" / "raesr_d_frontend.pt",
        help="RAESR-D frontend weights file",
    )
    parser.add_argument("--model", choices=("raesr", "raesr-d"), default="raesr")
    parser.add_argument("--dinov3-root", type=Path)
    parser.add_argument("--dinov3-weights", type=Path)
    args = parser.parse_args()

    result = {
        "decoder": verify_decoder(args.decoder),
        "frontend": verify_frontend(args.frontend),
        "runtime": runtime_record(),
        "runtime_verification": verify_runtime(),
    }
    if (args.dinov3_root is None) != (args.dinov3_weights is None):
        parser.error("--dinov3-root and --dinov3-weights must be supplied together")
    if args.dinov3_root is not None:
        weight_name = DINO_WEIGHT_NAME if args.model == "raesr" else DINO_S_WEIGHT_NAME
        result["dinov3"] = verify_dinov3(args.dinov3_root, args.dinov3_weights, weight_name)
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()