File size: 3,014 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | #!/usr/bin/env python3
"""Exercise package file contracts; optionally execute one real model forward."""
from __future__ import annotations
import argparse
import subprocess
import sys
import tempfile
from pathlib import Path
DELIVERY_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(DELIVERY_ROOT))
import numpy as np
from PIL import Image
from benchmark.validate_outputs import validate_outputs
from raesr.artifacts import verify_decoder, verify_frontend
from scripts.infer import collect_images, is_valid_output
def make_image(path: Path, size: tuple[int, int]) -> None:
y, x = np.indices((size[1], size[0]))
array = np.stack(((x * 3) % 256, (y * 5) % 256, ((x + y) * 7) % 256), axis=-1).astype(np.uint8)
Image.fromarray(array).save(path)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--decoder", type=Path)
parser.add_argument("--frontend", type=Path)
parser.add_argument("--model", choices=("raesr", "raesr-d"), default="raesr")
parser.add_argument("--dinov3-root", type=Path)
parser.add_argument("--dinov3-weights", type=Path)
parser.add_argument("--device", default="cuda")
parser.add_argument("--with-model", action="store_true")
args = parser.parse_args()
if args.with_model and (args.decoder is None or args.dinov3_root is None or args.dinov3_weights is None):
parser.error("--with-model requires --decoder, --dinov3-root, and --dinov3-weights")
if args.with_model and args.model == "raesr-d" and args.frontend is None:
parser.error("--with-model --model raesr-d requires --frontend")
with tempfile.TemporaryDirectory(prefix="raesr-smoke-") as temporary:
root = Path(temporary)
inputs, outputs = root / "inputs", root / "outputs"
inputs.mkdir()
outputs.mkdir()
make_image(inputs / "fixture.png", (63, 47))
assert len(collect_images(inputs)) == 1
make_image(outputs / "fixture.png", (512, 512))
assert is_valid_output(outputs / "fixture.png")
validate_outputs(inputs, outputs, expected_count=1)
if args.decoder:
verify_decoder(args.decoder)
if args.frontend:
verify_frontend(args.frontend)
if args.with_model:
outputs.mkdir(exist_ok=True)
command = [
sys.executable, str(DELIVERY_ROOT / "scripts" / "infer.py"),
"--input-dir", str(inputs), "--output-dir", str(outputs), "--model", args.model,
"--decoder", str(args.decoder),
"--dinov3-root", str(args.dinov3_root), "--dinov3-weights", str(args.dinov3_weights),
"--device", args.device, "--overwrite",
]
if args.frontend is not None:
command += ["--frontend", str(args.frontend)]
subprocess.run(command, check=True)
assert is_valid_output(outputs / "fixture.png")
print("smoke test passed")
if __name__ == "__main__":
main()
|