open-standards/wai/tools/wai_neural_audio_export_onnx.py

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