Skip to content

Support ONNX export for Demucs#10

Open
dhunstack wants to merge 6 commits intoadefossez:mainfrom
dhunstack:allchanges
Open

Support ONNX export for Demucs#10
dhunstack wants to merge 6 commits intoadefossez:mainfrom
dhunstack:allchanges

Conversation

@dhunstack
Copy link
Copy Markdown

@dhunstack dhunstack commented Oct 26, 2025

In this PR, I've modified the STFT and ISTFT layers to be ONNX exportable.
Details of the changes are documented in this blog - Mixxx Blog Link

The exported model measures closely on the MusDB benchmark.

Stem PyTorch Model (dB) ONNX Model (C++) (dB)
drums.wav 9.46 9.47
bass.wav 7.76 7.77
other.wav 4.69 4.65
vocals.wav 7.84 7.83
Overall 7.44 7.43.

Table: SI-SDR (dB) comparison for each stem and overall, using torchmetrics. Results are shown for the native PyTorch model and the exported ONNX model running in C++.

Add `onnx_exportable` flag to add a parallel computation path with
complex number expressed as a 2D tensor with real and imag components

Signed-off-by: Anmol Mishra <anmolmishra1997@gmail.com>
Add `STFT_Process` class that uses PyTorch's convolution to compute
STFT with real numbers, avoiding the use of complex numbers.

Signed-off-by: Anmol Mishra <anmolmishra1997@gmail.com>
Add `ISTFT_Process` class that computes ISTFT with real numbers,
avoiding any complex tensors

Signed-off-by: Anmol Mishra <anmolmishra1997@gmail.com>
Comment thread demucs/stft.py
KAISER = 'kaiser'

def __call__(self, window_length):
match self:
Copy link
Copy Markdown

@MrFletcheroni MrFletcheroni Dec 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The conda setup only supports python versions <3.10

@jurihock jurihock mentioned this pull request Feb 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants