fb103ac5ae
Every registered string now lives in one table (codecs/registry.rs) with its kind, status, media class, registering SPEC section, edition and registration date. `wai registry export` writes registry/capabilities.json, its JSON Schema, the generated table in REGISTRY.md and vocab v2; --check fails on drift. vocab v1 stays stable. A candidate capability (a registered string whose payload format no revision has pinned) must not be emitted. The C packers, the browser writers and `wai wrap` refuse one, and `codecs::is_dispatchable` is the single dispatch and selection predicate (SPEC §4): candidates, companion-only capabilities, records and derivations are not dispatchable, and an unregistered string is decoded only where the sink advertises it. The codec glossary is linked to the registry through a reviewed table and held to its schema by a library test; `org` becomes `body`, from a closed allowlist. The SPEC gains §5 sections on kinds and candidates, a §8 edition rule, corrected citations and rewritten appendices. The quantum-ops vectors are regenerated with neutral device ids. CI declares the feature-gated test targets with required-features, checks the registry export, builds wai-web for wasm32, and runs two public-text checks (comparative prose, and private name patterns) with a self-test. The specification, extensions, research notes, READMEs, glossary and source comments are rewritten to name no company or product and to make no comparisons. Co-Authored-By: Claude Opus 5.5 <[email protected]>
109 lines
4 KiB
Python
109 lines
4 KiB
Python
"""Export the DAC 44.1 kHz decoder to ONNX for the WAI browser runtime.
|
|
|
|
DAC's decode signature differs from EnCodec's: it takes
|
|
`audio_codes [B, n_codebooks, T]` directly (no chunking, no audio_scales).
|
|
We wrap it so the exported ONNX accepts the same `[1, 1, q, t]` int64 +
|
|
`[1]` float32 inputs that the shared WAI audio handler feeds for every
|
|
neural audio codec — the `audio_scales` input is accepted and ignored,
|
|
and the leading chunks dim is squeezed.
|
|
|
|
Usage:
|
|
python tools/wai_dac_export_onnx.py \
|
|
--out wai-web/demo/models/dac_44khz/decoder.onnx
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from transformers import DacModel
|
|
|
|
|
|
class WaiDacDecoder(torch.nn.Module):
|
|
"""Adapts DacModel.decode to the WAI neural-audio ONNX contract.
|
|
|
|
Inputs:
|
|
audio_codes int64 [1, 1, q, t] (chunks dim is unused; kept for ABI parity)
|
|
audio_scales float [1] (unused; DAC has no per-chunk scaling)
|
|
Output:
|
|
audio_values float [1, 1, n_samples]
|
|
"""
|
|
def __init__(self, model: DacModel):
|
|
super().__init__()
|
|
self.m = model
|
|
|
|
def forward(self, audio_codes: torch.Tensor,
|
|
audio_scales: torch.Tensor) -> torch.Tensor:
|
|
# Squeeze the WAI "chunks" dim → [B, q, t] which is what DAC wants.
|
|
codes = audio_codes.squeeze(1)
|
|
out = self.m.decode(audio_codes=codes, return_dict=True)
|
|
wav = out.audio_values # [B, samples] in this transformers version
|
|
if wav.dim() == 2:
|
|
wav = wav.unsqueeze(1) # → [B, 1, samples]
|
|
# Touch audio_scales so ONNX keeps it in the input signature for ABI
|
|
# consistency with the EnCodec/WavTokenizer exports. The expression collapses to 1.0.
|
|
keep_scales = audio_scales.sum() / audio_scales.sum().clamp(min=1e-9)
|
|
return wav * keep_scales
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--model-id", default="descript/dac_44khz", # allow-name: model-hub id the tool defaults to
|
|
help="HF model id or local path.")
|
|
ap.add_argument("--out", required=True, help="Output .onnx path.")
|
|
ap.add_argument("--opset", type=int, default=17)
|
|
args = ap.parse_args()
|
|
out = Path(args.out); out.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"loading DAC from {args.model_id}…")
|
|
m = DacModel.from_pretrained(args.model_id).eval()
|
|
print(f" n_codebooks={m.config.n_codebooks} "
|
|
f"sampling_rate={m.config.sampling_rate} "
|
|
f"hop_length={m.config.hop_length}")
|
|
wrap = WaiDacDecoder(m).eval()
|
|
|
|
q = m.config.n_codebooks
|
|
T = 128
|
|
dummy_codes = torch.zeros(1, 1, q, T, dtype=torch.long)
|
|
dummy_scales = torch.ones (1, dtype=torch.float32)
|
|
|
|
with torch.no_grad():
|
|
ref = wrap(dummy_codes, dummy_scales)
|
|
print(f"reference forward pass OK: {tuple(ref.shape)}")
|
|
|
|
print(f"exporting → {out} (opset {args.opset})…")
|
|
torch.onnx.export(
|
|
wrap,
|
|
(dummy_codes, dummy_scales),
|
|
out.as_posix(),
|
|
input_names=["audio_codes", "audio_scales"],
|
|
output_names=["audio_values"],
|
|
dynamic_axes={
|
|
"audio_codes": {3: "T"},
|
|
"audio_values": {2: "samples"},
|
|
},
|
|
opset_version=args.opset,
|
|
do_constant_folding=True,
|
|
dynamo=False, # legacy tracer; see tools/wai_encodec_export_onnx.py
|
|
)
|
|
sz = out.stat().st_size
|
|
print(f"wrote {sz:,} B")
|
|
|
|
import onnxruntime as ort, numpy as np
|
|
sess = ort.InferenceSession(out.as_posix(), providers=["CPUExecutionProvider"])
|
|
ort_out = sess.run(None, {
|
|
"audio_codes": dummy_codes.numpy(),
|
|
"audio_scales": dummy_scales.numpy(),
|
|
})
|
|
diff = float(np.max(np.abs(ort_out[0] - ref.numpy())))
|
|
print(f"onnxruntime sanity: shape={ort_out[0].shape} "
|
|
f"max |onnx-torch| = {diff:.3e}")
|
|
if diff > 1e-3:
|
|
print("WARNING: numerical mismatch > 1e-3", file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|