@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 +5 -3
- package/README.md +86 -10
- package/index.d.ts +10 -1
- package/package.json +2 -1
- package/scripts/compact.py +351 -0
- package/scripts/export-scnet.py +6 -6
- package/separate.js +90 -29
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
|
|
60
|
-
|
|
61
|
-
|
|
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 |
|
|
30
|
-
| `'scnet'` | SCNet: the same at half the width | 43 MB | MIT | 9.88 · 9.43 · 8.35 · 6.15 | 0.70 ·
|
|
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
|
|
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
|
|
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
|
|
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,
|
|
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
|
+
"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())
|
package/scripts/export-scnet.py
CHANGED
|
@@ -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
|
|
191
|
-
|
|
192
|
-
|
|
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
|
-
|
|
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
|
|
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,
|
|
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
|
-
//
|
|
375
|
-
//
|
|
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
|
-
//
|
|
409
|
-
//
|
|
410
|
-
|
|
411
|
-
|
|
412
|
-
|
|
413
|
-
let [
|
|
414
|
-
let
|
|
415
|
-
|
|
416
|
-
let
|
|
417
|
-
|
|
418
|
-
|
|
419
|
-
|
|
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
|
-
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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)
|