@audio/neural-separate 0.3.0 → 0.4.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/NOTICE CHANGED
@@ -56,9 +56,11 @@ The MIT License, for each notice above:
56
56
  OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
57
57
  DEALINGS IN THE SOFTWARE.
58
58
 
59
- No model weights are included. The export scripts write them to the local cache; their terms are
60
- their authors' (Open-Unmix umxhq: MIT; SCNet and SCNet-large: MIT, Copyright (c) 2024 starrytong,
61
- redistributable converted with attribution, github.com/starrytong/SCNet/issues/35; Demucs htdemucs
59
+ No model weights are included. The export scripts write them to the local cache, and the compact
60
+ files of SCNet-large, SCNet, MRX and TIGER are fetched from huggingface.co/audiojs (scnet-large, scnet,
61
+ mrx, tiger-dnr, each with its licence and model card); their terms are their authors' (Open-Unmix
62
+ umxhq: MIT; SCNet and SCNet-large: MIT, Copyright (c) 2024 starrytong, redistributable converted with
63
+ attribution, github.com/starrytong/SCNet/issues/35; Demucs htdemucs
62
64
  weights: see the README; MRX: MIT, Copyright 2023 Mitsubishi Electric Research Laboratories (MERL), the
63
65
  repository's .reuse/dep5 naming each checkpoint; TIGER's DnR model: Apache 2.0, its model card,
64
66
  huggingface.co/JusperLee/TIGER-DnR).
