TL;DR — This repo documents the conversion of the MIT HTDemucs v4 checkpoint into two ONNX models, along with a streaming audio separation pipeline validated on CPU (ONNX Runtime) and GPU (DirectML / RTX 5090). The work focuses exclusively on conversion, split export, numerical validation and the runtime pipeline. Not affiliated with facebookresearch/demucs.
HTDemucs v4 (Meta Research) is a state-of-the-art audio source separation model released under the MIT License. It separates a track into 4 stems: drums, bass, other, vocals.
This project solves two concrete problems for integrating HTDemucs into a .NET application (FrameShift):
-
ONNX export from the PyTorch checkpoint — the model uses operations (
torch.stft,nn.MultiheadAttention) that block naive ONNX export. The solutions are documented and fixed here. -
DirectML viability — the full ONNX graph causes total GPU saturation (~26 GB VRAM) due to 5 Conv1d layers with kernel=4096 (baked-in STFT/iSTFT). This repo introduces
htdemucs_split.onnx, a graph without STFT/iSTFT, which makes DirectML viable (+26 MB VRAM, 47 ms per inference, ~20× real-time on 30 min audio).
| Model | Size | Recommended EP | Description |
|---|---|---|---|
htdemucs.onnx |
289 MB | CPUExecutionProvider | Full export — STFT/iSTFT in-graph via real Conv1d. CPU validated. |
htdemucs_split.onnx |
161 MB | DirectMLExecutionProvider | Split export — STFT/iSTFT externalised. DirectML + CPU validated. |
Note: These files are not included in this repo (too large). See REPRODUCIBILITY.md to regenerate them from the checkpoint.
The PyTorch checkpoint 955717e8-8726e21a.th is the official MIT HTDemucs v4 model published by facebookresearch.
It is distributed separately under the MIT License — see NOTICE.md and MODEL_PROVENANCE.md.
⚠️ Do not use the CC-BY-NC ONNX model circulating on HuggingFace Hub. Its license is incompatible with commercial or product use. This repo uses exclusively the official MIT checkpoint.
Input : mix float32 (B, 2, 343980)
└── internal STFT (Conv1d kernel=4096)
└── encoders → cross-transformer → decoders
└── internal iSTFT (ConvTranspose1d kernel=4096)
Output : stems float32 (B, 4, 2, 343980) ← drums/bass/other/vocals
Simple to use: a single ONNX call, audio in and stems out. Incompatible with DirectML (5 giant Conv1d nodes).
Inputs:
mix float32 (B, 2, 343980) ← raw audio (time branch)
spec float32 (B, 2, 2048, 336, 2) ← spectrogram pre-computed by the host
[STFT handled by the host]
└── encoders → cross-transformer → decoders
[iSTFT handled by the host]
Outputs:
mask_spec float32 (B, 4, 2, 2048, 336, 2) ← spectral mask (pre-iSTFT)
stems_time float32 (B, 4, 2, 343980) ← time branch (ready to sum)
The host computes STFT(mix) → spec, runs inference, then iSTFT(mask_spec) + stems_time → stems.
Removes the 5 giant Conv1d nodes → DirectML works normally.
audio (WAV 44.1 kHz stereo)
│
├─[host CPU] STFT(mix_chunk) → spec (1, 2, 2048, 336, 2)
│
├─[DML GPU ] htdemucs_split.onnx({mix, spec})
│ → mask_spec (1, 4, 2, 2048, 336, 2)
│ → stems_time (1, 4, 2, 343980)
│
├─[host CPU] iSTFT(pad(mask_spec)) + stems_time → stems (4, 2, 343980)
│
└─[ring buffer OLA] → flush prefix → WAV streaming write
Overlap-Add (25%, triangular window) identical to demucs.apply.apply_model.
Fixed-size ring buffer: constant RAM regardless of audio duration.
| Pipeline | Model | Wall 30 min | RT factor | RAM peak | VRAM peak |
|---|---|---|---|---|---|
| V1 CPU | htdemucs.onnx |
343 s | 5.25× | 1 677 MB | n/a |
| V2 GPU (DML) | htdemucs_split.onnx |
90 s | 20.05× | 1 323 MB | +1 845 MB |
| Stem | SNR | OLA artefacts |
|---|---|---|
| drums | 82.1 dB | none (junction ratio 0.55) |
| bass | 84.0 dB | none (junction ratio 0.73) |
| other | 76.8 dB | none (junction ratio 0.69) |
| vocals | 83.4 dB | none (junction ratio 0.71) |
| instrumental | 88.4 dB | none (junction ratio 0.56) |
SNR ≥ 76.8 dB — inaudible on listening. V2 OLA ratios identical to V1.
79,579,428 frames = 1804.52 s exactly across all 5 stems × V1 and V2. No truncation.
.
├── README.md ← this file
├── LICENSE ← MIT for scripts
├── NOTICE.md ← Demucs / ORT attributions
├── MODEL_PROVENANCE.md ← provenance chain + SHA256
├── REPRODUCIBILITY.md ← steps to regenerate the ONNX files
│
├── stft_onnx.py ← host STFT/iSTFT (real Conv1d)
├── onnx_htdemucs.py ← V1 export wrapper
├── onnx_htdemucs_split.py ← V2 split export wrapper
├── export_demucs.py ← generates htdemucs.onnx
├── export_demucs_split.py ← generates htdemucs_split.onnx
│
├── separate_streaming.py ← V1 CPU production pipeline (CLI)
├── separate_streaming_gpu.py ← V2 GPU production pipeline (CLI)
├── separate_chunked.py ← in-memory pipeline (reference)
│
├── validate_chunking.py ← validation vs Demucs apply_model
├── validate_split_cpu.py ← V1 vs V2 split, 1 segment
├── compare_split_cpu_vs_dml.py ← CPU EP vs DML EP
├── compare_v1_vs_v2_stems.py ← full comparison + OLA scan
├── inspect_split_graph.py ← V1 vs V2 node inventory
├── test_split_directml.py ← guarded DML test (VRAM + time watchdog)
├── _split_dml_child.py ← DML child process
│
├── requirements.txt
│
└── docs/
├── REPORT_v1_streaming.md ← V1 streaming 30 min synth
├── REPORT_v2_split.md ← split graph: DML 1 segment
├── REPORT_v2_gpu_30min.md ← V2 GPU 30 min real audio
└── DEMUCS_FRAME_SHIFT_INTEGRATION_GUIDE.md ← .NET integration guide
Python 3.11
torch==2.1.2+cu121
onnx==1.21.0
onnxruntime==1.24.4 (or onnxruntime-directml for DML)
demucs==4.0.1
numpy==1.26.4
soundfile==0.13.1
einops>=0.6.0
psutil
See REPRODUCIBILITY.md for the full installation steps.
python separate_streaming.py \
--input /path/to/track.wav \
--out-dir ./stems_output \
--stems vocals,instrumentalpython separate_streaming_gpu.py \
--input /path/to/track.wav \
--out-dir ./stems_output \
--stems allThe scripts in this repo are released under the MIT License (see LICENSE).
The HTDemucs model weights (955717e8-8726e21a.th and the ONNX files derived from it)
are covered by the original MIT License of facebookresearch/demucs.
See NOTICE.md and MODEL_PROVENANCE.md for details.