"""Reproducible end-to-end validation harness for the elic + mlicpp decoders. elic (`wai.neural.elic`) and mlicpp (`wai.neural.mlicpp`) have no pretrained model in the CompressAI zoo, so we can't cross-check them against a reference codec the way bmshj2018 / mbt2018_mean are checked. Instead we build *minimal CompressAI-compatible* models (real EntropyBottleneck + GaussianConditional, random-but-non-trivial weights), run them through the SAME export + encode logic the real tools use, and emit the `.wai` envelope + a reference reconstruction. The Rust tests `elic_matches_reference` / `mlicpp_matches_reference` then decode the envelope and must match the reference within ORT-version g_s rounding. This makes those two validations reproducible from the repo (the artifacts the Rust tests consume otherwise live only in /tmp). Run before the gated tests: python tools/wai_validate_synthetic.py # writes /tmp artifacts cargo test --features neural --test neural_native_decode _matches_reference Needs torch + compressai + onnxruntime + numpy + Pillow. """ from __future__ import annotations import argparse import json import math import struct from pathlib import Path import numpy as np import torch import torch.nn as nn import onnxruntime as ort from PIL import Image from compressai.models import CompressionModel from compressai.entropy_models import GaussianConditional, EntropyBottleneck from compressai.ans import BufferedRansEncoder def scale_table(): return torch.exp(torch.linspace(math.log(0.11), math.log(256), 64)) def _convs(M, N): dn = lambda ci, co: nn.Conv2d(ci, co, 5, 2, 2) up = lambda ci, co: nn.ConvTranspose2d(ci, co, 5, 2, 2, output_padding=1) g_a = nn.Sequential(dn(3, M), nn.ReLU(), dn(M, M), nn.ReLU(), dn(M, M), nn.ReLU(), dn(M, M)) g_s = nn.Sequential(up(M, M), nn.ReLU(), up(M, M), nn.ReLU(), up(M, M), nn.ReLU(), up(M, 3)) h_a = nn.Sequential(dn(M, N), nn.ReLU(), dn(N, N)) h_s = nn.Sequential(up(N, M), nn.ReLU(), up(M, 2 * M)) return g_a, g_s, h_a, h_s class MiniElic(CompressionModel): def __init__(self, M=48, N=48): super().__init__() self.entropy_bottleneck = EntropyBottleneck(N) self.gaussian_conditional = GaussianConditional(None) self.g_a, self.g_s, self.h_a, self.h_s = _convs(M, N) self.context = nn.Conv2d(M, 2 * M, 5, 1, 2) class MiniMlic(CompressionModel): def __init__(self, M=48, N=48): super().__init__() self.entropy_bottleneck = EntropyBottleneck(N) self.gaussian_conditional = GaussianConditional(None) self.g_a, self.g_s, self.h_a, self.h_s = _convs(M, N) self.entropy = nn.Conv2d(3 * M, 2 * M, 5, 1, 2) def build(model_cls, seed): torch.manual_seed(seed) net = model_cls().eval() for p in net.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p, gain=3.0) # gain 3 → real latent entropy if hasattr(net, "gaussian_conditional"): net.gaussian_conditional.update_scale_table(scale_table()) net.update(force=True) return net def fblob(eb): cl = eb._cdf_length.int().tolist() return {"channels": len(cl), "cdf_length": cl, "offset": eb._offset.int().tolist(), "medians": eb._get_medians().detach().reshape(-1).tolist(), "cdf": [r[:cl[i]] for i, r in enumerate(eb._quantized_cdf.int().tolist())]} def gblob(g): cl = g._cdf_length.int().tolist() return {"scale_table": g.scale_table.reshape(-1).tolist(), "cdf_length": cl, "offset": g._offset.int().tolist(), "cdf": [r[:cl[i]] for i, r in enumerate(g._quantized_cdf.int().tolist())]} def export_onnx(mod, dummy, path): torch.onnx.export(mod, (dummy,), str(path), input_names=["i"], output_names=["o"], dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}}, opset_version=17, dynamo=False) def test_image(path, size): if path.exists(): return x, y = np.meshgrid(np.linspace(0, 1, size), np.linspace(0, 1, size)) r = (x * 255).astype(np.uint8); g = (y * 255).astype(np.uint8) b = (((np.sin(x * 12) * np.cos(y * 12)) * 0.5 + 0.5) * 255).astype(np.uint8) img = np.stack([r, g, b], -1) yy, xx = np.ogrid[:size, :size] img[(yy - size // 2) ** 2 + (xx - size // 2) ** 2 < (size // 5) ** 2] = [200, 40, 120] Image.fromarray(img).save(path) def write_wai(path, cap, kind, header, z, y, H, W): manifest = {"wai": "1.0", "media": "image", "intent": "replicate", "model_requirement": {"capability": cap}, "conditioning": {"kind": kind}, "target": {"w": W, "h": H}} mb = json.dumps(manifest, separators=(",", ":")).encode() payload = header + z + y path.write_bytes(b"WAI1" + struct.pack("> 16) << 32) * freq while state >= x_max: out.append(state & 0xFFFFFFFF); state >>= 32 state = ((state // freq) << 16) + (state % freq) + start words = [state & 0xFFFFFFFF, (state >> 32) & 0xFFFFFFFF] + list(reversed(out)) return b"".join(struct.pack(" int: ap = argparse.ArgumentParser() ap.add_argument("--size", type=int, default=256) args = ap.parse_args() img_path = Path("/tmp/wai_test.png") test_image(img_path, args.size) img = Image.open(img_path).convert("RGB").resize((args.size, args.size), Image.LANCZOS) x = torch.from_numpy(np.asarray(img, np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0) # Each harness is independent; a failure in one must not block the others # (e.g. a synthetic model that only traces at certain sizes). Report a # per-decoder PASS/FAIL summary and exit non-zero only if ALL failed. harnesses = [ ("elic", lambda: roundtrip_elic(build(MiniElic, 0), x, Path("/tmp/wai_elic"), args.size)), ("mlicpp", lambda: roundtrip_mlicpp(build(MiniMlic, 1), x, Path("/tmp/wai_mlic"), args.size)), ("snac", lambda: roundtrip_snac(Path("/tmp/wai_snac"))), ("dcvc", lambda: roundtrip_dcvc(build(MiniDCVC, 2), x, Path("/tmp/wai_dcvc"), args.size)), ("glc", lambda: roundtrip_glc(build(MiniGLC, 3), Path("/tmp/wai_glc"), args.size)), ] ok = 0 for name, fn in harnesses: try: fn() print(f" ✓ {name}") ok += 1 except Exception as e: # noqa: BLE001 — report and continue print(f" ✗ {name}: {type(e).__name__}: {e}") print(f"generated {ok}/{len(harnesses)} synthetic neural fixtures into /tmp") print("now: cargo test --features neural --test neural_native_decode -- _matches_reference") return 0 if ok else 1 if __name__ == "__main__": raise SystemExit(main())