package/README.md CHANGED
@@ -26,8 +26,8 @@ In the browser `@audio/neural-runtime` runs `onnxruntime-web` instead (`device:
26
26
  | `'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 |
27
27
  | `'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 |
28
28
  | `'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 |
29
- | `'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 |
30
- | `'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 |
29
+ | `'scnet-large'` | SCNet-large: band-split convolutions around dual-path LSTMs, on the complex spectrogram | 45 MB (export 169 MB) | MIT | 11.00 · 10.27 · 8.21 · 6.87 | 2.1 · 7.8 · 5.0 |
30
+ | `'scnet'` | SCNet: the same at half the width | 13 MB (export 43 MB) | MIT | 9.88 · 9.43 · 8.35 · 6.15 | 0.70 · 1.7 · 0.91 |
31
31
 
32
32
  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.
33
33
 
@@ -35,10 +35,10 @@ SDR as measured below (Verification), identical to the Python originals'. The pr
35
35
 
36
36
  A soundtrack's stems, Divide and Remaster's three (the cocktail fork problem: Petermann, Wichern, Wang, Le Roux, ICASSP 2022):
37
37
 
38
- | `model` | Architecture | Files | Weights | Targets | Real-time factor, Node CPU |
38
+ | `model` | Architecture | Files | Weights | Targets | Real-time factor: Node CPU · browser wasm · WebGPU |
39
39
  |---|---|---|---|---|---|
40
- | `'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 |
41
- | `'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 |
40
+ | `'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 | 31 MB (export 122 MB) | MIT | `dialogue`, `music`, `effects` | 0.21 · 0.31 · 0.75 |
41
+ | `'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 | 7.3 MB (export 29 MB) | Apache 2.0 | `dialogue`, `music`, `effects` | 7.5–23 · 9.2 · 3.1 |
42
42
 
43
43
  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.
44
44
 
@@ -53,11 +53,24 @@ Measured on Divide and Remaster v3's English test set (Watcharasupat, Wu, Orife,
53
53
 
54
54
  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.
55
55
 
56
+ The browser's, the shipped files (Size, below) on 30 s (a song; a Divide and Remaster clip), model load included:
57
+ onnxruntime-web 1.30 in headless Chromium 153, wasm on its default 4 threads, WebGPU on Metal, the faster of two runs
58
+ on a 14-core M4 Max shared with other jobs (load averages 110 to 300; the other run, at 9 to 140 with another job
59
+ training on the GPU, took 1.5 to 2 times as long), so upper bounds. Peak memory, the browser's processes
60
+ together (0.9 GB of it the browser idle): wasm `scnet` 3.0 GB, `scnet-large` 4.6, `mrx` 3.0, `tiger` 3.5; WebGPU 2.1,
61
+ 2.4, 3.4, 1.5. On WebGPU `tiger`'s session takes 13 s to make (38,817 nodes), a 60 s clip about 3 minutes, MRX's 45 s;
62
+ on wasm 9 and 0.3 minutes. Float16 compute does not pay: `tiger`'s on WebGPU is 13 % faster and its music 4.7 dB from
63
+ the float32 graph's.
64
+
56
65
  `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.
57
66
 
58
67
  ## Weights
59
68
 
60
- Nothing ships in this package and nothing is fetched by default. Export the weights once from the original checkpoints:
69
+ Nothing ships in this package. `scnet-large`, `scnet`, `mrx` and `tiger` read a compact file each, `<model>.int8.onnx`
70
+ (Size, below), hosted on Hugging Face (`audiojs/scnet-large`, `audiojs/scnet`, `audiojs/mrx`, `audiojs/tiger-dnr`, each
71
+ with its model card and licence): without `weights`, `https://huggingface.co/<repo>/resolve/<revision>/<file>` at the
72
+ revision `REVISIONS` pins, fetched once and checked against its SHA-256 (the browser keeps it in the Cache API, Node in
73
+ the neural cache, below). Every model can also be exported from its original checkpoint:
61
74
 
62
75
  ```sh
63
76
  pip install torch openunmix onnx onnxruntime onnxscript
@@ -74,9 +87,11 @@ python3 node_modules/@audio/neural-separate/scripts/export-mrx.py --verify
74
87
 
75
88
  pip install safetensors
76
89
  python3 node_modules/@audio/neural-separate/scripts/export-tiger.py --verify
90
+
91
+ python3 node_modules/@audio/neural-separate/scripts/compact.py --model tiger --verify # tiger.int8.onnx from tiger.onnx
77
92
  ```
78
93
 
79
- 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.
94
+ They land in `$AUDIO_NEURAL_CACHE/<model>/`, default `~/.cache/audiojs/neural/<model>/`, where a preset looks in Node: its compact file when it is the hosted one (its SHA-256), else its export (float32, the numbers the Python original gives), else the hosted file, fetched there once and written whole (offline, a compact file of another SHA-256, `compact.py`'s own build, is read as it is); a missing file throws, naming itself and the script. `weights` a directory: its files as they are, the compact one first. 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-scnet.py --fp16` float16 weights with float32 compute (`compact.py --as fp16`).
80
95
 
81
96
  - **`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).
82
97
  - **`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.
@@ -84,6 +99,63 @@ They land in `$AUDIO_NEURAL_CACHE/<model>/`, default `~/.cache/audiojs/neural/<m
84
99
  - **`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.
85
100
  - **`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.
86
101
 
102
+ ## Size
103
+
104
+ `scnet-large`, `scnet`, `mrx` and `tiger` each read one compact file, `<model>.int8.onnx`, that `scripts/compact.py`
105
+ makes from the export: the weights stored in 8 bits (symmetric, a scale per output channel; an LSTM's per gate row and
106
+ direction), the few whose rounding costs most in 16, computed in float32: onnxruntime folds each weight's Cast and scale
107
+ at load, on every backend, so the session holds float32 weights and runs as fast as the export. A variant is kept when
108
+ its cost is negligible: every stem's median within 0.05 dB of the export's (SDR on the 50 MUSDB18 test previews, SNR on
109
+ Divide and Remaster v3 test clips) and no song or clip worse by 0.5 dB or more; the smallest such is shipped.
110
+
111
+ | `model` | export (float32) | float16 weights | int8, every weight | **shipped**: int8, the costliest float16 | its gzip · brotli |
112
+ |---|---|---|---|---|---|
113
+ | `'scnet-large'` | 169.2 MB | 86.5 MB | 44.7 MB | **45.1 MB**: 83 of 88 weights int8 | 39.8 · 39.0 MB |
114
+ | `'scnet'` | 42.8 MB | 23.2 MB | 12.8 MB | **12.9 MB**: 77 of 82 | 10.7 · 10.2 MB |
115
+ | `'mrx'` | 122.3 MB | 61.5 MB | 31.3 MB | **31.3 MB**: all 39 | 27.5 · 26.6 MB |
116
+ | `'tiger'` | 28.6 MB (18.8 folded) | 10.9 MB | 7.3 MB | **7.3 MB**: all 444 | 5.2 · 2.6 MB |
117
+
118
+ Against the export, on the same songs or clips: the median per stem, then the largest loss on one song or clip, dB; and
119
+ the graph's own output against the export's on a calibration excerpt (two MUSDB18 training previews; a Divide and
120
+ Remaster v3 tuning clip), SNR.
121
+
122
+ | `model`, measure | export | float16 weights | int8, every weight | shipped |
123
+ |---|---|---|---|---|
124
+ | `'scnet-large'`, SDR: vocals · drums · bass · other | 11.00 · 10.27 · 8.21 · 6.87 | 10.99 · 10.27 · 8.21 · 6.87 | | **10.96 · 10.26 · 8.20 · 6.92** |
125
+ | … the song that lost most | | −0.004 | | −0.43 |
126
+ | … output, calibration | | 62.3 dB | 24.0 dB | 43.7 dB |
127
+ | `'scnet'`, SDR | 9.88 · 9.43 · 8.35 · 6.15 | 9.88 · 9.43 · 8.35 · 6.15 | 9.87 · 9.44 · 8.32 · 6.13 | **9.88 · 9.44 · 8.35 · 6.14** |
128
+ | … the song that lost most | | −0.004 | −2.61 | −0.29 |
129
+ | … output, calibration | | 60.6 dB | 23.4 dB | 42.4 dB |
130
+ | `'mrx'`, SNR, 30 clips: dialogue · music · effects | 10.92 · 5.17 · 5.72 | 10.92 · 5.17 · 5.72 | = shipped | **10.92 · 5.17 · 5.70** |
131
+ | … the clip that lost most | | −0.002 | | −0.04 |
132
+ | … output (masked spectra), calibration | | 73.8 dB | 40.4 dB | 40.4 dB |
133
+ | `'tiger'`, SNR, the same 30 clips | 12.68 · 10.24 · 8.06 | | = shipped | **12.66 · 10.23 · 8.03** |
134
+ | … the clip that lost most | | | | −0.10 |
135
+ | … output, calibration | | 73.9 dB | 41.4 dB | 41.4 dB |
136
+
137
+ Per song or clip the median change is within 0.015 dB everywhere; the largest losses fall where a stem is near silence:
138
+ SCNet's and SCNet-large's on PR - Happy Daze, whose vocals they separate at −2 dB SDR as exported (`scnet` with all 82
139
+ weights in int8: −2.6 dB there).
140
+ The remixes (`rebalance`'s, `scene`'s) move as little: vocals +6 dB 20.21 → 20.22 (`scnet-large`), dialogue +6 dB
141
+ 18.29 → 18.30 (`mrx`).
142
+
143
+ Which weights stay float16: each weight rounded alone to int8 on a calibration excerpt (two MUSDB18 training previews;
144
+ a Divide and Remaster v3 tuning clip), the others float32, its cost the error it adds over the output's power; the
145
+ weights go to int8 in order of that cost per byte saved while the costs' sum stays 40 dB under the output (separate
146
+ weights' costs add: the sum predicted the whole to 0.2 dB). SCNet keeps the five layers ending its decoder in float16,
147
+ 1 % of its values (`decoder.2.0`'s convolution, `decoder.1.1`'s three transposed convolutions, `decoder.2.1`'s first:
148
+ each alone at int8 26.8 to 39.0 dB); MRX (40.4 dB) and TIGER (41.4 dB) keep none.
149
+
150
+ Tried and left out: 4-bit weights (blocks of 32: SCNet's stems 7 dB from float32's; its LSTMs and linear layers alone,
151
+ 18 dB); int8 scales per block of 32 values rather than per channel (31.5 dB against 26.3: still the same decoder
152
+ layers); DequantizeLinear for the int8 weights (onnxruntime 1.30 fuses it with the MatMul it feeds into MatMulNBits,
153
+ which quantizes the activations to int8 on a CPU: MRX 38.4 dB from float32 so, 40.4 as stored).
154
+
155
+ Hosts send these files as they are: Hugging Face's CDN and raw.githubusercontent.com set no Content-Encoding on a binary
156
+ (both answer any origin, `Access-Control-Allow-Origin: *`). gzip would save 12 to 17 % of SCNet's and MRX's, 29 % of
157
+ TIGER's, brotli 14 to 21 % and 65 % (TIGER's graph, 38,817 nodes, is most of its file): not worth a decoder in the page.
158
+
87
159
  ## Algorithm
88
160
 
89
161
  **`'openunmix'`, `'mask'`** (spectrogram models):
@@ -111,12 +183,16 @@ TIGER's DnR model runs as `'complex'` with its own segments, those of `TIGERDNR.
111
183
  | 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 |
112
184
  | Pipeline vs `openunmix.Separator` in float64, 9 s reference mix (`test.js`) | 112–134 dB SNR per stem |
113
185
  | Pipeline vs `demucs.apply.apply_model`, same mix (`test.js`) | htdemucs 80–84 dB, htdemucs_ft 78–89 dB SNR per stem |
114
- | 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 |
186
+ | ONNX vs `SCNet.forward`, noise and tones (`export-scnet.py --verify`) | scnet-large max \|diff\| ≤ 2.8e-6 of max \|y\|, scnet ≤ 2.2e-6 |
115
187
  | 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 |
116
188
  | ONNX vs `MRX.forward`, noise and tones (`export-mrx.py --verify`) | max \|diff\| ≤ 1.2e-6 of max \|y\| |
117
189
  | 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 |
118
190
  | 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\| |
119
191
  | Pipeline vs `TIGERDNR.wav_chunk_inference`, its segments through the verified graph, a 20 s mix (`test.js`) | 95–140 dB SNR per stem |
192
+ | Compact file vs export, a calibration excerpt (`compact.py --verify`) | scnet-large 43.7 dB, scnet 42.4, mrx 40.4, tiger 41.4 SNR (the worst output); folding alone: bit for bit |
193
+ | Pipeline, compact file vs `reference.py`'s stems, the same 20 s mix (`test.js`) | each stem's error under the mixture: scnet-large 25.7 dB (its other and vocals, which trade the synthetic voice; drums and bass 51), scnet 47.0, mrx 45.6, tiger 58.9 |
194
+ | Hosted files: the pinned revisions' SHA-256 (`test.js`, online); a fetched file of another refused | as `models[name].sha256` |
195
+ | Browser (onnxruntime-web, Chromium) vs Node, the same compact file, 30 s | `rebalance` and `scene` through `audio`, 8 s, the files fetched from Hugging Face: the output 135 and 145 dB from Node's on wasm, 134 and 148 dB on WebGPU |
120
196
  | `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 π |
121
197
 
122
198
  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):
@@ -174,7 +250,7 @@ let vocals = complex.vocals.map(ch => istft(ch, { length: left.length }))
174
250
  | `sampleRate` | — | required unless `audio` is `{ channelData, sampleRate }` |
175
251
  | `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) |
176
252
  | `targets` | all | the targets to return |
177
- | `weights` | `$AUDIO_NEURAL_CACHE` or `~/.cache/audiojs/neural` (Node) | where a preset's files are: URL, or a directory in Node |
253
+ | `weights` | `$AUDIO_NEURAL_CACHE` or `~/.cache/audiojs/neural` (Node); a compact preset's hosted file (`REVISIONS`) | where a preset's files are: URL, or a directory in Node |
178
254
  | `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 |
179
255
  | `wiener` | `1` | EM iterations; `0` = raw masks. Ignored for `'hybrid'` and `'waveform'` |
180
256
  | `wienerWindow` | `300` | frames per EM window |
@@ -210,7 +286,7 @@ Audit any weight source yourself before shipping it — this table reflects what
210
286
 
211
287
  ## Reference
212
288
 
213
- 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).
289
+ 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, Zhu, Chen, Kang, Jiang, Li, Wu, Meng, "SCNet: Sparse Compression Network for Music Source Separation," *ICASSP* 2024 (arXiv:2401.13276). · 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).
214
290
 
215
291
  **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>
216
292
  **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 CHANGED
@@ -62,11 +62,20 @@ export interface ModelPreset {
62
62
  chunk?: number
63
63
  /** the export script that writes its weights */
64
64
  script?: string
65
+ /** its compact file (scripts/compact.py: int8 weights, a few float16, float32 compute), read before <name>.onnx */
66
+ file?: string
67
+ /** the Hugging Face repository hosting it */
68
+ repo?: string
69
+ /** the compact file's SHA-256, checked when it is fetched */
70
+ sha256?: string
65
71
  }
66
72
 
67
73
  /** Model presets by name */
68
74
  export const models: Record<ModelName, ModelPreset>
69
75
 
76
+ /** The hosted revision (commit) of each compact preset's repository: https://huggingface.co/<repo>/resolve/<revision>/<file> is read without opts.weights in the browser, and in Node when the cache holds neither the file nor its export. Empty: not hosted. */
77
+ export const REVISIONS: Record<'scnet-large' | 'scnet' | 'mrx' | 'tiger', string>
78
+
70
79
  /** A neural-runtime-shaped session — enough of it to drive separate() with a test double. */
71
80
  export interface Session {
72
81
  run(feeds: Record<string, { data: Float32Array; dims: number[]; type: string }>): Promise<Record<string, { data: Float32Array; dims: number[] }>>
@@ -84,7 +93,7 @@ export interface SeparateOptions {
84
93
  modelType?: ModelType
85
94
  /** subset of the model's targets to return; spectral models skip the other graphs, and one target runs Wiener EM against the residual */
86
95
  targets?: string[]
87
- /** where a preset's files are: URL, or a directory in Node (default $AUDIO_NEURAL_CACHE or ~/.cache/audiojs/neural) */
96
+ /** where a preset's files are: URL, or a directory in Node (default $AUDIO_NEURAL_CACHE or ~/.cache/audiojs/neural in Node; a compact preset's hosted file, REVISIONS) */
88
97
  weights?: string
89
98
  /** iterations of multichannel Wiener EM refinement (default 1); 0 = raw masks. Ignored for modelType 'hybrid' and 'waveform'. */
90
99
  wiener?: number
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@audio/neural-separate",
3
- "version": "0.3.0",
3
+ "version": "0.4.0",
4
4
  "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",
5
5
  "type": "module",
6
6
  "sideEffects": false,
@@ -22,6 +22,7 @@
22
22
  "scripts/export-mrx.py",
23
23
  "scripts/export-tiger.py",
24
24
  "scripts/reference.py",
25
+ "scripts/compact.py",
25
26
  "NOTICE"
26
27
  ],
27
28
  "dependencies": {
@@ -0,0 +1,351 @@
1
+ #!/usr/bin/env python3
2
+ """Compact an exported separation graph: its weights stored in fewer bits, its compute left in float32.
3
+
4
+ pip install onnx onnxruntime numpy
5
+ python3 compact.py --model scnet-large --verify # scnet-large.int8.onnx
6
+ python3 compact.py --model tiger --as fp16 --verify # tiger.fp16.onnx
7
+
8
+ Reads <dir>/<model>.onnx (an export-*.py output; <dir> the neural cache's <model>/) and writes <model>.<as>.onnx beside
9
+ it. The weights are the initializers a Conv, ConvTranspose, MatMul, Gemm, LSTM or GRU takes as its weight, of over 1024
10
+ values (biases, norms and constants stay float32):
11
+ int8 stored as int8 with a float32 scale per output channel (symmetric, the channel's largest magnitude to 127; an
12
+ LSTM's or GRU's per gate row of each direction), Cast and multiplied by its scale in the graph; all of them when
13
+ their rounding together costs --budget dB under the output or less (40; 0: no bound) on the calibration mix,
14
+ else int8 for the weights cheapest per byte saved (each rounded alone, the rest float32) while their summed
15
+ cost stays within the budget, float16 for the rest (separate weights' costs add: the sum predicts the whole
16
+ to a fraction of a dB)
17
+ fp16 stored as float16, Cast to float32 in the graph
18
+ fp32 as they are
19
+ onnxruntime folds the Cast and the product at load (both of constants): the session holds float32 weights and computes
20
+ as the float32 graph does, on every execution provider. DequantizeLinear would not: onnxruntime 1.30 fuses it with the
21
+ MatMul it feeds into MatMulNBits, whose default accuracy level quantizes the activations to int8 on a CPU (MRX's int8
22
+ graph 38.4 dB from float32 so, 40.4 dB as stored here).
23
+
24
+ Every variant is also folded and named short, neither of which changes a value. With each input's shape fixed (SCNet's
25
+ and TIGER's graphs; MRX's takes any length and is left as it is), every value computable from the initializers and the
26
+ shapes alone is computed once and stored as the graph computes it (SCNet's DFT matrices, made in float64 from a Range,
27
+ which onnxruntime-web's WebGPU session neither places nor folds), but for a fill (ConstantOfShape: SCNet's LSTM zero
28
+ states), kept with its shape stored. Identical initializers are stored once; node and value names become base-36
29
+ counters (TIGER's 56,652 nodes carried 11.7 MB of names), the graph's inputs and outputs keep theirs.
30
+
31
+ --verify runs the variant and the float32 graph on the calibration mix (default <dir>/test.f32, scripts/reference.py's;
32
+ --calibrate another: planar float32 stereo at 44.1 kHz), cut and transformed as the package does it, and prints each
33
+ output's largest difference over its peak and its SNR against the float32 graph's (MRX's masks as the spectra they
34
+ mask).
35
+ """
36
+ import argparse
37
+ import hashlib
38
+ import os
39
+ import sys
40
+ from pathlib import Path
41
+
42
+ import numpy as np
43
+
44
+ WEIGHT = {"Conv": (1,), "ConvTranspose": (1,), "MatMul": (1,), "Gemm": (1,), "LSTM": (1, 2), "GRU": (1, 2)}
45
+ MIN, SMALL = 1024, 1024
46
+
47
+
48
+ def users_of(g):
49
+ users = {}
50
+ for n in g.node:
51
+ for k, name in enumerate(n.input): users.setdefault(name, []).append((n, k))
52
+ return users
53
+
54
+
55
+ def weights(g):
56
+ """{name: node} of the float32 initializers of over MIN values that every user takes as its weight."""
57
+ from onnx import TensorProto
58
+
59
+ users, out = users_of(g), {}
60
+ for i in g.initializer:
61
+ uses = users.get(i.name, [])
62
+ if i.data_type != TensorProto.FLOAT or int(np.prod(i.dims)) <= MIN or not uses: continue
63
+ if all(n.op_type in WEIGHT and k in WEIGHT[n.op_type] for n, k in uses) and len({n.op_type for n, _ in uses}) == 1:
64
+ out[i.name] = uses[0][0]
65
+ return out
66
+
67
+
68
+ def q8(node, w):
69
+ """Symmetric int8 per output channel (an LSTM's or GRU's per gate row of each direction), the channel's largest
70
+ magnitude to 127: q in w's layout and its scale, shaped to broadcast against it."""
71
+ keep = (0, 1) if node.op_type in ("LSTM", "GRU") else ({"ConvTranspose": 1, "MatMul": w.ndim - 1}.get(node.op_type, 0),)
72
+ if node.op_type == "Gemm": keep = (0 if any(x.name == "transB" and x.i for x in node.attribute) else 1,)
73
+ scale = np.maximum(np.abs(w).max(axis=tuple(k for k in range(w.ndim) if k not in keep), keepdims=True), 1e-30) / 127
74
+ return np.clip(np.round(w / scale), -127, 127).astype(np.int8), scale.astype(np.float32)
75
+
76
+
77
+ def rounded(node, w):
78
+ """w through q8 and back, in float32: what the int8 graph computes with."""
79
+ q, scale = q8(node, w)
80
+ return q.astype(np.float32) * scale
81
+
82
+
83
+ def prepend(g, nodes):
84
+ body = list(g.node)
85
+ del g.node[:]
86
+ g.node.extend(nodes + body)
87
+
88
+
89
+ def store(m, kinds):
90
+ """Each weight as kinds[name] says: 'fp16' (Cast), 'int8' (Cast, Mul by its scale), else float32."""
91
+ from onnx import helper, numpy_helper, TensorProto
92
+
93
+ g = m.graph
94
+ inits = {i.name: i for i in g.initializer}
95
+ nodes = []
96
+ for name, node in weights(g).items():
97
+ kind = kinds.get(name)
98
+ if kind not in ("fp16", "int8"): continue
99
+ w = numpy_helper.to_array(inits[name])
100
+ g.initializer.remove(inits[name])
101
+ if kind == "fp16":
102
+ g.initializer.append(numpy_helper.from_array(w.astype(np.float16), name + "_h"))
103
+ nodes.append(helper.make_node("Cast", [name + "_h"], [name], to=TensorProto.FLOAT))
104
+ continue
105
+ q, scale = q8(node, w)
106
+ g.initializer.extend([numpy_helper.from_array(q, name + "_q"), numpy_helper.from_array(scale, name + "_s")])
107
+ nodes += [helper.make_node("Cast", [name + "_q"], [name + "_c"], to=TensorProto.FLOAT), helper.make_node("Mul", [name + "_c", name + "_s"], [name])]
108
+ prepend(g, nodes)
109
+ return m
110
+
111
+
112
+ def static(g):
113
+ return all(d.dim_value > 0 for i in g.input for d in i.type.tensor_type.shape.dim)
114
+
115
+
116
+ def fold(m):
117
+ """Values computable from the initializers and the input shapes alone: stored when small, else made at load."""
118
+ import onnx
119
+ import onnxruntime as ort
120
+ from onnx import helper, numpy_helper, shape_inference
121
+
122
+ g = m.graph
123
+ if not static(g): return dedupe(m)
124
+ const, cnodes = {i.name for i in g.initializer}, []
125
+ for n in g.node:
126
+ ins = [x for x in n.input if x]
127
+ if n.op_type in ("Constant", "Shape") or (ins and all(x in const for x in ins) and not n.op_type.startswith("Random")):
128
+ cnodes.append(n)
129
+ const.update(n.output)
130
+ computed = [o for n in cnodes for o in n.output]
131
+ types = {v.name: v.type.tensor_type.elem_type for v in shape_inference.infer_shapes(m).graph.value_info}
132
+ for n in cnodes:
133
+ if n.op_type == "Constant": types[n.output[0]] = next(a.t.data_type for a in n.attribute if a.name == "value")
134
+ probe = onnx.ModelProto()
135
+ probe.CopyFrom(m)
136
+ probe.graph.output.extend(helper.make_tensor_value_info(x, types[x], None) for x in computed)
137
+ s = session(probe.SerializeToString())
138
+ shapes = {i.name: [d.dim_value for d in i.type.tensor_type.shape.dim] for i in g.input}
139
+ vals = dict(zip(computed, s.run(computed, {k: np.zeros(v, np.float32) for k, v in shapes.items()})))
140
+ # the same from another input, or the graph holds a shape that depends on the data: left as it is
141
+ other = s.run(computed, {k: np.random.default_rng(0).standard_normal(v).astype(np.float32) for k, v in shapes.items()})
142
+ if any(not np.array_equal(vals[x], y, equal_nan=vals[x].dtype.kind == "f") for x, y in zip(computed, other)): return dedupe(m)
143
+ # what the rest of the graph reads, stored; a fill (ConstantOfShape) kept, its shape stored
144
+ producer = {o: n for n in cnodes for o in n.output}
145
+ inner = {id(n) for n in cnodes}
146
+ keep, stored = set(), {}
147
+ want = [x for n in g.node if id(n) not in inner for x in n.input] + [o.name for o in g.output]
148
+ while want:
149
+ x = want.pop()
150
+ if x not in producer or x in stored or id(producer[x]) in keep: continue
151
+ if producer[x].op_type == "ConstantOfShape" and vals[x].size > SMALL:
152
+ keep.add(id(producer[x]))
153
+ want += list(producer[x].input)
154
+ else: stored[x] = vals[x]
155
+ body = [n for n in g.node if id(n) not in inner or id(n) in keep]
156
+ del g.node[:]
157
+ g.node.extend(body)
158
+ g.initializer.extend(numpy_helper.from_array(np.asarray(v), x) for x, v in stored.items())
159
+ return dedupe(m)
160
+
161
+
162
+ def dedupe(m):
163
+ """Identical initializers stored once; those no node reads dropped."""
164
+ from onnx import numpy_helper
165
+
166
+ g = m.graph
167
+ seen, alias = {}, {}
168
+ for i in list(g.initializer):
169
+ a = numpy_helper.to_array(i)
170
+ key = hashlib.sha256(str((a.dtype, a.shape)).encode() + a.tobytes()).hexdigest()
171
+ if key in seen:
172
+ alias[i.name] = seen[key]
173
+ g.initializer.remove(i)
174
+ else: seen[key] = i.name
175
+ for n in g.node:
176
+ for k, x in enumerate(n.input):
177
+ if x in alias: n.input[k] = alias[x]
178
+ read = {x for n in g.node for x in n.input} | {o.name for o in g.output}
179
+ for i in [i for i in g.initializer if i.name not in read]: g.initializer.remove(i)
180
+ return m
181
+
182
+
183
+ def shorten(m):
184
+ """Node and value names as base-36 counters; the inputs and outputs keep theirs."""
185
+ g = m.graph
186
+ keep = {v.name for v in list(g.input) + list(g.output)}
187
+ names = {}
188
+
189
+ def base36(k):
190
+ s = ""
191
+ while True:
192
+ s, k = "0123456789abcdefghijklmnopqrstuvwxyz"[k % 36] + s, k // 36
193
+ if not k: return s
194
+
195
+ def short(n):
196
+ if n == "" or n in keep: return n
197
+ if n not in names: names[n] = "_" + base36(len(names))
198
+ return names[n]
199
+
200
+ for i in g.initializer: i.name = short(i.name)
201
+ for k, n in enumerate(g.node):
202
+ for j, x in enumerate(n.input): n.input[j] = short(x)
203
+ for j, x in enumerate(n.output): n.output[j] = short(x)
204
+ n.name, n.doc_string = base36(k), "" # unique: onnxruntime-web's partitioner reports nodes by name
205
+ del g.value_info[:]
206
+ m.doc_string = ""
207
+ return m
208
+
209
+
210
+ # ------------------------------------------------------------------ calibration
211
+
212
+ def session(model, fold=False):
213
+ import onnxruntime as ort
214
+
215
+ o = ort.SessionOptions()
216
+ o.enable_cpu_mem_arena = False
217
+ o.enable_mem_pattern = False
218
+ o.log_severity_level = 3 # the overridable weights' warnings
219
+ if not fold: o.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
220
+ return ort.InferenceSession(model if isinstance(model, bytes) else str(model), o, providers=["CPUExecutionProvider"])
221
+
222
+
223
+ def feeds(g, mix):
224
+ """The graph's inputs from a mixture [2, samples] at 44.1 kHz, as separate.js cuts and transforms it: SCNet one
225
+ 11 s segment (no window, 1/sqrt(n)), TIGER one 12 s segment of the first channel (Hann), MRX 10 s of each
226
+ channel at three resolutions (Hann, 1/sqrt(n), magnitudes)."""
227
+ ins = {i.name: [d.dim_value or None for d in i.type.tensor_type.shape.dim] for i in g.input}
228
+
229
+ def spec(x, n, hop, win, scale):
230
+ pad = np.pad(x, ((0, 0), (n // 2, n // 2)), mode="reflect")
231
+ T = 1 + (pad.shape[1] - n) // hop
232
+ fr = np.stack([pad[:, t * hop:t * hop + n] for t in range(T)], 1) * win
233
+ return np.fft.rfft(fr, axis=-1).transpose(0, 2, 1) * scale # [C, F, T]
234
+
235
+ hann = lambda n: 0.5 - 0.5 * np.cos(2 * np.pi * np.arange(n) / n)
236
+ if list(ins) == ["mix_spec"]:
237
+ _, C2, F, T = ins["mix_spec"]
238
+ n = 2 * (F - 1)
239
+ tiger = n == 2048
240
+ x = mix[:C2 // 2, :(T - 1) * (n // 4)]
241
+ X = spec(x, n, n // 4, hann(n) if tiger else 1.0, 1.0 if tiger else n ** -0.5)[..., :T]
242
+ X = np.pad(X, ((0, 0), (0, 0), (0, T - X.shape[2])))
243
+ return {"mix_spec": np.stack([X.real, X.imag], 1).reshape(1, C2, F, T).astype(np.float32)}
244
+ x = mix[:, :441000]
245
+ return {k: np.abs(spec(x, int(k.split("_")[1]), 256, hann(int(k.split("_")[1])), int(k.split("_")[1]) ** -0.5)).astype(np.float32) for k in ins}
246
+
247
+
248
+ def outputs(s, feed):
249
+ """The outputs as the stems see them: a mask times the magnitudes it masks."""
250
+ ys = dict(zip([o.name for o in s.get_outputs()], s.run(None, feed)))
251
+ return {k: (y * feed["mag_" + k[5:]][:, None] if k.startswith("mask_") else y).astype(np.float64) for k, y in ys.items()}
252
+
253
+
254
+ def noise(y, y0):
255
+ """Error power over signal power, the worst output's."""
256
+ return max(((y[k] - y0[k]) ** 2).sum() / max((y0[k] ** 2).sum(), 1e-300) for k in y0)
257
+
258
+
259
+ def costs(m, feed, W):
260
+ """The cost at int8 of all the weights rounded together, and a function giving each one's alone (the others
261
+ float32): noise power over output power. The weights are fed as overridable initializers, one session for all."""
262
+ import onnx
263
+ from onnx import helper, numpy_helper
264
+
265
+ inits = {i.name: i for i in m.graph.initializer}
266
+ probe = onnx.ModelProto()
267
+ probe.CopyFrom(m)
268
+ probe.graph.input.extend(helper.make_tensor_value_info(n, inits[n].data_type, list(inits[n].dims)) for n in W)
269
+ s = session(probe.SerializeToString())
270
+ y0 = outputs(s, feed)
271
+ q = {n: rounded(node, numpy_helper.to_array(inits[n])) for n, node in W.items()}
272
+ return noise(outputs(s, {**feed, **q}), y0), lambda n: noise(outputs(s, {**feed, n: q[n]}), y0)
273
+
274
+
275
+ def int8(m, feed, budget):
276
+ """{weight: 'int8' | 'fp16'}: every weight int8 when their rounding together costs budget dB under the output or
277
+ less (no budget: always); else int8 for the weights cheapest per byte saved, each rounded alone, while their summed
278
+ cost stays within the budget (costs of separate weights add), float16 for the rest."""
279
+ W = weights(m.graph)
280
+ if budget is None: return {n: "int8" for n in W}
281
+ whole, alone = costs(m, feed, W)
282
+ print(f" every weight int8: SNR {-10 * np.log10(whole):.1f} dB against float32")
283
+ if whole <= 10 ** (-budget / 10): return {n: "int8" for n in W}
284
+ size = {i.name: int(np.prod(i.dims)) for i in m.graph.initializer}
285
+ cost = {}
286
+ for k, n in enumerate(W):
287
+ cost[n] = alone(n)
288
+ print(f"\r cost {k + 1}/{len(W)}", end="", file=sys.stderr)
289
+ print(file=sys.stderr)
290
+ kinds, total = {}, 0.0
291
+ for n in sorted(W, key=lambda n: cost[n] / size[n]):
292
+ if total + cost[n] <= 10 ** (-budget / 10): kinds[n], total = "int8", total + cost[n]
293
+ else: kinds[n] = "fp16"
294
+ for n in sorted((n for n in kinds if kinds[n] == "fp16"), key=lambda n: -cost[n]):
295
+ print(f" float16: {W[n].name or n}, {size[n]} values, alone at int8 {-10 * np.log10(max(cost[n], 1e-30)):.1f} dB")
296
+ print(f" int8: {sum(k == 'int8' for k in kinds.values())} of {len(kinds)} weights, "
297
+ f"{sum(size[n] for n in kinds if kinds[n] == 'int8') / sum(size[n] for n in W):.1%} of their values; predicted SNR {-10 * np.log10(max(total, 1e-30)):.1f} dB")
298
+ return kinds
299
+
300
+
301
+ def compact(path, out, kind, calibration=None, budget=40.0):
302
+ import onnx
303
+
304
+ m = onnx.load(str(path))
305
+ # the trained weights, chosen before folding (never a folded constant)
306
+ kinds = int8(m, feeds(m.graph, calibration), budget) if kind == "int8" else {n: kind for n in weights(m.graph)}
307
+ m = shorten(store(fold(m), kinds))
308
+ m.metadata_props.append(onnx.StringStringEntryProto(key="compact", value=kind + (f", budget {budget:g} dB" if kind == "int8" and budget else "")))
309
+ onnx.checker.check_model(m)
310
+ onnx.save(m, str(out))
311
+ return out
312
+
313
+
314
+ def verify(path, base, mix):
315
+ import onnx
316
+
317
+ a, b = session(base, True), session(path, True)
318
+ feed = feeds(onnx.load(str(base), load_external_data=False).graph, mix)
319
+ y0, y = outputs(a, feed), outputs(b, feed)
320
+ for k in y0:
321
+ d = y[k] - y0[k]
322
+ rel = np.abs(d).max() / np.abs(y0[k]).max()
323
+ print(f" {path.name} {k}: max|diff|/max|y| = {rel:.2e}, SNR {10 * np.log10((y0[k] ** 2).sum() / max((d ** 2).sum(), 1e-300)):.1f} dB against float32")
324
+
325
+
326
+ def main():
327
+ ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
328
+ ap.add_argument("--model", required=True, help="scnet-large | scnet | mrx | tiger (any <dir>/<model>.onnx)")
329
+ ap.add_argument("--as", dest="kind", default="int8", choices=["fp32", "fp16", "int8"])
330
+ ap.add_argument("--budget", type=float, default=40, help="int8: the rounding's cost, dB under the output (40); 0: every weight int8")
331
+ ap.add_argument("--dir", default=None, help="default: $AUDIO_NEURAL_CACHE/<model> or ~/.cache/audiojs/neural/<model>")
332
+ ap.add_argument("--out", default=None, help="default: <dir>/<model>.<as>.onnx")
333
+ ap.add_argument("--calibrate", default=None, help="planar float32 stereo at 44.1 kHz (default <dir>/test.f32)")
334
+ ap.add_argument("--verify", action="store_true", help="the variant against the float32 graph on the calibration mix")
335
+ args = ap.parse_args()
336
+
337
+ cache = Path(os.environ.get("AUDIO_NEURAL_CACHE") or Path.home() / ".cache" / "audiojs" / "neural")
338
+ d = Path(args.dir) if args.dir else cache / args.model
339
+ path = d / f"{args.model}.onnx"
340
+ out = Path(args.out) if args.out else d / f"{args.model}.{args.kind}.onnx"
341
+ cal = Path(args.calibrate) if args.calibrate else d / "test.f32"
342
+ mix = np.fromfile(cal, np.float32).reshape(2, -1) if cal.exists() else None
343
+ budget = args.budget or None
344
+ if mix is None and ((args.kind == "int8" and budget) or args.verify): raise SystemExit(f"{cal}: no calibration mix (export with --verify writes test.f32)")
345
+ compact(path, out, args.kind, mix, budget)
346
+ print(f"wrote {out} ({out.stat().st_size / 1e6:.2f} MB, from {path.stat().st_size / 1e6:.2f})")
347
+ if args.verify: verify(out, path, mix)
348
+
349
+
350
+ if __name__ == "__main__":
351
+ sys.exit(main())
@@ -25,7 +25,7 @@ The STFT is SCNet's own: n_fft 4096, hop 1024, no window (torch.stft's default,
25
25
  ones), normalized (1/sqrt(4096)), centered with reflect padding.
26
26
 
27
27
  Writes <out-dir>/<model>.onnx (scnet-large 169 MB, scnet 43 MB) and, with --fp16,
28
- <model>.fp16.onnx (half that), weights embedded. Default out-dir:
28
+ <model>.fp16.onnx (half that: float16 weights, float32 compute, scripts/compact.py's), weights embedded. Default out-dir:
29
29
  $AUDIO_NEURAL_CACHE/<model>, else ~/.cache/audiojs/neural/<model>. --verify compares the stems rebuilt from the ONNX
30
30
  output with SCNet.forward on the same segment, and writes the reference
31
31
  separation test.js compares the JS pipeline against (scripts/reference.py),
@@ -187,11 +187,11 @@ def export(model, path, opset):
187
187
 
188
188
 
189
189
  def fp16(path, out):
190
- """float16 weights, float32 inputs and outputs (onnxruntime's converter, as export-openunmix.py's)."""
191
- import onnx
192
- from onnxruntime.transformers.float16 import convert_float_to_float16
190
+ """float16 weights, float32 compute (scripts/compact.py --as fp16)."""
191
+ sys.path.insert(0, str(Path(__file__).parent))
192
+ import compact
193
193
 
194
- onnx.save(convert_float_to_float16(onnx.load(str(path)), keep_io_types=True), str(out))
194
+ compact.compact(path, out, "fp16")
195
195
  print(f"wrote {out} ({out.stat().st_size / 1e6:.1f} MB)")
196
196
 
197
197
 
@@ -245,7 +245,7 @@ def main():
245
245
  ap.add_argument("--model", default="scnet-large", choices=list(MODELS))
246
246
  ap.add_argument("--out-dir", default=None, help="default: $AUDIO_NEURAL_CACHE/<model> or ~/.cache/audiojs/neural/<model>")
247
247
  ap.add_argument("--opset", type=int, default=17)
248
- ap.add_argument("--fp16", action="store_true", help="also write <model>.fp16.onnx (float16 weights, float32 I/O)")
248
+ ap.add_argument("--fp16", action="store_true", help="also write <model>.fp16.onnx (float16 weights, float32 compute)")
249
249
  ap.add_argument("--verify", action="store_true", help="compare with SCNet.forward; write the reference separation")
250
250
  args = ap.parse_args()
251
251
 
package/separate.js CHANGED
@@ -12,7 +12,7 @@
12
12
  // of sevagh/demucs.onnx. 7.8 s segments overlap-added as demucs.apply.apply_model does; each
13
13
  // runs as waveform + complex-as-channels spectrogram in, frequency- and time-branch estimates
14
14
  // out, summed after the iSTFT.
15
- // 'complex': complex spectrogram models (SCNet: Tong, Zhang, Liu, Li, Yu, "SCNet: Sparse Compression
15
+ // 'complex': complex spectrogram models (SCNet: Tong, Zhu, Chen, Kang, Jiang, Li, Wu, Meng, "SCNet: Sparse Compression
16
16
  // Network for Music Source Separation", ICASSP 2024; the Band-Split RoFormer family has the same
17
17
  // contract) with their STFT and iSTFT out of the graph: each channel's STFT in, re and im as
18
18
  // channels, each source's out. 11 s segments overlap-added under linear fades, as
@@ -27,7 +27,7 @@
27
27
 
28
28
  import { fft, ifft } from 'fourier-transform'
29
29
  import resampleSinc from '@audio/resample-sinc'
30
- import { load as neuralLoad, tensor } from '@audio/neural-runtime'
30
+ import { load as neuralLoad, tensor, fetchModel } from '@audio/neural-runtime'
31
31
 
32
32
  const PI2 = Math.PI * 2
33
33
  const isNode = typeof process !== 'undefined' && !!process.versions?.node
@@ -371,20 +371,41 @@ const TIGER = {
371
371
  standardize: false, mono: true, script: 'export-tiger.py',
372
372
  }
373
373
 
374
- // Presets: files are <weights>/<name>/<target>.onnx (one graph per target) or
375
- // <weights>/<name>/<name>.onnx, as scripts/export-*.py write them.
374
+ // The compact files (scripts/compact.py: int8 weights, a few float16, float32 compute), each within 0.05 dB of its
375
+ // export's median SDR or SNR per stem (README, Size); their Hugging Face repositories, and SHA-256 checked when fetched
376
+ const FILES = {
377
+ 'scnet-large': { file: 'scnet-large.int8.onnx', repo: 'audiojs/scnet-large', sha256: 'b2dc586a1e0e6c0afe4915e9057ea29111397b7bbf589edc3d85de91de79cd72' },
378
+ scnet: { file: 'scnet.int8.onnx', repo: 'audiojs/scnet', sha256: '98228931494151762a1c4ab1ec7899a894b1f81fd4a509921a9d2f9bacc50845' },
379
+ mrx: { file: 'mrx.int8.onnx', repo: 'audiojs/mrx', sha256: 'd876c92d224f5d8cd2207d528b23369da895a6de71fb24132416fcb16f9f8424' },
380
+ tiger: { file: 'tiger.int8.onnx', repo: 'audiojs/tiger-dnr', sha256: '2189859ac7e89cf4d9d284c6f393e737aaa2b3233ed7ee424ad622012bce7bca' },
381
+ }
382
+
383
+ // The hosted revision of each compact file's repository (a commit), its URL https://huggingface.co/<repo>/resolve/<revision>/
384
+ // <file>: what a compact preset reads without opts.weights, in the browser, and in Node fetched once into the neural cache
385
+ // (unless the cache holds it or the export). Empty: not hosted.
386
+ export const REVISIONS = {
387
+ 'scnet-large': '95f22d9e68f6cab78aef358f1a0de248d4edf805',
388
+ scnet: 'e71bcdb42bce1d62b7fae25ba90fcbcf90f73f46',
389
+ mrx: 'dedc823a0bea2ed085d19e8b7f32589d35169ee8',
390
+ tiger: '1c959ab7ff8794dc5663b104ba301687ab95d2ab',
391
+ }
392
+
393
+ const hosted = name => models[name].repo && REVISIONS[name] ? `https://huggingface.co/${models[name].repo}/resolve/${REVISIONS[name]}/${models[name].file}` : null
394
+
395
+ // Presets: files are <weights>/<name>/<target>.onnx (one graph per target) or <weights>/<name>/<name>.onnx, as
396
+ // scripts/export-*.py write them; a compact preset's <weights>/<name>/<file> first, as scripts/compact.py writes it.
376
397
  export const models = {
377
398
  umxhq: { modelType: 'openunmix', sampleRate: 44100, targets: ['vocals', 'drums', 'bass', 'other'], perTarget: true },
378
399
  htdemucs: { modelType: 'hybrid', sampleRate: 44100, targets: DEMUCS },
379
400
  htdemucs_ft: { modelType: 'hybrid', sampleRate: 44100, targets: DEMUCS, perTarget: true },
380
401
  // SCNet-large and SCNet (starrytong/SCNet, MIT weights) as scripts/export-scnet.py exports them
381
- 'scnet-large': SCNET,
382
- scnet: SCNET,
402
+ 'scnet-large': { ...SCNET, ...FILES['scnet-large'] },
403
+ scnet: { ...SCNET, ...FILES.scnet },
383
404
  // MRX (merlresearch/cocktail-fork-separation, MIT weights) as scripts/export-mrx.py exports it: dialogue, music
384
405
  // and effects of a soundtrack (Divide and Remaster)
385
- mrx: MRX,
406
+ mrx: { ...MRX, ...FILES.mrx },
386
407
  // TIGER (JusperLee/TIGER, Apache-2.0 weights): the same three stems, three band-split models, about 50 times MRX's time
387
- tiger: TIGER,
408
+ tiger: { ...TIGER, ...FILES.tiger },
388
409
  }
389
410
 
390
411
  async function cacheDir() {
@@ -405,18 +426,40 @@ async function weightsBase(opts, name) {
405
426
  return pathToFileURL(w.endsWith('/') ? w : w + '/').href
406
427
  }
407
428
 
408
- // Weights produced locally must exist before a session loads them: name the missing file and
409
- // how to make it, rather than surfacing ENOENT from deep inside the runtime.
410
- async function checkLocal(graphs, name) {
411
- let files = graphs.map(g => g.spec).filter(s => typeof s === 'string' && s.startsWith('file:'))
412
- if (!files.length) return
413
- let [{ existsSync }, { fileURLToPath }] = await Promise.all([import('node:fs'), import('node:url')])
414
- let missing = files.map(f => fileURLToPath(f)).filter(f => !existsSync(f))
415
- if (!missing.length) return
416
- let script = models[name].script ?? { hybrid: 'export-htdemucs.py', complex: 'export-scnet.py' }[models[name].modelType] ?? 'export-openunmix.py'
417
- throw new Error(`neural-separate: ${name} weights not found: ${missing.join(', ')}. ` +
418
- `Produce them with python3 node_modules/@audio/neural-separate/scripts/${script} --model ${name}, ` +
419
- `or pass opts.weights: a URL or directory holding ${name}/<file>.onnx`)
429
+ // A preset's files in a local directory; a file missing named with how to make it, rather than ENOENT from the runtime.
430
+ // A compact preset reads its compact file, else its export (float32). From the neural cache (no opts.weights) its compact
431
+ // file only as hosted (its SHA-256), else the export, else the hosted file fetched into the cache, checked, written whole;
432
+ // offline, a compact file of another SHA-256 (scripts/compact.py's own) as it is.
433
+ async function local(graphs, name, opts) {
434
+ let [fs, { fileURLToPath }] = await Promise.all([import('node:fs'), import('node:url')])
435
+ let file = g => typeof g.spec === 'string' && g.spec.startsWith('file:') ? fileURLToPath(g.spec) : null
436
+ let here = g => !file(g) || fs.existsSync(file(g)), p = models[name], why = ''
437
+ let whole = g => ({ ...g, spec: g.spec.replace(/[^/]*$/, `${name}.onnx`), sha256: null })
438
+ if (p.file) graphs = await Promise.all(graphs.map(async g => {
439
+ if (!file(g)) return g
440
+ let cached = opts.weights == null
441
+ if (here(g) && (!cached || await sha256(fs.readFileSync(file(g))) === p.sha256)) return { ...g, sha256: null }
442
+ if (here(whole(g))) return whole(g)
443
+ if (cached && hosted(name)) try { return { ...g, spec: await fetched(name, file(g)), sha256: null } } catch (e) { why = ` (${e.message})` }
444
+ return { ...g, sha256: null }
445
+ }))
446
+ let missing = graphs.filter(g => !here(g)).map(file)
447
+ if (!missing.length) return graphs
448
+ let script = p.script ?? { hybrid: 'export-htdemucs.py', complex: 'export-scnet.py' }[p.modelType] ?? 'export-openunmix.py'
449
+ throw new Error(`neural-separate: ${name} weights not found: ${missing.join(', ')}${why}. ` +
450
+ `Produce them with python3 node_modules/@audio/neural-separate/scripts/${script} --model ${name}` +
451
+ (p.file ? ` (and scripts/compact.py --model ${name} for ${p.file})` : '') +
452
+ `, or pass opts.weights: a URL or directory holding ${name}/<file>.onnx`)
453
+ }
454
+
455
+ // The hosted compact file into the cache, checked against its SHA-256, written whole or not at all → its file URL
456
+ async function fetched(name, dest) {
457
+ let [fs, path, { pathToFileURL }] = await Promise.all([import('node:fs'), import('node:path'), import('node:url')])
458
+ let bytes = await checked(hosted(name), models[name].sha256, { cache: false })
459
+ fs.mkdirSync(path.dirname(dest), { recursive: true })
460
+ fs.writeFileSync(dest + '.part', bytes)
461
+ fs.renameSync(dest + '.part', dest)
462
+ return pathToFileURL(dest).href
420
463
  }
421
464
 
422
465
  // opts.model → { graphs: [{ spec, sources, keep }], modelType, sampleRate, targets }. A graph's
@@ -429,14 +472,15 @@ async function resolveModel(opts) {
429
472
  let pick = names => want ? names.filter(n => want.includes(n)) : names
430
473
 
431
474
  if (typeof model === 'string' && Object.hasOwn(models, model)) {
432
- let p = models[model], base = await weightsBase(opts, model)
475
+ // the browser without opts.weights: the hosted file
476
+ let p = models[model], url = !isNode && opts.weights == null && hosted(model), base = url ? '' : await weightsBase(opts, model)
433
477
  if (want) for (let t of want) if (!p.targets.includes(t)) throw new Error(`neural-separate: ${model} has no target '${t}' (has ${p.targets.join(', ')})`)
434
478
  let targets = pick(p.targets)
435
479
  let sources = p.sources ?? (p.modelType === 'hybrid' || p.modelType === 'complex' ? DEMUCS : null)
436
480
  let graphs = p.perTarget
437
481
  ? targets.map(t => ({ spec: `${base}${model}/${t}.onnx`, sources: sources ?? [t], keep: [t] }))
438
- : [{ spec: `${base}${model}/${model}.onnx`, sources: sources ?? p.targets, keep: targets }]
439
- if (isNode) await checkLocal(graphs, model)
482
+ : [{ spec: url || `${base}${model}/${p.file ?? `${model}.onnx`}`, sources: sources ?? p.targets, keep: targets, sha256: p.sha256 }]
483
+ if (isNode) graphs = await local(graphs, model, opts)
440
484
  return { graphs, modelType: p.modelType, sampleRate: p.sampleRate, targets, preset: p }
441
485
  }
442
486
 
@@ -459,8 +503,25 @@ async function resolveModel(opts) {
459
503
  // No pooled arena: one htdemucs segment peaks at 2.7 GB RSS in onnxruntime-node instead of 3.2 GB.
460
504
  const sessionOptions = { enableCpuMemArena: false, enableMemPattern: false }
461
505
 
506
+ // A graph's session; a fetched file with a SHA-256 checked against it first
462
507
  function loader(opts) {
463
- return spec => opts.session ? opts.session(spec, opts) : neuralLoad(spec, { backend: opts.device, sessionOptions })
508
+ return async g => {
509
+ if (opts.session) return opts.session(g.spec, opts)
510
+ let spec = g.sha256 && !g.spec.startsWith('file:') ? await checked(g.spec, g.sha256) : g.spec
511
+ return neuralLoad(spec, { backend: opts.device, sessionOptions })
512
+ }
513
+ }
514
+
515
+ // A fetched file's bytes, refused unless of its SHA-256
516
+ async function checked(url, hash, opts) {
517
+ let bytes = await fetchModel(url, opts), hex = await sha256(bytes)
518
+ if (hex !== hash) throw new Error(`neural-separate: ${url} has SHA-256 ${hex}, expected ${hash}: a partial or altered download (clear the 'audio-neural' cache), or another build of the file`)
519
+ return bytes
520
+ }
521
+
522
+ async function sha256(bytes) {
523
+ let subtle = globalThis.crypto?.subtle ?? (await import('node:crypto')).webcrypto.subtle
524
+ return Array.from(new Uint8Array(await subtle.digest('SHA-256', bytes)), b => b.toString(16).padStart(2, '0')).join('')
464
525
  }
465
526
 
466
527
  function outputOf(session, out, i = 0) {
@@ -549,7 +610,7 @@ async function separateSpectral(channels, rate, opts, graphs, modelType, load) {
549
610
  let wopts = { iterations, softmask: opts.softmask ?? false, eps: opts.eps, residual }
550
611
  let sessions = []
551
612
  try {
552
- for (let g of graphs) sessions.push(await load(g.spec))
613
+ for (let g of graphs) sessions.push(await load(g))
553
614
  let C = channels.length, N = channels[0].length
554
615
  let { starts, size, weight } = crossfadeChunks(N, rate, opts)
555
616
  return await overlapAdd(N, C, starts, size, weight, async (start, len) => {
@@ -643,7 +704,7 @@ async function separateHybrid(channels, rate, opts, graphs, load) {
643
704
 
644
705
  let result = {}
645
706
  for (let g of graphs) {
646
- let session = await load(g.spec)
707
+ let session = await load(g)
647
708
  try {
648
709
  let L = session.inputs?.[0]?.dims?.[2]
649
710
  L = typeof L === 'number' && L > 0 ? L : opts.segment ?? 343980 // htdemucs: int(7.8 · 44100)
@@ -718,7 +779,7 @@ async function separateComplex(channels, rate, opts, graphs, load, p) {
718
779
  let { norm, scale: back } = opts.standardize ?? p?.standardize ?? true ? normalize(channels) : { norm: channels, scale: c => c }
719
780
  let result = {}
720
781
  for (let g of graphs) {
721
- let session = await load(g.spec)
782
+ let session = await load(g)
722
783
  try {
723
784
  let T = session.inputs?.[0]?.dims?.[3]
724
785
  T = typeof T === 'number' && T > 0 ? T : opts.frames ?? p?.frames
@@ -764,7 +825,7 @@ async function separateMultires(channels, rate, opts, graphs, load, p) {
764
825
  let windows = opts.windows ?? p?.windows, hop = opts.hop ?? p?.hop ?? 256
765
826
  let C = channels.length, N = channels[0].length, result = {}
766
827
  for (let g of graphs) {
767
- let session = await load(g.spec)
828
+ let session = await load(g)
768
829
  try {
769
830
  let { starts, size, weight } = crossfadeChunks(N, rate, { ...opts, chunk: opts.chunk ?? p?.chunk })
770
831
  Object.assign(result, await overlapAdd(N, C, starts, size, weight, async (start, len) => {
@@ -847,7 +908,7 @@ async function separateWaveform(channels, rate, opts, graphs, load) {
847
908
  let { starts, size, weight } = crossfadeChunks(N, rate, opts)
848
909
  let result = {}
849
910
  for (let g of graphs) {
850
- let session = await load(g.spec)
911
+ let session = await load(g)
851
912
  try {
852
913
  let out = await overlapAdd(N, C, starts, size, weight, async (start, len) => {
853
914
  let data = new Float32Array(C * len)