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))