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
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]>
441 lines
22 KiB
Python
441 lines
22 KiB
Python
"""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())
|