audiojs/mel-roformer

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): 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): 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 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 at ac9b0614ab3cd7f77219e18ba494dfd93956c348 (MIT).
• Config and model code: 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, MIT).

Licence and attribution

MIT (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.

@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.

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