[Feature request] ONNX Export

[Puiching-Memory]

I noticed that experimental support for ONNX export has already been added. However, when I attempted to write my own export script, I still encountered some shape matching issues. It would be greatly appreciated if you could provide more details about this.

I would like to track the progress of this task on this page.

[Puiching-Memory]

I have conducted preliminary experiments in my fork:
https://github.com/Puiching-Memory/focalcodec

Running this script will export two ONNX files, one for the Encoder and one for the Decoder.

MSE difference comparison results:

Running ONNX model validation
ONNX Encoder Token difference count: 0
ONNX Encoder Codes MSE: 0.000000e+00
ONNX Decoder Reconstruction MSE: 5.708673e-04
Calling onnxslim to optimize models
Running validation on onnxslim optimized models
ONNX Slim Encoder Token difference count: 0
ONNX Slim Encoder Codes MSE: 0.000000e+00
ONNX Slim Decoder Reconstruction MSE: 5.708673e-04

If you'd like to hear the difference, here is a sample audio:
reconstruction_onnx_slim.wav

To achieve this goal, I modified two places in the code:
Puiching-Memory@fc0ca9f#diff-577b3d341e2cf1e538c4649ad39c0bbc9ae445ef603acf83dfb527c8e8825773L281-R297

Puiching-Memory@fc0ca9f#diff-1b3760692af40eec0d387d960b261a5db305b4b67912b1c0eb7606296ccdda15L466-R516

[lucadellalib]

Hi @Puiching-Memory,

each module includes a test_onnx function to verify ONNX exportability. For example:

focalcodec/focalcodec/focalnet.py

Line 1479 in e2e7521

def test_onnx() -> "None":

You can check that function to see how ONNX export is handled. It uses the tracing interface, so I had to ensure that every module is JIT-traceable. The only module where the test currently fails is vocos.py, since the inverse STFT operation isn’t supported in the current opset.

To export the full codec, there's a helper method that takes care of all the export configuration details:

focalcodec/focalcodec/codec.py

Line 1165 in e2e7521

