fb103ac5ae
Every registered string now lives in one table (codecs/registry.rs) with its kind, status, media class, registering SPEC section, edition and registration date. `wai registry export` writes registry/capabilities.json, its JSON Schema, the generated table in REGISTRY.md and vocab v2; --check fails on drift. vocab v1 stays stable. A candidate capability (a registered string whose payload format no revision has pinned) must not be emitted. The C packers, the browser writers and `wai wrap` refuse one, and `codecs::is_dispatchable` is the single dispatch and selection predicate (SPEC §4): candidates, companion-only capabilities, records and derivations are not dispatchable, and an unregistered string is decoded only where the sink advertises it. The codec glossary is linked to the registry through a reviewed table and held to its schema by a library test; `org` becomes `body`, from a closed allowlist. The SPEC gains §5 sections on kinds and candidates, a §8 edition rule, corrected citations and rewritten appendices. The quantum-ops vectors are regenerated with neutral device ids. CI declares the feature-gated test targets with required-features, checks the registry export, builds wai-web for wasm32, and runs two public-text checks (comparative prose, and private name patterns) with a self-test. The specification, extensions, research notes, READMEs, glossary and source comments are rewritten to name no company or product and to make no comparisons. Co-Authored-By: Claude Opus 5.5 <[email protected]>
376 lines
22 KiB
Python
376 lines
22 KiB
Python
#!/usr/bin/env python3
|
||
"""Export the pretrained mbt2018_mean mean-scale hyperprior (CompressAI zoo,
|
||
quality 6, MSE) as an INTEGER-exact decoder and measure what that costs:
|
||
per-channel int12 weights + fixed-point activations + integer IGDN, all i64 —
|
||
byte-identical on every machine (NeuralIntegerExact). The tool reports the
|
||
integer and float PSNR on the given image; tests/int_hyper_synth.rs requires the
|
||
integer decode within 1 dB of the float decode.
|
||
|
||
python tools/wai_sota_int_synth_export.py <image.png> <out-dir>
|
||
|
||
Normally run through tools/wai_sota_vectors.sh, which pins the source image and
|
||
the Python stack and produces the whole vector directory the Rust int_hyper tests
|
||
read.
|
||
"""
|
||
import sys, os, io, json, hashlib
|
||
import numpy as np, torch, warnings; warnings.filterwarnings('ignore')
|
||
from compressai.zoo import image as zoo_image
|
||
from compressai.zoo.pretrained import load_pretrained
|
||
|
||
AFRAC = 12; WBITS = 12; SHIFT = 24; GF = 16
|
||
|
||
# The pretrained checkpoint the model-derived vectors come from (all but the source
|
||
# image, source.png and source.rgb), pinned by its full SHA-256.
|
||
# CompressAI's own `mbt2018_mean(pretrained=True)` fetches it through torch.hub
|
||
# without checking any hash (and unpickles it with weights_only=False), so a changed
|
||
# object at the URL, or a stale file in the torch.hub cache, would silently yield
|
||
# different vectors that the Rust tests — comparing against Python output built
|
||
# from the same checkpoint — would still pass. So the bytes are checked here, and
|
||
# only then loaded, as tensors only (weights_only=True).
|
||
CKPT_URL = "https://compressai.s3.amazonaws.com/models/v1/mbt2018-mean-6-a19628ab.pth.tar"
|
||
CKPT_SHA256 = "a19628aba19732125e02caf5c38e233edbe2a641145a5b1779c51f4527f51dd6"
|
||
|
||
|
||
def load_pinned_mbt2018_mean():
|
||
"""mbt2018_mean quality 6 (MSE), built from checkpoint bytes whose SHA-256 is
|
||
CKPT_SHA256 — or exit non-zero. Downloads into the torch.hub cache if absent."""
|
||
zoo_url = zoo_image.model_urls["mbt2018-mean"]["mse"][6]
|
||
if zoo_url != CKPT_URL:
|
||
sys.exit(f"error: the installed compressai points mbt2018_mean q6 at {zoo_url}, "
|
||
f"not the pinned {CKPT_URL}; review the new checkpoint before re-pinning")
|
||
path = os.path.join(torch.hub.get_dir(), "checkpoints", os.path.basename(CKPT_URL))
|
||
if not os.path.exists(path):
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
torch.hub.download_url_to_file(CKPT_URL, path, progress=False)
|
||
with open(path, "rb") as f:
|
||
blob = f.read()
|
||
got = hashlib.sha256(blob).hexdigest()
|
||
if got != CKPT_SHA256:
|
||
sys.exit(f"error: {path} has sha256 {got}, pinned {CKPT_SHA256}. "
|
||
"Delete the file to re-download; if the published checkpoint itself "
|
||
"changed, re-pin only after reviewing it.")
|
||
print(f"checkpoint {os.path.basename(path)}: sha256 {got} (pinned, verified)")
|
||
state = torch.load(io.BytesIO(blob), map_location="cpu", weights_only=True)
|
||
return zoo_image.model_architectures["mbt2018-mean"].from_state_dict(load_pretrained(state))
|
||
|
||
def q_weight_pc(w): # ConvTranspose [in,out,k,k] → per-OUT-channel int12 + scale
|
||
s = np.abs(w).max(axis=(0,2,3), keepdims=True)/(2**(WBITS-1)-1) + 1e-12
|
||
wi = np.round(w/s).clip(-(2**(WBITS-1)), 2**(WBITS-1)-1).astype(np.int64)
|
||
return wi, s.reshape(-1)
|
||
|
||
def q_weight_conv(w): # Conv2d [out,in,k,k] → per-OUT-channel int12 + scale
|
||
s = np.abs(w).max(axis=(1,2,3), keepdims=True)/(2**(WBITS-1)-1) + 1e-12
|
||
wi = np.round(w/s).clip(-(2**(WBITS-1)), 2**(WBITS-1)-1).astype(np.int64)
|
||
return wi, s.reshape(-1)
|
||
|
||
def conv_int(xi, wi, mult, bias, stride, pad, shift):
|
||
"""Integer regular conv (stride/pad) with per-out-channel (mult,shift) requant —
|
||
the EXACT arithmetic of Rust `int_hyper_synth` 'conv'. `wi` is [out,in,k,k]."""
|
||
co, ci, k, _ = wi.shape; H, W = xi.shape[1:]
|
||
Ho = (H + 2*pad - k)//stride + 1; Wo = (W + 2*pad - k)//stride + 1
|
||
xp = np.zeros((ci, H+2*pad, W+2*pad), dtype=np.int64); xp[:, pad:pad+H, pad:pad+W] = xi
|
||
raw = np.zeros((co, Ho, Wo), dtype=np.int64)
|
||
for ky in range(k):
|
||
for kx in range(k):
|
||
sub = xp[:, ky:ky+stride*Ho:stride, kx:kx+stride*Wo:stride] # [ci,Ho,Wo]
|
||
raw += np.tensordot(wi[:, :, ky, kx], sub, axes=([1], [0])) # [co,Ho,Wo]
|
||
return ((raw.astype(object)*mult[:,None,None]) >> shift).astype(np.int64) + bias[:,None,None]
|
||
|
||
def conv_t_int(x_int, wi, scale_oc, bias, stride, pad, opad):
|
||
cin, co, k, _ = wi.shape; H, W = x_int.shape[1:]
|
||
Ho = (H-1)*stride + k + opad - 2*pad; Wo = (W-1)*stride + k + opad - 2*pad
|
||
raw = np.zeros((co, Ho+2*pad, Wo+2*pad), dtype=np.int64)
|
||
for ic in range(cin):
|
||
xi = x_int[ic]
|
||
for (iy, ix) in np.argwhere(xi != 0):
|
||
raw[:, iy*stride:iy*stride+k, ix*stride:ix*stride+k] += wi[ic] * xi[iy, ix]
|
||
raw = raw[:, pad:pad+Ho, pad:pad+Wo]
|
||
mult = np.round(scale_oc*(2**SHIFT)).astype(np.int64)
|
||
bias_int = np.round(bias*(2**AFRAC)).astype(np.int64)
|
||
return ((raw*mult[:,None,None]) >> SHIFT) + bias_int[:,None,None]
|
||
|
||
def isqrt_vec(n): # exact floor(sqrt(n)) for int64 ≥0 — matches Rust u128::isqrt
|
||
s = np.maximum(np.floor(np.sqrt(n.astype(np.float64))).astype(np.int64), 0)
|
||
s = np.where((s+1)*(s+1) <= n, s+1, s)
|
||
s = np.where(s*s > n, s-1, s)
|
||
return s
|
||
|
||
def igdn_int(x_int, gi, bi): # gi [C,C] int@2^GF, bi [C] int@2^(2·AFRAC)
|
||
C = x_int.shape[0]; xf = x_int.reshape(C,-1).astype(np.int64); x2 = xf*xf
|
||
out = np.zeros_like(xf)
|
||
for i in range(C):
|
||
norm = np.maximum(bi[i] + ((gi[i][:,None]*x2) >> GF).sum(0), 0)
|
||
out[i] = (xf[i]*isqrt_vec(norm)) >> AFRAC
|
||
return out.reshape(x_int.shape)
|
||
|
||
def conv_t_raw(xi, wi, mult, bias, stride, pad, opad, shift):
|
||
"""Integer transposed conv with per-out-channel (mult,shift) requant — the
|
||
EXACT arithmetic of Rust `int_hyper_synth::conv_t` (uses the exported int mult,
|
||
not a recomputed float scale), so the Python expected == the Rust decode."""
|
||
cin, co, k, _ = wi.shape; H, W = xi.shape[1:]
|
||
Ho = (H-1)*stride + k + opad - 2*pad; Wo = (W-1)*stride + k + opad - 2*pad
|
||
raw = np.zeros((co, Ho+2*pad, Wo+2*pad), dtype=np.int64)
|
||
for ic in range(cin):
|
||
xc = xi[ic]
|
||
for (iy, ix) in np.argwhere(xc != 0):
|
||
raw[:, iy*stride:iy*stride+k, ix*stride:ix*stride+k] += wi[ic] * xc[iy, ix]
|
||
raw = raw[:, pad:pad+Ho, pad:pad+Wo]
|
||
return ((raw.astype(object)*mult[:,None,None]) >> shift).astype(np.int64) + bias[:,None,None]
|
||
|
||
def run_int_net(model, xi):
|
||
"""Run an exported integer model dict (the same JSON Rust `run_net` consumes)
|
||
on `xi` ([c,h,w] int64 @ 2^afrac). Mirrors Rust byte-for-byte."""
|
||
shift = model['shift']
|
||
for L in model['layers']:
|
||
t = L['type']
|
||
if t == 'convT':
|
||
cin, co, k = L['in_ch'], L['out_ch'], L['k']
|
||
wi = np.array(L['weight'], dtype=np.int64).reshape(cin, co, k, k)
|
||
mult = np.array(L['mult'], dtype=np.int64); bias = np.array(L['bias'], dtype=np.int64)
|
||
xi = conv_t_raw(xi, wi, mult, bias, L['stride'], L['pad'], L['opad'], shift)
|
||
elif t == 'conv':
|
||
cin, co, k = L['in_ch'], L['out_ch'], L['k']
|
||
wi = np.array(L['weight'], dtype=np.int64).reshape(co, cin, k, k)
|
||
mult = np.array(L['mult'], dtype=np.int64); bias = np.array(L['bias'], dtype=np.int64)
|
||
xi = conv_int(xi, wi, mult, bias, L['stride'], L['pad'], shift)
|
||
elif t == 'igdn':
|
||
C = L['c']; gi = np.array(L['gamma'], dtype=np.int64).reshape(C, C); bi = np.array(L['beta'], dtype=np.int64)
|
||
xi = igdn_int(xi, gi, bi)
|
||
elif t == 'relu':
|
||
xi = np.maximum(xi, 0)
|
||
elif t == 'leaky':
|
||
xi = np.where(xi >= 0, xi, (xi*L['slope_q']) >> L['slope_frac'])
|
||
return xi
|
||
|
||
def det_build_index_vec(sigma_q, edges):
|
||
"""Integer σ→bucket, vectorized — the EXACT logic of Rust
|
||
`int_entropy::det_build_index` (and CompressAI `build_indexes`, but on the
|
||
integer σ so encoder and Rust decoder agree by construction)."""
|
||
n = len(edges); sq = np.maximum(sigma_q, edges[0])
|
||
idx = np.full(sq.shape, n-1, dtype=np.int64)
|
||
for e in edges[:n-1]:
|
||
idx -= (sq <= e).astype(np.int64)
|
||
return idx
|
||
|
||
def fnv1a(buf):
|
||
h = 14695981039346656037
|
||
for byte in buf:
|
||
h = ((h ^ byte) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
|
||
return h
|
||
|
||
def model_to_bin(model):
|
||
"""Pack an integer model dict into the compact binary `IntModel::from_bytes`
|
||
reads: magic WHM1, afrac/gf/shift/n_layers (u32), then per layer a u8 tag with
|
||
int12 weights as i16 and IGDN gamma as i32 (the rest i64) — roughly halves the
|
||
on-disk/wire size vs JSON decimal, so the in-browser model download is practical."""
|
||
import struct
|
||
out = bytearray(b'WHM1')
|
||
out += struct.pack('<IIII', model['afrac'], model['gf'], model['shift'], len(model['layers']))
|
||
for L in model['layers']:
|
||
t = L['type']
|
||
if t == 'convT':
|
||
out += struct.pack('<BIIIIII', 0, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'], L['opad'])
|
||
w = np.array(L['weight'], dtype=np.int64); assert np.abs(w).max() < 2**15, 'convT weight > i16'
|
||
out += w.astype('<i2').tobytes()
|
||
out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
|
||
elif t == 'conv':
|
||
out += struct.pack('<BIIIII', 1, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'])
|
||
w = np.array(L['weight'], dtype=np.int64); assert np.abs(w).max() < 2**15, 'conv weight > i16'
|
||
out += w.astype('<i2').tobytes()
|
||
out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
|
||
elif t == 'igdn':
|
||
out += struct.pack('<BI', 2, L['c'])
|
||
g = np.array(L['gamma'], dtype=np.int64); assert np.abs(g).max() < 2**31, 'gamma > i32'
|
||
out += g.astype('<i4').tobytes() + np.array(L['beta'], dtype='<i8').tobytes()
|
||
elif t == 'relu':
|
||
out += struct.pack('<B', 3)
|
||
elif t == 'leaky':
|
||
out += struct.pack('<B', 4) + np.array([L['slope_q']], dtype='<i8').tobytes() + struct.pack('<I', L['slope_frac'])
|
||
return bytes(out)
|
||
|
||
def main(imgpath, outdir):
|
||
from PIL import Image
|
||
# The ENCODER side (g_a, h_a, the float h_s used for the reference) is float, so
|
||
# its latent can depend on kernel choice and reduction order. Both are pinned
|
||
# to make two runs of this tool on one machine emit identical vectors. That
|
||
# held for repeated runs on an arm64 macOS machine, but two runs on an x86_64
|
||
# Linux runner wrote different cond.json and int_synth_sota.json, most likely
|
||
# in psnr_float (`base` below), which the Rust tests read only as a quality
|
||
# threshold; the latent and both rgb_hash values matched. The decode being
|
||
# tested is integer and needs no such pinning.
|
||
torch.use_deterministic_algorithms(True)
|
||
torch.set_num_threads(1)
|
||
# Verify the checkpoint before writing anything, so a refused checkpoint leaves
|
||
# no partial vector directory behind.
|
||
m = load_pinned_mbt2018_mean().eval()
|
||
os.makedirs(outdir, exist_ok=True)
|
||
img = Image.open(imgpath).convert('RGB')
|
||
src8 = np.array(img)
|
||
x = torch.from_numpy(src8.astype(np.float32)/255.).permute(2,0,1).unsqueeze(0)
|
||
h,w = x.shape[2]//64*64, x.shape[3]//64*64; x = x[:,:,:h,:w]
|
||
# The exact pixels the codec saw (HWC u8), so a decoder can measure its own
|
||
# reconstruction quality instead of trusting the PSNR this tool reports.
|
||
np.ascontiguousarray(src8[:h, :w]).tofile(f'{outdir}/source.rgb')
|
||
def psnr(a,b): return float(-10*np.log10(((a-b)**2).mean()+1e-12))
|
||
with torch.no_grad():
|
||
y = m.g_a(x); z = m.h_a(y); z_hat = torch.round(z)
|
||
gp = m.h_s(z_hat); _, means = gp.chunk(2,1)
|
||
y_hat_t = torch.round(y - means) + means
|
||
xr_float = m.g_s(y_hat_t).clamp(0,1)[0].numpy()
|
||
y_hat = y_hat_t[0].numpy()
|
||
y_np = y[0].numpy() # raw float latent — the conditional encoder codes round(y - mu)
|
||
base = psnr(x[0].numpy(), xr_float)
|
||
|
||
xi = np.round(y_hat*(2**AFRAC)).astype(np.int64)
|
||
model = {'afrac':AFRAC,'gf':GF,'shift':SHIFT,'layers':[]}
|
||
inp = {'c':int(xi.shape[0]),'h':int(xi.shape[1]),'w':int(xi.shape[2]),'data':xi.reshape(-1).tolist()}
|
||
for layer in m.g_s.children():
|
||
t = type(layer).__name__
|
||
if t == 'ConvTranspose2d':
|
||
wi, sc = q_weight_pc(layer.weight.detach().numpy())
|
||
mult = np.round(sc*(2**SHIFT)).astype(np.int64)
|
||
bias_int = np.round(layer.bias.detach().numpy()*(2**AFRAC)).astype(np.int64)
|
||
model['layers'].append({'type':'convT','in_ch':int(wi.shape[0]),'out_ch':int(wi.shape[1]),
|
||
'k':int(wi.shape[2]),'stride':int(layer.stride[0]),'pad':int(layer.padding[0]),
|
||
'opad':int(layer.output_padding[0]),'weight':wi.reshape(-1).tolist(),
|
||
'mult':mult.tolist(),'bias':bias_int.tolist()})
|
||
xi = conv_t_int(xi, wi, sc, layer.bias.detach().numpy(),
|
||
layer.stride[0], layer.padding[0], layer.output_padding[0])
|
||
elif t == 'GDN':
|
||
with torch.no_grad():
|
||
beta = layer.beta_reparam(layer.beta).numpy(); C = beta.shape[0]
|
||
gamma = layer.gamma_reparam(layer.gamma).reshape(C, C).numpy()
|
||
gi = np.round(gamma*(2**GF)).astype(np.int64); bi = np.round(beta*(2**(2*AFRAC))).astype(np.int64)
|
||
model['layers'].append({'type':'igdn','c':int(C),'gamma':gi.reshape(-1).tolist(),'beta':bi.tolist()})
|
||
xi = igdn_int(xi, gi, bi)
|
||
# INTEGER pixel quantization (deterministic, no float) — what the Rust port reproduces.
|
||
out_c = np.clip(xi, 0, 1 << AFRAC)
|
||
pix = (((out_c * 255 + (1 << (AFRAC-1))) >> AFRAC)).clip(0, 255).astype(np.uint8) # [3,H,W] round-half-up
|
||
rec = out_c.astype(np.float64) / (2**AFRAC)
|
||
p_int = psnr(x[0].numpy(), rec); p_vf = psnr(xr_float, rec)
|
||
print(f'FLOAT baseline {base:.2f} dB | INTEGER synthesis {p_int:.2f} dB (vs source) | {p_vf:.2f} dB (vs float synth)')
|
||
rgb8 = np.ascontiguousarray(pix.transpose(1, 2, 0)) # HWC uint8
|
||
# FNV-1a 64-bit over the pixel bytes — a dependency-free checksum the Rust
|
||
# port reproduces with identical integer arithmetic (cross-language exact).
|
||
h = 14695981039346656037
|
||
for byte in rgb8.tobytes():
|
||
h = ((h ^ byte) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
|
||
digest = format(h, '016x')
|
||
json.dump({'afrac':AFRAC,'wbits':WBITS,'shift':SHIFT,'psnr_int':p_int,'psnr_float':base,'rgb_hash':digest,
|
||
'out_c':int(rec.shape[0]),'out_h':int(rec.shape[1]),'out_w':int(rec.shape[2])},
|
||
open(f'{outdir}/int_synth_sota.json','w'))
|
||
json.dump(model, open(f'{outdir}/int_synth_sota_model.json','w'))
|
||
open(f'{outdir}/int_synth_sota_model.bin','wb').write(model_to_bin(model))
|
||
json.dump(inp, open(f'{outdir}/int_synth_sota_input.json','w'))
|
||
print(f'wrote model ({len(model["layers"])} layers) + input + conformance (rgb_hash {digest[:16]}…) to {outdir}')
|
||
|
||
# ---- integer hyper-decoder h_s: z_hat -> (sigma, mu), per-channel int12 + ReLU ----
|
||
# z is entropy-coded with the FACTORIZED prior (CompressAI EntropyBottleneck) so the
|
||
# whole bitstream is integer-entropy-coded, not part-raw. The channel median is folded
|
||
# into a fixed-point constant median_q = round(median·2^AFRAC), so the decoder rebuilds
|
||
# z_q = s·2^AFRAC + median_q with NO float — the integer-entropy contract on z too.
|
||
from compressai.ans import RansEncoder as _RE
|
||
eb = m.entropy_bottleneck; eb.update()
|
||
with torch.no_grad():
|
||
z = m.h_a(m.g_a(x))
|
||
z_np = z[0].numpy()
|
||
median = eb._get_medians().reshape(-1).detach().numpy().astype(np.float64) # [Cz]
|
||
median_q = np.round(median * (2**AFRAC)).astype(np.int64)
|
||
z_sym = np.round(z_np - median[:,None,None]).astype(np.int64) # factorized symbols s
|
||
z_cdfs = eb._quantized_cdf.tolist()
|
||
z_cdf_lengths = eb._cdf_length.reshape(-1).int().tolist()
|
||
z_offsets = eb._offset.reshape(-1).int().tolist()
|
||
Cz = z_sym.shape[0]
|
||
z_indexes = np.repeat(np.arange(Cz), z_sym.shape[1]*z_sym.shape[2]) # per-channel CDF
|
||
z_bits = _RE().encode_with_indexes(z_sym.reshape(-1).tolist(), z_indexes.tolist(),
|
||
z_cdfs, z_cdf_lengths, z_offsets)
|
||
zi = z_sym * (2**AFRAC) + median_q[:,None,None] # z_q: INTEGER latent @ 2^AFRAC (what the decoder rebuilds)
|
||
hmodel = {'afrac':AFRAC,'gf':GF,'shift':SHIFT,'layers':[]}
|
||
hin = {'c':int(zi.shape[0]),'h':int(zi.shape[1]),'w':int(zi.shape[2]),'data':zi.reshape(-1).tolist()}
|
||
xi = zi.copy()
|
||
for layer in m.h_s.children():
|
||
t = type(layer).__name__
|
||
if t == 'ConvTranspose2d':
|
||
wi, sc = q_weight_pc(layer.weight.detach().numpy())
|
||
mult = np.round(sc*(2**SHIFT)).astype(np.int64)
|
||
bias_int = np.round(layer.bias.detach().numpy()*(2**AFRAC)).astype(np.int64)
|
||
hmodel['layers'].append({'type':'convT','in_ch':int(wi.shape[0]),'out_ch':int(wi.shape[1]),
|
||
'k':int(wi.shape[2]),'stride':int(layer.stride[0]),'pad':int(layer.padding[0]),
|
||
'opad':int(layer.output_padding[0]),'weight':wi.reshape(-1).tolist(),
|
||
'mult':mult.tolist(),'bias':bias_int.tolist()})
|
||
xi = conv_t_int(xi, wi, sc, layer.bias.detach().numpy(), layer.stride[0], layer.padding[0], layer.output_padding[0])
|
||
elif t == 'Conv2d': # the final stride-1 conv → σ‖μ (M*3//2 → M*2)
|
||
wi, sc = q_weight_conv(layer.weight.detach().numpy())
|
||
mult = np.round(sc*(2**SHIFT)).astype(np.int64)
|
||
bias_int = np.round(layer.bias.detach().numpy()*(2**AFRAC)).astype(np.int64)
|
||
hmodel['layers'].append({'type':'conv','in_ch':int(wi.shape[1]),'out_ch':int(wi.shape[0]),
|
||
'k':int(wi.shape[2]),'stride':int(layer.stride[0]),'pad':int(layer.padding[0]),
|
||
'weight':wi.reshape(-1).tolist(),'mult':mult.tolist(),'bias':bias_int.tolist()})
|
||
xi = conv_int(xi, wi, mult, bias_int, layer.stride[0], layer.padding[0], SHIFT)
|
||
elif t == 'ReLU':
|
||
hmodel['layers'].append({'type':'relu'})
|
||
xi = np.maximum(xi, 0)
|
||
elif t == 'LeakyReLU': # slope 0.01 → fixed-point 655/2^16
|
||
hmodel['layers'].append({'type':'leaky','slope_q':655,'slope_frac':16})
|
||
xi = np.where(xi >= 0, xi, (xi*655) >> 16)
|
||
sigmamu_int = xi.copy() # [2M, sh, sw] integer σ‖μ @ 2^AFRAC, from the INTEGER h_s
|
||
# FNV-1a over the integer sigma||mu output (LE i64 bytes) — matches Rust run_net
|
||
hh = fnv1a(xi.astype('<i8').tobytes())
|
||
json.dump(hmodel, open(f'{outdir}/int_hs_model.json','w'))
|
||
open(f'{outdir}/int_hs_model.bin','wb').write(model_to_bin(hmodel))
|
||
json.dump(hin, open(f'{outdir}/int_hs_input.json','w'))
|
||
json.dump({'afrac':AFRAC,'out_c':int(xi.shape[0]),'out_h':int(xi.shape[1]),'out_w':int(xi.shape[2]),
|
||
'params_hash':format(hh,'016x')}, open(f'{outdir}/int_hs.json','w'))
|
||
print(f'wrote h_s ({len(hmodel["layers"])} layers): z_hat {zi.shape} -> sigma||mu {xi.shape}, params_hash {format(hh,"016x")}')
|
||
|
||
# ---- the conditional ENTROPY layer: code y against integer σ (the bitstream) ----
|
||
# Everything the DECODER needs is integer: z_hat (stored) → integer h_s → σ,μ_int;
|
||
# σ_int → bucket (int_entropy); rANS-decode residual r; y_hat_int = r·2^AFRAC + μ_int.
|
||
# The encoder runs the SAME integer h_s so it agrees on every σ-bucket — the whole
|
||
# point of the integer-entropy contract. CDFs/offsets come from a real CompressAI
|
||
# GaussianConditional; the rANS bytes from CompressAI's encoder (the Rust RansDecoder
|
||
# is a byte-exact port of it, bypass path included), so no float touches the decode.
|
||
from compressai.entropy_models import GaussianConditional
|
||
from compressai.ans import RansEncoder
|
||
M = sigmamu_int.shape[0] // 2
|
||
sigma_int = sigmamu_int[:M]; mu_int = sigmamu_int[M:] # each [M, sh, sw] @ 2^AFRAC
|
||
scale_table = list(torch.exp(torch.linspace(np.log(0.11), np.log(256.0), 64)).numpy().astype(float))
|
||
edges = np.floor(np.array(scale_table) * (2**AFRAC)).astype(np.int64) # σ buckets @ 2^AFRAC
|
||
gc = GaussianConditional(scale_table); gc.update()
|
||
cdfs = gc._quantized_cdf.tolist()
|
||
cdf_lengths = gc._cdf_length.reshape(-1).int().tolist()
|
||
offsets = gc._offset.reshape(-1).int().tolist()
|
||
# residual r = round(y - μ_int/2^AFRAC); decoder rebuilds y_hat_int = r·2^AFRAC + μ_int
|
||
resid = np.round(y_np - mu_int.astype(np.float64) / (2**AFRAC)).astype(np.int64)
|
||
indexes = det_build_index_vec(sigma_int, edges) # integer σ→bucket, per element
|
||
enc = RansEncoder()
|
||
bits = enc.encode_with_indexes(resid.reshape(-1).tolist(), indexes.reshape(-1).tolist(),
|
||
cdfs, cdf_lengths, offsets)
|
||
# expected reconstruction the DECODER must reproduce: integer g_s on y_hat_int
|
||
y_hat_int = (resid << AFRAC) + mu_int
|
||
rec_int = run_int_net(model, y_hat_int)
|
||
out_c2 = np.clip(rec_int, 0, 1 << AFRAC)
|
||
pix2 = (((out_c2 * 255 + (1 << (AFRAC-1))) >> AFRAC)).clip(0, 255).astype(np.uint8)
|
||
rgb2 = np.ascontiguousarray(pix2.transpose(1, 2, 0))
|
||
cond_hash = format(fnv1a(rgb2.tobytes()), '016x')
|
||
p_cond = psnr(x[0].numpy(), out_c2.astype(np.float64) / (2**AFRAC))
|
||
# entropy tables: the y Gaussian-conditional CDFs + the z FACTORIZED CDFs (+ median_q,
|
||
# the fixed-point channel medians) — all integer, both bitstreams decode with no float.
|
||
json.dump({'afrac':AFRAC,'frac':AFRAC,'scale_edges':edges.tolist(),
|
||
'cdfs':cdfs,'cdf_lengths':cdf_lengths,'offsets':offsets,
|
||
'z_cdfs':z_cdfs,'z_cdf_lengths':z_cdf_lengths,'z_offsets':z_offsets,
|
||
'z_median_q':median_q.tolist()},
|
||
open(f'{outdir}/cond_entropy.json','w'))
|
||
json.dump({'m':int(M),'sh':int(sigma_int.shape[1]),'sw':int(sigma_int.shape[2]),
|
||
'n_z':int(Cz),'zh':int(z_sym.shape[1]),'zw':int(z_sym.shape[2]),
|
||
'z_rans':list(z_bits),'rans':list(bits)},
|
||
open(f'{outdir}/cond_bitstream.json','w'))
|
||
json.dump({'rgb_hash':cond_hash,'psnr_cond':p_cond,'psnr_float':base,
|
||
'out_c':int(rec_int.shape[0]),'out_h':int(rec_int.shape[1]),'out_w':int(rec_int.shape[2])},
|
||
open(f'{outdir}/cond.json','w'))
|
||
print(f'wrote conditional bitstream: {len(z_bits)} z + {len(bits)} y rANS bytes '
|
||
f'-> {p_cond:.2f} dB, rgb_hash {cond_hash}')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main(sys.argv[1], sys.argv[2])
|