ce755f2fa9
int_mlicpp was run_net_checked's last production caller. run_net_checked, check_layer, igdn_check and the values' failure constructor move into the test-only int_hyper_synth::legacy (cfg(any(test, feature = "legacy_int"))), beside the direct decodes that call them; the direct decodes of splat, audio and int_mlicpp, the guard's own unit tests and the integration tests that hold the guard and the check ranges to each other import them from there. Production docs that named the guard now name net_shape or the lowering. run_net and decode (the world renderer's) stay for M6. The export tool's note now names the check ranges. CHANGELOG. Co-Authored-By: Claude Opus 5.5 <[email protected]>
482 lines
24 KiB
Python
482 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""Export a small `wai.neural.int_mlicpp` conformance vector.
|
|
|
|
`wai.neural.int_mlicpp` is the MLIC++ entropy topology made integer end to end:
|
|
the latent `y` is split into K channel slices decoded in order, each slice in two
|
|
checkerboard passes (anchors, where (row + col) is even, then non-anchors), and
|
|
before every pass an entropy network over concat[y decoded so far, hyper] gives
|
|
the scale sigma and the mean mu of every element. The scale picks the element's
|
|
distribution by integer comparison (`int_entropy::det_build_index`), the rANS
|
|
coder decodes the residual r, and y = r * 2^afrac + mu. Every network on the
|
|
decode side is an integer network (per-output-channel int12 weights,
|
|
(mult, shift) requantisation, integer IGDN), so the decoder computes exactly the
|
|
sigma the encoder computed, on every platform and at every thread count.
|
|
|
|
The weights are SYNTHETIC. They are drawn from a fixed seed and quantised the way
|
|
a trained model's would be; no trained MLIC++ checkpoint is used. The vector
|
|
therefore shows decode-equivalence -- the Rust decoder reproduces this tool's
|
|
integer reference byte for byte -- and says nothing about compression quality.
|
|
|
|
python tools/wai_int_mlicpp_export.py <out-dir> [--stats]
|
|
|
|
Writes into <out-dir>:
|
|
model.wmb the decode-side model bundle: magic `WMB1`, then four sections,
|
|
each a u32 length and its bytes, in order h_s, entropy, g_s (WHM1
|
|
integer networks, `IntModel::from_bytes`) and the entropy tables
|
|
(JSON, the fields of `int_hyper_synth::CondEntropy`).
|
|
stream.wai a WAI1 envelope over `wai.neural.int_mlicpp`, intent `replicate`,
|
|
pinning model.wmb by SHA-256 as an integer-exact prior. Its payload
|
|
is the `WIM1` codestream (see `wai-rs/src/int_mlicpp.rs`).
|
|
model.sha256, the set-form twin (tools/wai_int_mlicpp_set_form.py): the listing of
|
|
stream_set.wai the one-file set, and stream.wai pinned through it (`wai.prior.set`).
|
|
expected.json what the decode must reproduce: the RGB hashes (FNV-1a and
|
|
SHA-256), and FNV-1a over the y stream's bucket indices (u16 LE)
|
|
and symbols (i32 LE) in coding order.
|
|
|
|
How the expected values are made. The weights are drawn in float from the seed and
|
|
quantised; from the quantised weights on, everything is integer. The encoder's
|
|
analysis transforms (g_a, h_a) are integer networks too, so the latents, symbols
|
|
and reconstruction are integer functions of the quantised weights and the
|
|
(integer) source pattern. Two runs on the development machine wrote identical
|
|
files; whether another machine's float draws and CompressAI tables come out the
|
|
same was not checked, and the conformance check does not depend on it -- the
|
|
committed files are the vector. The encoder computes each sigma and mu with the
|
|
same integer networks as the decoder; the reference reconstruction runs the
|
|
integer synthesis in numpy int64 arithmetic, under explicit range checks that keep
|
|
every intermediate inside int64 (numpy wraps silently, so the checks are what make
|
|
this reference exact). Before anything is written, the rANS bytes are decoded back
|
|
in Python, pass by pass, with CompressAI's decoder and the integer networks, and
|
|
must give the encoder's latent exactly. The rANS bytes come from CompressAI's
|
|
encoder, whose bitstream the Rust `RansDecoder` reads (bypass-coded escapes
|
|
included). The Gaussian tables are CompressAI's `GaussianConditional` tables over
|
|
a 16-level scale table; the hyper-latent's per-channel tables come from
|
|
`pmf_to_quantized_cdf` over discretised Gaussians.
|
|
"""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from wai_sota_int_synth_export import ( # noqa: E402 (the shared integer ops)
|
|
conv_int, conv_t_raw, fnv1a, igdn_int, isqrt_vec, model_to_bin, q_weight_conv, q_weight_pc,
|
|
det_build_index_vec,
|
|
)
|
|
|
|
AFRAC = 12 # activations at 2^12 (the int_hyper family's fixed point)
|
|
GF = 16 # IGDN gamma at 2^16
|
|
SHIFT = 24 # requantisation shift
|
|
SEED = 20260929
|
|
|
|
H = W = 64 # image
|
|
CY = 8 # latent channels (M)
|
|
CZ = 4 # hyper-latent channels (N)
|
|
K = 4 # channel slices, CY / K channels each
|
|
SCALE_MIN, SCALE_MAX, SCALE_LEVELS = 0.11, 16.0, 16
|
|
# Gains of the synthetic networks' weight draws, and the scale head's per-channel
|
|
# base scales: chosen so the latent spans several integers, the scales spread over
|
|
# the table, and a few residuals fall outside their table (escapes).
|
|
GA_GAIN, HA_GAIN, HS_GAIN, E1_GAIN, E2_GAIN, GS_OUT_GAIN = 4.0, 1.5, 1.4, 1.0, 0.6, 0.1
|
|
SIGMA_BASE = (0.7, 10.0)
|
|
Z_WIDTH = 1.0 # hyper-latent table half-width, in units of the channel's spread
|
|
I64_MAX = (1 << 63) - 1
|
|
|
|
CAPABILITY = "wai.neural.int_mlicpp"
|
|
PRIOR_ID = "int_mlicpp/synthetic-s20260929"
|
|
DETERMINISM = "int-rans/p16+int-mlicpp/i64"
|
|
|
|
|
|
# ---------------------------------------------------------------- synthetic float model
|
|
|
|
def xavier(rng, shape, fan_in, fan_out, gain):
|
|
a = gain * np.sqrt(6.0 / (fan_in + fan_out))
|
|
return rng.uniform(-a, a, size=shape)
|
|
|
|
|
|
def f_conv(rng, cin, cout, k, stride, pad, gain, bias=0.05):
|
|
return {'kind': 'conv', 'w': xavier(rng, (cout, cin, k, k), cin * k * k, cout * k * k, gain),
|
|
'b': rng.uniform(-bias, bias, size=cout), 'stride': stride, 'pad': pad}
|
|
|
|
|
|
def f_convT(rng, cin, cout, k, stride, pad, opad, gain, bias=0.05):
|
|
return {'kind': 'convT', 'w': xavier(rng, (cin, cout, k, k), cin * k * k, cout * k * k, gain),
|
|
'b': rng.uniform(-bias, bias, size=cout), 'stride': stride, 'pad': pad, 'opad': opad}
|
|
|
|
|
|
def f_igdn(rng, c):
|
|
gamma = 0.1 * np.eye(c) + rng.uniform(0.0, 0.02, size=(c, c)) # non-negative
|
|
beta = 1.0 + rng.uniform(0.0, 0.2, size=c) # positive
|
|
return {'kind': 'igdn', 'gamma': gamma, 'beta': beta}
|
|
|
|
|
|
def ptq(layers):
|
|
"""Quantise a synthetic float network the way `wai_sota_int_synth_export.py`
|
|
quantises a trained one: per-output-channel int12 weights, mult = round(scale *
|
|
2^SHIFT), bias at 2^AFRAC, IGDN gamma at 2^GF and beta at 2^(2*AFRAC). Returns
|
|
the integer model dict `IntModel::from_json` / `model_to_bin` take."""
|
|
out = {'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': []}
|
|
for L in layers:
|
|
k = L['kind']
|
|
if k == 'convT':
|
|
wi, sc = q_weight_pc(L['w'])
|
|
out['layers'].append({'type': 'convT', 'in_ch': int(wi.shape[0]), 'out_ch': int(wi.shape[1]),
|
|
'k': int(wi.shape[2]), 'stride': L['stride'], 'pad': L['pad'], 'opad': L['opad'],
|
|
'weight': wi.reshape(-1).tolist(),
|
|
'mult': np.round(sc * 2 ** SHIFT).astype(np.int64).tolist(),
|
|
'bias': np.round(L['b'] * 2 ** AFRAC).astype(np.int64).tolist()})
|
|
elif k == 'conv':
|
|
wi, sc = q_weight_conv(L['w'])
|
|
out['layers'].append({'type': 'conv', 'in_ch': int(wi.shape[1]), 'out_ch': int(wi.shape[0]),
|
|
'k': int(wi.shape[2]), 'stride': L['stride'], 'pad': L['pad'],
|
|
'weight': wi.reshape(-1).tolist(),
|
|
'mult': np.round(sc * 2 ** SHIFT).astype(np.int64).tolist(),
|
|
'bias': np.round(L['b'] * 2 ** AFRAC).astype(np.int64).tolist()})
|
|
elif k == 'igdn':
|
|
c = L['gamma'].shape[0]
|
|
out['layers'].append({'type': 'igdn', 'c': int(c),
|
|
'gamma': np.round(L['gamma'] * 2 ** GF).astype(np.int64).reshape(-1).tolist(),
|
|
'beta': np.round(L['beta'] * 2 ** (2 * AFRAC)).astype(np.int64).tolist()})
|
|
elif k == 'relu':
|
|
out['layers'].append({'type': 'relu'})
|
|
elif k == 'leaky': # slope 0.01 as 655 / 2^16, as the int_hyper export writes it
|
|
out['layers'].append({'type': 'leaky', 'slope_q': 655, 'slope_frac': 16})
|
|
else:
|
|
raise ValueError(k)
|
|
return out
|
|
|
|
|
|
def build_model(rng):
|
|
"""The synthetic networks. Topology mirrors the float `wai.neural.mlicpp`
|
|
validation model (`tools/wai_validate_synthetic.py`, MiniMlic) at a smaller
|
|
scale: 3 stride-2 stages instead of 4 (a 64x64 image -> an 8x8 latent), 8 latent
|
|
channels, a 2-layer entropy network whose first layer is MiniMlic's 5x5 conv over
|
|
concat[y, hyper], and IGDN in the synthesis."""
|
|
M, N = CY, CZ
|
|
g_a = [f_conv(rng, 3, M, 5, 2, 2, GA_GAIN), {'kind': 'relu'},
|
|
f_conv(rng, M, M, 5, 2, 2, GA_GAIN), {'kind': 'relu'},
|
|
f_conv(rng, M, M, 5, 2, 2, GA_GAIN)]
|
|
h_a = [f_conv(rng, M, N, 5, 2, 2, HA_GAIN), {'kind': 'relu'},
|
|
f_conv(rng, N, N, 5, 2, 2, HA_GAIN)]
|
|
h_s = [f_convT(rng, N, M, 5, 2, 2, 1, HS_GAIN), {'kind': 'relu'},
|
|
f_convT(rng, M, 2 * M, 5, 2, 2, 1, HS_GAIN)]
|
|
e1 = f_conv(rng, 3 * M, 2 * M, 5, 1, 2, E1_GAIN)
|
|
e2 = f_conv(rng, 2 * M, 2 * M, 1, 1, 0, E2_GAIN)
|
|
# The sigma half starts positive with a different base scale per channel, as a
|
|
# trained scale head's would: some channels carry more entropy than others.
|
|
e2['b'][:M] += np.geomspace(SIGMA_BASE[0], SIGMA_BASE[1], M)
|
|
entropy = [e1, {'kind': 'leaky'}, e2]
|
|
g_s = [f_convT(rng, M, M, 5, 2, 2, 1, 1.0), f_igdn(rng, M),
|
|
f_convT(rng, M, M, 5, 2, 2, 1, 1.0), f_igdn(rng, M),
|
|
f_convT(rng, M, 3, 5, 2, 2, 1, GS_OUT_GAIN)]
|
|
g_s[-1]['b'] += 0.5 # pixels centred in [0, 1]
|
|
return {name: ptq(layers) for name, layers in
|
|
(('g_a', g_a), ('h_a', h_a), ('h_s', h_s), ('entropy', entropy), ('g_s', g_s))}
|
|
|
|
|
|
# ---------------------------------------------------------------- checked integer reference
|
|
|
|
def max_abs(x):
|
|
return int(np.max(np.abs(x.astype(object)))) if x.size else 0
|
|
|
|
|
|
def run_checked(model, x, c):
|
|
"""Run an integer model on x ([c, h, w] int64 at 2^AFRAC) layer by layer with the
|
|
ops `run_int_net` uses, after checking before every layer that no intermediate
|
|
can leave int64 -- numpy int64 wraps silently, so these checks are what make the
|
|
result exact. They are at least as strict as the `check` ranges of the Rust
|
|
decoder's lowered networks (integer-operators §13.2, which accumulates IGDN in
|
|
i128), so a vector that passes here decodes there without refusal."""
|
|
for i, L in enumerate(model['layers']):
|
|
t = L['type']
|
|
xm = max_abs(x)
|
|
if t in ('conv', 'convT'):
|
|
cin, co, k = L['in_ch'], L['out_ch'], L['k']
|
|
assert cin == c, f'layer {i}: in_ch {cin} != running channels {c}'
|
|
wm = max(abs(v) for v in L['weight'])
|
|
acc = cin * k * k * wm * xm
|
|
assert acc <= I64_MAX, f'layer {i}: accumulation bound {acc} > i64'
|
|
q = ((acc * max(abs(v) for v in L['mult'])) >> model['shift']) + 1
|
|
assert q + max(abs(v) for v in L['bias']) <= I64_MAX, f'layer {i}: requant bound'
|
|
w = np.array(L['weight'], dtype=np.int64)
|
|
mult = np.array(L['mult'], dtype=np.int64)
|
|
bias = np.array(L['bias'], dtype=np.int64)
|
|
if t == 'convT':
|
|
x = conv_t_raw(x, w.reshape(cin, co, k, k), mult, bias, L['stride'], L['pad'], L['opad'],
|
|
model['shift'])
|
|
else:
|
|
x = conv_int(x, w.reshape(co, cin, k, k), mult, bias, L['stride'], L['pad'], model['shift'])
|
|
c = co
|
|
elif t == 'igdn':
|
|
C = L['c']
|
|
assert C == c
|
|
gm = max(abs(v) for v in L['gamma'])
|
|
bm = max(abs(v) for v in L['beta'])
|
|
x2 = xm * xm
|
|
norm = bm + C * (((gm * x2) >> model['gf']) + 1)
|
|
assert x2 <= I64_MAX and gm * x2 <= I64_MAX and norm <= I64_MAX, f'layer {i}: igdn bound'
|
|
s = int(isqrt_vec(np.array([norm], dtype=np.int64))[0])
|
|
assert xm * s <= I64_MAX, f'layer {i}: igdn product bound'
|
|
assert model['afrac'] == AFRAC and model['gf'] == GF # igdn_int reads the module constants
|
|
x = igdn_int(x, np.array(L['gamma'], dtype=np.int64).reshape(C, C),
|
|
np.array(L['beta'], dtype=np.int64))
|
|
elif t == 'relu':
|
|
x = np.maximum(x, 0)
|
|
elif t == 'leaky':
|
|
assert xm * abs(L['slope_q']) <= I64_MAX
|
|
x = np.where(x >= 0, x, (x * L['slope_q']) >> L['slope_frac'])
|
|
else:
|
|
raise ValueError(t)
|
|
return x, c
|
|
|
|
|
|
def to_rgb8(x):
|
|
"""int_hyper_synth::to_rgb8: clamp to [0, 2^afrac], round half up to 8 bits, CHW -> HWC."""
|
|
one = 1 << AFRAC
|
|
pix = ((np.clip(x, 0, one) * 255 + (1 << (AFRAC - 1))) >> AFRAC).clip(0, 255).astype(np.uint8)
|
|
return np.ascontiguousarray(pix.transpose(1, 2, 0))
|
|
|
|
|
|
def is_anchor(r, c):
|
|
return (r + c) % 2 == 0
|
|
|
|
|
|
def coding_order(cy, yh, yw, k):
|
|
"""(slice, anchor, [flat indices]) in the order the y stream codes them."""
|
|
sw = cy // k
|
|
for ki in range(k):
|
|
for anchor in (True, False):
|
|
idx = [ch * yh * yw + r * yw + c for ch in range(ki * sw, (ki + 1) * sw)
|
|
for r in range(yh) for c in range(yw) if is_anchor(r, c) == anchor]
|
|
yield ki, anchor, idx
|
|
|
|
|
|
# ---------------------------------------------------------------- entropy tables
|
|
|
|
def gaussian_tables():
|
|
import torch
|
|
from compressai.entropy_models import GaussianConditional
|
|
scale_table = torch.exp(torch.linspace(np.log(SCALE_MIN), np.log(SCALE_MAX), SCALE_LEVELS))
|
|
gc = GaussianConditional(scale_table.tolist())
|
|
gc.update()
|
|
lengths = gc._cdf_length.reshape(-1).int().tolist()
|
|
cdfs_padded = gc._quantized_cdf.int().tolist()
|
|
cdfs = [row[:lengths[i]] for i, row in enumerate(cdfs_padded)]
|
|
offsets = gc._offset.reshape(-1).int().tolist()
|
|
edges = np.floor(scale_table.numpy().astype(np.float64) * 2 ** AFRAC).astype(np.int64)
|
|
return edges, cdfs_padded, cdfs, lengths, offsets
|
|
|
|
|
|
def z_tables(z_sym):
|
|
"""One discretised-Gaussian table per hyper-latent channel, over symbols
|
|
-half..half with offset -half; the width follows the channel's spread, and a
|
|
symbol outside it is coded by the rANS escape (bypass)."""
|
|
from compressai._CXX import pmf_to_quantized_cdf
|
|
cdfs, lengths, offsets = [], [], []
|
|
for ch in range(z_sym.shape[0]):
|
|
spread = max(0.6, float(np.std(z_sym[ch].astype(np.float64))))
|
|
half = max(1, int(np.ceil(Z_WIDTH * spread)))
|
|
v = np.arange(-half, half + 1, dtype=np.float64)
|
|
pmf = np.exp(-(v * v) / (2.0 * spread * spread))
|
|
pmf = pmf / pmf.sum() * (1.0 - 1e-3)
|
|
cdf = list(pmf_to_quantized_cdf(pmf.tolist() + [1e-3], 16))
|
|
cdfs.append(cdf)
|
|
lengths.append(len(cdf))
|
|
offsets.append(-half)
|
|
return cdfs, lengths, offsets
|
|
|
|
|
|
def check_cdf(cdf, n):
|
|
assert n >= 2 and len(cdf) >= n and cdf[0] == 0 and cdf[n - 1] == 1 << 16
|
|
assert all(cdf[i] < cdf[i + 1] for i in range(n - 1)), 'CDF must be strictly increasing'
|
|
|
|
|
|
# ---------------------------------------------------------------- encode
|
|
|
|
def source_pattern():
|
|
"""A 64x64 RGB test pattern built from integers only (so it is the same bytes on
|
|
every machine): two ramps, an xor texture and a disc."""
|
|
yy, xx = np.mgrid[0:H, 0:W]
|
|
img = np.stack([xx * 255 // (W - 1), yy * 255 // (H - 1), ((xx ^ yy) * 4) & 255], -1)
|
|
img[(yy - H // 2) ** 2 + (xx - W // 2) ** 2 < (H // 5) ** 2] = [200, 40, 120]
|
|
return img.astype(np.uint8)
|
|
|
|
|
|
def encode(models, stats=False):
|
|
from compressai.ans import RansDecoder, RansEncoder
|
|
one = 1 << AFRAC
|
|
half = 1 << (AFRAC - 1)
|
|
img = source_pattern()
|
|
x = (img.transpose(2, 0, 1).astype(np.int64) << AFRAC) // 255 # [3, H, W] at 2^AFRAC
|
|
|
|
y_q, c = run_checked(models['g_a'], x, 3) # encoder-side analysis
|
|
assert c == CY
|
|
z_q, c = run_checked(models['h_a'], y_q, CY)
|
|
assert c == CZ
|
|
_, yh, yw = y_q.shape
|
|
_, zh, zw = z_q.shape
|
|
|
|
# hyper-latent: per-channel median (a multiple of 2^(AFRAC-2)), symbols, tables
|
|
median_q = np.array([int(np.median(z_q[ch])) >> (AFRAC - 2) << (AFRAC - 2) for ch in range(CZ)],
|
|
dtype=np.int64)
|
|
z_sym = (z_q - median_q[:, None, None] + half) >> AFRAC # round half up
|
|
z_cdfs, z_lengths, z_offsets = z_tables(z_sym)
|
|
for cdf, n in zip(z_cdfs, z_lengths):
|
|
check_cdf(cdf, n)
|
|
z_int = z_sym * one + median_q[:, None, None] # what the decoder rebuilds
|
|
z_idx = np.repeat(np.arange(CZ), zh * zw)
|
|
z_bits = RansEncoder().encode_with_indexes(z_sym.reshape(-1).tolist(), z_idx.tolist(),
|
|
z_cdfs, z_lengths, z_offsets)
|
|
|
|
hyper, c = run_checked(models['h_s'], z_int, CZ) # decode-side from here on
|
|
assert hyper.shape == (2 * CY, yh, yw), hyper.shape
|
|
|
|
edges, cdfs_padded, cdfs, lengths, offsets = gaussian_tables()
|
|
for cdf, n in zip(cdfs, lengths):
|
|
check_cdf(cdf, n)
|
|
|
|
n = CY * yh * yw
|
|
yq = y_q.reshape(-1)
|
|
y_int = np.zeros(n, dtype=np.int64)
|
|
hyper_flat = hyper.reshape(-1)
|
|
syms, idxs, sig_all = [], [], []
|
|
for _, _, pos in coding_order(CY, yh, yw, K):
|
|
ctx = np.concatenate([y_int, hyper_flat]).reshape(3 * CY, yh, yw)
|
|
params, c = run_checked(models['entropy'], ctx, 3 * CY)
|
|
assert params.shape == (2 * CY, yh, yw)
|
|
params = params.reshape(-1)
|
|
sigma, mu = params[:n], params[n:]
|
|
p = np.array(pos)
|
|
b = det_build_index_vec(sigma[p], edges)
|
|
r = (yq[p] - mu[p] + half) >> AFRAC # round half up
|
|
y_int[p] = r * one + mu[p]
|
|
syms.extend(int(v) for v in r)
|
|
idxs.extend(int(v) for v in b)
|
|
sig_all.extend(int(v) for v in sigma[p])
|
|
y_bits = RansEncoder().encode_with_indexes(syms, idxs, cdfs_padded, lengths, offsets)
|
|
|
|
# ---- decode the rANS bytes back, pass by pass, as the decoder does
|
|
dz = RansDecoder()
|
|
dz.set_stream(z_bits)
|
|
z_dec = np.array(dz.decode_stream(z_idx.tolist(), z_cdfs, z_lengths, z_offsets), dtype=np.int64)
|
|
assert np.array_equal(z_dec, z_sym.reshape(-1)), 'z round trip'
|
|
hyper2, _ = run_checked(models['h_s'], z_dec.reshape(CZ, zh, zw) * one + median_q[:, None, None], CZ)
|
|
assert np.array_equal(hyper2, hyper)
|
|
dy = RansDecoder()
|
|
dy.set_stream(y_bits)
|
|
y_dec = np.zeros(n, dtype=np.int64)
|
|
for _, _, pos in coding_order(CY, yh, yw, K):
|
|
ctx = np.concatenate([y_dec, hyper2.reshape(-1)]).reshape(3 * CY, yh, yw)
|
|
params, _ = run_checked(models['entropy'], ctx, 3 * CY)
|
|
params = params.reshape(-1)
|
|
p = np.array(pos)
|
|
b = det_build_index_vec(params[:n][p], edges)
|
|
vals = np.array(dy.decode_stream(b.tolist(), cdfs_padded, lengths, offsets), dtype=np.int64)
|
|
y_dec[p] = vals * one + params[n:][p]
|
|
assert np.array_equal(y_dec, y_int), 'y round trip: the Python decode must give the encoder latent'
|
|
|
|
rec, c = run_checked(models['g_s'], y_int.reshape(CY, yh, yw), CY)
|
|
assert rec.shape == (3, H, W), rec.shape
|
|
rgb = to_rgb8(rec)
|
|
|
|
# escapes: a symbol the table cannot hold is coded by the rANS bypass
|
|
y_escape = sum(1 for s, b in zip(syms, idxs) if not (0 <= s - offsets[b] < lengths[b] - 2))
|
|
z_escape = sum(1 for s, ch in zip(z_sym.reshape(-1).tolist(), z_idx.tolist())
|
|
if not (0 <= s - z_offsets[ch] < z_lengths[ch] - 2))
|
|
if stats:
|
|
sg = np.array(sig_all) / one
|
|
print(f'y/2^A: std {np.std(y_q / one):.2f}, max {np.max(np.abs(y_q)) / one:.1f}; '
|
|
f'z/2^A: std {np.std(z_q / one):.2f}; hyper max {np.max(np.abs(hyper)) / one:.1f}')
|
|
print(f'sigma: min {sg.min():.3f} median {np.median(sg):.3f} max {sg.max():.3f}; '
|
|
f'buckets used {len(set(idxs))}/{SCALE_LEVELS}: {np.bincount(idxs, minlength=SCALE_LEVELS).tolist()}')
|
|
print(f'|r| max {max(abs(s) for s in syms)}; escapes y {y_escape}/{len(syms)}, z {z_escape}/{z_sym.size}')
|
|
print(f'pixels: min {rgb.min()} max {rgb.max()} mean {rgb.mean():.1f}; '
|
|
f'clamped {int(np.sum((rgb == 0) | (rgb == 255)))}/{rgb.size}')
|
|
tables = {'afrac': AFRAC, 'scale_edges': edges.tolist(), 'cdfs': cdfs, 'cdf_lengths': lengths,
|
|
'offsets': offsets, 'z_cdfs': z_cdfs, 'z_cdf_lengths': z_lengths, 'z_offsets': z_offsets,
|
|
'z_median_q': median_q.tolist()}
|
|
return {
|
|
'tables': tables, 'z_bits': z_bits, 'y_bits': y_bits, 'rgb': rgb,
|
|
'dims': (zh, zw, yh, yw), 'syms': syms, 'idxs': idxs,
|
|
'counts': {'y_symbols': len(syms), 'z_symbols': int(z_sym.size), 'y_escapes': y_escape,
|
|
'z_escapes': z_escape, 'buckets_used': len(set(idxs))},
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------- containers
|
|
|
|
def bundle(models, tables):
|
|
"""model.wmb: magic WMB1, then h_s, entropy, g_s (WHM1) and the tables (JSON),
|
|
each as a u32 length and its bytes."""
|
|
out = bytearray(b'WMB1')
|
|
tj = json.dumps(tables, separators=(',', ':')).encode()
|
|
for part in (model_to_bin(models['h_s']), model_to_bin(models['entropy']), model_to_bin(models['g_s']), tj):
|
|
out += struct.pack('<I', len(part)) + part
|
|
return bytes(out)
|
|
|
|
|
|
def payload(enc):
|
|
zh, zw, yh, yw = enc['dims']
|
|
# 36 bytes: magic, u32 H, u32 W, u16 Cz, Zh, Zw, Cy, Yh, Yw, K, reserved (0),
|
|
# u32 Lz, u32 Ly -- then the z stream and the y stream.
|
|
hdr = b'WIM1' + struct.pack('<II8HII', H, W, CZ, zh, zw, CY, yh, yw, K, 0,
|
|
len(enc['z_bits']), len(enc['y_bits']))
|
|
assert len(hdr) == 36
|
|
return hdr + enc['z_bits'] + enc['y_bits']
|
|
|
|
|
|
def envelope(pl, model_sha):
|
|
manifest = {'wai': '1.0', 'media': 'image', 'intent': 'replicate',
|
|
'model_requirement': {'capability': CAPABILITY},
|
|
'conditioning': {'kind': 'int_mlicpp',
|
|
'prior': {'id': PRIOR_ID, 'sha256': model_sha, 'determinism': DETERMINISM}},
|
|
'target': {'w': W, 'h': H}}
|
|
mb = json.dumps(manifest, separators=(',', ':')).encode()
|
|
return b'WAI1' + struct.pack('<I', len(mb)) + mb + struct.pack('<I', len(pl)) + pl
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('outdir')
|
|
ap.add_argument('--stats', action='store_true', help='print latent, scale and escape statistics')
|
|
a = ap.parse_args()
|
|
rng = np.random.default_rng(SEED)
|
|
models = build_model(rng)
|
|
enc = encode(models, stats=a.stats)
|
|
wmb = bundle(models, enc['tables'])
|
|
model_sha = hashlib.sha256(wmb).hexdigest()
|
|
pl = payload(enc)
|
|
env = envelope(pl, model_sha)
|
|
rgb = enc['rgb']
|
|
expected = {
|
|
'capability': CAPABILITY, 'width': W, 'height': H,
|
|
'rgb_fnv1a': format(fnv1a(rgb.tobytes()), '016x'),
|
|
'rgb_sha256': hashlib.sha256(rgb.tobytes()).hexdigest(),
|
|
'bucket_fnv1a': format(fnv1a(np.array(enc['idxs'], dtype='<u2').tobytes()), '016x'),
|
|
'symbol_fnv1a': format(fnv1a(np.array(enc['syms'], dtype='<i4').tobytes()), '016x'),
|
|
'model_sha256': model_sha,
|
|
'payload_sha256': hashlib.sha256(pl).hexdigest(),
|
|
**enc['counts'],
|
|
'seed': SEED,
|
|
'weights': 'synthetic (seeded), not a trained model: decode-equivalence only, no quality claim',
|
|
}
|
|
os.makedirs(a.outdir, exist_ok=True)
|
|
with open(os.path.join(a.outdir, 'model.wmb'), 'wb') as f:
|
|
f.write(wmb)
|
|
with open(os.path.join(a.outdir, 'stream.wai'), 'wb') as f:
|
|
f.write(env)
|
|
with open(os.path.join(a.outdir, 'expected.json'), 'w') as f:
|
|
json.dump(expected, f, indent=1)
|
|
f.write('\n')
|
|
from wai_int_mlicpp_set_form import write_set_form
|
|
write_set_form(a.outdir)
|
|
print(f'int_mlicpp vector: {W}x{H}, y {CY}x{enc["dims"][2]}x{enc["dims"][3]} in {K} slices, '
|
|
f'{len(enc["z_bits"])} z + {len(enc["y_bits"])} y rANS bytes, model {len(wmb)} B, '
|
|
f'rgb_fnv1a {expected["rgb_fnv1a"]}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|