def onnx(

And here's the corresponding test:

focalcodec/focalcodec/codec.py

Line 1707 in e2e7521

def test_onnx(config: "Optional[str]" = None) -> "None":

You can also find an example of how to use the exported ONNX model for streaming inference here:

focalcodec/demo.py

Line 295 in e2e7521

def infer_onnx(

It works with both the OpenVINOExecutionProvider and CPUExecutionProvider backends, but I couldn't get it to run with the CUDAExecutionProvider. Another issue is that, although the offline tests pass, for some checkpoints the resynthesized audio quality during ONNX streaming inference can be degraded. I've debugged this a bit and found that the issue isn't systematic, but rather checkpoint-dependent.

[Puiching-Memory]

Based on the discussion here, there seems to be unofficial implementations of the istft operator, but after I tried replacing it, the audio quality reconstructed by torch inference has significantly degraded.

gemelo-ai/vocos#38
onnx/onnx#4777

https://github.com/biendltb/torch-istft-onnx

[lucadellalib]

The current implementation already uses a custom istft, but the real issue is the call to torch.fft.irfft, which expects a complex-typed tensor as input, while ONNX does not yet support complex dtypes. A possible workaround is to reimplement irfft in a way that represents the complex input using two real channels, rather than relying on the complex dtype.

Is there a specific reason you need to export Vocos to ONNX? The causal checkpoints use WaveNeXt, which is fully ONNX-exportable, and I recommend using those for ONNX export instead.

[lucadellalib]

@Puiching-Memory I updated the code to improve ONNX support.

I was able to export the Vocos-based version by re-implementing irfft using 2 real channels instead of complex input.

A few caveats:

• Custom ISTFT implementation is not as optimized as PyTorch native one, so you may observe a slight slowdown in inference speed.
• To be able to ONNX export, you need to use a fixed chunk size (for example 10 seconds, zero-pad or split if your input is shorter/longer). This limitation comes from nn.functional.fold and there's no workaround for that.

Here is the script to reproduce (onnx==1.17.0 onnxruntime==1.21.1 torch==2.8.0 torchaudio==2.8.0). Hope this helps!

#!/usr/bin/env python3

# onnx==1.17.0 onnxruntime==1.21.1 torch==2.8.0 torchaudio==2.8.0

import io
import math

import numpy as np
import onnxruntime as ort
import torch
import torchaudio
from torch import Tensor, nn


class ONNXExportableISTFT(nn.Module):
    """Custom implementation of inverse STFT with support for non-causal and causal padding.

    Parameters
    ----------
    n_fft:
        Size of Fourier transform.
    hop_length:
        Distance between neighboring sliding window frames.
    win_length:
        Size of window frame and STFT filter.
    causal:
        Whether the module should be causal.

    """

    def __init__(
        self,
        n_fft: "int" = 1024,
        hop_length: "int" = 320,
        win_length: "int" = 1024,
        causal: "bool" = False,
    ) -> "None":
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.causal = causal
        self.non_causal_pad = (win_length - hop_length) // 2

        # Buffers (JIT compilable)
        window = torch.hann_window(win_length)
        self.register_buffer("window", window, persistent=False)
        self.register_buffer("window_sq", window**2, persistent=False)

    def forward(self, input: "Tensor") -> "Tensor":
        """Forward pass.

        Parameters
        ----------
        input:
            Input tensor of shape (batch_size, seq_length, n_fft // 2 + 1, 2).

        Returns
        -------
            Output tensor of shape (batch_size, seq_length * hop_length).

        """
        # Inverse FFT
        T = input.shape[1]
        # input = torch.view_as_complex(input).permute(0, 2, 1)
        # ifft = torch.fft.irfft(input, self.n_fft, dim=1, norm="backward")
        ifft = self._irfft(input, self.n_fft, norm="backward").permute(0, 2, 1)

        if self.causal:
            output = ifft[:, -self.hop_length :].permute(0, 2, 1)
            output = output.flatten(start_dim=1)
            return output

        # Overlap-add
        ifft = ifft * self.window[None, :, None]
        output_size = int((T - 1) * self.hop_length + self.win_length)
        output = nn.functional.fold(
            ifft,
            output_size=(1, output_size),
            kernel_size=(1, self.win_length),
            stride=(1, self.hop_length),
        )[:, 0, 0, self.non_causal_pad : -self.non_causal_pad]

        # Window envelope
        window_sq = self.window_sq.expand(1, T, -1).permute(0, 2, 1)
        window_envelope = nn.functional.fold(
            window_sq,
            output_size=(1, output_size),
            kernel_size=(1, self.win_length),
            stride=(1, self.hop_length),
        )[0, 0, 0, self.non_causal_pad : -self.non_causal_pad]

        # Normalize
        output /= window_envelope

        return output

    def _irfft(
        self,
        input: "Tensor",
        n: "int",
        norm: "str" = "backward",
    ) -> "Tensor":
        B, T, F, _ = input.shape
        device = input.device
        dtype = input.dtype

        # Split real/imag parts: [..., F]
        re = input[..., 0]
        im = input[..., 1]

        # Indices
        n_idx = torch.arange(n, device=device, dtype=dtype)  # [N]
        k_idx = torch.arange(
            1, F - 1, device=device, dtype=dtype
        )  # [F-2] for inner sum

        # Angles: shape [F-2, N]
        angles = 2.0 * math.pi * k_idx[:, None] * (n_idx[None, :] / n)

        cos_table = torch.cos(angles)  # [F-2, N]
        sin_table = torch.sin(angles)  # [F-2, N]

        # Middle frequencies k = 1..F-2
        re_mid = re[..., 1:-1]  # [..., F-2]
        im_mid = im[..., 1:-1]  # [..., F-2]

        # Broadcasted inner sums over k:
        # 2 * sum_k R[k] cos(2πkn/N)
        # 2 * sum_k I[k] sin(2πkn/N)
        # Using einsum: "...f, f n -> ... n"
        two = torch.tensor(2.0, device=device, dtype=dtype)
        sum_re = two * torch.einsum("...f,fn->...n", re_mid, cos_table)  # [..., N]
        sum_im = two * torch.einsum("...f,fn->...n", im_mid, sin_table)  # [..., N]

        # DC and Nyquist bins
        dc = re[..., 0][..., None]  # [..., 1]
        nyquist = re[..., -1][..., None]  # [..., 1]  (k = N/2)

        # (-1)^n term: cos(pi * n) = (-1)^n
        minus_one_power_n = torch.cos(math.pi * n_idx)[None, ...]  # [1, N] -> broadcast

        # Base signal before normalization:
        # x[n] = R[0] + R[N/2] * (-1)^n + sum_re - sum_im
        x = dc + nyquist * minus_one_power_n  # [..., N]
        x = x + sum_re - sum_im

        # Normalization (match torch.fft.irfft semantics)
        if norm == "backward":
            scale = 1.0 / n
        elif norm == "forward":
            scale = 1.0
        elif norm == "ortho":
            scale = 1.0 / math.sqrt(n)
        else:
            raise ValueError(f"Unsupported norm: {norm}")

        x = x * scale
        return x

    def __repr__(self) -> "str":
        return (
            f"{self.__class__.__name__}("
            f"n_fft={self.n_fft}, "
            f"hop_length={self.hop_length}, "
            f"win_length={self.win_length})"
        )


class ONNXExportableISTFTHead(nn.Module):
    """Inverse STFT head.

    Parameters
    ----------
    dim:
        Number of input channels.
    n_fft:
        Size of Fourier transform.
    hop_length:
        Distance between neighboring sliding window frames.
    causal:
        Whether the module should be causal.

    """

    def __init__(
        self,
        dim: "int" = 512,
        n_fft: "int" = 1024,
        hop_length: "int" = 320,
        causal: "bool" = False,
    ) -> "None":
        super().__init__()
        self.dim = dim
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.causal = causal

        # Modules
        self.proj = nn.Linear(dim, n_fft + 2)
        self.istft = ONNXExportableISTFT(n_fft, hop_length, n_fft, causal)

    def forward(self, input: "Tensor") -> "Tensor":
        """Forward pass.

        Parameters
        ----------
        input:
            Input tensor of shape (batch_size, seq_length, dim).

        Returns
        -------
            Output tensor of shape (batch_size, seq_length * hop_length).

        """
        output = self.proj(input)
        mag, phase = output.chunk(2, dim=-1)
        mag = mag.exp()
        # Safeguard to prevent excessively large magnitudes
        mag = mag.clamp(max=1e2)
        # Real and imaginary value
        # JIT compilable
        stft = mag[..., None] * torch.stack([phase.cos(), phase.sin()], dim=-1)
        output = self.istft(stft)
        return output


if __name__ == "__main__":
    # Load FocalCodec model
    codec = torch.hub.load(
        repo_or_dir="lucadellalib/focalcodec",
        model="focalcodec",
        config="lucadellalib/focalcodec_50hz",
        force_reload=True,  # Fetch the latest FocalCodec version from Torch Hub
    )

    # Replace original Vocos head with ONNX exportable one
    istft_head_state_dict = codec.decoder.head.state_dict()
    codec.decoder.head = ONNXExportableISTFTHead(
        codec.decoder.head.dim,
        codec.decoder.head.n_fft,
        codec.decoder.head.hop_length,
        codec.decoder.head.causal,
    )
    codec.decoder.head.load_state_dict(istft_head_state_dict)
    codec.eval().requires_grad_(False)

    # Load and preprocess the input audio
    audio_file = "audios/librispeech-dev-clean/251-118436-0003.wav"
    sig, sample_rate = torchaudio.load(audio_file)
    sig = torchaudio.functional.resample(sig, sample_rate, codec.sample_rate_input)

    # IMPORTANT: nn.functional.fold requires a fixed chunk size to be ONNX-exportable
    f = io.BytesIO()
    codec.onnx(f, chunk_size=sig.shape[1])
    onnx_bytes = f.getvalue()

    # Run ONNX inference
    session = ort.InferenceSession(onnx_bytes)
    inputs_ort = {}
    inputs_ort["sig"] = sig.numpy().astype(np.float32)
    outputs_ort = session.run([x.name for x in session.get_outputs()], inputs_ort)

    # Encode audio into tokens
    toks = outputs_ort[0]  # Shape: (batch, time)
    print(toks.shape)
    print(toks)

    # Convert tokens to their corresponding binary spherical codes
    codes = outputs_ort[1]  # Shape: (batch, code_time, log2 codebook_size)
    print(codes.shape)
    print(codes)

    # Decode tokens back into a waveform
    rec_sig = outputs_ort[2]

    # Save the reconstructed audio
    rec_sig = torchaudio.functional.resample(
        torch.as_tensor(rec_sig), codec.sample_rate_output, sample_rate
    )
    torchaudio.save("reconstruction_onnx.wav", rec_sig, sample_rate)

[Puiching-Memory]

Thank you very much for your assistance! Your method can export all models to ONNX format.

---

Regarding inputs with dynamic lengths, I have an idea to use scatter_add_

It has not been verified in detail at present, and inputting audio with a length > 10 seconds (512 frames) on the cacusal model will result in errors.

output = nn.functional.fold(
    ifft,
    output_size=(1, output_size),
    kernel_size=(1, self.win_length),
    stride=(1, self.hop_length),
)[:, 0, 0, self.non_causal_pad : -self.non_causal_pad]

output_size = int((T - 1) * self.hop_length + self.win_length)
signal = ifft.new_zeros((B, output_size))

frame_offsets = torch.arange(T, device=ifft.device, dtype=torch.long) * self.hop_length
sample_offsets = torch.arange(self.win_length, device=ifft.device, dtype=torch.long)

frame_indices = frame_offsets[:, None] + sample_offsets[None, :]
frame_indices = frame_indices.reshape(1, -1).expand(B, -1)

frame_values = ifft.permute(0, 2, 1).reshape(B, -1)
signal.scatter_add_(1, frame_indices, frame_values)

all code

class ONNXExportableISTFT(nn.Module):

    def __init__(
        self,
        n_fft: int = 1024,
        hop_length: int = 320,
        win_length: int = 1024,
        causal: bool = False,
    ) -> None:
        super().__init__()
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.win_length = win_length
        self.causal = causal
        self.non_causal_pad = (win_length - hop_length) // 2

        self.window: Tensor
        self.window_sq: Tensor
        window = torch.hann_window(win_length)
        self.register_buffer("window", window, persistent=False)
        self.register_buffer("window_sq", window**2, persistent=False)

    def forward(self, input: Tensor) -> Tensor:
        # input: (batch, seq, n_fft // 2 + 1, 2)
        batch = input.shape[0]
        time_steps = input.shape[1]
        ifft = self._irfft(input, self.n_fft, norm="backward").permute(0, 2, 1)

        if self.causal:
            output = ifft[:, -self.hop_length :].permute(0, 2, 1)
            return output.flatten(start_dim=1)

        windowed = ifft * self.window[None, :, None]
        frames = windowed.permute(0, 2, 1)  # (batch, frames, win_length)

        output_size = (time_steps - 1) * self.hop_length + self.win_length
        signal = windowed.new_zeros((batch, output_size))

        frame_offsets = torch.arange(time_steps, device=frames.device, dtype=torch.long) * self.hop_length
        sample_offsets = torch.arange(self.win_length, device=frames.device, dtype=torch.long)
        frame_indices = frame_offsets[:, None] + sample_offsets[None, :]
        frame_indices = frame_indices.reshape(1, -1).expand(batch, -1)

        frame_values = frames.reshape(batch, -1)
        signal.scatter_add_(1, frame_indices, frame_values)

        envelope = self.window_sq.view(1, 1, -1).expand(batch, time_steps, self.win_length)
        envelope = envelope.reshape(batch, -1)
        norm = windowed.new_zeros((batch, output_size))
        norm.scatter_add_(1, frame_indices, envelope)

        signal = signal[:, self.non_causal_pad : output_size - self.non_causal_pad]
        norm = norm[:, self.non_causal_pad : output_size - self.non_causal_pad]

        signal = signal / norm.clamp_min(1e-8)
        return signal

    def _irfft(self, input: Tensor, n: int, norm: str = "backward") -> Tensor:
        _, _, freq_bins, _ = input.shape
        device = input.device
        dtype = input.dtype

        real = input[..., 0]
        imag = input[..., 1]

        n_idx = torch.arange(n, device=device, dtype=dtype)
        k_idx = torch.arange(1, freq_bins - 1, device=device, dtype=dtype)

        if k_idx.numel() == 0:
            cos_table = torch.zeros(0, n, device=device, dtype=dtype)
            sin_table = cos_table
        else:
            angles = 2.0 * math.pi * k_idx[:, None] * (n_idx[None, :] / n)
            cos_table = torch.cos(angles)
            sin_table = torch.sin(angles)

        two = torch.tensor(2.0, device=device, dtype=dtype)
        real_mid = real[..., 1:-1]
        imag_mid = imag[..., 1:-1]
        sum_real = two * torch.einsum("...f,fn->...n", real_mid, cos_table)
        sum_imag = two * torch.einsum("...f,fn->...n", imag_mid, sin_table)

        dc = real[..., 0][..., None]
        nyquist = real[..., -1][..., None]
        minus_one_power_n = torch.cos(math.pi * n_idx)[None, ...]
        signal_time = dc + nyquist * minus_one_power_n + sum_real - sum_imag

        if norm == "backward":
            scale = 1.0 / n
        elif norm == "ortho":
            scale = 1.0 / math.sqrt(n)
        else:
            scale = 1.0

        signal_time = signal_time * scale
        return signal_time


class ONNXExportableISTFTHead(nn.Module):

    def __init__(self, dim: int, n_fft: int, hop_length: int, causal: bool) -> None:
        super().__init__()
        self.dim = dim
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.causal = causal
        self.proj = nn.Linear(dim, n_fft + 2)
        self.istft = ONNXExportableISTFT(n_fft, hop_length, n_fft, causal)

    def forward(self, input: Tensor) -> Tensor:
        proj = self.proj(input)
        mag, phase = proj.chunk(2, dim=-1)
        mag = mag.exp().clamp(max=1e2)
        stft = mag[..., None] * torch.stack([phase.cos(), phase.sin()], dim=-1)
        return self.istft(stft)