@audio/neural-separate 0.1.0 → 0.3.0

package/NOTICE

@@ -14,6 +14,30 @@ htdemucs.py's spectrogram framing, into separate.js. Under the MIT License:
 
     Copyright (c) Meta Platforms, Inc. and affiliates.
 
+This product includes a JavaScript translation of parts of Music-Source-Separation-Training
+(https://github.com/ZFTurbo/Music-Source-Separation-Training): utils/model_utils.py's demix(),
+its segments, fades and reflection, into separate.js. Under the MIT License:
+
+    Copyright (c) 2024 Roman Solovyev (ZFTurbo)
+
+This product includes a JavaScript translation of parts of cocktail-fork-separation
+(https://github.com/merlresearch/cocktail-fork-separation): mrx.py's three STFTs, masks and iSTFTs, and
+separate.py's input loudness, into separate.js. Under the MIT License:
+
+    Copyright (c) 2023 Mitsubishi Electric Research Laboratories (MERL)
+
+This product includes a JavaScript translation of parts of TIGER (https://github.com/JusperLee/TIGER):
+look2hear/models/tiger_dnr.py's TIGERDNR.wav_chunk_inference, its segments and their sum, into separate.js.
+Under the MIT License:
+
+    Copyright (c) 2026 Kai Li
+
+This product includes a JavaScript translation of parts of pyloudnorm
+(https://github.com/csteinmetz1/pyloudnorm): meter.py's integrated loudness and iirfilter.py's K-weighting
+coefficients, into separate.js. Under the MIT License:
+
+    Copyright (c) 2018 Christian Steinmetz
+
 The MIT License, for each notice above:
 
     Permission is hereby granted, free of charge, to any person obtaining a copy of this
@@ -33,4 +57,8 @@ The MIT License, for each notice above:
     DEALINGS IN THE SOFTWARE.
 
 No model weights are included. The export scripts write them to the local cache; their terms are
-their authors' (Open-Unmix umxhq: MIT; Demucs htdemucs weights: see the README).
+their authors' (Open-Unmix umxhq: MIT; SCNet and SCNet-large: MIT, Copyright (c) 2024 starrytong,
+redistributable converted with attribution, github.com/starrytong/SCNet/issues/35; Demucs htdemucs
+weights: see the README; MRX: MIT, Copyright 2023 Mitsubishi Electric Research Laboratories (MERL), the
+repository's .reuse/dep5 naming each checkpoint; TIGER's DnR model: Apache 2.0, its model card,
+huggingface.co/JusperLee/TIGER-DnR).

package/README.md

@@ -1,13 +1,13 @@
 # @audio/neural-separate
 
-> Source separation (stems): Open-Unmix and Hybrid Transformer Demucs through ONNX, matching the original PyTorch implementations.
+> Source separation (stems): a song's by SCNet, Hybrid Transformer Demucs and Open-Unmix, a soundtrack's dialogue, music and effects by MRX and TIGER, through ONNX, matching the original PyTorch implementations.
 
 The ML upgrade to [`@audio/vocals`](https://github.com/audiojs/vocals)'s classical center-cancel: instead of one M/S trick, a trained model estimates each target (vocals, drums, bass, other), and the pipeline around it (STFT, Wiener refinement, chunking, overlap-add) is ported from the reference implementations and checked against them.
 
 ```js
 import separate from '@audio/neural-separate'
 
-let { stems, residual } = await separate([left, right], { sampleRate: 44100, model: 'umxhq' })
+let { stems, residual } = await separate([left, right], { sampleRate: 44100, model: 'scnet-large' })
 stems.vocals   // Float32Array[2]
 ```
 
@@ -25,11 +25,33 @@ In the browser `@audio/neural-runtime` runs `onnxruntime-web` instead (`device:
 |---|---|---|---|---|---|
 | `'umxhq'` | Open-Unmix: bi-LSTM on magnitude spectrograms, multichannel Wiener EM | 4 × 35.6 MB (fp16: 4 × 17.8 MB) | MIT | 6.75 · 6.11 · 5.00 · 3.36 | 0.19 (vocals only: 0.11) · 0.12 · 0.29 |
 | `'htdemucs'` | Hybrid Transformer Demucs: waveform and spectrogram U-Nets joined by a cross-domain transformer | 174 MB | research only | 8.86 · 9.55 · 9.30 · 5.69 | 1.4 (vocals only: 1.2) · 0.70 · 0.21 |
-| `'htdemucs_ft'` | four fine-tuned HTDemucs, one per source | 4 × 174 MB | research only | not measured | vocals only: 1.1 · not measured · not measured |
+| `'htdemucs_ft'` | four fine-tuned HTDemucs, one per source | 4 × 174 MB | research only | 8.67 · 9.45 · 9.68 · 5.58 | vocals only: 1.1 · not measured · not measured |
+| `'scnet-large'` | SCNet-large: band-split convolutions around dual-path LSTMs, on the complex spectrogram | 169 MB (fp16: 85 MB) | MIT | 11.00 · 10.27 · 8.21 · 6.87 | 2.1 · not measured · not measured |
+| `'scnet'` | SCNet: the same at half the width | 43 MB | MIT | 9.88 · 9.43 · 8.35 · 6.15 | 0.70 · not measured · not measured |
 
 SDR as measured below (Verification), identical to the Python originals'. The previews are 7 s excerpts, so these numbers sit apart from full-track results in scale, not in order: the Demucs README reports overall SDR 9.0 for fine-tuned HT Demucs against 5.3 for Open-Unmix on the MUSDB18-HQ test set.
 
-Real-time factor: processing time over duration, lower is faster. 27 s of music; onnxruntime-node and onnxruntime-web 1.30, the web one in headless Chromium (wasm, Metal WebGPU); the best of two to four runs spread over a day on a 14-core M4 Max shared with other jobs (load averages 45 to 180), so upper bounds. The Python originals' best there: 0.43 (open-unmix) and 2.8 (demucs). umxhq's time is mostly JS (ONNX Runtime takes 6 to 12 %: STFT, Wiener EM, iSTFT around it), and it runs faster on wasm than on WebGPU; htdemucs's is about 80 % inference, where WebGPU pays. Memory: umxhq 1.3 GB, hybrid models 3 GB at peak in Node.
+`scnet-large` is ahead of iZotope RX 12 Music Rebalance at its Best quality on drums and other and on most songs for every stem, behind on bass by the median (Measured). Its time is the segments' overlap: each sample is heard in four 11 s segments (`step` 0.25); `step: 0.5` takes half the time (1.0) for 0.07 to 0.17 dB less (27 s of four MUSDB18 training previews end to end, six of them), and a song shorter than a segment takes one. The SCNet times are of one run, a busy machine's (load averages 53 to 93), beside htdemucs's 0.24 and RX 12 Music Rebalance's 1.22 at Best in the same run. Real-time factor: processing time over duration, lower is faster. 27 s of music; onnxruntime-node and onnxruntime-web 1.30, the web one in headless Chromium (wasm, Metal WebGPU); the best of two to four runs spread over a day on a 14-core M4 Max shared with other jobs (load averages 45 to 180), so upper bounds. The Python originals' best there: 0.43 (open-unmix) and 2.8 (demucs). umxhq's time is mostly JS (ONNX Runtime takes 6 to 12 %: STFT, Wiener EM, iSTFT around it), and it runs faster on wasm than on WebGPU; htdemucs's is about 80 % inference, where WebGPU pays. Memory: umxhq 1.3 GB, hybrid models 3 GB at peak in Node.
+
+A soundtrack's stems, Divide and Remaster's three (the cocktail fork problem: Petermann, Wichern, Wang, Le Roux, ICASSP 2022):
+
+| `model` | Architecture | Files | Weights | Targets | Real-time factor, Node CPU |
+|---|---|---|---|---|---|
+| `'mrx'` | MRX: magnitudes at three STFT resolutions (Hann 1024, 2048, 8192, hop 256) into one hidden layer, a BLSTM per source, their mean, a real mask per source and resolution | 122 MB | MIT | `dialogue`, `music`, `effects` | 0.21 |
+| `'tiger'` | TIGER: three band-split models (57 bands, multi-scale convolutions and frame and frequency attention, 1.4 M parameters each), a complex mask each, one source kept from each | 29 MB | Apache 2.0 | `dialogue`, `music`, `effects` | 7.5–23 |
+
+Measured on Divide and Remaster v3's English test set (Watcharasupat, Wu, Orife, 2024; CC BY-SA 4.0, built from sources that allow commercial use), 150 of its 1200 clips of 60 s, every 8th ([`audio`](https://github.com/audiojs/audio)'s `bench/rx/scene.mjs`): each stem against its reference over the whole clip, SNR the median over clips (the measure Bandit v2's paper reports on this set) and SI-SDR the mean (MRX's on v2); remixes, a stem 6 dB up or down as the input plus (g − 1) times it, against the true remix, SNR, the median.
+
+| | dialogue | music | effects | dialogue +6 dB | music −6 dB | effects −6 dB |
+|---|---|---|---|---|---|---|
+| the mixture as it is | | | | 7.17 | 13.35 | 12.84 |
+| `mrx`, SNR · SI-SDR | 11.82 · 10.13 | 5.32 · 2.54 | 6.03 · 3.66 | 18.91 | 18.66 | 17.93 |
+| Bandit v2, its paper (all 1200 clips; weights CC BY-SA 4.0) | 15.6 | 10.4 | 9.9 | | | |
+| 30 of the clips (every 40th): the mixture as it is | | | | 7.75 | 13.38 | 11.93 |
+| `mrx` | 10.92 · 9.81 | 5.17 · 2.80 | 5.72 · 3.27 | 18.65 | 18.29 | 17.72 |
+| `tiger` | 12.68 · 12.35 | 10.24 · 8.06 | 8.06 · 7.50 | 20.98 | 21.46 | 21.15 |
+
+TIGER, at about 50 times MRX's time (RTF 10.6 against 0.21 here), is ahead on 29, 28 and 29 of the 30 clips (dialogue, music, effects), by a median 2.3, 3.7 and 3.4 dB. MRX's own README reports 12.5 · 4.2 · 5.7 dB SI-SDR on DnR v2, the set it trained on, whose music holds singing; v3's music holds none, and its loudness and languages differ (v3's paper). DnR v2's test split comes only inside a 116 GB gzip of the whole set; v3's clips are fetched one by one. The real-time factors are of one run on a 14-core M4 Max shared with other jobs (load averages 25 to 60), so upper bounds: `mrx` a 60 s clip in 13 s; `tiger` three models of 15,000 small operators each, every sample in three 12 s segments.
 
 `targets` picks a subset: `targets: ['vocals']` runs only the vocals graph of `umxhq` (against the residual, see Algorithm: 2.4× faster, vocals SDR 6.50 instead of 6.75) and of `htdemucs_ft`, and skips the other sources' iSTFT for `htdemucs`. Presets resample to their 44.1 kHz and back.
 
@@ -43,11 +65,23 @@ python3 node_modules/@audio/neural-separate/scripts/export-openunmix.py --model
 
 pip install demucs
 python3 node_modules/@audio/neural-separate/scripts/export-htdemucs.py --model htdemucs --verify
+
+pip install pyyaml
+python3 node_modules/@audio/neural-separate/scripts/export-scnet.py --model scnet-large --verify
+
+pip install pyloudnorm
+python3 node_modules/@audio/neural-separate/scripts/export-mrx.py --verify
+
+pip install safetensors
+python3 node_modules/@audio/neural-separate/scripts/export-tiger.py --verify
 ```
 
 They land in `$AUDIO_NEURAL_CACHE/<model>/`, default `~/.cache/audiojs/neural/<model>/`, where a preset looks in Node; a missing file throws, naming itself and the script. Elsewhere, serve that directory and pass its URL: `separate(audio, { model: 'umxhq', weights: 'https://…/' })`; `@audio/neural-runtime` fetches each file once and caches it (Cache API in the browser). `--verify` compares the ONNX graphs with the PyTorch modules and writes the reference separation `test.js` checks the JS pipeline against. `export-openunmix.py --fp16` also writes float16 weights with float32 I/O, half the size.
 
 - **`export-openunmix.py`** exports each target's `OpenUnmix` through a wrapper that takes the frame count from the traced tensor (upstream reads it off `x.data.shape`, which export freezes), with the TorchScript exporter (the `torch.export` one bakes the LSTM's sequence length into a reshape). Magnitude `[1, C, F, T]` in and out, T dynamic. `umxhq`'s bandwidth restriction (`max_bin`, 1487 of 2049 bins for `n_fft=4096`@44.1kHz) crops the network's *input* below 16 kHz for efficiency; the final dense layer regresses the **full** bin range from that reduced representation (a learned extrapolation, not a literal zero-fill).
+- **`export-scnet.py`** fetches SCNet's code from ZFTurbo/Music-Source-Separation-Training at a pinned commit and the author's MUSDB18-HQ checkpoint from that repository's releases, and exports the network between its STFT and iSTFT (the same split as demucs.onnx's): `mix_spec [1, 4, 2049, 476]` in, `stems_spec [1, 16, 2049, 476]` out, one 11 s segment. Two changes to the graph, neither to its function: the rFFT and irFFT over time in its separation network become products with cosine and sine matrices (torch 2.14 exports no `fft_rfft`), and each `GroupNorm(1, C)` reduces its mean and variance one axis at a time: exported whole, onnxruntime summed a 10-million-value group in float32 and the stems came out 52 to 55 dB from PyTorch's, against 110 dB now (PyTorch's own float32 against float64: 110 dB). `--fp16` also writes float16 weights.
+- **`export-mrx.py`** fetches MRX's code and MERL's checkpoint (`default_mrx_pre_trained_weights.pth`, trained with the SNR loss so its stems keep the mixture's level; the three others are scale-invariant or trained on loudness- or EQ-adapted stems, and their stems come out at another level: on Divide and Remaster v3 their SNR is negative where their SI-SDR is up to 1 dB higher) from merlresearch/cocktail-fork-separation at a pinned commit, and exports the network between its STFTs and iSTFTs: each channel's three magnitude spectrograms `mag_1024 [B, 513, T]`, `mag_2048 [B, 1025, T]`, `mag_8192 [B, 4097, T]` in (scaled 1/√n, torch.stft's `normalized=True`), the masks `mask_<n> [B, 3, F, T]` out, music, speech and sfx; B and T free.
+- **`export-tiger.py`** fetches TIGER's model code (JusperLee/TIGER, `look2hear/models/tiger_dnr.py` and the two layer modules it reads) at a pinned commit and its Divide and Remaster weights (huggingface.co/JusperLee/TIGER-DnR) at a pinned revision, and exports TIGERDNR's three models between their STFT and iSTFT as one graph, each its kept source (the first's dialogue, the second's effects, the third's music): `mix_spec [1, 2, 1025, 1034]` (one channel's STFT, n 2048, hop 512, Hann, unnormalized) in, `stems_spec [1, 6, 1025, 1034]` (dialogue, music, effects, re and im) out, one 12 s segment. Three changes to the graph, none to its function: `F.adaptive_avg_pool1d`, which the exporter takes only for output sizes dividing the input's, becomes a cumulative sum read at torch's windows; UConvBlock's global sum starts at its first term, not at a zeros tensor exported as a constant the size of the features; each `GroupNorm(1, C)` reduces axis by axis, as SCNet's. ONNX Runtime runs the graph in 12 s segments at about twice real time on a CPU: three models of 15,000 small operators each, every sample in three segments.
 - **`export-htdemucs.py`** follows [sevagh/demucs.onnx](https://github.com/sevagh/demucs.onnx): the STFT and iSTFT, which ONNX export cannot carry, move out of the graph. Rather than a vendored copy of `htdemucs.py`, it runs upstream's own `forward` with its four transform methods swapped on the instance. One graph per 7.8 s segment: `mix [1, 2, 343980]` and `mix_spec [1, 4, 2048, 336]` (complex as channels) in, `stems_spec [1, 4, 4, 2048, 336]` and `stems_wave [1, 4, 2, 343980]` out; a source is the iSTFT of the first plus the second.
 
 ## Algorithm
@@ -62,6 +96,12 @@ They land in `$AUDIO_NEURAL_CACHE/<model>/`, default `~/.cache/audiojs/neural/<m
 
 **`'hybrid'`** (Hybrid Transformer Demucs): `demucs.apply.apply_model` with `split=True, overlap=0.25, shifts=0`, and `demucs.api`'s normalization by the whole input's mean and standard deviation. Segments of 7.8 s start every 5.85 s; each is centered in its window with the neighbouring input as context, zeros past the ends. Per segment: HTDemucs's framing (reflect re-padding by ¾ hop, `normalized=True`, the Nyquist bin and two edge frames dropped) → ONNX → frequency branch through the iSTFT plus time branch. Segments overlap-add under a triangular window. `shifts` (averaging randomly time-shifted runs, "up to 0.2 points" of SDR per demucs's docstring) is not implemented; the CLI's default single shift averages nothing.
 
+**`'complex'`** (SCNet; any model taking each channel's STFT, re and im as channels, and giving each source's): Music-Source-Separation-Training's `demix()`, its generic mode. The input is normalized by its mean and deviation; segments of 11 s start every `step` of a segment (0.25: every sample in four), each weighted by a window that fades in and out linearly over a tenth of it; an input longer than two segments less a step is first reflected out by that much on each side, and a segment running past the end is reflected out when more than half of it is input, zero-padded when not. Per segment: zero-padded to 476 frames as `SCNet.forward` pads, the model's STFT (n 4096, hop 1024, no window, normalized, centered), ONNX, iSTFT. Departures from `demix()`: the segments stop at the first to reach the end (it runs on while one starts before the end: a 7 s song took three segments, now one), the first and last fades are set per segment (it sets them per batch of 8, so its first segment fades in), and the mean goes back to no source (its callers add it to every one), so the stems sum to the input less its DC.
+
+TIGER's DnR model runs as `'complex'` with its own segments, those of `TIGERDNR.wav_chunk_inference`: 12 s every 4 s (`step` 1/3), unweighted (`fade` 0), the input zero-padded by 8 s at each end and a segment past the end with zeros (`pad: 'zero'`), nothing normalized (`standardize: false`), each channel apart (`mono`); every sample is heard in three segments, their sum divided by three.
+
+**`'multires'`** (MRX): upstream's `separate.separate_soundtrack`. The input is set to -27 LUFS (BS.1770-4 integrated loudness as pyloudnorm measures it, its K-weighting designed at the model's rate) and its stems scaled back by the same gain. Each channel is heard apart, so mono stays mono: its STFT at each window (Hann, hop 256, centered, reflect padding), the magnitudes into the graph, and each source the sum over the three resolutions of the iSTFT of its mask times that resolution's spectrogram. MRX's LSTMs hear the whole input at once; here it runs in chunks of 20 s (`chunk`) crossfaded over 2 s (`overlap`), which hold Node's peak at 1.9 GB on a 60 s clip (4.4 GB whole, 1.3 GB in 10 s chunks): on the 18 tune clips of Divide and Remaster v3 that costs 0.09, 0.02 and 0.24 dB of SNR (median, dialogue, music, effects) against the whole clip at once (`chunk: 3600`), and the remixes stay within 0.15 dB; 10 s chunks cost 0.47 dB of the dialogue.
+
 `modelType: 'waveform'` (Demucs v2-class) skips steps 1–4 entirely: chunked raw audio in `[1, C, N]`, stacked stems out `[1, S, C, N]`, no STFT or Wiener step — Demucs v2 operates in the time domain by design.
 
 ## Verification
@@ -71,6 +111,12 @@ They land in `$AUDIO_NEURAL_CACHE/<model>/`, default `~/.cache/audiojs/neural/<m
 | ONNX vs PyTorch module, random input (`--verify`) | umxhq max \|diff\| ≤ 5.7e-6 of max \|y\| (fp16: ≤ 5.3e-3); stems of htdemucs ≤ 1.6e-4, of htdemucs_ft ≤ 2.5e-4 |
 | Pipeline vs `openunmix.Separator` in float64, 9 s reference mix (`test.js`) | 112–134 dB SNR per stem |
 | Pipeline vs `demucs.apply.apply_model`, same mix (`test.js`) | htdemucs 80–84 dB, htdemucs_ft 78–89 dB SNR per stem |
+| ONNX vs `SCNet.forward`, noise and tones (`export-scnet.py --verify`) | scnet-large max \|diff\| ≤ 2.8e-6 of max \|y\| (fp16: ≤ 6.7e-3), scnet ≤ 2.2e-6 |
+| Pipeline vs `SCNet.forward` on `reference.py`'s segments, a 20 s mix (`test.js`) | scnet-large 114–134 dB, scnet 123–133 dB SNR per stem |
+| ONNX vs `MRX.forward`, noise and tones (`export-mrx.py --verify`) | max \|diff\| ≤ 1.2e-6 of max \|y\| |
+| Pipeline vs `separate.separate_soundtrack` (pyloudnorm -27 LUFS in and back), a 20 s mix (`test.js`) | 101–134 dB SNR per stem, in one chunk; in 8 s chunks 12–50 dB |
+| ONNX vs the three `TIGER.forward`, one 12 s segment of tones and noise (`export-tiger.py --verify`) | max \|diff\| ≤ 4.2e-7 of max \|y\| |
+| Pipeline vs `TIGERDNR.wav_chunk_inference`, its segments through the verified graph, a 20 s mix (`test.js`) | 95–140 dB SNR per stem |
 | `wienerFilter` vs open-unmix's `wiener()`, float64 fixture (`test.js`) | ratio mask: 3e-12; mixture phase: 1.7e-7, upstream's `atan2` adds a float32 π |
 
 The 50 test tracks of the MUSDB18 7 s previews (`musdb.DB(download=True)`; its terms are educational and non-commercial, so it serves measurement only), SDR by [museval](https://github.com/sigsep/sigsep-mus-eval) (BSSEval v4, 1 s windows, median over windows, then over tracks: the SiSEC 2018 aggregation):
@@ -84,6 +130,26 @@ The 50 test tracks of the MUSDB18 7 s previews (`musdb.DB(download=True)`; its t
 
 Per track, the SDRs differ by at most 0.003 dB.
 
+## Measured against iZotope RX 12
+
+RX 12 Advanced Music Rebalance (VST3 hosted by Pedalboard; [`audio`](https://github.com/audiojs/audio)'s `bench/rx/separate.mjs`, October 2026) beside this package on the same 50 MUSDB18 test previews, the mixture stream in, the stems as references. RX at each quality from its defaults (sensitivity 50 %), a stem soloed (`vocal_solo` and the like); ours at its defaults. Stems: BSSEval v4 SDR as above (museval 0.4.1), the median over tracks, then the songs where each is ahead (paired). Remixes, RX Music Rebalance's own use: vocals 6 dB up, vocals 6 dB down, drums 6 dB down, against the true remix (the mixture with that stem's own part scaled, so the coding difference between MUSDB18's mixture and its stems, 22 to 31 dB, stays as it is), the same SDR of one source; RX's remix renders equal the input plus (g − 1) times its solo to 126 dB and more, so its remixes are made from its solos. Harm: every gain 0 dB, the output against the input. RTF: one render over the preview's 6.8 s.
+
+| 50 songs | vocals | drums | bass | other | vocals +6 | vocals −6 | drums −6 | harm | RTF |
+|---|---|---|---|---|---|---|---|---|---|
+| the mixture as it is | | | | | 8.39 | 10.99 | 10.67 | | |
+| RX 12, Good / Real-time | 6.68 | 6.52 | 4.43 | 3.15 | 15.58 | 18.60 | 17.50 | 128.7 dB | 0.08 |
+| RX 12, Better / Offline | 9.29 | 8.93 | 6.89 | 5.44 | 18.53 | 20.96 | 19.98 | 130.2 dB | 0.53 |
+| RX 12, Best / Offline | **10.89** | 9.85 | **9.66** | 6.45 | 19.46 | **22.26** | 20.85 | 130.2 dB | 2.48 |
+| `umxhq` | 6.74 | 6.11 | 5.00 | 3.36 | 15.09 | 18.35 | 17.32 | bit for bit | 0.16 |
+| `htdemucs` | 8.86 | 9.55 | 9.30 | 5.69 | 17.99 | 20.76 | 20.76 | bit for bit | 1.74 |
+| `htdemucs_ft` | 8.67 | 9.45 | 9.68 | 5.58 | 17.59 | 20.33 | 21.08 | bit for bit | 4.70 |
+| `scnet` | 9.88 | 9.43 | 8.35 | 6.15 | 18.23 | 21.05 | 20.91 | bit for bit | 0.45 |
+| `scnet-large` | **11.00** | **10.27** | 8.21 | **6.87** | **20.20** | 22.11 | **21.88** | bit for bit | 0.96 |
+| … ahead of RX Best on | 33 of 50 | 40 | 34 | 32 | 34 | 30 | 41 | | |
+| … paired median over RX Best | +0.31 | +0.68 | +0.97 | +0.33 | +0.31 | +0.18 | +0.68 | | |
+
+Our remixes are the input plus (g − 1) times the stems, as `audio`'s `rebalance` makes them, so every gain at 0 dB returns the input (RX's, to 128 dB); `rebalance` itself (its solo is the input less the other three stems) scores 10.75 · 10.30 · 8.17 · 6.94 and the same remixes. Against `htdemucs`, the best model here before, RX at Best was ahead on 46 of the 50 songs for vocals (1.62 dB, the paired median). On vocals `scnet-large`'s median is 0.11 dB over RX's and `rebalance`'s 0.14 under it (ahead on 33 songs either way, the means 10.29 and 10.16 against 9.99). On bass it is 1.5 dB under RX's median while ahead on 34 songs, the means level (8.68 against 8.67): its bass falls under 3 dB on 8 songs (RX's on 6), mostly electronic, where it hears the synth bass as other (PR - Happy Daze: none of it, 0.0 dB, against RX's 18.3); SCNet trained on MUSDB18's 100 songs alone, RX on iZotope's own. Averaging `scnet` in (vocals 10.56, bass 8.21), test-time augmentation (the channels swapped, the polarity inverted: 3× the time, +0.01 to +0.02 dB paired) and multichannel Wiener EM over its stems (−0.5 to −2.4 dB on 13 of the songs: it takes the mixture's phase where SCNet estimated its own) were tried and left out. The previews are 7 s, one 11 s segment each: on whole songs each sample is heard in four.
+
 Upstream runs its Wiener EM in float32. On the reference mix that alone moves its stems by 71 to 93 dB SNR against a float64 run of the same code, which this port matches to 112 to 134 dB; the float64 run is the reference.
 
 ## Wiener filter
@@ -106,41 +172,48 @@ let vocals = complex.vocals.map(ch => istft(ch, { length: left.length }))
 | Option | Default | |
 |---|---|---|
 | `sampleRate` | — | required unless `audio` is `{ channelData, sampleRate }` |
-| `model` | — | required. Preset `'umxhq' \| 'htdemucs' \| 'htdemucs_ft'` · `url \| bytes` (single target, named `'stem'`) · `{ target: url \| { url, targets }, ... }` (one graph per target, Open-Unmix's own layout; a multi-source graph contributes its own target) · `{ url, targets: [...] }` (one multi-target graph, stacks a target axis) |
+| `model` | — | required. Preset `'scnet-large' \| 'scnet' \| 'umxhq' \| 'htdemucs' \| 'htdemucs_ft' \| 'mrx' \| 'tiger'` · `url \| bytes` (single target, named `'stem'`) · `{ target: url \| { url, targets }, ... }` (one graph per target, Open-Unmix's own layout; a multi-source graph contributes its own target) · `{ url, targets: [...] }` (one multi-target graph, stacks a target axis) |
 | `targets` | all | the targets to return |
 | `weights` | `$AUDIO_NEURAL_CACHE` or `~/.cache/audiojs/neural` (Node) | where a preset's files are: URL, or a directory in Node |
-| `modelType` | `'openunmix'` | `'openunmix'` (magnitude out) · `'mask'` (`[0,1]` mask out, multiplied by mixture magnitude) · `'hybrid'` (demucs.onnx contract) · `'waveform'` (Demucs v2-class); presets know theirs |
+| `modelType` | `'openunmix'` | `'openunmix'` (magnitude out) · `'mask'` (`[0,1]` mask out, multiplied by mixture magnitude) · `'hybrid'` (demucs.onnx contract) · `'complex'` (complex spectrogram in and out, SCNet's) · `'multires'` (magnitudes at several resolutions in, masks out, MRX's) · `'waveform'` (Demucs v2-class); presets know theirs |
 | `wiener` | `1` | EM iterations; `0` = raw masks. Ignored for `'hybrid'` and `'waveform'` |
 | `wienerWindow` | `300` | frames per EM window |
-| `chunk` / `overlap` | `30` / `2` (seconds) | `overlap` must be `< chunk`; not used by `'hybrid'`, which segments as demucs does |
-| `segment` | `343980` | `'hybrid'`: segment length in samples when the graph does not declare its input length |
+| `chunk` / `overlap` | `30` / `2` (seconds; `mrx` 20) | `overlap` must be `< chunk`; not used by `'hybrid'`, which segments as demucs does |
+| `segment` | `343980` | `'hybrid'`: segment length in samples when the graph does not declare its input length; `'complex'`: the preset's (485100) |
+| `step` | `0.25` | `'complex'`: segments start every `step` of a segment (each sample heard in 1/`step` of them; `tiger` 1/3) |
+| `fade` / `pad` / `standardize` | `0.1` / `'reflect'` / `true` | `'complex'`: each segment's linear fade, a fraction of it (`tiger` 0); the ends reflected or zero-padded (`tiger` `'zero'`); the input normalized by its mean and deviation (`tiger` not) |
+| `window` / `normalized` / `frames` | the preset's | `'complex'`: the model's STFT window (`'ones'`, torch.stft's window=None, or `'hann'`), its 1/√n scaling, its frames per segment when the graph does not declare them |
 | `targetRate` | the preset's rate, else the input rate | resample to the model's rate for inference; stems are resampled back to the input rate |
 | `device` | — | passed through to `@audio/neural-runtime`'s `load()` as `backend` |
 | `dtype` | `'float32'` | only `'float32'` tensor marshalling is implemented; anything else throws |
 | `progress` | — | `({ chunk, totalChunks }) => {}` |
 | `session` | — | overrides `@audio/neural-runtime`'s `load()` — for tests, or a custom ORT setup |
 
-Mono input is duplicated to stereo internally (matching `openunmix.utils.preprocess`'s own "if we have mono, we duplicate it to get stereo"), so stems always come back stereo. `separate()` returns `{ stems, sampleRate, residual }` — `residual` is the mixture minus the sum of all stems, per channel, at the input rate.
+Mono input is duplicated to stereo internally (matching `openunmix.utils.preprocess`'s own "if we have mono, we duplicate it to get stereo"), so stems come back stereo; `mrx` and `tiger` hear each channel apart and give mono back mono. `separate()` returns `{ stems, sampleRate, residual }` — `residual` is the mixture minus the sum of all stems, per channel, at the input rate.
 
 ## Precedence and licenses
 
-Open-source stem separation has three lineages; this package runs the first two.
+Open-source stem separation has these lineages; this package runs SCNet, Open-Unmix and Demucs.
 
 | Project | Code | Weights | |
 |---|---|---|---|
 | **Open-Unmix** (`umx`, `umxhq`) | MIT | MIT ([Zenodo](https://zenodo.org/records/3370489)-declared) | trained on MUSDB18(-HQ); this package's primary target |
 | **Open-Unmix** (`umxl`) | MIT | **CC BY-NC-SA 4.0 — non-commercial only** | despite being the `openunmix` package's own default variant name; trained on a private stems dataset (see the project [README](https://github.com/sigsep/open-unmix-pytorch#pre-trained-models)) |
-| **Demucs** (Meta; `htdemucs`, `htdemucs_ft`) | MIT | **research only**: "The model weights are not covered by the MIT license, and are provided only for scientific purposes" (the author in [#327](https://github.com/facebookresearch/demucs/issues/327), 2022; again in [#508](https://github.com/facebookresearch/demucs/issues/508), 2023) | trained on MUSDB18-HQ plus 800 songs; the repository is archived, maintained at [adefossez/demucs](https://github.com/adefossez/demucs) |
+| **Demucs** (Meta; `htdemucs`, `htdemucs_ft`) | MIT | **research only**: "The model weights are not covered by the MIT license, and are provided only for scientific purposes" (the author in [#327](https://github.com/facebookresearch/demucs/issues/327), 2022; again in [#508](https://github.com/facebookresearch/demucs/issues/508), 2023; the Hugging Face copies he published in 2026 had their `license: mit` tag removed) | trained on MUSDB18-HQ plus 800 songs; the repository is archived, maintained at [adefossez/demucs](https://github.com/adefossez/demucs) |
+| **SCNet** (`scnet`, `scnet-large`; starrytong) | MIT | MIT: "I confirm that the released SCNet and SCNet-large pretrained weights are distributed under the MIT License, consistent with the source code. You are welcome to redistribute the original checkpoints and format-converted versions, including ONNX exports, as part of your MIT-licensed tool, with appropriate attribution." (the author in [starrytong/SCNet#35](https://github.com/starrytong/SCNet/issues/35), 2026) | trained on MUSDB18-HQ; hosted with its configs by [ZFTurbo/Music-Source-Separation-Training](https://github.com/ZFTurbo/Music-Source-Separation-Training) (MIT), whose own MUSDB18-HQ checkpoints (SCNet XL, BS-RoFormer) carry its author's MIT statement only for some files ([#254](https://github.com/ZFTurbo/Music-Source-Separation-Training/issues/254)) |
+| **MRX** (MERL; `mrx`) | MIT | MIT: the repository's [`.reuse/dep5`](https://github.com/merlresearch/cocktail-fork-separation/blob/main/.reuse/dep5) names the four checkpoints, "Copyright: 2023 Mitsubishi Electric Research Laboratories (MERL) / License: MIT" | trained on Divide and Remaster v2, whose music and effects come from FMA and FSD50K clips each under its own licence, some non-commercial (the DnR v3 paper's account) |
+| **TIGER** (Tsinghua; `tiger`) | MIT (repository LICENSE; its README badge says Apache 2.0) | Apache 2.0: the [model card](https://huggingface.co/JusperLee/TIGER-DnR) states `license: apache-2.0` | trained on Divide and Remaster v1 (652 test mixtures in its paper), sourced as v2 is |
+| **Bandit v2** (DnR v3) | Apache 2.0 | **CC BY-SA 4.0** ([Zenodo](https://zenodo.org/records/12701995)): a converted copy is an adaptation under the same terms | not shipped; BandIt v1's DnR checkpoints are CC BY-NC 4.0 ([Zenodo](https://zenodo.org/records/10160698)) |
 | **Spleeter** (Deezer) | MIT | **undocumented** | the README licenses only "the code of Spleeter"; the pretrained weights' license is an open, unresolved question ([deezer/spleeter#898](https://github.com/deezer/spleeter/issues/898)) — do not assume MIT |
 
-Audit any weight source yourself before shipping it — this table reflects what each project states as of this writing, not a guarantee. `scripts/export-openunmix.py` defaults to `umxhq` (not `umxl`) for exactly this reason.
+Audit any weight source yourself before shipping it — this table reflects what each project states as of this writing, not a guarantee. `scripts/export-openunmix.py` defaults to `umxhq` (not `umxl`) for exactly this reason. MUSDB18, which every model here trained on, is licensed for educational use; whether that reaches the weights no project has settled (Défossez in [demucs#384](https://github.com/facebookresearch/demucs/issues/384): "a grey area").
 
 ## Reference
 
-Stöter, Uhlich, Liutkus, Mitsufuji, "Open-Unmix - A Reference Implementation for Music Source Separation," *JOSS* 4(41), 2019. · Rouard, Massa, Défossez, "Hybrid Transformers for Music Source Separation," *ICASSP* 2023. · Duong, Vincent, Gribonval, "Under-determined reverberant audio source separation using a full-rank spatial covariance model," *IEEE TASLP* 18(7), 2010. · Rafii, Liutkus, Stöter, Mimilakis, Bittner, "The MUSDB18 corpus for music separation," 2017. · [norbert](https://github.com/sigsep/norbert) (Liutkus & Stöter) · [open-unmix-pytorch](https://github.com/sigsep/open-unmix-pytorch) · [Demucs](https://github.com/facebookresearch/demucs) · [demucs.onnx](https://github.com/sevagh/demucs.onnx) · [museval](https://github.com/sigsep/sigsep-mus-eval) · [Spleeter](https://github.com/deezer/spleeter).
+Stöter, Uhlich, Liutkus, Mitsufuji, "Open-Unmix - A Reference Implementation for Music Source Separation," *JOSS* 4(41), 2019. · Rouard, Massa, Défossez, "Hybrid Transformers for Music Source Separation," *ICASSP* 2023. · Xu, Li, Chen, Hu, "TIGER: Time-frequency Interleaved Gain Extraction and Reconstruction for Efficient Speech Separation," *ICLR* 2025. · Petermann, Wichern, Wang, Le Roux, "The Cocktail Fork Problem: Three-Stem Audio Separation for Real-World Soundtracks," *ICASSP* 2022. · Watcharasupat, Wu, Orife, "Remastering Divide and Remaster: A Cinematic Audio Source Separation Dataset with Multilingual Support," 2024. · Tong, Zhang, Liu, Li, Yu, "SCNet: Sparse Compression Network for Music Source Separation," *ICASSP* 2024. · Duong, Vincent, Gribonval, "Under-determined reverberant audio source separation using a full-rank spatial covariance model," *IEEE TASLP* 18(7), 2010. · Rafii, Liutkus, Stöter, Mimilakis, Bittner, "The MUSDB18 corpus for music separation," 2017. · [norbert](https://github.com/sigsep/norbert) (Liutkus & Stöter) · [open-unmix-pytorch](https://github.com/sigsep/open-unmix-pytorch) · [Demucs](https://github.com/facebookresearch/demucs) · [demucs.onnx](https://github.com/sevagh/demucs.onnx) · [SCNet](https://github.com/starrytong/SCNet) · [cocktail-fork-separation](https://github.com/merlresearch/cocktail-fork-separation) · [TIGER](https://github.com/JusperLee/TIGER) · [pyloudnorm](https://github.com/csteinmetz1/pyloudnorm) · [Music-Source-Separation-Training](https://github.com/ZFTurbo/Music-Source-Separation-Training) · [museval](https://github.com/sigsep/sigsep-mus-eval) · [Spleeter](https://github.com/deezer/spleeter).
 
-**Use when:** you have (or can license) an ONNX-exported spectrogram-mask, Hybrid-Demucs or waveform separation model and want to run it — with proper multichannel Wiener refinement, not just the raw mask — dependency-free, in Node or the browser; `umxhq` where the weights must be MIT, `htdemucs` for the higher SDR under research-only terms.<br>
-**Not for:** the classical, model-free case — reach for [`@audio/vocals`](https://github.com/audiojs/vocals) when a center-panned M/S trick is all you need; training a model (this is inference-only); real-time streaming (neither the bi-LSTM Open-Unmix nor Hybrid Transformer Demucs is causal — offline/chunked only, same as upstream).
+**Use when:** you have (or can license) an ONNX-exported spectrogram-mask, Hybrid-Demucs or waveform separation model and want to run it — with proper multichannel Wiener refinement, not just the raw mask — dependency-free, in Node or the browser; `scnet-large` for the highest SDR here under MIT weights, `scnet` for a quarter of its size and a third of its time, `htdemucs` under research-only terms.<br>
+**Not for:** the classical, model-free case — reach for [`@audio/vocals`](https://github.com/audiojs/vocals) when a center-panned M/S trick is all you need; training a model (this is inference-only); real-time streaming (none of Open-Unmix's bi-LSTM, Hybrid Transformer Demucs or SCNet's dual-path LSTMs is causal: offline and chunked only, as upstream).
 
 ---
 

package/index.d.ts

@@ -1,17 +1,18 @@
 /**
  * Source separation (stems) through @audio/neural-runtime's ONNX adapter: Open-Unmix-class
  * spectrogram models (Stöter, Uhlich, Liutkus, Mitsufuji, JOSS 2019), Hybrid Transformer Demucs
- * (Rouard, Massa, Défossez, ICASSP 2023) with STFT and iSTFT outside the graph, and waveform graphs.
+ * (Rouard, Massa, Défossez, ICASSP 2023), SCNet (Tong et al., ICASSP 2024) and MRX (Petermann et al., ICASSP 2022)
+ * with STFT and iSTFT outside the graph, and waveform graphs.
  */
 
-/** Mono duplicates internally; multichannel arrays must share the same length. */
+/** Mono duplicates internally (a model hearing channels apart, mrx and tiger, keeps it mono); multichannel arrays must share the same length. */
 export type AudioInput = Float32Array[] | { channelData: Float32Array[]; sampleRate: number }
 
 /** A model spec resolvable by @audio/neural-runtime's load(): URL string or raw ONNX bytes. */
 export type ModelSpec = string | Uint8Array
 
 /** Presets whose files the export scripts write: <weights>/<name>/<target>.onnx or <weights>/<name>/<name>.onnx */
-export type ModelName = 'umxhq' | 'htdemucs' | 'htdemucs_ft'
+export type ModelName = 'umxhq' | 'htdemucs' | 'htdemucs_ft' | 'scnet-large' | 'scnet' | 'mrx' | 'tiger'
 
 /** One graph whose output stacks several sources on its S axis */
 export type MultiGraph = { url: ModelSpec; targets: string[] }
@@ -26,6 +27,8 @@ export type ModelType =
 	| 'openunmix' // target model outputs the estimated magnitude directly
 	| 'mask' // target model outputs a [0,1] mask; multiplied by the mixture magnitude
 	| 'hybrid' // demucs.onnx contract: mix [1,C,L] + CaC spectrogram [1,2C,F,T] in; [1,S,2C,F,T] + [1,S,C,L] out
+	| 'complex' // SCNet-class: CaC spectrogram [1,2C,F,T] in; [1,2SC,F,T] out (source, channel, re/im)
+	| 'multires' // MRX-class: one channel's magnitudes mag_<n> [1,F,T] at each window n in; real masks mask_<n> [1,S,F,T] out
 	| 'waveform' // Demucs v2-class: [1,C,N] waveform in, [1,S,C,N] stacked waveforms out
 
 export interface ModelPreset {
@@ -35,6 +38,30 @@ export interface ModelPreset {
 	targets: string[]
 	/** one graph per target (<target>.onnx), else one graph for all (<name>.onnx) */
 	perTarget?: boolean
+	/** 'complex': the model's STFT (size, hop, window, normalized) and segment (samples, frames) */
+	n?: number
+	hop?: number
+	window?: 'ones' | 'hann'
+	normalized?: boolean
+	segment?: number
+	frames?: number
+	/** the graph's source order, when not the targets' (or Demucs's for 'hybrid' and 'complex') */
+	sources?: string[]
+	/** 'multires': its windows, one hop */
+	windows?: number[]
+	/** integrated loudness (LUFS) the input is set to, its stems scaled back */
+	loudness?: number
+	/** channels heard apart: mono is not duplicated */
+	mono?: boolean
+	/** 'complex': segments' step and fade (fractions of a segment), their padding past the ends, the input normalized or not */
+	step?: number
+	fade?: number
+	pad?: 'reflect' | 'zero'
+	standardize?: boolean
+	/** chunk length (s) when not 30 */
+	chunk?: number
+	/** the export script that writes its weights */
+	script?: string
 }
 
 /** Model presets by name */
@@ -43,7 +70,7 @@ export const models: Record<ModelName, ModelPreset>
 /** A neural-runtime-shaped session — enough of it to drive separate() with a test double. */
 export interface Session {
 	run(feeds: Record<string, { data: Float32Array; dims: number[]; type: string }>): Promise<Record<string, { data: Float32Array; dims: number[] }>>
-	/** 'hybrid' reads its segment length off inputs[0].dims[2] when declared */
+	/** 'hybrid' reads its segment length off inputs[0].dims[2] when declared, 'complex' its frames off inputs[0].dims[3] */
 	inputs?: { name: string; dims?: number[] }[]
 	outputs?: { name: string }[]
 	free?(): void
@@ -67,9 +94,23 @@ export interface SeparateOptions {
 	eps?: number
 	/** frames per Wiener EM window (default 300, open-unmix Separator's wiener_win_len) */
 	wienerWindow?: number
-	/** 'hybrid': segment length in samples when the graph does not declare it (default 343980, htdemucs's 7.8 s) */
+	/** 'hybrid': segment length in samples when the graph does not declare it (default 343980, htdemucs's 7.8 s); 'complex': the preset's (SCNet 485100, 11 s) */
 	segment?: number
-	/** 'openunmix' | 'mask' | 'waveform': chunk length in seconds (default 30) */
+	/** 'complex': segments start every `step` of a segment (default 0.25: each sample in four, as Music-Source-Separation-Training's num_overlap 4) */
+	step?: number
+	/** 'complex': each segment's linear fade in and out, a fraction of it (default 0.1, demix's; tiger 0) */
+	fade?: number
+	/** 'complex': past the input's ends, reflected (default) or zeros (tiger) */
+	pad?: 'reflect' | 'zero'
+	/** 'complex': the input normalized by its mean and deviation (default true; tiger false) */
+	standardize?: boolean
+	/** 'complex': frames per segment when the graph does not declare them (inputs[0].dims[3]) */
+	frames?: number
+	/** 'complex': the STFT window, 'ones' (torch.stft's window=None, SCNet's) or 'hann' (default: the preset's, else 'hann') */
+	window?: 'ones' | 'hann'
+	/** 'complex': the STFT scaled by 1/√n and back, torch.stft's normalized=True (default: the preset's, else true) */
+	normalized?: boolean
+	/** 'openunmix' | 'mask' | 'waveform' | 'multires': chunk length in seconds (default 30; mrx 20) */
 	chunk?: number
 	/** crossfade overlap in seconds between chunks (default 2); must be < chunk */
 	overlap?: number

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@audio/neural-separate",
-  "version": "0.1.0",
-  "description": "Source separation (stems): Open-Unmix-class spectrogram masking + multichannel Wiener refinement, Hybrid Transformer Demucs, model-agnostic ONNX",
+  "version": "0.3.0",
+  "description": "Source separation (stems): a song's by SCNet (MIT weights), Hybrid Transformer Demucs, Open-Unmix with multichannel Wiener refinement; a soundtrack's dialogue, music and effects by MRX (MIT weights) and TIGER (Apache-2.0); model-agnostic ONNX",
   "type": "module",
   "sideEffects": false,
   "main": "separate.js",
@@ -18,6 +18,9 @@
     "index.d.ts",
     "scripts/export-openunmix.py",
     "scripts/export-htdemucs.py",
+    "scripts/export-scnet.py",
+    "scripts/export-mrx.py",
+    "scripts/export-tiger.py",
     "scripts/reference.py",
     "NOTICE"
   ],
@@ -41,6 +44,11 @@
     "stems",
     "open-unmix",
     "demucs",
+    "scnet",
+    "mrx",
+    "tiger",
+    "cinematic",
+    "dialogue",
     "wiener",
     "onnx"
   ],

package/scripts/export-mrx.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+"""Export MRX (Petermann, Wichern, Wang, Le Roux, "The Cocktail Fork Problem: Three-Stem Audio Separation for
+Real-World Soundtracks", ICASSP 2022) to the ONNX contract @audio/neural-separate expects for modelType
+'multires': the network between its three STFTs and their inverses, which JS computes.
+
+    pip install torch onnx onnxruntime pyloudnorm
+    python3 export-mrx.py --verify
+
+The model is MERL's own, trained on Divide and Remaster with the SNR loss so its stems keep the mixture's level
+(merlresearch/cocktail-fork-separation, checkpoints/default_mrx_pre_trained_weights.pth, 30.5 M parameters; the
+repository's .reuse/dep5 licenses the four checkpoints MIT, Copyright 2023 Mitsubishi Electric Research Laboratories);
+its code and checkpoint are fetched at a pinned commit. Its README, DnR test, average SI-SDR: speech 12.5, music 4.2,
+sound effects 5.7 dB (the mixture 1.0, -6.8, -5.0).
+
+The graph sees each channel's magnitude spectrogram at three resolutions, scaled as torch.stft(normalized=True)
+scales it (1/sqrt(n)): Hann windows of 1024, 2048 and 8192 samples, hop 256, centered with reflect padding, at
+44.1 kHz. B (channels, encoded apart) and T (frames) are free:
+  inputs   mag_1024 [B, 513, T]   mag_2048 [B, 1025, T]   mag_8192 [B, 4097, T]
+  outputs  mask_1024 [B, 3, 513, T]   mask_2048 [B, 3, 1025, T]   mask_8192 [B, 3, 4097, T]
+the sources in MRX's order: music, speech, sfx. A source is the sum over resolutions of the iSTFT of its mask times
+that resolution's spectrogram (the scale 1/sqrt(n) and its inverse cancel: the masks are real).
+
+Writes <out-dir>/mrx.onnx (122 MB), weights embedded. Default out-dir: $AUDIO_NEURAL_CACHE/mrx, else
+~/.cache/audiojs/neural/mrx. --verify compares the sources rebuilt from the ONNX output with MRX.forward, and writes
+the reference separation test.js compares the JS pipeline against: separate.separate_soundtrack, as upstream's
+separate.py runs it (the input at -27 LUFS by pyloudnorm, the stems scaled back), on scripts/reference.py's mix.
+Last run: torch 2.14, onnxruntime 1.30.
+"""
+import argparse
+import os
+import sys
+import urllib.request
+from pathlib import Path
+
+COMMIT = "19b3de827ebc4bfb014570cf92dd32b4ee3b6921"
+REPO = f"https://raw.githubusercontent.com/merlresearch/cocktail-fork-separation/{COMMIT}/"
+LFS = f"https://media.githubusercontent.com/media/merlresearch/cocktail-fork-separation/{COMMIT}/"
+CKPT = "checkpoints/default_mrx_pre_trained_weights.pth"
+WINDOWS, HOP, RATE = (1024, 2048, 8192), 256, 44100
+
+
+def fetch(url, path):
+    if not path.exists():
+        print(f"fetching {url}")
+        urllib.request.urlretrieve(url, str(path) + ".part")
+        os.replace(str(path) + ".part", path)
+    return path
+
+
+def load(src):
+    """MRX with its weights, and upstream's separate module: code and checkpoint fetched into `src` once."""
+    import torch
+
+    for f in ("mrx.py", "separate.py", "consistency.py", "dnr_dataset.py"):
+        fetch(REPO + f, src / f)
+    sys.path.insert(0, str(src))
+    sys.dont_write_bytecode = True
+    from mrx import MRX
+
+    model = MRX().eval()
+    model.load_state_dict(torch.load(fetch(LFS + CKPT, src / Path(CKPT).name), map_location="cpu", weights_only=True))
+    return model
+
+
+def mags(x):
+    """x (B, N) -> the three normalized magnitude spectrograms (B, F, T) and complex spectrograms MRX._stft makes"""
+    import mrx
+
+    specs = [mrx._stft(x, n, HOP) for n in WINDOWS]
+    return [s.abs() for s in specs], specs
+
+
+class Core:
+    """MRX.forward between its STFTs and its iSTFTs: magnitudes in, masks out."""
+
+    @staticmethod
+    def build(model):
+        import torch
+        import torch.nn as nn
+
+        class C(nn.Module):
+            def __init__(self, m):
+                super().__init__()
+                self.m = m
+
+            def forward(self, m1, m2, m3):
+                m = self.m
+                enc = []
+                for mag, layer in zip((m1, m2, m3), m._encoders):
+                    B, F, T = mag.shape
+                    e = layer.layer(mag.transpose(1, 2).reshape(-1, F))           # _EncoderBlock, its abs() done
+                    enc.append(torch.tanh(e.reshape(B, T, -1)))
+                cross = m._cross_net(enc)
+                out = []
+                for r, mag in enumerate((m1, m2, m3)):
+                    B, F, T = mag.shape
+                    masks = [dec[r](cross).transpose(-1, -2) for dec in m._decoders]  # (B, F, T) per source
+                    out.append(torch.stack(masks, 1))
+                return tuple(out)
+
+        return C(model).eval()
+
+
+def export(model, path, opset):
+    import torch
+
+    m, _ = mags(0.1 * torch.randn(2, RATE * 4))
+    names = [f"mag_{n}" for n in WINDOWS]
+    outs = [f"mask_{n}" for n in WINDOWS]
+    axes = {**{k: {0: "B", 2: "T"} for k in names}, **{k: {0: "B", 3: "T"} for k in outs}}
+    with torch.no_grad():
+        torch.onnx.export(Core.build(model), tuple(m), str(path), input_names=names, output_names=outs,
+                          dynamic_axes=axes, opset_version=opset, dynamo=False)
+    print(f"wrote {path} ({path.stat().st_size / 1e6:.1f} MB)")
+
+
+def session(path):
+    import onnxruntime as ort
+
+    o = ort.SessionOptions()
+    o.enable_cpu_mem_arena = False
+    o.enable_mem_pattern = False
+    return ort.InferenceSession(str(path), o, providers=["CPUExecutionProvider"])
+
+
+def verify(model, path, tol):
+    """Sources rebuilt from onnxruntime's masks against MRX.forward: noise and tones, mono and stereo."""
+    import math
+    import numpy as np
+    import torch
+    import mrx
+
+    s = session(path)
+    t = torch.arange(RATE * 3) / RATE
+    tone = sum(torch.sin(2 * math.pi * f * t) / k for k, f in enumerate((110, 220, 440, 880, 1760), 1))
+    for name, x in (("noise", 0.1 * torch.randn(1, RATE * 3)), ("tones", 0.2 * torch.stack([tone, 0.8 * tone]))):
+        with torch.no_grad():
+            want = model(x).numpy()
+            m, specs = mags(x)
+            masks = s.run(None, {f"mag_{n}": v.numpy() for n, v in zip(WINDOWS, m)})
+            got = sum(mrx._istft(torch.from_numpy(k) * z[:, None], HOP, x.shape[-1]) for k, z in zip(masks, specs))
+            got = got.transpose(0, 1).numpy()
+        rel = np.abs(got - want).max() / max(np.abs(want).max(), 1e-12)
+        print(f"  verify {path.name} {name}: max|diff|/max|y| = {rel:.2e} " + ("OK" if rel < tol else "FAIL"))
+        if rel >= tol:
+            raise SystemExit(f"{path}: verification failed ({rel:.2e} >= {tol})")
+
+
+def reference(model, out_dir):
+    """The reference separation: upstream's separate_soundtrack (-27 LUFS in, back out) on reference.py's mix, 20 s
+    (longer than the package's chunks are not: MRX hears the whole input at once, as upstream)."""
+    import torch
+
+    sys.path.insert(0, str(Path(__file__).parent))
+    import reference as ref
+    import separate
+
+    def run(mix):
+        with torch.no_grad():
+            y = separate.separate_soundtrack(torch.from_numpy(mix.copy()), separation_model=model)
+        return {"dialogue": y["speech"].numpy(), "music": y["music"].numpy(), "effects": y["sfx"].numpy()}
+
+    ref.write(out_dir, run, 20.0)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--out-dir", default=None, help="default: $AUDIO_NEURAL_CACHE/mrx or ~/.cache/audiojs/neural/mrx")
+    ap.add_argument("--opset", type=int, default=17)
+    ap.add_argument("--verify", action="store_true", help="compare with MRX.forward; write the reference separation")
+    args = ap.parse_args()
+
+    cache = Path(os.environ.get("AUDIO_NEURAL_CACHE") or Path.home() / ".cache" / "audiojs" / "neural")
+    out_dir = Path(args.out_dir) if args.out_dir else cache / "mrx"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    src = out_dir / "src"
+    src.mkdir(exist_ok=True)
+    try:
+        model = load(src)
+    except ImportError as e:
+        raise SystemExit(f"torch is required: pip install torch onnx onnxruntime pyloudnorm ({e})")
+    path = out_dir / "mrx.onnx"
+    export(model, path, args.opset)
+    if args.verify:
+        verify(model, path, 1e-4)
+        reference(model, out_dir)
+    print(f"done: mrx in {out_dir}")
+
+
+if __name__ == "__main__":
+    sys.exit(main())