MRX, ONNX (int8 weights)
MRX (Petermann, Wichern, Wang, Le Roux, "The Cocktail Fork Problem: Three-Stem Audio Separation for Real-World
Soundtracks", ICASSP 2022, arXiv:2110.09958): magnitudes at three STFT resolutions
into one hidden layer, a BLSTM per source, a real mask per source and resolution; 30.5 M parameters, MERL's own
checkpoint trained on Divide and Remaster with the SNR loss. It splits a soundtrack into three stems: dialogue, music,
effects. This is the network between its STFTs and iSTFTs, as
@audio/neural-separate runs it (model: 'mrx').
| File | Size | SHA-256 |
|---|---|---|
mrx.int8.onnx | 31.3 MB (31,287,081 bytes) | d876c92d224f5d8cd2207d528b23369da895a6de71fb24132416fcb16f9f8424 |
The float32 export it is made from is 122.3 MB; float16 weights would be 61.5 MB.
Source
- Model, code and checkpoint: merlresearch/cocktail-fork-separation
at
19b3de827ebc4bfb014570cf92dd32b4ee3b6921:mrx.py,checkpoints/default_mrx_pre_trained_weights.pth.
Licence and attribution
MIT, Copyright (c) 2023 Mitsubishi Electric Research Laboratories (MERL) (LICENSE). The repository's
.reuse/dep5
names the checkpoints:
Files: checkpoints/default_mrx_pre_trained_weights.pth [and its three other checkpoints] Copyright: 2023 Mitsubishi Electric Research Laboratories (MERL) License: MIT
MRX by Darius Petermann, Gordon Wichern, Zhong-Qiu Wang and Jonathan Le Roux (MERL); ONNX export and compaction by audiojs. Trained on Divide and Remaster v2, whose music and effects come from FMA and FSD50K clips, each under its own licence, some non-commercial.
@inproceedings{petermann2022cocktail,
title = {The Cocktail Fork Problem: Three-Stem Audio Separation for Real-World Soundtracks},
author = {Petermann, Darius and Wichern, Gordon and Wang, Zhong-Qiu and Le Roux, Jonathan},
booktitle = {IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
year = {2022},
eprint = {2110.09958},
archivePrefix = {arXiv}
}
Graph
Each channel apart (B), any length (T frames):
| name | shape | ||
|---|---|---|---|
| input | mag_1024, mag_2048, mag_8192 | [B, 513, T], [B, 1025, T], [B, 4097, T] | magnitude STFTs: Hann windows of 1024, 2048, 8192, hop 256, scaled by 1/√n, centered with reflect padding, at 44.1 kHz |
| output | mask_1024, mask_2048, mask_8192 | [B, 3, F, T] | a real mask per source (music, speech, sfx) and resolution |
A source is the sum over resolutions of the iSTFT of its mask times that resolution's spectrogram; the input is set to
-27 LUFS first and the stems scaled back (upstream's separate.py); @audio/neural-separate runs it in 20 s chunks.
import separate from '@audio/neural-separate'
let { stems } = await separate([mono], { sampleRate: 48000, model: 'mrx' }) // stems.dialogue, .music, .effects
Export, compaction, verification
scripts/export-mrx.py --verify exports the network and compares the sources rebuilt from the graph with
MRX.forward on noise and tones: max |diff| ≤ 1.2e-6 of max |y|; the package's pipeline matches upstream's
separate_soundtrack to 101–134 dB SNR per stem. scripts/compact.py --model mrx --calibrate <a Divide and Remaster v3 tuning clip> makes this file from it:
- Weights: all 39 stored in int8 (symmetric, a scale per output channel, an LSTM's per gate row and direction): rounded together they move the masked spectra 40.4 dB under their power on the calibration clip, within the 40 dB the script allows before keeping any in float16.
- Computed in float32 on every backend: each weight is Cast and multiplied by its scale in the graph, which onnxruntime folds at load (the session holds float32 weights).
- Node and value names shortened (its input length is free, so nothing is folded).
- Against the export on the calibration clip: max |diff| 9.9e-3 of max |y|, SNR 40.4 dB (the worst resolution).
Quality
30 clips of Divide and Remaster v3's English test set (every 40th of its 1200 clips of 60 s; Watcharasupat, Wu, Orife, 2024, CC BY-SA 4.0): each stem against its reference over the whole clip, SNR, the median over clips, dB:
| dialogue | music | effects | |
|---|---|---|---|
| export (float32) | 10.92 | 5.17 | 5.72 |
| this file | 10.92 | 5.17 | 5.70 |
| change per clip: median · the clip that lost most | +0.00 · −0.01 | +0.00 · −0.04 | +0.00 · −0.02 |
SI-SDR (mean over clips) 9.819 · 2.804 · 3.275 against the export's 9.815 · 2.805 · 3.273. Remixes (a stem 6 dB up or down, against the true remix): dialogue +6 dB, music −6 dB, effects −6 dB each within 0.03 dB of the export's median, no clip lower by more than 0.04 dB. Float16 weights (61.5 MB) change no median by more than 0.001 dB.