c151bdb819
The cargo-class push deliberately excluded the neural/semantic work to avoid reverting origin; origin never independently changed neural, so this brings it forward cleanly and wires it up. - neural/mod.rs declares the full decoder family; rans.rs gains the categorical decoder it needs (glc/mlicpp/elic/hyperprior/int_synth/ cmic/dcvc/snac + the JPEG-AI entropy/codestream/σ/value/filter port). - runtime.rs uses the shared deterministic ort_session loader. - Restores the semantic/generative path (semantic.rs, min_sim/sim_metric) behind the `semantic` feature, with the bins/tests that construct ModelRequirement updated accordingly. - Re-adds byte-exact-conformance and the tools/ export/encode scripts. Green together: wai-rs all features (75 lib + integration tests) + clippy, byte-exact-conformance (8), wai-web wasm, wai-spz-gen. Co-Authored-By: Claude Opus 4.8 <[email protected]>
101 lines
4.3 KiB
Python
101 lines
4.3 KiB
Python
"""Quantize a synthesis transform to the integer `int_synth.json` spec decoded
|
|
bit-exactly by `wai-rs/src/neural/int_synth.rs` — the integer-exact `replicate`
|
|
path.
|
|
|
|
`int_synth` runs an all-integer transposed-conv stack (no float, no ONNX), so
|
|
its decode is identical on every machine. This tool turns a trained float
|
|
synthesis (ConvTranspose2d + ReLU layers) into the integer weights + per-layer
|
|
requantize shifts the Rust decoder needs.
|
|
|
|
IMPORTANT on quality: integer fixed-point inference is bit-exact regardless, but
|
|
to match the float model's *quality* the model should be **quantization-aware
|
|
trained** (QAT) for these shifts. A post-hoc round of an arbitrary float model
|
|
will decode deterministically but may look worse — the determinism is free, the
|
|
quality is the training step. This tool does the export; QAT is upstream.
|
|
|
|
Spec written (matches int_synth.rs):
|
|
{ "layers": [ { in_ch, out_ch, k, stride, pad, out_pad,
|
|
weight:[i32; in*out*k*k] # layout ic,oc,ky,kx
|
|
bias:[i32; out_ch], shift:u32, relu:bool } ] }
|
|
|
|
Run (needs torch):
|
|
python tools/wai_int_synth_export.py --weights synth.pt \
|
|
--out wai-web/demo/models/int_synth/int_synth.json
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
def quantize_layer(conv: nn.ConvTranspose2d, w_bits: int, shift: int, relu: bool) -> dict:
|
|
"""Per-tensor symmetric integer quantization of a ConvTranspose2d.
|
|
|
|
weight_q = round(weight / w_scale), bias_q = round(bias / (w_scale * in_scale))
|
|
folded so the integer accumulate matches; the decoder applies `>> shift` to
|
|
requantize. We pick `w_scale` from the weight's max-abs to fill the int range.
|
|
"""
|
|
w = conv.weight.detach().cpu() # [in_ch, out_ch, kH, kW]
|
|
qmax = (1 << (w_bits - 1)) - 1
|
|
w_scale = float(w.abs().max()) / qmax if float(w.abs().max()) > 0 else 1.0
|
|
w_q = torch.clamp(torch.round(w / w_scale), -qmax - 1, qmax).to(torch.int64)
|
|
|
|
bias = conv.bias.detach().cpu() if conv.bias is not None else torch.zeros(conv.out_channels)
|
|
# Bias lives in the accumulator domain (before the >> shift).
|
|
bias_q = torch.round(bias / w_scale * (1 << shift)).to(torch.int64)
|
|
|
|
in_ch, out_ch, kh, kw = w.shape
|
|
assert kh == kw, "int_synth assumes square kernels"
|
|
sH, sW = conv.stride
|
|
pH, pW = conv.padding
|
|
opH, opW = conv.output_padding
|
|
assert sH == sW and pH == pW and opH == opW, "int_synth assumes symmetric stride/pad/out_pad"
|
|
|
|
return {
|
|
"in_ch": in_ch, "out_ch": out_ch, "k": kh,
|
|
"stride": sH, "pad": pH, "out_pad": opH,
|
|
"weight": w_q.reshape(-1).tolist(),
|
|
"bias": bias_q.reshape(-1).tolist(),
|
|
"shift": shift,
|
|
"relu": relu,
|
|
"_w_scale": w_scale, # informational
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--weights", required=True,
|
|
help="torch checkpoint: a Sequential of ConvTranspose2d (+ReLU between).")
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--w-bits", type=int, default=8)
|
|
ap.add_argument("--shift", type=int, default=8,
|
|
help="per-layer requantize right-shift; tune with QAT for quality.")
|
|
args = ap.parse_args()
|
|
|
|
model = torch.load(args.weights, map_location="cpu", weights_only=False)
|
|
model.eval()
|
|
|
|
convs = [m for m in model.modules() if isinstance(m, nn.ConvTranspose2d)]
|
|
if not convs:
|
|
raise SystemExit("no ConvTranspose2d layers found in the checkpoint")
|
|
# ReLU after every layer except the last (final layer emits the image).
|
|
layers = [quantize_layer(c, args.w_bits, args.shift, relu=(i < len(convs) - 1))
|
|
for i, c in enumerate(convs)]
|
|
if layers[-1]["out_ch"] != 3:
|
|
print(f"WARNING: final layer out_ch={layers[-1]['out_ch']} != 3 (RGB)")
|
|
|
|
out = Path(args.out); out.parent.mkdir(parents=True, exist_ok=True)
|
|
out.write_text(json.dumps({"layers": layers}, separators=(",", ":")))
|
|
print(f"wrote {len(layers)}-layer integer synthesis → {out} "
|
|
f"(w_bits={args.w_bits}, shift={args.shift})")
|
|
print("determinism: bit-exact on any machine by construction; "
|
|
"quality depends on QAT for these shifts.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|