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]>
78 lines
3 KiB
Python
78 lines
3 KiB
Python
"""Export the REAL pretrained SNAC model to the ONNX `wai.neural.snac` decoder
|
||
(`wai-rs/src/neural/snac.rs`), via the model surgery needed to make it both
|
||
exportable AND deterministic.
|
||
|
||
Two upstream blockers, both patched here (monkey-patch, no fork):
|
||
1. NoiseBlock injects `torch.randn` — that makes the decode STOCHASTIC (a codec
|
||
must be reproducible) AND bakes a bad shape on ONNX export. We zero the
|
||
noise: `x + 0*linear(x)` → deterministic, exportable. (Quality cost is the
|
||
noise texture only.)
|
||
2. Snake1d's `@torch.jit.script snake` does a reshape-roundtrip that ONNX
|
||
mis-traces. We replace it with the identical 3D formula, no reshape.
|
||
|
||
Verified: snac_24khz exports clean and the ONNX matches the (now deterministic)
|
||
torch decode to ~9e-7; `wai-rs` then decodes the resulting envelope to ~9e-7 of
|
||
that reference (`snac_real_model_matches_reference`).
|
||
|
||
Run (needs torch + snac + onnxruntime):
|
||
python tools/wai_snac_export.py --model hubertsiuzdak/snac_24khz \
|
||
--out-dir wai-web/demo/models/snac_24khz
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
|
||
def patch_for_export():
|
||
import snac.layers
|
||
# 1. deterministic, export-clean noise.
|
||
def noise_fwd(self, x):
|
||
return x + 0.0 * self.linear(x)
|
||
snac.layers.NoiseBlock.forward = noise_fwd
|
||
# 2. reshape-free snake.
|
||
def snake1d_fwd(self, x):
|
||
return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x) ** 2
|
||
snac.layers.Snake1d.forward = snake1d_fwd
|
||
|
||
|
||
class _Dec(nn.Module):
|
||
def __init__(s, m): super().__init__(); s.m = m
|
||
def forward(s, *codes):
|
||
return s.m.decoder(s.m.quantizer.from_codes(list(codes)))
|
||
|
||
|
||
def main() -> int:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--model", default="hubertsiuzdak/snac_24khz")
|
||
ap.add_argument("--out-dir", required=True)
|
||
ap.add_argument("--opset", type=int, default=17)
|
||
args = ap.parse_args()
|
||
|
||
patch_for_export()
|
||
from snac import SNAC
|
||
m = SNAC.from_pretrained(args.model).eval()
|
||
out = Path(args.out_dir); out.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Probe the codebook hierarchy (N levels, their strides) from a dummy encode.
|
||
dummy_audio = torch.zeros(1, 1, m.sampling_rate)
|
||
with torch.no_grad():
|
||
codes = m.encode(dummy_audio)
|
||
n = len(codes)
|
||
names = [f"codes_{i}" for i in range(n)]
|
||
torch.onnx.export(
|
||
_Dec(m).eval(), tuple(codes), str(out / "decoder.onnx"),
|
||
input_names=names, output_names=["audio_values"],
|
||
dynamic_axes={**{f"codes_{i}": {1: f"t{i}"} for i in range(n)}, "audio_values": {2: "s"}},
|
||
opset_version=args.opset, dynamo=False)
|
||
print(f"exported {args.model}: {n} code levels @ {m.sampling_rate} Hz → {out}/decoder.onnx")
|
||
print("decode is deterministic (zero-noise). Encode tool: pack <III> sr,n_samples,n_levels "
|
||
"+ n×u32 lengths + per-level u16 codes (see snac.rs).")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|