46c5688b86
Commit the WAI standard source: a content + capability-dispatch standard
whose thesis is "carries whatever it can reconstruct, verifies whatever
it carries." Source only — large vendored/model assets and build output
are gitignored.
Reconstruct classes (each: spec + conformance corpus + in-tab browser
sink that re-derives the corpus hash, one wai.det.fixed64 fixed-point
floor, one *-equivalence apiece):
- interactive worlds (replay-equivalence) + 3 sim registrants,
multiplayer over a relay, session receipts, glTF + real SPZ assets
- spatial audio (mixdown-equivalence)
- haptics (sample-equivalence)
- state-feeds (replay-equivalence, non-physics sim)
- avatars / digital humans (pose-equivalence) + telepresence over the relay
- volumetric / 4D splats (frame-state-equivalence) + WebGPU EWA render,
real-SPZ-derived clip
Verify half (lever 4):
- asset provenance: a signed, energy-accounted lineage for opaque
captured content WAI does not reconstruct; tamper-evident, in-tab.
Cross-cutting: JWP receipt profiles per class; a native energy meter
(IOReport via macmon, no sudo) emitting receipts with REAL measured
joules; the learned-image/neural codec family and JPEG AI entropy port
from earlier work.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
395 lines
20 KiB
Python
395 lines
20 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 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 = ort.InferenceSession(str(mdir / "hyper_synthesis.onnx"))
|
||
cx = ort.InferenceSession(str(mdir / "context.onnx"))
|
||
gs = ort.InferenceSession(str(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 = ort.InferenceSession(str(mdir / "hyper_synthesis.onnx"))
|
||
en = ort.InferenceSession(str(mdir / "entropy.onnx"))
|
||
gs = ort.InferenceSession(str(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 = ort.InferenceSession(str(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 = ort.InferenceSession(str(mdir / "hyper_synthesis.onnx"))
|
||
gs = ort.InferenceSession(str(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 = ort.InferenceSession(str(mdir / "hyper_entropy.onnx")); ce = ort.InferenceSession(str(mdir / "context_entropy.onnx"))
|
||
gs = ort.InferenceSession(str(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)
|
||
|
||
ny, ns = roundtrip_elic(build(MiniElic, 0), x, Path("/tmp/wai_elic"), args.size)
|
||
print(f"elic → /tmp/test.elic.wai (y={ny}B, {ns} symbols)")
|
||
ny, ns = roundtrip_mlicpp(build(MiniMlic, 1), x, Path("/tmp/wai_mlic"), args.size)
|
||
print(f"mlicpp → /tmp/test.mlic.wai (y={ny}B, {ns} symbols)")
|
||
nsamp = roundtrip_snac(Path("/tmp/wai_snac"))
|
||
print(f"snac → /tmp/test.snac.wai ({nsamp} samples, 3 levels)")
|
||
roundtrip_dcvc(build(MiniDCVC, 2), x, Path("/tmp/wai_dcvc"), args.size)
|
||
print("dcvc → /tmp/test.dcvc.wai (3 inter-frames, feature propagation)")
|
||
n = roundtrip_glc(build(MiniGLC, 3), Path("/tmp/wai_glc"), args.size)
|
||
print(f"glc → /tmp/test.glc.wai ({n} categorical indices)")
|
||
print("artifacts written; now: cargo test --features neural --test neural_native_decode _matches_reference")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|