Skip to content

Commit 20f3274

Browse files
dhunstackdhunstack
committed
[blog] Add Demucs ONNX export blog
Blog documenting the "Converting Demucs v4 (Hybrid Transformer) AI model to ONNX format" GSOC project. Signed-off-by: Anmol Mishra <anmolmishra1997@gmail.com>
1 parent b962977 commit 20f3274

4 files changed

Lines changed: 288 additions & 0 deletions

File tree

‎content/images/avatars/dhunstack.png‎

172 KB
Loading

‎content/images/news/demucs-to-onnx/htdemucs.png‎

450 KB
Loading

‎content/news/XXXX-XX-XX-gsoc2025-demucs-to-onnx-dhunstack.md‎

Lines changed: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
1+
title: GSOC 2025 - Converting Demucs v4 (Hybrid Transformer) AI model to ONNX format
2+
authors: Anmol Mishra
3+
status: draft
4+
tags: gsoc, gsoc-2025, stems
5+
math: yes
6+
comments: yes
7+
8+
Disclaimer: *This blog post primarily serves as the documentation for the [Google Summer of Code](https://summerofcode.withgoogle.com/programs/2025/projects/lRQpeA7K) 2025 project: "Converting Demucs v4 (Hybrid Transformer) AI model to ONNX format".*
9+
10+
Mixxx 2.6 will support playback of [stem files]( https://www.stems-music.com/), which can be created from the original track stems recorded in a studio (or DAW), or by using AI based stem separation models on the final mixed track. [Demucs v4](https://github.com/adefossez/demucs) is a state-of-the-art open-source music source separation model developed by Antoine Défossez. While it provides exceptional quality in separating audio into stems, it is currently implemented in Python and therefore cannot be used directly in C++ applications or run efficiently on hardware accelerators through [ONNX Runtime](https://onnxruntime.ai/).
11+
12+
This project is a step towards supporting real time stem separation within Mixxx, by exporting the Python based Demucs model to [ONNX](https://onnx.ai/) (Open Neural Network Exchange format). The key contributions of this project are -
13+
14+
- Preparing the Demucs code for ONNX export by rewriting non-exportable operations
15+
- Validating all the modified parts with numerical tests to ensure export of existing model weights without the need for retraining.
16+
- Scripts for exporting ONNX and ORT formats of the model.
17+
- Example scripts for deployment of ONNX model with C++ using ONNX Runtime.
18+
- Benchmarking the exported model for performance and separation quality.
19+
- Upstream PR with our modifications to the Demucs repository
20+
- [Conference talk](https://conference.audio.dev/session/2025/converting-source-separation-models-to-onnx-for-real-time-usage-in-dj-software/) to be delivered at the ADC 2025
21+
22+
---
23+
24+
## Introduction
25+
26+
Source separation models are typically developed using Python-based libraries such as PyTorch or TensorFlow, which require a Python runtime to execute machine learning models. This dependency can make it challenging to use these models in environments where Python is not available or practical, such as audio plugins, which are often written in C++. To address this, ONNX (Open Neural Network Exchange) has become a widely adopted serialization format, enabling language-independent deployment of machine learning models.
27+
28+
While both PyTorch and TensorFlow provide straightforward ONNX export for simple models, complex audio models like Demucs often present significant hurdles. This is primarily due to their use of complex tensors and custom implementations of operations like STFT and ISTFT, which are not natively supported by ONNX. As a result, previous attempts to export Demucs to ONNX have typically bypassed these layers, requiring developers to reimplement STFT/ISTFT in every target language, a process that is both error prone and time consuming.
29+
30+
In this project, we successfully exported Demucs v4 as a fully self-contained ONNX model, including the STFT and ISTFT layers. This means the model can be used directly in any language that supports ONNX, greatly simplifying integration for audio developers and enabling broader adoption of source separation technology.
31+
32+
A key part of our approach was to diagnose and address each export issue step by step, developing custom numerical tests to ensure that our ONNX-compatible replacements matched the original PyTorch calculations as closely as possible. This was crucial, as retraining the model was not a viable option: Demucs was originally trained on a large private dataset over weeks using multiple GPUs, and reproducing those results without the dataset and computational resources would be extremely difficult. By focusing on numerical fidelity, we ensured that the exported model retained the high performance of the original, without the need for retraining.
33+
34+
---
35+
36+
## Understanding the Math Behind Audio
37+
38+
Before diving into transforms and neural networks, it helps to review how audio is represented mathematically and why **complex numbers** and **Fourier transforms** are so central to signal processing.
39+
40+
#### Real Numbers (Time-Domain Audio)
41+
42+
Digital audio is simply a sequence of **real numbers** — samples taken at regular time intervals.
43+
44+
Example: a 44.1 kHz stereo track stores **44,100 real values per second per channel**.
45+
Each value (e.g. `0.25`, `-0.67`) represents instantaneous amplitude (air pressure or voltage).
46+
47+
#### Complex Numbers (Frequency Representation)
48+
49+
To analyze frequencies, we extend real numbers to **complex numbers**.
50+
51+
$$
52+
z = a + bi
53+
$$
54+
55+
Where:
56+
57+
- $a$: real part, $b$: imaginary part, $i^2 = -1$
58+
59+
Complex numbers are often represented in **polar form**:
60+
61+
$$
62+
z = r e^{i\theta}
63+
$$
64+
65+
Where:
66+
67+
- $r = \sqrt{a^2 + b^2}$ (magnitude)
68+
- $\theta = \operatorname{atan2}(b, a)$ (phase angle)
69+
70+
**Conversions:**
71+
72+
- $r = \sqrt{a^2 + b^2}$, $\theta = \operatorname{atan2}(b, a)$, $a = r \cos\theta$, $b = r \sin\theta$
73+
74+
**In signal processing:**
75+
76+
- *Magnitude* → strength of each frequency component.
77+
- *Phase* → timing/offset of that frequency’s oscillation.
78+
79+
Complex numbers allow us to represent **both amplitude and phase**, making them perfect for describing sound in the **frequency domain**.
80+
81+
---
82+
83+
## From Time to Frequency: FFT, STFT, and ISTFT
84+
85+
#### Fourier Transform (FT)
86+
87+
The **Fourier Transform** expresses a signal as a sum of sine and cosine waves of different frequencies.
88+
89+
For digital (sampled) signals, we use the **Discrete Fourier Transform (DFT)**:
90+
91+
$$
92+
X[k] = \sum_{n=0}^{N-1} x[n] \, e^{-2\pi i \frac{kn}{N}}
93+
$$
94+
95+
Its inverse reconstructs the signal:
96+
97+
$$
98+
x[n] = \frac{1}{N} \sum_{k=0}^{N-1} X[k] \, e^{2\pi i \frac{kn}{N}}
99+
$$
100+
101+
Where:
102+
103+
- $x[n]$ is the time-domain signal (length $N$)
104+
- $X[k]$ is the frequency-domain representation (DFT coefficients)
105+
- $n$ is the time index, $k$ is the frequency bin
106+
107+
For example, with a 4-sample input:
108+
109+
- **Input:** $[1, 2, 3, 4]$
110+
- **DFT:** $[10, -2+2i, -2, -2-2i]$ (frequency components)
111+
- **IDFT:** Reconstructs the original signal from its DFT.
112+
113+
#### Short-Time Fourier Transform (STFT)
114+
115+
A single DFT gives the *overall* frequency content but ignores *when* things happen. To capture **time-varying** frequency information, we use the **Short-Time Fourier Transform**.
116+
117+
Steps:
118+
119+
1. Split the signal into short, overlapping frames (e.g. 2048 samples).
120+
2. Apply a **window** (Hann/Hamming) to each frame.
121+
3. Compute the FFT of each windowed frame.
122+
123+
Mathematically:
124+
125+
$X(m,\omega) = \sum_n x[n] \, w[n - m] \, e^{-j \omega n}$
126+
127+
- $m$: frame index (time)
128+
- $\omega$: frequency bin
129+
- Output: **complex spectrogram** showing how frequency content evolves over time.
130+
131+
A spectrogram is a 2-D matrix (time × frequency) of complex values - often visualized by plotting magnitude in dB.
132+
133+
#### Inverse STFT (ISTFT)
134+
135+
To go back to the time domain:
136+
137+
1. Compute **inverse FFT** for each frame.
138+
2. Multiply by a synthesis window if required.
139+
3. **Overlap-add** frames to reconstruct the waveform.
140+
141+
Errors here cause audible artifacts (clicks, smearing, or phase shifts).
142+
143+
---
144+
145+
## Understanding HTDemucs and the Challenges with ONNX Export
146+
147+
To understand the export process, it’s important to look closely at how **HTDemucs (Hybrid Transformer Demucs)** is structured, and why its architecture introduces challenges when converting to the **ONNX** format.
148+
149+
![htdemucs]({static}/images/news/demucs-to-onnx/htdemucs.png)
150+
151+
#### HTDemucs Explained
152+
153+
HTDemucs is a *hybrid* model — it combines two parallel processing branches:
154+
155+
##### **Time-Domain Branch**
156+
157+
- Operates directly on the raw waveform (real-valued input).
158+
- Uses temporal convolutions and Transformer blocks to learn how each instrument evolves over time.
159+
160+
##### **Time-Frequency Domain Branch**
161+
162+
- Operates on the signal’s **spectrogram**, obtained by applying a **Short-Time Fourier Transform (STFT)**.
163+
- The STFT produces a **complex-valued** tensor, representing both magnitude and phase across time and frequency.
164+
- The model immediately converts this to a **magnitude spectrogram**, a purely **real-valued** representation, which is then processed by convolutional and Transformer layers.
165+
- At the end of this branch, an **Inverse STFT (ISTFT)** is applied to reconstruct the real-valued waveform.
166+
167+
These two branches are then **fused** to produce the final stem-separated outputs (vocals, drums, bass, and others).
168+
169+
#### The ONNX Export Problem
170+
171+
The key difficulty in exporting HTDemucs lies in ONNX’s lack of support for complex tensors.
172+
173+
- In the **time-domain branch**, all operations are real-valued and export cleanly to ONNX.
174+
- In the **time-frequency branch**, two specific operations, **STFT** and **ISTFT**, involve complex numbers.
175+
176+
The rest of the model (all convolutional, transformer, and linear layers) operates purely on **real numbers**, and therefore poses **no export issue**.
177+
178+
ONNX, as of current opset versions, does not fully support:
179+
180+
- Complex-valued tensors as native types - [PyTorch Issue](https://github.com/pytorch/pytorch/issues/126972)
181+
- Complex FFT operations (`torch.stft`, `torch.istft`) from PyTorch - [PyTorch Issue](https://github.com/pytorch/pytorch/issues/65666)
182+
183+
As a result, a direct `torch.onnx.export()` call on HTDemucs fails, since the exporter encounters unsupported complex operations.
184+
185+
---
186+
187+
## Our Solution
188+
189+
#### Real-Valued STFT and ISTFT Rewrites
190+
191+
Because complex values appear **only** in the initial STFT and final ISTFT layers, our solution was to **reimplement these operations** in a real-valued form.
192+
193+
##### 1. **STFT Rewrite**
194+
195+
Instead of using `torch.stft()`, we expressed the Fourier transform as a set of **1-D convolutions** with precomputed sine and cosine kernels.
196+
197+
This allows us to compute:
198+
$\text{Re}(X) = x * \cos(\omega)$, $\text{Im}(X) = -x * \sin(\omega)$
199+
200+
Thus, we can store the **real and imaginary parts** as separate real-valued tensors, fully ONNX-compatible.
201+
202+
##### 2. **ISTFT Rewrite**
203+
204+
Similarly, for inverse reconstruction:
205+
206+
- The original ISTFT combines complex values through real + imaginary synthesis.
207+
- We reconstructed the time-domain signal by performing the same series of **overlap-add** and **cosine/sine inverse convolutions**, again using only real-valued tensors.
208+
209+
By carefully ensuring numerical equivalence to PyTorch’s implementation, we achieved perfect parity (MSE < 1e-4) between the original and rewritten layers.
210+
211+
---
212+
213+
#### Final Export
214+
215+
The **only** problematic parts were the **STFT** and **ISTFT** layers, and once these were rewritten with real-valued math, the model exported **fully and cleanly** to ONNX — no retraining required.
216+
217+
This approach made it possible to:
218+
219+
- Preserve the original model weights.
220+
- Ensure perfect numerical equivalence.
221+
- Achieve seamless ONNX export compatible with **ONNX Runtime (ORT)** and **C++ deployment**.
222+
223+
#### Unified Benchmarking: Timing and SI-SDR
224+
225+
To evaluate the exported ONNX model, we developed a unified benchmarking script (`benchmark.py`) that supports both PyTorch and ONNX/ORT backends. Key features:
226+
227+
- **Separation Backend:**
228+
- PyTorch: Runs separation using the original Demucs model.
229+
- ONNX: Runs separation by invoking the ONNX CLI tool via subprocess, matching the workflow of the provided example scripts.
230+
- **Timing:**
231+
- The script records detailed timing for each track, printing and saving the results.
232+
- **SI-SDR Evaluation:**
233+
- Computes the SI-SDR (Scale-Invariant Signal-to-Distortion Ratio) for each separated track using the `torchmetrics` library.
234+
- **Output:**
235+
- Results are saved in timestamped folders, with timing and SI-SDR metrics output to JSON files named according to the backend (e.g., `results_onnx.json`).
236+
237+
#### How to Run the Benchmark
238+
239+
1. Place your test tracks in the input directory.
240+
2. Run the benchmark script:
241+
242+
```python
243+
python benchmark.py --backend onnx --input_dir ./testfiles --output_dir ./outputs
244+
```
245+
246+
3. Review the JSON results for timing and SI-SDR metrics.
247+
248+
## Benchmark Results
249+
250+
The quality of the model is expected to be equal or slightly worse when exported to ONNX. While there are plenty of ways of measuring the benchmarks models (another blog post incoming), we've chosen to measure our models with `SI-SDR` metric, Scale Invariant Signal To Distortion Ratio, on the MusDB dataset. This is the standard metric on which researchers report their source separation model's performance.
251+
After export, the model's performance is nearly identical.
252+
253+
| Stem | PyTorch Model (dB) | ONNX Model (C++) (dB) |
254+
|--------------|--------------------|-----------------------|
255+
| drums.wav | 9.53 | 9.53 |
256+
| bass.wav | 13.24 | 13.20 |
257+
| other.wav | 12.61 | 12.64 |
258+
| vocals.wav | 7.42 | 7.41 |
259+
| **Overall** | **10.70** | **10.69** |
260+
261+
*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++.*
262+
263+
Now that we have a running platform independent high quality ONNX Demucs model that can utilize hardware acceleration and be deployed with C++, we plan to integrate this into Mixxx DJ for future.
264+
We've prepared example scripts for running the exported Demucs model, which can be used following the instructions documented in our READMEs.
265+
266+
## Plans for future integration into Mixxx
267+
268+
Here's a list of all the PRs that document our incremental changes towards exporting Demucs:
269+
270+
| Task | PR Status | PR Link |
271+
|-----------------------|---------------------|----------------------------------------------------------|
272+
| ONNX Computation Path | Merged | [PR #1](https://github.com/mixxxdj/demucs/pull/1) |
273+
| STFT Rewrite | Merged | [PR #2](https://github.com/mixxxdj/demucs/pull/3) |
274+
| Inverse STFT Rewrite | Merged | [PR #4](https://github.com/mixxxdj/demucs/pull/4) |
275+
| ONNX Export Scripts | Merged | [PR #5](https://github.com/mixxxdj/demucs/pull/5) |
276+
| CI for model export | Merged | [PR #6](https://github.com/mixxxdj/demucs/pull/6) |
277+
| Example C++ Scripts | In Review | [PR #7](https://github.com/mixxxdj/demucs/pull/7) |
278+
| Benchmarking Scripts | In Review | [PR #9](https://github.com/mixxxdj/demucs/pull/8) |
279+
280+
We've raised a PR with all our Demucs changes to the upstream here - [ADD LINK].
281+
A talk will be presented on this project at the [Audio Developers Conference 2025](https://conference.audio.dev/session/2025/converting-source-separation-models-to-onnx-for-real-time-usage-in-dj-software/) in Bristol.
282+
283+
Finally, we've created this [EPIC](https://github.com/mixxxdj/mixxx/issues/15495) to track all issues related to merging Demucs to Mixxx. Feel free to join the discussion and get involved with supporting real time stems separation inside Mixxx.

‎pelicanconf.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,11 @@ def __init__(self, url, title, context, css="", children=()):
302302
"github": "jclsn",
303303
"tagline": "Mixxx Contributor",
304304
},
305+
"Anmol Mishra": {
306+
"github": "dhunstack",
307+
"discourse": "dhunstack",
308+
"tagline": "GSoC 2025 Contributor",
309+
},
305310
}
306311

307312
# Needed for Jinja2 markdown filter

0 commit comments

Comments
 (0)