StemSplitio/htdemucs-6s-onnx

HT-Demucs 6-stem — ONNX (with guitar + piano)

The first ONNX export of the 6-stem `htdemucs_6s` variant on the Hugging Face Hub. Adds guitar and piano stems on top of the standard 4 (drums / bass / other / vocals). Runs in onnxruntime on CPU out of the box, and on CoreML / CUDA / DirectML with a one-line provider change. No PyTorch required at inference.

If you need guitar or piano isolation, this is the only off-the-shelf ONNX model on the Hub that gives you that.

TL;DR

pip install onnxruntime numpy soundfile

# 258 MB fp32 model — all 6 stems:
python infer.py your-song.mp3 ./out/

# 136 MB fp16weights variant (same runtime cost):
python infer.py your-song.mp3 ./out/ --small

# Just the guitar stem:
python infer.py your-song.mp3 ./out/ --stems guitar

The repo contains:

• —htdemucs_6s.onnx — 258 MB, opset 17, parity-verified vs PyTorch fp32.
• —htdemucs_6s_fp16weights.onnx — 136 MB, fp16-stored weights, same runtime memory / latency.
• —infer.py — pure-numpy reference inference (~200 lines, no torch).
• —requirements.txt — three small packages, no PyTorch.

What stems do I get?

SOURCES = ("drums", "bass", "other", "vocals", "guitar", "piano")

Output tensor: stems[1, 6, 2, 343980] in that exact stem order. The 6-stem variant overlaps with the 4-stem on the first 4 stems but with slightly different separation behavior — the extra guitar and piano heads change what "other" learns to keep.

Quality

Parity vs PyTorch fp32 (random input, 7.8 s segment):

• —htdemucs_6s.onnx max abs diff: 2.42 × 10⁻⁴
• —htdemucs_6s_fp16weights.onnx max abs diff (vs fp32 weights): 1.06 × 10⁻⁴

Both well within the 1e-3 publish threshold.

Stem-specific SDR (informal; the official paper covers in-depth eval):

If you care about absolute drums/vocals SDR, prefer `htdemucs-ft-onnx`. If you specifically need guitar or piano isolation, this is the model.

Performance

Single 7.8 s segment, Apple M4 Pro CPU:

CUDA / DirectML / CoreML EPs are typically ≥ 5× faster on real GPUs.

Quick start

Python

import soundfile as sf
import infer

audio, sr = sf.read("your-song.mp3", dtype="float32", always_2d=True)
stems = infer.separate(audio.T, sr,
                       model_path=infer.DEFAULT_MODEL,
                       providers=["CPUExecutionProvider"])
sf.write("guitar.wav", stems["guitar"].T, sr)
sf.write("piano.wav",  stems["piano"].T,  sr)

CLI

python infer.py your-song.mp3 ./out/                          # all 6 stems
python infer.py your-song.mp3 ./out/ --stems guitar piano     # guitar + piano only
python infer.py your-song.mp3 ./out/ --providers coreml       # macOS arm64
python infer.py your-song.mp3 ./out/ --providers cuda         # Linux + NVIDIA
python infer.py your-song.mp3 ./out/ --small                  # 136 MB variant

Mobile / Web

// iOS / Swift — 258 MB or 136 MB bundled
import onnxruntime_objc
let session = try ORTSession(env: env,
    modelPath: Bundle.main.path(forResource: "htdemucs_6s_fp16weights",
                                 ofType: "onnx")!,
    sessionOptions: opts)

// Browser
import * as ort from "onnxruntime-web";
const sess = await ort.InferenceSession.create(
  "htdemucs_6s_fp16weights.onnx",
  { executionProviders: ["wasm"] },
);
const t = new ort.Tensor("float32", audioBuffer, [1, 2, 343980]);
const out = await sess.run({ mix: t });   // out.stems is (1, 6, 2, 343980)

For a turnkey browser demo with file-picker + chunked overlap-add, see `demucs-onnx browser-demo`.

Input / output spec

For longer audio, chunk with overlap-add — see infer.py::separate.

Tooling — demucs-onnx Python package

This model can be run via the open-source `demucs-onnx` Python package on PyPI. It auto-downloads from this repo on first use.

pip install demucs-onnx

# 6-stem mode — all 6 stems, single session:
demucs-onnx separate song.mp3 stems/ --model htdemucs_6s

# Just guitar + piano:
demucs-onnx separate song.mp3 stems/ --model htdemucs_6s --stems guitar piano

# Python API:
python -c "from demucs_onnx import separate_stem; \
  guitar = separate_stem('song.mp3', 'guitar')"

To re-export your own fine-tune:

pip install 'demucs-onnx[export]'
demucs-onnx export htdemucs_6s out/htdemucs_6s.onnx

How it was built

The export pipeline lives in the open-source `demucs-onnx` package at `demucs_onnx/export/`. It applies the same four patches that make htdemucs_ft exportable:

1. 1.Complex-typed torch.stft outputs → Conv1d with sin/cos kernels.
2. 2.model.segment fractions.Fraction → plain float.
3. 3.random.randrange in transformer pos-embedding → hardcoded shift=0.
4. 4.aten::_native_multi_head_attention (no ONNX symbolic) → drop-in nn.MultiheadAttention.forward built from Linear/bmm/softmax.

The 6-stem head is wider than the 4-stem one but the surgery is identical — no new blockers. Parity at 2.42 × 10⁻⁴ on first try.

Related work

Sibling ONNX repos from the same export pipeline:

Full benchmark across every popular open-source separator: StemSplitio/stem-separation-benchmark-2026.

Skip the infrastructure — use the StemSplit API

Don't want to ship a 258 MB model in your app, manage a GPU pool, or write overlap-add chunking? Use the [StemSplit API](https://stemsplit.io/developers) instead — same model under the hood, hosted for you, with credits.

• —🌐 stemsplit.io
• —📘 Developer docs
• —🔌 API reference

License & attribution

This repo is MIT-licensed, matching the original HT-Demucs.

@inproceedings{rouard2023hybrid,
  title     = {Hybrid Transformers for Music Source Separation},
  author    = {Rouard, Simon and Massa, Francisco and D{\'e}fossez, Alexandre},
  booktitle = {ICASSP},
  year      = {2023}
}

• —Original PyTorch model: `facebookresearch/demucs`
• —ONNX export, parity verification, and packaging by StemSplit
• —Search keywords: htdemucs 6 stem onnx, htdemucs_6s onnx, guitar isolation onnx, piano isolation onnx, demucs 6-stem mobile, stem separation guitar onnx.