open-standards/wai/tools/wai_elic_export.py
Transaction Science 46c5688b86 wai: the WAI open standard — full cargo-class surface + provenance
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]>
2026-06-13 22:28:51 -04:00

74 lines
3.1 KiB
Python

"""Export sink artifacts for the `wai.neural.elic` checkerboard context codec
decoded by `wai-rs/src/neural/context.rs`.
context.rs needs (next to the registered model path):
hyper_synthesis.onnx z_hat -> [1, 2*Cy, Sy, Sy] (scales ‖ means, base)
context.onnx y_anchors -> [1, 2*Cy, Sy, Sy] (refined params for non-anchors)
synthesis.onnx y_hat -> [1, 3, H, W]
hyper_cdfs_q<q>.json factorized z CDF (EntropyBottleneck)
gauss_cdfs_q<q>.json { scale_table, cdf[], cdf_length[], offset[] } (GaussianConditional)
It wraps an ELIC-style checkpoint exposing submodules `h_s` (hyper synthesis),
`context` (the checkerboard context net: anchor-filled latent → params), `g_s`
(synthesis), `entropy_bottleneck`, `gaussian_conditional`. CDF extraction is
identical to the hyperprior tool (the entropy coders are the same).
Run (needs torch + onnxruntime; bring your own ELIC checkpoint):
python tools/wai_elic_export.py --ckpt elic.pt --quality 3 --size 256 \
--out-dir wai-web/demo/models/elic
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import torch
# Reuse the CDF extractors from the hyperprior tool (same entropy coders).
from wai_mbt2018_mean_export import factorized_cdf_blob, gaussian_cdf_blob # noqa: E402
def export_onnx(module, dummy, path: Path, in_name, out_name, opset):
torch.onnx.export(
module, (dummy,), path.as_posix(),
input_names=[in_name], output_names=[out_name],
dynamic_axes={in_name: {2: "h", 3: "w"}, out_name: {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="ELIC checkpoint exposing h_s/context/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}]")
export_onnx(net.h_s, torch.zeros(1, N, Sz, Sz), out / "hyper_synthesis.onnx", "z", "params", args.opset)
export_onnx(net.context, torch.zeros(1, M, Sy, Sy), out / "context.onnx", "y", "params", args.opset)
export_onnx(net.g_s, torch.zeros(1, M, Sy, Sy), out / "synthesis.onnx", "y", "image", 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())