22928147f3
SPEC §3.1 defines the pin a manifest carries (conditioning.prior, and a WAI2 rendition's prior): id, sha256, determinism, and the optional form, tier and supersedes. The digest covers the whole parameter set. A set of several files is pinned through its listing, the new registered prior form wai.prior.set. A sink verifies the pin before it decodes, reports a refusal with one of an ordered set of snake_case reasons, and decodes only from the set it verified. Tier is a closed vocabulary; supersedes is informative. A pin that does not resolve makes its capability unsupported for that envelope: SPEC §4 dispatch continues at the fallback, §7.1 selection at the next rendition, and an inert envelope's error names the reasons. A rendition entry may carry id, sha256 and, on a companion, companion_of_id. SPEC §7 gains a Digests subsection, and §8 says how a MINOR revision adds an optional member. In wai-rs: the manifest types are #[non_exhaustive] with builders; a new pin module verifies pins in memory and on disk; Wai::dispatch_verified, WaiMulti::pick and primary check pins and component digests; every neural decoder reads through ParamSource; a codec-free `ffi` feature carries the C envelope surface with wai_envelope_select and wai_neural_decode_ex. wai-web wraps wai::container. New vectors: the int_mlicpp set-form twin and the pin-conformance corpus with an independent Python verifier, plus sha256.json sidecars for the vectors of capabilities that decode against a parameter set. Squashed from wai/stream-a-parameter-set-identity onto link 12. Conflicts are resolved to keep both sides; the integration commits that follow make the two streams agree. 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 Rust decoder's
|
|
`int_hyper_synth::run_net_checked` (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()
|