| |
| """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() |
|
|
