46c5688b86
Commit the WAI standard source: a content + capability-dispatch standard
whose thesis is "carries whatever it can reconstruct, verifies whatever
it carries." Source only — large vendored/model assets and build output
are gitignored.
Reconstruct classes (each: spec + conformance corpus + in-tab browser
sink that re-derives the corpus hash, one wai.det.fixed64 fixed-point
floor, one *-equivalence apiece):
- interactive worlds (replay-equivalence) + 3 sim registrants,
multiplayer over a relay, session receipts, glTF + real SPZ assets
- spatial audio (mixdown-equivalence)
- haptics (sample-equivalence)
- state-feeds (replay-equivalence, non-physics sim)
- avatars / digital humans (pose-equivalence) + telepresence over the relay
- volumetric / 4D splats (frame-state-equivalence) + WebGPU EWA render,
real-SPZ-derived clip
Verify half (lever 4):
- asset provenance: a signed, energy-accounted lineage for opaque
captured content WAI does not reconstruct; tamper-evident, in-tab.
Cross-cutting: JWP receipt profiles per class; a native energy meter
(IOReport via macmon, no sudo) emitting receipts with REAL measured
joules; the learned-image/neural codec family and JPEG AI entropy port
from earlier work.
Co-Authored-By: Claude Opus 4.8 <[email protected]>
75 lines
3.1 KiB
Python
75 lines
3.1 KiB
Python
"""Export sink artifacts for the `wai.neural.mlicpp` channel-autoregressive
|
|
codec decoded by `wai-rs/src/neural/mlicpp.rs`.
|
|
|
|
mlicpp.rs needs (next to the registered model path):
|
|
hyper_synthesis.onnx z_hat -> [1, 2*Cy, Sy, Sy] (base scales ‖ means)
|
|
entropy.onnx concat[y_sofar, hyper] -> [1, 2*Cy, Sy, Sy] (refined params)
|
|
input is [1, 3*Cy, Sy, Sy] (y is Cy ch, hyper is 2*Cy)
|
|
synthesis.onnx y_hat -> [1, 3, H, W]
|
|
hyper_cdfs_q<q>.json factorized z CDF (EntropyBottleneck)
|
|
gauss_cdfs_q<q>.json GaussianConditional CDF table
|
|
|
|
It wraps a checkpoint exposing `h_s`, `entropy` (the multi-reference net in the
|
|
WAI-native single-net form: concat[y_sofar, hyper] → params), `g_s`,
|
|
`entropy_bottleneck`, `gaussian_conditional`. This is the MLIC++ *class*
|
|
(channel-autoregressive multi-reference); the attention internals live inside
|
|
`entropy.onnx`. CDF extraction matches the hyperprior tool (same coders).
|
|
|
|
Run (needs torch + onnxruntime; bring your own mlicpp-class checkpoint):
|
|
python tools/wai_mlicpp_export.py --ckpt mlic.pt --quality 3 --size 256 \
|
|
--out-dir wai-web/demo/models/mlicpp
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from wai_mbt2018_mean_export import factorized_cdf_blob, gaussian_cdf_blob # noqa: E402
|
|
|
|
|
|
def export_onnx(module, dummy, path: Path, opset):
|
|
torch.onnx.export(
|
|
module, (dummy,), path.as_posix(),
|
|
input_names=["i"], output_names=["o"],
|
|
dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}},
|
|
opset_version=opset, do_constant_folding=True, dynamo=False)
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--ckpt", required=True, help="mlicpp-class ckpt: h_s/entropy/g_s + entropy coders")
|
|
ap.add_argument("--quality", type=int, default=3)
|
|
ap.add_argument("--size", type=int, default=256)
|
|
ap.add_argument("--out-dir", required=True)
|
|
ap.add_argument("--opset", type=int, default=17)
|
|
args = ap.parse_args()
|
|
out = Path(args.out_dir); out.mkdir(parents=True, exist_ok=True)
|
|
q = args.quality
|
|
|
|
net = torch.load(args.ckpt, map_location="cpu", weights_only=False).eval()
|
|
net.update(force=True)
|
|
|
|
with torch.no_grad():
|
|
y = net.g_a(torch.zeros(1, 3, args.size, args.size))
|
|
z = net.h_a(y)
|
|
_, M, Sy, _ = y.shape
|
|
_, N, Sz, _ = z.shape
|
|
print(f" y[1,{M},{Sy},{Sy}] z[1,{N},{Sz},{Sz}] entropy-in=[1,{3 * M},{Sy},{Sy}]")
|
|
|
|
export_onnx(net.h_s, torch.zeros(1, N, Sz, Sz), out / "hyper_synthesis.onnx", args.opset)
|
|
export_onnx(net.entropy, torch.zeros(1, 3 * M, Sy, Sy), out / "entropy.onnx", args.opset)
|
|
export_onnx(net.g_s, torch.zeros(1, M, Sy, Sy), out / "synthesis.onnx", args.opset)
|
|
|
|
(out / f"hyper_cdfs_q{q}.json").write_text(
|
|
json.dumps(factorized_cdf_blob(net.entropy_bottleneck, q), separators=(",", ":")))
|
|
(out / f"gauss_cdfs_q{q}.json").write_text(
|
|
json.dumps(gaussian_cdf_blob(net.gaussian_conditional, q), separators=(",", ":")))
|
|
print(f"wrote 3 ONNX + 2 CDF tables → {out}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|