cc83d1445e
Extracts the IcciNet weights from the two eICCI wrappers into the JSON the icci calibrations read. Heads and scales come out by name; the trunk convolutions are anonymous in the graph and are assigned by file order — conv_first, then each residual block's 1x3 tap followed by its 3x1 tap. These wrappers are BN-pre-folded, unlike post_filters/model_*.onnx, so there is nothing to fold. The structure is asymmetric in two ways worth stating: the luma path has 2 residual blocks where chroma has 4, and the blocks are factorized into 1x3 and 3x1 taps rather than a single 3x3. Which pair of wrappers applies is signalled per picture — icci_model_signalled_idx indexes a shortlist, which for the 256x256 conformance stream selects Y = model_6 and UV = model_3. Having this as a mode rather than a one-off script means the weights are regenerable; the previous extraction was lost with its scratch directory. Co-Authored-By: Claude Opus 4.8 <[email protected]>
283 lines
11 KiB
Python
283 lines
11 KiB
Python
"""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 _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 <y_wrapper> <uv_wrapper> <out>
|
|
else:
|
|
print(f"unknown mode {mode!r}; expected list | ptq | int | icci")
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv))
|