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]>
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())
|