open-standards/wai/tools/wai_validate_synthetic.py
Transaction Science 2579c9c42d
All checks were successful
standards workspaces build + test / jouleclaw-rs (pull_request) Successful in 1m38s
standards workspaces build + test / aac-rs (pull_request) Successful in 32s
standards workspaces build + test / joule-code-rs (pull_request) Successful in 34s
standards workspaces build + test / joulecontract-rs (pull_request) Successful in 32s
standards workspaces build + test / eoc-rs (pull_request) Successful in 1m48s
standards workspaces build + test / map-rs (pull_request) Successful in 30s
standards workspaces build + test / mesh-rs (pull_request) Successful in 26s
standards workspaces build + test / jcp-rs (pull_request) Successful in 1m17s
standards workspaces build + test / openpay (pull_request) Successful in 1m19s
standards workspaces build + test / proof-rs (pull_request) Successful in 1m11s
standards workspaces build + test / sandbox-rs (pull_request) Successful in 2m8s
standards workspaces build + test / arl-rs (pull_request) Successful in 2m24s
standards workspaces build + test / smart-byte-rs (pull_request) Successful in 2m37s
WAI neural decoders end-to-end / synthetic fixtures → Rust decode → reference (±4) (pull_request) Successful in 1m5s
standards workspaces build + test / joule-ui-rs (pull_request) Successful in 2m44s
standards workspaces build + test / jouledesk-rs (pull_request) Successful in 3m15s
standards workspaces build + test / joulehook-rs (pull_request) Successful in 3m17s
standards workspaces build + test / sense-rs (pull_request) Successful in 3m29s
standards workspaces build + test / macOS energy backend (jcp-energy) (pull_request) Successful in 2m29s
wai-neural-e2e: make the job run, and fix the decoder mismatch it was hiding
This job has failed at ~1 minute on every run, so its reference comparisons have
never executed in CI. Two root causes, both now established from evidence — the
first from the job's actual log.

1. The Rust test could not link on the runner.

The job reaches its last step and fails there. From the job log (forge 16.0.5
exposes `GET /repos/{o}/{r}/actions/jobs/{id}/logs`):

    rust-lld: error: undefined symbol: __isoc23_strtoull
      >>> referenced by parser.cc ... in libort_sys-....rlib
    rust-lld: error: undefined symbol:
      std::__cxx11::basic_string<wchar_t>::_M_replace_cold(...)

`ort`'s `download-binaries` feature links a prebuilt static onnxruntime (1.28.0
for ort 2.0.0-rc.13) built against glibc >= 2.38 and GCC 13's libstdc++. The
runner's default job image, catthehacker/ubuntu:act-22.04, has glibc 2.35 and
GCC 11. All five undefined symbols were checked in both images: act-22.04 lacks
all five, act-24.04 (glibc 2.39, g++ 13) provides all five:

    __isoc23_strtol  __isoc23_strtoll  __isoc23_strtoull      glibc 2.38+
    basic_string<char|wchar_t>::_M_replace_cold               libstdc++ GCC 13+

The job now runs in catthehacker/ubuntu:act-24.04. No Rust or Cargo change, so
nothing here touches the files the landing chain edits. It ran in ~1 minute
because the runner's Python caches are warm; the fast failure was never early.

A correction: an earlier version of this commit replaced actions/setup-python
with uv, on the theory — reached by elimination, without the log — that
setup-python was the fault. The log refutes that: setup-python worked in every
run ("Create Python 3.11.16 folder"), and the old runs fail at the identical
link error. The workflow is back on setup-python exactly as main has it. My
elimination had verified the Python side in Linux containers but built the Rust
side natively on macOS, where the prebuilt is a Mac binary and glibc never
enters — so it cleared everything except the one step it never reproduced.

2. Encode the fixtures the way the decoder evaluates them.

With the link fixed, the comparisons would have exposed a real mismatch. The
Rust decoder opens every ONNX model via `ort_session::load_deterministic` — one
intra-op thread, one inter-op thread, full optimisation, on purpose. The fixture
generator used onnxruntime's defaults, which use every core; a convolution's
reduction is split across threads and its float output shifts.

