From 963410a8cc95a1c9e41c4b28c8644a16b61d7e7e Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Sat, 18 Jul 2026 10:12:42 -0400 Subject: [PATCH] =?UTF-8?q?wai/tools:=20add=20wai=5Fjpegai=5Fexport=5Fweig?= =?UTF-8?q?hts.py=20=E2=80=94=20reproducible=20weight=20export=20for=20the?= =?UTF-8?q?=20decoder?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The byte-exact JPEG-AI decoder in wai-rs is driven by weights lifted from the reference ONNX assets, but the export step lived only in scratch and had to be rewritten each time it was lost. This commits it. Two regimes, and choosing wrongly yields plausible-but-wrong output: ptq float nets (synthesis model_y/model_uv, hyper_decoder, MCM stages) -> 16-bit per-tensor symmetric PTQ (conv -> int16 + mult=round(scale*2^31), shift=31; bias/LayerNorm/temperature -> Q16), the layout Weights::from_json expects. int hyper_scale_decoder -> verbatim. That net is integer end to end, so quantizing it would corrupt an exactly-representable network; IntNet::from_json reads this. Dependency-free: only the ONNX initializer table is needed, so it reads the protobuf directly (no onnx/protobuf/torch import). Verified to reproduce the weights the merged calibration tests were validated against: synthesis_y 36 conv + 28 scalar (3,341,769 params, worst round-trip 6.6e-05), MCM stage0 3 conv (76,800 params, 6.9e-05), hyper_scale 9 integer tensors. Re-running jpegai_hyperscale_calibration against freshly exported weights still reproduces the reference scale indices 40960/40960 exactly. Documents the eICCI caveat: post_filters/model_*.onnx export the model's `forward`, which the decoder never calls, and their initializers are BN-folded and incomplete (no BNY, scaleU/scaleV, or per-block scale) — take the reference state_dict instead. Co-Authored-By: Claude Opus 4.8 --- wai/tools/wai_jpegai_export_weights.py | 226 +++++++++++++++++++++++++ 1 file changed, 226 insertions(+) create mode 100644 wai/tools/wai_jpegai_export_weights.py diff --git a/wai/tools/wai_jpegai_export_weights.py b/wai/tools/wai_jpegai_export_weights.py new file mode 100644 index 000000000..7e3f043bf --- /dev/null +++ b/wai/tools/wai_jpegai_export_weights.py @@ -0,0 +1,226 @@ +"""Export JPEG-AI decoder weights from the ISO reference ONNX assets. + +The byte-exact decoder in `wai-rs` (`jpegai_synthesis`, `jpegai_pipeline`, +`jpegai_icci`) is driven by weights lifted straight out of the reference +`jpegai-vm-assets` ONNX files. Two regimes, and using the wrong one silently +produces plausible-but-wrong output: + + float nets — synthesis (model_y / model_uv), hyper_decoder, MCM stages. + Exported as 16-bit per-tensor symmetric PTQ: + conv weight [O, I/g, kh, kw] -> int16 + per-tensor + mult = round(scale * 2^31), shift = 31 + bias / LayerNorm / temperature -> Q16 (round(v * 2^16)) + This is what `jpegai_synthesis::Weights::from_json` expects. + + integer net — hyper_scale_decoder. Already integer end to end (int32 conv + weights, integer-valued Pow exponents); exported VERBATIM with + no quantization, for `jpegai_synthesis::IntNet::from_json`. + Quantizing it would corrupt an exactly-representable net. + +Reading the file is dependency-free: ONNX is protobuf, and we only need the +initializer table, so there is no `onnx` / `protobuf` / `torch` requirement. + + ModelProto.graph = 7 + GraphProto.initializer = 5 + TensorProto{ dims=1, data_type=2, name=8, raw_data=9 } + data_type: 1 = float32, 6 = int32, 7 = int64 + +Usage +----- + # float net -> Q16 PTQ (synthesis / hyper_decoder / MCM stage) + python wai_jpegai_export_weights.py ptq synthesis16.onnx synth16_y.json + + # integer net -> verbatim + python wai_jpegai_export_weights.py int hyper_scale_decoder.onnx hs_y.json + + # just look at what is inside + python wai_jpegai_export_weights.py list model_15.onnx + +Caveat worth knowing: the eICCI post-filter models under `post_filters/` export +the model's `forward`, which the reference decoder never calls (it calls +`process_dwt444_2` + `process_comp`), and their initializers are BatchNorm-folded +and incomplete — no `BNY`, no `scaleU`/`scaleV`, no per-residual-block `scale`. +For eICCI take the reference's PyTorch `state_dict` instead; `list` is still +useful there for inspection. +""" + +import json +import struct +import sys +from collections import OrderedDict + +SHIFT = 31 # conv requant shift, matches int_transformer::conv_grouped +Q16 = 1 << 16 # bias / LN / temperature fixed-point scale +INT16_MAX = 32767 + + +# ---- dependency-free protobuf reader ---- + +def _varint(buf, i): + shift = value = 0 + while True: + byte = buf[i] + i += 1 + value |= (byte & 0x7F) << shift + if not byte & 0x80: + return value, i + shift += 7 + + +def _fields(buf): + """Yield (field_number, wire_type, payload) for one protobuf message.""" + i = 0 + while i < len(buf): + tag, i = _varint(buf, i) + field, wire = tag >> 3, tag & 7 + if wire == 0: + value, i = _varint(buf, i) + yield field, 0, value + elif wire == 2: + length, i = _varint(buf, i) + yield field, 2, buf[i:i + length] + i += length + elif wire == 1: + yield field, 1, buf[i:i + 8] + i += 8 + elif wire == 5: + yield field, 5, buf[i:i + 4] + i += 4 + else: + raise ValueError(f"unsupported wire type {wire}") + + +def _parse_tensor(buf): + dims, name, dtype, raw = [], "", 0, b"" + for field, wire, payload in _fields(buf): + if field == 1: + if wire == 0: + dims.append(payload) + elif wire == 2: # packed dims + j = 0 + while j < len(payload): + d, j = _varint(payload, j) + dims.append(d) + elif field == 2 and wire == 0: + dtype = payload + elif field == 8 and wire == 2: + name = payload.decode("utf-8", "replace") + elif field == 9 and wire == 2: + raw = payload + return name, dims, dtype, raw + + +def initializers(path): + """name -> (dims, data_type, raw_bytes) for every initializer, in file order.""" + out = OrderedDict() + model = open(path, "rb").read() + for field, wire, payload in _fields(model): + if field == 7 and wire == 2: # ModelProto.graph + for gfield, gwire, gpayload in _fields(payload): + if gfield == 5 and gwire == 2: # GraphProto.initializer + name, dims, dtype, raw = _parse_tensor(gpayload) + out[name] = (dims, dtype, raw) + return out + + +def _prod(xs): + n = 1 + for x in xs: + n *= x + return n + + +def _decode(dims, dtype, raw): + """raw bytes -> list of numbers, or None if the tensor is empty/unsupported.""" + n = _prod(dims) + if not raw: + return None + if dtype == 1 and len(raw) == 4 * n: + return list(struct.unpack(f"<{n}f", raw)) + if dtype == 6 and len(raw) == 4 * n: + return list(struct.unpack(f"<{n}i", raw)) + if dtype == 7 and len(raw) == 8 * n: + return list(struct.unpack(f"<{n}q", raw)) + return None + + +# ---- export modes ---- + +def export_ptq(src, dst, strip_prefix=True): + """Float net -> 16-bit symmetric PTQ, the layout `Weights::from_json` reads.""" + import re + result = OrderedDict() + n_conv = n_scalar = total = 0 + worst = 0.0 + for name, (dims, dtype, raw) in initializers(src).items(): + values = _decode(dims, dtype, raw) + if values is None or dtype != 1: + continue + total += len(values) + # the reference wraps modules as `hop_.inst.`; the decoder + # addresses them by the bare module path + key = re.sub(r"^hop_\w+\.inst\.", "", name) if strip_prefix else name + if len(dims) == 4 and dims[0] > 1: # conv weight [O, I/g, kh, kw] + amax = max((abs(v) for v in values), default=0.0) or 1e-12 + scale = amax / INT16_MAX + quant = [max(-32768, min(INT16_MAX, round(v / scale))) for v in values] + worst = max(worst, max(abs(v - round(v / scale) * scale) for v in values)) + result[key] = {"kind": "conv", "dims": dims, "shift": SHIFT, + "mult": round(scale * (1 << SHIFT)), "w": quant} + n_conv += 1 + else: # bias / LN gamma,beta / temperature + result[key] = {"kind": "q16", "dims": dims, + "v": [round(v * Q16) for v in values]} + n_scalar += 1 + json.dump(result, open(dst, "w")) + print(f"{n_conv} conv + {n_scalar} scalar tensors ({total:,} params) -> {dst}") + print(f"conv round-trip worst |error| = {worst:.3e} " + f"({'lossless at 16 bits' if worst < 1e-3 else 'CHECK'})") + + +def export_int(src, dst): + """Integer net -> verbatim. Pow exponents are integer-valued floats.""" + result = OrderedDict() + for name, (dims, dtype, raw) in initializers(src).items(): + values = _decode(dims, dtype, raw) + if values is None: + continue + if dtype == 1: # integer-valued float (Pow exponent) + values = [int(round(v)) for v in values] + result[name] = {"dims": dims, "v": values} + json.dump(result, open(dst, "w")) + print(f"{len(result)} integer tensors -> {dst}") + for name, entry in result.items(): + lo, hi = min(entry["v"]), max(entry["v"]) + print(f" {name:24s} dims={entry['dims']} range=[{lo}, {hi}]") + + +def list_tensors(src): + dtypes = {1: "f32", 6: "i32", 7: "i64"} + total = 0 + for name, (dims, dtype, raw) in initializers(src).items(): + n = _prod(dims) + total += n + print(f" {name:38s} {dtypes.get(dtype, dtype):4s} dims={dims}") + print(f"{total:,} params") + + +def main(argv): + if len(argv) < 3: + print(__doc__) + return 1 + mode, src = argv[1], argv[2] + if mode == "list": + list_tensors(src) + elif mode == "ptq": + export_ptq(src, argv[3]) + elif mode == "int": + export_int(src, argv[3]) + else: + print(f"unknown mode {mode!r}; expected list | ptq | int") + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) -- 2.54.0 (Apple Git-157)