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_<x>.inst.<name>`; 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))