For MLIC++ that is fatal. Traced symbol by symbol, encoder and decoder agree on
the first 416 of 12,288 symbols — bypass-coded ones included — then diverge
inside the first pass, before anything is fed back:

    element (ch 3, row 4, col 0)     scale         CDF
    encoder, all cores               0.78787208    16
    decoder, single thread           0.78788924    17
    scale_table[16]                  0.78788900

One scale straddles one threshold and the arithmetic-coded stream desyncs:
11,841 symbols wrong, max|delta|=255. On main the garbage overflows the bypass
loop and panics at rans.rs:77; the landing chain bounds that loop, so there it
decodes silently to a wrong image.

Ruled out, each by experiment: the ML stack and numpy 2 -> 1.26 downgrade
(reproduced on linux/amd64); onnxruntime version drift (encoder pinned to the
decoder's 1.28.0 — still broken); the bypass decode (first 64 symbols match);
the hyper-latent path (decoded z bit-identical, FNV over 768 values);
`build_index` (scale_table[0] is exactly 0.11).

All 12 session sites now go through `_session`, mirroring load_deterministic.
MLIC++ max|delta| 255 -> 0; DCVC 1 -> 0. ELIC's residual max|delta|=1 does NOT
change, so it has another source (likely a final 8-bit rounding tie); it is
inside the +/-4 tolerance and is not claimed as fixed.

3. A correction to this PR's own guard.

The filter `_matches_reference` selects six tests, but
`snac_real_model_matches_reference` needs real pretrained SNAC weights this
synthetic job never builds, so it would skip on every run and keep the job red
with every decoder correct. I had verified the guard fires with no fixtures, but
never ran it with fixtures present. It is now excluded by name, with the reason,
so the guard stays strict for the five in scope.

Also corrects main()'s comment, which still described the old
"exit non-zero only if ALL failed" rule.

Verified on main's decoder with the job's exact commands: 5/5 fixtures (exit 0),
5 passed / 0 failed, guard clean. Falsified: removing only the threading pin
brings the MLIC++ panic back; deleting one in-scope fixture makes the guard fail
the step. The container change is verified by this PR's CI run.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-09-29 09:46:15 -04:00

441 lines
22 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""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("<I", len(mb)) + mb + struct.pack("<I", len(payload)) + payload)
def _session(path):
"""Open an ONNX model exactly as the Rust decoder does.
`wai-rs/src/ort_session.rs::load_deterministic` evaluates every model with one
intra-op thread, one inter-op thread and full graph optimisation, because float
ONNX inference is only reproducible that way. The fixtures must be ENCODED under
the same settings. With onnxruntime's defaults (every core) a convolution's
reduction is split across threads, the summation order changes and the float
output shifts. Where a scale then straddles a threshold of the Gaussian scale
table, encoder and decoder select adjacent CDFs, and the arithmetic-coded stream
desyncs from that symbol onward.
Measured, not hypothetical: under default sessions the MLIC++ fixture diverged at
symbol 416 of 12,288 -- encoded with scale 0.78787208 against CDF 16, decoded
with 0.78788924 against CDF 17, either side of scale_table[16] = 0.78788900 --
and decoded to max|delta|=255. Under these settings it decodes exactly, and
DCVC's residual max|delta|=1 goes to 0.
"""
so = ort.SessionOptions()
so.intra_op_num_threads = 1
so.inter_op_num_threads = 1
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
return ort.InferenceSession(str(path), sess_options=so, providers=["CPUExecutionProvider"])
def onx(sess, t):
return torch.from_numpy(sess.run(None, {sess.get_inputs()[0].name: t.numpy().astype(np.float32)})[0])
def roundtrip_elic(net, x, mdir, size):
mdir.mkdir(parents=True, exist_ok=True)
gc = net.gaussian_conditional
export_onnx(net.h_s, torch.zeros(1, 48, 4, 4), mdir / "hyper_synthesis.onnx")
export_onnx(net.context, torch.zeros(1, 48, 16, 16), mdir / "context.onnx")
export_onnx(net.g_s, torch.zeros(1, 48, 16, 16), mdir / "synthesis.onnx")
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
(mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
hs = _session(mdir / "hyper_synthesis.onnx")
cx = _session(mdir / "context.onnx")
gs = _session(mdir / "synthesis.onnx")
with torch.no_grad():
y = net.g_a(x); z = net.h_a(y)
zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
base = onx(hs, z_hat); sb, mb = base.chunk(2, 1)
_, Cy, Sy, _ = y.shape
yy, xx = torch.meshgrid(torch.arange(Sy), torch.arange(Sy), indexing="ij")
anc = ((yy + xx) % 2 == 0).numpy()
syb = torch.round(y - mb).to(torch.int32)
ya = ((syb.float() + mb) * torch.from_numpy(anc)).to(y.dtype)
ctx = onx(cx, ya); sc, mc = ctx.chunk(2, 1); syc = torch.round(y - mc).to(torch.int32)
ib = gc.build_indexes(sb)[0]; ic = gc.build_indexes(sc)[0]; syb, syc = syb[0], syc[0]
syms, idxs = [], []
for isa in (True, False):
st, it = (syb, ib) if isa else (syc, ic)
for ch in range(Cy):
for r in range(Sy):
for c in range(Sy):
if bool(anc[r, c]) == isa:
syms.append(int(st[ch, r, c])); idxs.append(int(it[ch, r, c]))
enc = BufferedRansEncoder()
enc.encode_with_indexes(syms, idxs, gc._quantized_cdf.tolist(), gc._cdf_length.tolist(), gc._offset.tolist())
yb = enc.flush()
yhat = (syb.float() + mb) * torch.from_numpy(anc) + (syc.float() + mc) * torch.from_numpy(~anc)
img = onx(gs, yhat if yhat.dim() == 4 else yhat.unsqueeze(0))
_, Cz, Sz, _ = z.shape
hdr = struct.pack("<IIIIHHBBH", size, size, len(zs[0]), len(yb), Cz, Cy, 2, 4, 3)
write_wai(Path("/tmp/test.elic.wai"), "wai.neural.elic", "elic", hdr, zs[0], yb, size, size)
(np.clip(img[0].numpy(), 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8).tofile("/tmp/ref_elic.rgb")
return len(yb), len(syms)
def roundtrip_mlicpp(net, x, mdir, size, K=4):
mdir.mkdir(parents=True, exist_ok=True)
gc = net.gaussian_conditional
export_onnx(net.h_s, torch.zeros(1, 48, 4, 4), mdir / "hyper_synthesis.onnx")
export_onnx(net.entropy, torch.zeros(1, 144, 16, 16), mdir / "entropy.onnx")
export_onnx(net.g_s, torch.zeros(1, 48, 16, 16), mdir / "synthesis.onnx")
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
(mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
hs = _session(mdir / "hyper_synthesis.onnx")
en = _session(mdir / "entropy.onnx")
gs = _session(mdir / "synthesis.onnx")
with torch.no_grad():
y = net.g_a(x); z = net.h_a(y)
zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
hyper = onx(hs, z_hat); _, Cy, Sy, _ = y.shape; sw = Cy // K
yy, xx = torch.meshgrid(torch.arange(Sy), torch.arange(Sy), indexing="ij")
anc = ((yy + xx) % 2 == 0).numpy()
y_dec = torch.zeros(1, Cy, Sy, Sy); y0 = y[0]; syms, idxs = [], []
for k in range(K):
c0, c1 = k * sw, (k + 1) * sw
for isa in (True, False):
params = onx(en, torch.cat([y_dec, hyper], 1)); sc, mc = params.chunk(2, 1)
idx = gc.build_indexes(sc)[0]; m0 = mc[0]
for ch in range(c0, c1):
for r in range(Sy):
for c in range(Sy):
if bool(anc[r, c]) == isa:
sym = int(round(float(y0[ch, r, c] - m0[ch, r, c])))
syms.append(sym); idxs.append(int(idx[ch, r, c]))
y_dec[0, ch, r, c] = sym + m0[ch, r, c]
enc = BufferedRansEncoder()
enc.encode_with_indexes(syms, idxs, gc._quantized_cdf.tolist(), gc._cdf_length.tolist(), gc._offset.tolist())
yb = enc.flush(); img = onx(gs, y_dec)
_, Cz, Sz, _ = z.shape
hdr = struct.pack("<IIIIHHBBHH", size, size, len(zs[0]), len(yb), Cz, Cy, 2, 4, 3, K)
write_wai(Path("/tmp/test.mlic.wai"), "wai.neural.mlicpp", "mlicpp", hdr, zs[0], yb, size, size)
(np.clip(img[0].numpy(), 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8).tofile("/tmp/ref_mlic.rgb")
return len(yb), len(syms)
class SynthSnac(nn.Module):
"""SNAC's multi-rate-sum decode interface (3 levels, ×4/×2/×1 → finest),
with clean ONNX ops only — the real SNAC decoder's noise/attention blocks
don't export (a sink-side concern). Validates snac.rs's multi-level wire
parse + dynamic codes_0..codes_{N-1} dispatch."""
HOP = 512
def forward(self, c0, c1, c2):
f0 = c0.float().repeat_interleave(4, dim=1)
f1 = c1.float().repeat_interleave(2, dim=1)
z = (f0 + f1 + c2.float()) / 12000.0
return z.unsqueeze(1).repeat_interleave(self.HOP, dim=2)
def roundtrip_snac(mdir: Path):
mdir.mkdir(parents=True, exist_ok=True)
sr, lengths = 24_000, [12, 24, 48]
nsamp = lengths[-1] * SynthSnac.HOP
rng = np.random.RandomState(7)
codes = [torch.from_numpy(rng.randint(0, 4096, size=(1, L)).astype(np.int64)) for L in lengths]
names = [f"codes_{i}" for i in range(3)]
torch.onnx.export(SynthSnac().eval(), tuple(codes), str(mdir / "decoder.onnx"),
input_names=names, output_names=["audio_values"],
dynamic_axes={**{f"codes_{i}": {1: f"t{i}"} for i in range(3)},
"audio_values": {2: "s"}}, opset_version=17, dynamo=False)
sess = _session(mdir / "decoder.onnx")
ref = sess.run(None, {f"codes_{i}": codes[i].numpy() for i in range(3)})[0]
hdr = struct.pack("<III", sr, nsamp, 3) + b"".join(struct.pack("<I", L) for L in lengths)
body = b"".join(codes[i][0].numpy().astype(np.uint16).tobytes() for i in range(3))
man = (b'{"wai":"1.0","media":"audio","intent":"replicate","model_requirement":'
b'{"capability":"wai.neural.snac"},"conditioning":{"kind":"snac"}}')
payload = hdr + body
Path("/tmp/test.snac.wai").write_bytes(
b"WAI1" + struct.pack("<I", len(man)) + man + struct.pack("<I", len(payload)) + payload)
ref.astype(np.float32).tofile("/tmp/ref_snac.f32")
return nsamp
class MiniDCVC(CompressionModel):
"""Hyperprior entropy + a temporal decoder (latents, ref_feature) →
(image, feature). Validates the inter-frame frame-loop orchestration."""
def __init__(self, M=48, N=48, Cf=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.t_feat = nn.Conv2d(M + Cf, Cf, 3, 1, 1)
up = lambda ci, co: nn.ConvTranspose2d(ci, co, 5, 2, 2, output_padding=1)
self.t_img = nn.Sequential(up(M + Cf, M), nn.ReLU(), up(M, M), nn.ReLU(),
up(M, M), nn.ReLU(), up(M, 3))
class _TDec(nn.Module):
def __init__(s, m): super().__init__(); s.m = m
def forward(s, latents, ref_feature):
x = torch.cat([latents, ref_feature], 1)
return s.m.t_img(x), torch.relu(s.m.t_feat(x))
def roundtrip_dcvc(net, x, mdir: Path, size, Cf=48):
mdir.mkdir(parents=True, exist_ok=True)
gc = net.gaussian_conditional
export_onnx(net.h_s, torch.zeros(1, 48, 4, 4), mdir / "hyper_synthesis.onnx")
torch.onnx.export(_TDec(net).eval(), (torch.zeros(1, 48, 16, 16), torch.zeros(1, Cf, 16, 16)),
str(mdir / "temporal_decoder.onnx"), input_names=["latents", "ref_feature"],
output_names=["image", "feature"],
dynamic_axes={"latents": {2: "h", 3: "w"}, "ref_feature": {2: "h", 3: "w"},
"image": {2: "h", 3: "w"}, "feature": {2: "h", 3: "w"}},
opset_version=17, dynamo=False)
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
(mdir / "gauss_cdfs_q3.json").write_text(json.dumps(gblob(gc), separators=(",", ":")))
hs = _session(mdir / "hyper_synthesis.onnx")
gs = _session(mdir / "temporal_decoder.onnx")
frames = [np.roll(np.asarray(x[0].permute(1, 2, 0)), shift=8 * k, axis=1) for k in range(3)]
F = torch.zeros(1, Cf, 16, 16); fcodes = []; ref = []
for fr in frames:
xi = torch.from_numpy(np.ascontiguousarray(fr)).permute(2, 0, 1).unsqueeze(0)
with torch.no_grad():
y = net.g_a(xi); z = net.h_a(y); zs = net.entropy_bottleneck.compress(z)
z_hat = net.entropy_bottleneck.decompress(zs, z.size()[-2:])
params = onx(hs, z_hat); sc, mn = params.chunk(2, 1)
ys = gc.compress(y, gc.build_indexes(sc), means=mn); y_hat = gc.decompress(ys, gc.build_indexes(sc), means=mn)
img, F = onx(gs, [y_hat, F]) if False else [torch.from_numpy(t) for t in gs.run(None, {"latents": y_hat.numpy(), "ref_feature": F.numpy()})]
fcodes.append((zs[0], ys[0])); ref.append((np.clip(img[0].numpy(), 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8))
_, Cz, Sz, _ = z.shape; _, Cy, Sy, _ = y.shape
hdr = struct.pack("<IIIIHHBBHH", size, size, 3, 30000, Cy, Cz, 4, 2, 3, Cf)
table = b"".join(struct.pack("<II", len(zb), len(yb)) for zb, yb in fcodes)
body = b"".join(zb + yb for zb, yb in fcodes)
write_wai(Path("/tmp/test.dcvc.wai"), "wai.neural.dcvc_rt", "dcvc_rt", hdr, b"", table + body, size, size)
np.concatenate([r.reshape(-1) for r in ref]).astype(np.uint8).tofile("/tmp/ref_dcvc.rgb")
return 3
from compressai.entropy_models.entropy_models import pmf_to_quantized_cdf # noqa: E402
_RL = 1 << 31
def _rans_cat_encode(symbols, cdfs):
"""Bypass-free categorical rANS matching wai-rs decode_categorical."""
state, out = _RL, []
for sym, cdf in zip(reversed(symbols), reversed(cdfs)):
start, freq = cdf[sym], cdf[sym + 1] - cdf[sym]
x_max = ((_RL >> 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("<I", w) for w in words)
class MiniGLC(CompressionModel):
"""Categorical/VQ entropy: hyper + checkerboard context predict per-position
logits over a K-codebook; synthesis maps the index grid → image. Logits are
ROUNDED in-graph so encoder/decoder ORT builds agree (categorical CDFs are
razor-sensitive — the determinism contract, sharper than Gaussian)."""
def __init__(self, K=16, Cz=32):
super().__init__()
self.entropy_bottleneck = EntropyBottleneck(Cz)
up = lambda ci, co: nn.ConvTranspose2d(ci, co, 5, 2, 2, output_padding=1)
self.hyper_entropy = nn.Sequential(up(Cz, 64), nn.ReLU(), up(64, K))
self.context_entropy = nn.Conv2d(1, K, 5, 1, 2)
self.synthesis = nn.Sequential(up(1, 32), nn.ReLU(), up(32, 32), nn.ReLU(),
up(32, 32), nn.ReLU(), up(32, 3))
class _Round(nn.Module):
def __init__(s, m): super().__init__(); s.m = m
def forward(s, x): return torch.round(s.m(x))
def roundtrip_glc(net, mdir: Path, size, K=16, Cz=32, Sz=4, Sy=16):
mdir.mkdir(parents=True, exist_ok=True)
torch.onnx.export(_Round(net.hyper_entropy), (torch.zeros(1, Cz, Sz, Sz),), str(mdir / "hyper_entropy.onnx"),
input_names=["i"], output_names=["o"], dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}}, opset_version=17, dynamo=False)
torch.onnx.export(_Round(net.context_entropy), (torch.zeros(1, 1, Sy, Sy),), str(mdir / "context_entropy.onnx"),
input_names=["i"], output_names=["o"], dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}}, opset_version=17, dynamo=False)
export_onnx(net.synthesis, torch.zeros(1, 1, Sy, Sy), mdir / "synthesis.onnx")
(mdir / "hyper_cdfs_q3.json").write_text(json.dumps(fblob(net.entropy_bottleneck), separators=(",", ":")))
he = _session(mdir / "hyper_entropy.onnx"); ce = _session(mdir / "context_entropy.onnx")
gs = _session(mdir / "synthesis.onnx")
rng = np.random.RandomState(7)
sm = lambda l: (lambda e: e / e.sum())(np.exp(l - l.max()))
cdf_of = lambda l: pmf_to_quantized_cdf(torch.tensor(sm(l), dtype=torch.float32), 16).tolist()
z = torch.randn(1, Cz, Sz, Sz); zs = net.entropy_bottleneck.compress(z); z_hat = net.entropy_bottleneck.decompress(zs, (Sz, Sz))
base = he.run(None, {"i": z_hat.numpy()})[0][0]
idx = rng.randint(0, K, size=(Sy, Sy)).astype(np.int64)
anc = np.array([[(y + x) % 2 == 0 for x in range(Sy)] for y in range(Sy)])
ctx = ce.run(None, {"i": np.where(anc, idx, 0).astype(np.float32)[None, None]})[0][0]
syms, cdfs = [], []
for isa in (True, False):
for y in range(Sy):
for x in range(Sy):
if anc[y, x] == isa:
lg = base[:, y, x] if isa else base[:, y, x] + ctx[:, y, x]
syms.append(int(idx[y, x])); cdfs.append(cdf_of(lg))
yb = _rans_cat_encode(syms, cdfs)
ref = gs.run(None, {"i": idx.astype(np.float32)[None, None]})[0]
(np.clip(ref[0], 0, 1).transpose(1, 2, 0) * 255).round().astype(np.uint8).tofile("/tmp/ref_glc.rgb")
hdr = struct.pack("<IIIIHHBBH", size, size, len(zs[0]), len(yb), Cz, K, 2, 4, 3)
write_wai(Path("/tmp/test.glc.wai"), "wai.neural.glc", "glc", hdr, zs[0], yb, size, size)
return len(syms)
def main() -> 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 unless EVERY harness succeeded.
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")
# Every harness must succeed, not merely one of them.
#
# `0 if ok else 1` reported success whenever a single harness worked, and the
# loop above deliberately catches per-harness exceptions to keep going. So a
# 1-of-5 run exited 0, the CI step went green, and the four decoders whose
# fixtures were never written then SKIPPED in the Rust suite — where a skip is
# reported as a pass. The job came out fully green having checked one decoder
# out of five, and nothing in the output said so.
#
# Partial generation is still useful interactively: the per-harness ✓/✗ lines
# are printed either way, and a caller who wants one fixture can ignore the
# exit code. What is not acceptable is CI reading a partial run as a pass.
return 0 if ok == len(harnesses) else 1
if __name__ == "__main__":
raise SystemExit(main())