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