---
license: mit
library_name: onnx
pipeline_tag: audio-to-audio
base_model: KimberleyJSN/melbandroformer
base_model_relation: quantized
tags:
- audio
- source-separation
- music-source-separation
- vocals
- onnx
- webgpu
---

# Mel-Band RoFormer (vocals), ONNX for the GPU (float16 weights)

Mel-Band RoFormer (Wang, Lu, Won, Chen, Song, "Mel-Band RoFormer for Music Source Separation", ISMIR 2023 LBD,
[arXiv:2310.01809](https://arxiv.org/abs/2310.01809)): axial transformers over time and 60 overlapping mel bands of the
complex spectrogram, a mask estimated per band. This is Kimberley Jensen's vocals model
([KimberleyJSN/melbandroformer](https://huggingface.co/KimberleyJSN/melbandroformer)): it takes a song's vocals out;
the rest is the song less them. The network between its STFT and iSTFT, as
[@audio/neural-separate](https://github.com/audiojs/neural/tree/main/packages/neural-separate) runs it
(`model: 'mel-roformer'`), written for the GPU.

| File | Size | SHA-256 |
|---|---|---|
| `mel-roformer.fp16.onnx` | 464.1 MB (464,050,097 bytes) | `adcbbb826f12453fbb86652b503242d565a9c81cdef1bca2fe0e54752eba91ef` |
| `mel-roformer.engine.bin` | 466.1 MB (466,082,640 bytes) | `140fedcd9d03a8f918b235688c69f4dfdd86b08f9d430277659c095688a026ff` |
| `mel-roformer.engine.json` | 16 KB (16,282 bytes) | `3588f4da4a862d939d3e9e314d04c9b4c79706a54838fead43a73a14ca78f10b` |

The ONNX file: float16 weights, float32 compute (onnxruntime casts them at load); the float32 export it is made from is
927.5 MB. The engine files: the same tensors (weights float16, tables float32, indices int32) for the package's own WebGPU
engine, `roformer.js`, which a browser whose GPU offers subgroup matrices (Chromium on Metal) runs instead of onnxruntime:
the `.json` names each tensor, its type, shape and offset in the `.bin`. A browser downloads one or the other.

## Source

- Checkpoint: `MelBandRoformer.ckpt` (SHA-256 `87201f4d31afb5bc79993230fc49446918425574db48c01c405e44f365c7559e`),
  [KimberleyJSN/melbandroformer](https://huggingface.co/KimberleyJSN/melbandroformer) at
  `ac9b0614ab3cd7f77219e18ba494dfd93956c348` (MIT).
- Config and model code: [ZFTurbo/Music-Source-Separation-Training](https://github.com/ZFTurbo/Music-Source-Separation-Training)
  (MIT) at `84b1eac0887756b4f1a9d7a1ff49105939749ed2`: `configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml`,
  `models/bs_roformer/mel_band_roformer.py` (from Phil Wang's [BS-RoFormer](https://github.com/lucidrains/BS-RoFormer), MIT).

## Licence and attribution

MIT ([LICENSE](LICENSE)), as the checkpoint's repository states. The weights by Kimberley Jensen; the architecture by
Wang, Lu, Won, Chen and Song (ByteDance); the code it was written in by Phil Wang and Roman Solovyev (ZFTurbo); ONNX
export for the GPU and float16 weights by audiojs. What the model was trained on its repository does not say.

```bibtex
@inproceedings{wang2023melroformer,
  title     = {Mel-Band RoFormer for Music Source Separation},
  author    = {Wang, Ju-Chiang and Lu, Wei-Tsung and Won, Minz and Chen, Kuan-Yi and Song, Xuchen},
  booktitle = {ISMIR Late-Breaking Demo},
  year      = {2023},
  eprint    = {2310.01809},
  archivePrefix = {arXiv}
}
```

## Graph

One 8 s segment (352,800 samples at 44.1 kHz, 801 frames) per run:

| | name | shape | |
|---|---|---|---|
| input | `mix_spec` | [1, 4, 1025, 801] | STFT: n 2048, hop 441, Hann of 2048, unnormalized, centered with reflect padding; L re, L im, R re, R im |
| output | `stems_spec` | [1, 4, 1025, 801] | the vocals: L re, L im, R re, R im |

Segments every 4 s, faded over 0.8 s, the input reflected out at the ends, as Music-Source-Separation-Training's
`demix()` runs this model (`num_overlap: 2`); nothing normalized.

```js
import separate from '@audio/neural-separate'
let { stems } = await separate([left, right], { sampleRate: 44100, model: 'mel-roformer' })   // stems.vocals
```

## Written for the GPU

`scripts/export-roformer.py` exports what `MelBandRoformer.forward` computes between its STFT and iSTFT, rearranged so a
GPU runs it in few, large operations (826 of them); none of it changes the function:

- the 60 bands' input projections and mask MLPs as batched matrix products, bands of near widths grouped and
  zero-padded (10 groups);
- each RMSNorm's scale folded into the weights it feeds, the attention's scale into the queries, its gates' projection
  joined to q, k and v;
- the rotary embedding's interleaved pairs turned to halves (the same permutation of q's and k's features), so it is
  two products with cos and sin tables and one with a fixed 64×64 matrix;
- the overlapping bands' masks averaged by two gathers, not a scatter;
- no tensor past 100 MB: each transformer runs on slices of its sequences (onnxruntime 1.30's WebGPU Softmax over a
  tensor past WebGPU's default 128 MiB storage binding came out wrong one run in two).

`--verify`: the float32 graph against `MelBandRoformer.forward` on noise and tones, max |diff| 1e-5 to 2e-4 of max |y|
(PyTorch's own float32 stands 5e-3 from its float64 on the tones); the package's pipeline against the forward on its
segments, 89 to 92 dB SNR. This file against the float32 graph: 48 dB SNR on a segment of noise-floored tones.

On onnxruntime 1.30's WebGPU (Metal, Apple M4 Max), 2.4 s a segment for the float32 graph, 3.0 for this one; on its CPU,
9 s. SCNet-large's LSTMs take 19 of its 26 s a segment on the same WebGPU. The engine, in Chromium on Metal under the same
load as onnxruntime-web's WebGPU there: 1.5 to 1.7 s a segment against 2.6 to 3.3, 47.8 dB from onnxruntime's CPU on the
float32 graph (the float16 weights' rounding, as this ONNX file's).

## Quality

The 50 MUSDB18 test previews, BSSEval v4 SDR (museval), the median over songs, dB (the trained-on data unknown, so these
may include songs it heard):

| | vocals | vocals +6 dB remix | vocals −6 dB remix |
|---|---|---|---|
| float32 export | 12.08 | 20.33 | 23.38 |
| this file | 12.08 | 20.33 | 23.38 |
| SCNet-large (`audiojs/scnet-large`) | 11.00 | 20.20 | 22.11 |
| iZotope RX 12 Music Rebalance, Best | 10.89 | 19.46 | 22.26 |