open-standards/wai/tools/wai_int_mlicpp_export.py

#!/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`).
  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')
    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()