"""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 _conv_node(w_dims, w_raw, b_raw): """One `Conv` as `jpegai_icci`'s `conv_of` reads it. Padding is kernel//2, which keeps the spatial size for both the 3x3 stems and the factorized 1x3 / 3x1 taps.""" cout, cin, kh, kw = w_dims return {"w": _decode(w_dims, 1, w_raw), "b": _decode([len(b_raw) // 4], 1, b_raw), "cin": cin, "cout": cout, "kh": kh, "kw": kw, "ph": kh // 2, "pw": kw // 2} def _icci_path(inits, n_blocks, conv_last_name, scale_name, hidden_prefix): """One eICCI plane-path out of a wrapper's initializers. Heads and scales are found by NAME; the trunk convolutions are anonymous (`onnx::Conv_*`) and are assigned by FILE ORDER — conv_first, then each residual block's 1x3 tap followed by its 3x1 tap. The wrappers are BN-pre-folded, so nothing is folded here. """ named = {n: (d, r) for n, (d, _dt, r) in inits.items()} anon = [(d, r) for n, (d, _dt, r) in inits.items() if n.startswith("onnx::Conv")] convs = [(anon[i][0], anon[i][1], anon[i + 1][1]) for i in range(0, len(anon), 2)] assert len(convs) == 1 + 2 * n_blocks, \ f"{len(convs)} convs for {n_blocks} blocks in {conv_last_name}" scalar = lambda name: _decode(named[name][0], 1, named[name][1])[0] blocks = [{"c1": _conv_node(*convs[1 + 2 * b]), "c2": _conv_node(*convs[2 + 2 * b]), "scale": scalar(f"{hidden_prefix}.{b}.scale")} for b in range(n_blocks)] cl_d, cl_r = named[conv_last_name + ".weight"] return {"conv_first": _conv_node(*convs[0]), "blocks": blocks, "conv_last": _conv_node(cl_d, cl_r, named[conv_last_name + ".bias"][1]), "scale": scalar(scale_name)} def export_icci(y_src, uv_src, dst): """The two eICCI model wrappers -> the `IcciNet` weight JSON. Which wrappers is signalled per picture (`icci_model_signalled_idx` indexes a shortlist); for the 256x256 conformance stream that is Y = model_6, UV = model_3. The luma path has 2 residual blocks and the chroma path 4, and U and V carry separate residual scales. """ luma = _icci_path(initializers(y_src), 2, "comp_model.conv_last_Y", "comp_model.scaleY", "comp_model.hidden_layer_Y") ui = initializers(uv_src) chroma = _icci_path(ui, 4, "comp_model.conv_last_u", "comp_model.scaleU", "comp_model.hidden_layer_UV") clv = ui["comp_model.conv_last_v.weight"] out = {"luma": luma, "chroma": chroma, "conv_last_v": _conv_node(clv[0], clv[2], ui["comp_model.conv_last_v.bias"][2]), "scale_v": _decode(ui["comp_model.scaleV"][0], 1, ui["comp_model.scaleV"][2])[0]} json.dump(out, open(dst, "w")) print(f"eICCI: luma {len(luma['blocks'])} blocks, chroma {len(chroma['blocks'])} blocks; " f"scaleY={luma['scale']:.6f} scaleU={chroma['scale']:.6f} " f"scaleV={out['scale_v']:.6f} -> {dst}") 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]) elif mode == "icci": export_icci(argv[2], argv[3], argv[4]) # icci else: print(f"unknown mode {mode!r}; expected list | ptq | int | icci") return 1 return 0 if __name__ == "__main__": sys.exit(main(sys.argv))