"""Export the decoder of EnCodec 24 kHz, EnCodec 48 kHz, DAC 16 kHz or DAC 24 kHz to the one ONNX file a sink holds for `wai.neural.encodec24`, `wai.neural.encodec48`, `wai.neural.dac16` or `wai.neural.dac24` (SPEC ยง5, *Payloads of the neural audio codecs*). The decoder takes the model's full codebook count Q for every operating point: inputs: audio_codes int64 [1, 1, Q, T] row i < q: codebook i's codes; a row the payload does not carry holds codebook_size (1024) audio_scales float [1] the chunk's scale (1.0 for a model with no chunks) output: audio_values float [1, C, T*320 - d] C = 2 for EnCodec 48 kHz, else 1; d = 8 for DAC 16 and 24 kHz (their upsampling by 5 gives 5L - 1 samples for L), else 0 A row holding codebook_size adds nothing to the quantized latent, which is what the model's own decode does with fewer codebooks (its residual quantizer sums the codebooks it is given), so one file decodes every bitrate. The route in wai-rs (neural/audio.rs) checks the payload, fills the rows and, for EnCodec 48 kHz, decodes chunk by chunk and joins the chunks by the model's linear overlap-add. The export uses the checkpoints of the model hub through `transformers`, traces with the legacy tracer (the codebook loop unrolls at Q; T stays dynamic), and then checks itself: for each operating point it decodes seeded codes through onnxruntime and through the model's own decode with the first q codebooks, and fails if they differ by more than 1e-4. python tools/wai_neural_audio_export_onnx.py --model encodec24 --out models/encodec24/decoder.onnx Needs torch, transformers, onnx and onnxruntime; the weights are downloaded by `transformers`. EnCodec's code and weights are MIT; DAC's are MIT. """ from __future__ import annotations import argparse import sys from pathlib import Path MODELS = { # name: (model-hub id, kind, sample rate, channels, codebooks, operating points, d) "encodec24": ("facebook/encodec_24khz", "encodec", 24_000, 1, 32, [2, 4, 8, 16, 32], 0), # allow-name: model-hub id the tool defaults to "encodec48": ("facebook/encodec_48khz", "encodec", 48_000, 2, 16, [2, 4, 8, 16], 0), # allow-name: model-hub id the tool defaults to "dac16": ("descript/dac_16khz", "dac", 16_000, 1, 12, [1, 6, 12], 8), # allow-name: model-hub id the tool defaults to "dac24": ("descript/dac_24khz", "dac", 24_000, 1, 32, [1, 16, 32], 8), # allow-name: model-hub id the tool defaults to } CODEBOOK_SIZE = 1024 HOP = 320 def wrapper(kind, model, q_full): import torch class Encodec(torch.nn.Module): """EnCodec's residual-quantizer decode over the rows that hold a code, then its decoder.""" def __init__(self, m): super().__init__() self.m = m def forward(self, audio_codes, audio_scales): codes = audio_codes[0] # [1, Q, T] emb = None for i in range(q_full): idx = codes[:, i, :] present = (idx < CODEBOOK_SIZE).to(torch.float32).unsqueeze(-1) # [1, T, 1] safe = torch.where(idx < CODEBOOK_SIZE, idx, torch.zeros_like(idx)) e = torch.nn.functional.embedding(safe, self.m.quantizer.layers[i].codebook.embed) * present e = e.permute(0, 2, 1) # [1, D, T] emb = e if emb is None else emb + e return self.m.decoder(emb) * audio_scales.view(-1, 1, 1) class Dac(torch.nn.Module): """DAC's from_codes over the rows that hold a code (each codebook's out_proj, bias included, only where it is used), then its decoder.""" def __init__(self, m): super().__init__() self.m = m def forward(self, audio_codes, audio_scales): codes = audio_codes[0] # [1, Q, T] z = None for i in range(q_full): idx = codes[:, i, :] present = (idx < CODEBOOK_SIZE).to(torch.float32).unsqueeze(1) # [1, 1, T] safe = torch.where(idx < CODEBOOK_SIZE, idx, torch.zeros_like(idx)) qz = self.m.quantizer.quantizers[i] zi = qz.out_proj(qz.codebook(safe).transpose(1, 2)) * present z = zi if z is None else z + zi wav = self.m.decoder(z) if wav.dim() == 2: wav = wav.unsqueeze(1) return wav * audio_scales.view(-1, 1, 1) return (Encodec if kind == "encodec" else Dac)(model).eval() def reference(kind, model, codes_q, scale): """The model's own decode of one chunk with the codebooks it is given.""" import torch with torch.no_grad(): if kind == "encodec": # EncodecModel._decode_frame: codes [1, q, T] -> [1, C, samples], times the scale. emb = model.quantizer.decode(codes_q.transpose(0, 1)) return model.decoder(emb) * scale.view(-1, 1, 1) wav = model.decode(audio_codes=codes_q, return_dict=True).audio_values return (wav.unsqueeze(1) if wav.dim() == 2 else wav) * scale.view(-1, 1, 1) def main() -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--model", required=True, choices=sorted(MODELS)) ap.add_argument("--out", required=True, help="output .onnx path") ap.add_argument("--source", help="a local checkpoint directory instead of the model hub id") ap.add_argument("--opset", type=int, default=17) ap.add_argument("--frames", type=int, default=40, help="frames of the self-check's codes") args = ap.parse_args() import numpy as np import onnxruntime as ort import torch from transformers import DacModel, EncodecModel hub, kind, sr, channels, q_full, points, shrink = MODELS[args.model] src = args.source or hub model = (EncodecModel if kind == "encodec" else DacModel).from_pretrained(src).eval() cfg = model.config got = (cfg.sampling_rate, cfg.codebook_size) if got != (sr, CODEBOOK_SIZE): sys.exit(f"{src}: sampling rate and codebook size {got}, expected {(sr, CODEBOOK_SIZE)}") held = len(model.quantizer.layers) if kind == "encodec" else len(model.quantizer.quantizers) if held != q_full: sys.exit(f"{src}: {held} codebooks, expected {q_full}") wrap = wrapper(kind, model, q_full) out = Path(args.out) out.parent.mkdir(parents=True, exist_ok=True) dummy = torch.full((1, 1, q_full, 64), CODEBOOK_SIZE, dtype=torch.long) dummy[0, 0, :2] = 0 torch.onnx.export( wrap, (dummy, torch.ones(1)), 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, ) print(f"{args.model}: wrote {out} ({out.stat().st_size:,} B)") sess = ort.InferenceSession(out.as_posix(), providers=["CPUExecutionProvider"]) rng = np.random.default_rng(20261109) worst = 0.0 for q in points: codes = rng.integers(0, CODEBOOK_SIZE, size=(1, q, args.frames), dtype=np.int64) padded = np.full((1, 1, q_full, args.frames), CODEBOOK_SIZE, dtype=np.int64) padded[0, 0, :q] = codes[0] scale = np.array([0.5 if kind == "encodec" and channels == 2 else 1.0], dtype=np.float32) got = sess.run(None, {"audio_codes": padded, "audio_scales": scale})[0] want = reference(kind, model, torch.from_numpy(codes), torch.from_numpy(scale)).numpy() length = args.frames * HOP - shrink if got.shape != (1, channels, length) or got.shape != want.shape: sys.exit(f"{args.model} q={q}: shapes {got.shape} and {want.shape}, expected (1, {channels}, {length})") d = float(np.max(np.abs(got - want))) worst = max(worst, d) print(f" q={q:2d}: max |onnx - model| = {d:.2e}") if worst > 1e-4: sys.exit(f"{args.model}: the export differs from the model's decode by {worst:.2e}") return 0 if __name__ == "__main__": raise SystemExit(main())