@audio/neural-separate 0.1.0

package/NOTICE

@@ -0,0 +1,36 @@
+@audio/neural-separate
+Copyright 2026 Dmitry Iv. and audiojs contributors
+
+This product includes a JavaScript translation of parts of open-unmix-pytorch
+(https://github.com/sigsep/open-unmix-pytorch) and norbert (https://github.com/sigsep/norbert):
+filtering.py's multichannel Wiener filter and its expectation-maximization, with their defaults,
+and model.Separator's 300-frame windowing, into separate.js. Both under the MIT License:
+
+    Copyright (c) 2019 Inria (Fabian-Robert Stöter, Antoine Liutkus)
+
+This product includes a JavaScript translation of parts of Demucs
+(https://github.com/facebookresearch/demucs): apply.apply_model's segment overlap-add and
+htdemucs.py's spectrogram framing, into separate.js. Under the MIT License:
+
+    Copyright (c) Meta Platforms, Inc. and affiliates.
+
+The MIT License, for each notice above:
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy of this
+    software and associated documentation files (the "Software"), to deal in the Software
+    without restriction, including without limitation the rights to use, copy, modify, merge,
+    publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons
+    to whom the Software is furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all copies or
+    substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
+    INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
+    PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
+    FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+    OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+    DEALINGS IN THE SOFTWARE.
+
+No model weights are included. The export scripts write them to the local cache; their terms are
+their authors' (Open-Unmix umxhq: MIT; Demucs htdemucs weights: see the README).

package/README.md

@@ -0,0 +1,149 @@
+# @audio/neural-separate
+
+> Source separation (stems): Open-Unmix and Hybrid Transformer Demucs through ONNX, matching the original PyTorch implementations.
+
+The ML upgrade to [`@audio/vocals`](https://github.com/audiojs/vocals)'s classical center-cancel: instead of one M/S trick, a trained model estimates each target (vocals, drums, bass, other), and the pipeline around it (STFT, Wiener refinement, chunking, overlap-add) is ported from the reference implementations and checked against them.
+
+```js
+import separate from '@audio/neural-separate'
+
+let { stems, residual } = await separate([left, right], { sampleRate: 44100, model: 'umxhq' })
+stems.vocals   // Float32Array[2]
+```
+
+## Install
+
+```
+npm install @audio/neural-separate onnxruntime-node
+```
+
+In the browser `@audio/neural-runtime` runs `onnxruntime-web` instead (`device: 'wasm'` or `'webgpu'`).
+
+## Models
+
+| `model` | Architecture | Files | Weights | SDR, MUSDB18 previews: vocals · drums · bass · other | Real-time factor: Node CPU · browser wasm · WebGPU |
+|---|---|---|---|---|---|
+| `'umxhq'` | Open-Unmix: bi-LSTM on magnitude spectrograms, multichannel Wiener EM | 4 × 35.6 MB (fp16: 4 × 17.8 MB) | MIT | 6.75 · 6.11 · 5.00 · 3.36 | 0.19 (vocals only: 0.11) · 0.12 · 0.29 |
+| `'htdemucs'` | Hybrid Transformer Demucs: waveform and spectrogram U-Nets joined by a cross-domain transformer | 174 MB | research only | 8.86 · 9.55 · 9.30 · 5.69 | 1.4 (vocals only: 1.2) · 0.70 · 0.21 |
+| `'htdemucs_ft'` | four fine-tuned HTDemucs, one per source | 4 × 174 MB | research only | not measured | vocals only: 1.1 · not measured · not measured |
+
+SDR as measured below (Verification), identical to the Python originals'. The previews are 7 s excerpts, so these numbers sit apart from full-track results in scale, not in order: the Demucs README reports overall SDR 9.0 for fine-tuned HT Demucs against 5.3 for Open-Unmix on the MUSDB18-HQ test set.
+
+Real-time factor: processing time over duration, lower is faster. 27 s of music; onnxruntime-node and onnxruntime-web 1.30, the web one in headless Chromium (wasm, Metal WebGPU); the best of two to four runs spread over a day on a 14-core M4 Max shared with other jobs (load averages 45 to 180), so upper bounds. The Python originals' best there: 0.43 (open-unmix) and 2.8 (demucs). umxhq's time is mostly JS (ONNX Runtime takes 6 to 12 %: STFT, Wiener EM, iSTFT around it), and it runs faster on wasm than on WebGPU; htdemucs's is about 80 % inference, where WebGPU pays. Memory: umxhq 1.3 GB, hybrid models 3 GB at peak in Node.
+
+`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.
+
+## Weights
+
+Nothing ships in this package and nothing is fetched by default. Export the weights once from the original checkpoints:
+
+```sh
+pip install torch openunmix onnx onnxruntime onnxscript
+python3 node_modules/@audio/neural-separate/scripts/export-openunmix.py --model umxhq --verify
+
+pip install demucs
+python3 node_modules/@audio/neural-separate/scripts/export-htdemucs.py --model htdemucs --verify
+```
+
+They land in `$AUDIO_NEURAL_CACHE/<model>/`, default `~/.cache/audiojs/neural/<model>/`, where a preset looks in Node; a missing file throws, naming itself and the script. Elsewhere, serve that directory and pass its URL: `separate(audio, { model: 'umxhq', weights: 'https://…/' })`; `@audio/neural-runtime` fetches each file once and caches it (Cache API in the browser). `--verify` compares the ONNX graphs with the PyTorch modules and writes the reference separation `test.js` checks the JS pipeline against. `export-openunmix.py --fp16` also writes float16 weights with float32 I/O, half the size.
+
+- **`export-openunmix.py`** exports each target's `OpenUnmix` through a wrapper that takes the frame count from the traced tensor (upstream reads it off `x.data.shape`, which export freezes), with the TorchScript exporter (the `torch.export` one bakes the LSTM's sequence length into a reshape). Magnitude `[1, C, F, T]` in and out, T dynamic. `umxhq`'s bandwidth restriction (`max_bin`, 1487 of 2049 bins for `n_fft=4096`@44.1kHz) crops the network's *input* below 16 kHz for efficiency; the final dense layer regresses the **full** bin range from that reduced representation (a learned extrapolation, not a literal zero-fill).
+- **`export-htdemucs.py`** follows [sevagh/demucs.onnx](https://github.com/sevagh/demucs.onnx): the STFT and iSTFT, which ONNX export cannot carry, move out of the graph. Rather than a vendored copy of `htdemucs.py`, it runs upstream's own `forward` with its four transform methods swapped on the instance. One graph per 7.8 s segment: `mix [1, 2, 343980]` and `mix_spec [1, 4, 2048, 336]` (complex as channels) in, `stems_spec [1, 4, 4, 2048, 336]` and `stems_wave [1, 4, 2, 343980]` out; a source is the iSTFT of the first plus the second.
+
+## Algorithm
+
+**`'openunmix'`, `'mask'`** (spectrogram models):
+
+1. **STFT** each channel — `n_fft` 4096, hop 1024, periodic Hann, `center=True` (torch.stft-compatible reflect padding) — own implementation on top of [`fourier-transform`](https://github.com/scijs/fourier-transform)'s raw FFT (not its bundled `stft` submodule, which zero-pads instead of centering — Open-Unmix's own filterbank is `torch.stft(..., center=True, pad_mode="reflect")`, so this package matches that framing exactly rather than reusing a differently-conventioned STFT).
+2. **Magnitude** per channel → **one ONNX run per target** (or one multi-target graph) → estimated magnitude (`modelType: 'openunmix'`) or a `[0,1]` mask multiplied by the mixture magnitude (`modelType: 'mask'`).
+3. **Multichannel Wiener EM** (`wiener` option, default 1 iteration) refines the per-target magnitude estimates into full complex spectra, using the mixture's spatial (inter-channel) structure — see below. It runs in windows of 300 frames, as open-unmix's `Separator` does (`wiener_win_len=300`), each window estimating its own spatial covariances; a single target runs against a residual (open-unmix's `residual=True`), since EM needs two sources.
+4. **iSTFT** each target's channel spectra back to waveforms.
+5. Long files are **chunked** (`chunk` seconds, default 30, with `overlap` seconds crossfade, default 2) so a bi-LSTM-class model's memory stays bounded; chunk stems are stitched back with a linear crossfade. Trade-off: the model only sees `chunk` seconds of context at once, same reasoning as Open-Unmix's own reference `Separator` batching its Wiener step into `wiener_win_len=300`-frame windows (`openunmix/model.py`) — coarser here (whole pipeline per chunk, not just the Wiener step), chosen for a simple, uniform memory bound across both spectral and waveform model families.
+
+**`'hybrid'`** (Hybrid Transformer Demucs): `demucs.apply.apply_model` with `split=True, overlap=0.25, shifts=0`, and `demucs.api`'s normalization by the whole input's mean and standard deviation. Segments of 7.8 s start every 5.85 s; each is centered in its window with the neighbouring input as context, zeros past the ends. Per segment: HTDemucs's framing (reflect re-padding by ¾ hop, `normalized=True`, the Nyquist bin and two edge frames dropped) → ONNX → frequency branch through the iSTFT plus time branch. Segments overlap-add under a triangular window. `shifts` (averaging randomly time-shifted runs, "up to 0.2 points" of SDR per demucs's docstring) is not implemented; the CLI's default single shift averages nothing.
+
+`modelType: 'waveform'` (Demucs v2-class) skips steps 1–4 entirely: chunked raw audio in `[1, C, N]`, stacked stems out `[1, S, C, N]`, no STFT or Wiener step — Demucs v2 operates in the time domain by design.
+
+## Verification
+
+| Check | Result |
+|---|---|
+| ONNX vs PyTorch module, random input (`--verify`) | umxhq max \|diff\| ≤ 5.7e-6 of max \|y\| (fp16: ≤ 5.3e-3); stems of htdemucs ≤ 1.6e-4, of htdemucs_ft ≤ 2.5e-4 |
+| Pipeline vs `openunmix.Separator` in float64, 9 s reference mix (`test.js`) | 112–134 dB SNR per stem |
+| Pipeline vs `demucs.apply.apply_model`, same mix (`test.js`) | htdemucs 80–84 dB, htdemucs_ft 78–89 dB SNR per stem |
+| `wienerFilter` vs open-unmix's `wiener()`, float64 fixture (`test.js`) | ratio mask: 3e-12; mixture phase: 1.7e-7, upstream's `atan2` adds a float32 π |
+
+The 50 test tracks of the MUSDB18 7 s previews (`musdb.DB(download=True)`; its terms are educational and non-commercial, so it serves measurement only), SDR by [museval](https://github.com/sigsep/sigsep-mus-eval) (BSSEval v4, 1 s windows, median over windows, then over tracks: the SiSEC 2018 aggregation):
+
+| | vocals | drums | bass | other | this package vs Python, waveform SNR per track (median / lowest) |
+|---|---|---|---|---|---|
+| `umxhq`, open-unmix `Separator` | 6.75 | 6.11 | 5.00 | 3.36 | |
+| `umxhq`, this package | 6.75 | 6.11 | 5.00 | 3.36 | 116–125 / 72–83 dB |
+| `htdemucs`, demucs `apply_model` | 8.86 | 9.55 | 9.30 | 5.69 | |
+| `htdemucs`, this package | 8.86 | 9.55 | 9.30 | 5.69 | 78–88 / 56–75 dB |
+
+Per track, the SDRs differ by at most 0.003 dB.
+
+Upstream runs its Wiener EM in float32. On the reference mix that alone moves its stems by 71 to 93 dB SNR against a float64 run of the same code, which this port matches to 112 to 134 dB; the float64 run is the reference.
+
+## Wiener filter
+
+`wienerFilter(mixStft, estimates, opts)` is a from-scratch JS port of the algorithm in [norbert](https://github.com/sigsep/norbert) (Liutkus & Stöter), with [open-unmix-pytorch](https://github.com/sigsep/open-unmix-pytorch)'s `openunmix/filtering.py` defaults (`eps=1e-10`, `softmask=False`, `scale_factor=10`) — softmask off by default, matching upstream's own recommendation ("`softmask=False` is recommended... once the model estimates are themselves good"). The core update — re-estimate each source's power spectral density and spatial covariance matrix, rebuild the modelled mixture covariance, apply the resulting multichannel Wiener gain, iterate — is the local Gaussian model from **Duong, Vincent, Gribonval, "Under-determined reverberant audio source separation using a full-rank spatial covariance model," IEEE TASLP 18(7), 2010**. The gain applies as `v_j R_j (Cxx⁻¹ x)`: one inverse times the mixture per bin, then one C×C product per source. The covariances are Hermitian, so stereo takes the closed-form 2×2 inverse, other channel counts a Gauss-Jordan solver.
+
+```js
+import { stft, wienerFilter, istft } from '@audio/neural-separate'
+
+let mixStft = [stft(left, {}), stft(right, {})]                    // per channel
+let estimates = { vocals: vocalsMag, drums: drumsMag }             // per target: magnitude[channel][frame] (Float64Array(bins))
+let complex = wienerFilter(mixStft, estimates, { iterations: 1 })  // per target: complex STFT per channel
+let vocals = complex.vocals.map(ch => istft(ch, { length: left.length }))
+```
+
+`iterations: 0` returns the initial estimate untouched ("raw masks"). `softmask: true` uses a ratio mask that sums to the mixture exactly, by construction; `residual: true` appends a `'residual'` target (mixture minus the other targets) computed before EM. **Known limitation**: multichannel Wiener EM assumes genuine inter-channel diversity — mono content duplicated to stereo (identical L≡R) gives a rank-1 spatial covariance matrix per bin, a degenerate case where EM iterations can occasionally underperform `iterations: 0` on an already-good separation (real stereo mixes, with any actual left/right difference, don't hit this — verified in `test.js`).
+
+## `opts`
+
+| Option | Default | |
+|---|---|---|
+| `sampleRate` | — | required unless `audio` is `{ channelData, sampleRate }` |
+| `model` | — | required. Preset `'umxhq' \| 'htdemucs' \| 'htdemucs_ft'` · `url \| bytes` (single target, named `'stem'`) · `{ target: url \| { url, targets }, ... }` (one graph per target, Open-Unmix's own layout; a multi-source graph contributes its own target) · `{ url, targets: [...] }` (one multi-target graph, stacks a target axis) |
+| `targets` | all | the targets to return |
+| `weights` | `$AUDIO_NEURAL_CACHE` or `~/.cache/audiojs/neural` (Node) | where a preset's files are: URL, or a directory in Node |
+| `modelType` | `'openunmix'` | `'openunmix'` (magnitude out) · `'mask'` (`[0,1]` mask out, multiplied by mixture magnitude) · `'hybrid'` (demucs.onnx contract) · `'waveform'` (Demucs v2-class); presets know theirs |
+| `wiener` | `1` | EM iterations; `0` = raw masks. Ignored for `'hybrid'` and `'waveform'` |
+| `wienerWindow` | `300` | frames per EM window |
+| `chunk` / `overlap` | `30` / `2` (seconds) | `overlap` must be `< chunk`; not used by `'hybrid'`, which segments as demucs does |
+| `segment` | `343980` | `'hybrid'`: segment length in samples when the graph does not declare its input length |
+| `targetRate` | the preset's rate, else the input rate | resample to the model's rate for inference; stems are resampled back to the input rate |
+| `device` | — | passed through to `@audio/neural-runtime`'s `load()` as `backend` |
+| `dtype` | `'float32'` | only `'float32'` tensor marshalling is implemented; anything else throws |
+| `progress` | — | `({ chunk, totalChunks }) => {}` |
+| `session` | — | overrides `@audio/neural-runtime`'s `load()` — for tests, or a custom ORT setup |
+
+Mono input is duplicated to stereo internally (matching `openunmix.utils.preprocess`'s own "if we have mono, we duplicate it to get stereo"), so stems always come back stereo. `separate()` returns `{ stems, sampleRate, residual }` — `residual` is the mixture minus the sum of all stems, per channel, at the input rate.
+
+## Precedence and licenses
+
+Open-source stem separation has three lineages; this package runs the first two.
+
+| Project | Code | Weights | |
+|---|---|---|---|
+| **Open-Unmix** (`umx`, `umxhq`) | MIT | MIT ([Zenodo](https://zenodo.org/records/3370489)-declared) | trained on MUSDB18(-HQ); this package's primary target |
+| **Open-Unmix** (`umxl`) | MIT | **CC BY-NC-SA 4.0 — non-commercial only** | despite being the `openunmix` package's own default variant name; trained on a private stems dataset (see the project [README](https://github.com/sigsep/open-unmix-pytorch#pre-trained-models)) |
+| **Demucs** (Meta; `htdemucs`, `htdemucs_ft`) | MIT | **research only**: "The model weights are not covered by the MIT license, and are provided only for scientific purposes" (the author in [#327](https://github.com/facebookresearch/demucs/issues/327), 2022; again in [#508](https://github.com/facebookresearch/demucs/issues/508), 2023) | trained on MUSDB18-HQ plus 800 songs; the repository is archived, maintained at [adefossez/demucs](https://github.com/adefossez/demucs) |
+| **Spleeter** (Deezer) | MIT | **undocumented** | the README licenses only "the code of Spleeter"; the pretrained weights' license is an open, unresolved question ([deezer/spleeter#898](https://github.com/deezer/spleeter/issues/898)) — do not assume MIT |
+
+Audit any weight source yourself before shipping it — this table reflects what each project states as of this writing, not a guarantee. `scripts/export-openunmix.py` defaults to `umxhq` (not `umxl`) for exactly this reason.
+
+## Reference
+
+Stöter, Uhlich, Liutkus, Mitsufuji, "Open-Unmix - A Reference Implementation for Music Source Separation," *JOSS* 4(41), 2019. · Rouard, Massa, Défossez, "Hybrid Transformers for Music Source Separation," *ICASSP* 2023. · Duong, Vincent, Gribonval, "Under-determined reverberant audio source separation using a full-rank spatial covariance model," *IEEE TASLP* 18(7), 2010. · Rafii, Liutkus, Stöter, Mimilakis, Bittner, "The MUSDB18 corpus for music separation," 2017. · [norbert](https://github.com/sigsep/norbert) (Liutkus & Stöter) · [open-unmix-pytorch](https://github.com/sigsep/open-unmix-pytorch) · [Demucs](https://github.com/facebookresearch/demucs) · [demucs.onnx](https://github.com/sevagh/demucs.onnx) · [museval](https://github.com/sigsep/sigsep-mus-eval) · [Spleeter](https://github.com/deezer/spleeter).
+
+**Use when:** you have (or can license) an ONNX-exported spectrogram-mask, Hybrid-Demucs or waveform separation model and want to run it — with proper multichannel Wiener refinement, not just the raw mask — dependency-free, in Node or the browser; `umxhq` where the weights must be MIT, `htdemucs` for the higher SDR under research-only terms.<br>
+**Not for:** the classical, model-free case — reach for [`@audio/vocals`](https://github.com/audiojs/vocals) when a center-panned M/S trick is all you need; training a model (this is inference-only); real-time streaming (neither the bi-LSTM Open-Unmix nor Hybrid Transformer Demucs is causal — offline/chunked only, same as upstream).
+
+---
+
+Part of the [@audio/neural](https://github.com/audiojs/neural) lane.
+
+MIT © [audiojs](https://github.com/audiojs)

package/index.d.ts

@@ -0,0 +1,163 @@
+/**
+ * Source separation (stems) through @audio/neural-runtime's ONNX adapter: Open-Unmix-class
+ * spectrogram models (Stöter, Uhlich, Liutkus, Mitsufuji, JOSS 2019), Hybrid Transformer Demucs
+ * (Rouard, Massa, Défossez, ICASSP 2023) with STFT and iSTFT outside the graph, and waveform graphs.
+ */
+
+/** Mono duplicates internally; multichannel arrays must share the same length. */
+export type AudioInput = Float32Array[] | { channelData: Float32Array[]; sampleRate: number }
+
+/** A model spec resolvable by @audio/neural-runtime's load(): URL string or raw ONNX bytes. */
+export type ModelSpec = string | Uint8Array
+
+/** Presets whose files the export scripts write: <weights>/<name>/<target>.onnx or <weights>/<name>/<name>.onnx */
+export type ModelName = 'umxhq' | 'htdemucs' | 'htdemucs_ft'
+
+/** One graph whose output stacks several sources on its S axis */
+export type MultiGraph = { url: ModelSpec; targets: string[] }
+
+export type ModelOption =
+	| ModelName
+	| ModelSpec // single target, named 'stem'
+	| Record<string, ModelSpec | MultiGraph> // one graph per target (Open-Unmix's layout); a multi-source graph contributes its own target
+	| MultiGraph
+
+export type ModelType =
+	| 'openunmix' // target model outputs the estimated magnitude directly
+	| 'mask' // target model outputs a [0,1] mask; multiplied by the mixture magnitude
+	| 'hybrid' // demucs.onnx contract: mix [1,C,L] + CaC spectrogram [1,2C,F,T] in; [1,S,2C,F,T] + [1,S,C,L] out
+	| 'waveform' // Demucs v2-class: [1,C,N] waveform in, [1,S,C,N] stacked waveforms out
+
+export interface ModelPreset {
+	modelType: ModelType
+	sampleRate: number
+	/** in the order stems come back */
+	targets: string[]
+	/** one graph per target (<target>.onnx), else one graph for all (<name>.onnx) */
+	perTarget?: boolean
+}
+
+/** Model presets by name */
+export const models: Record<ModelName, ModelPreset>
+
+/** A neural-runtime-shaped session — enough of it to drive separate() with a test double. */
+export interface Session {
+	run(feeds: Record<string, { data: Float32Array; dims: number[]; type: string }>): Promise<Record<string, { data: Float32Array; dims: number[] }>>
+	/** 'hybrid' reads its segment length off inputs[0].dims[2] when declared */
+	inputs?: { name: string; dims?: number[] }[]
+	outputs?: { name: string }[]
+	free?(): void
+}
+
+export interface SeparateOptions {
+	/** required unless audio is the { channelData, sampleRate } form */
+	sampleRate?: number
+	model: ModelOption
+	/** ignored for presets, which know theirs (default 'openunmix') */
+	modelType?: ModelType
+	/** subset of the model's targets to return; spectral models skip the other graphs, and one target runs Wiener EM against the residual */
+	targets?: string[]
+	/** where a preset's files are: URL, or a directory in Node (default $AUDIO_NEURAL_CACHE or ~/.cache/audiojs/neural) */
+	weights?: string
+	/** iterations of multichannel Wiener EM refinement (default 1); 0 = raw masks. Ignored for modelType 'hybrid' and 'waveform'. */
+	wiener?: number
+	/** wienerFilter's softmask option (default false, matching open-unmix-pytorch) */
+	softmask?: boolean
+	/** wienerFilter's eps (default 1e-10) */
+	eps?: number
+	/** frames per Wiener EM window (default 300, open-unmix Separator's wiener_win_len) */
+	wienerWindow?: number
+	/** 'hybrid': segment length in samples when the graph does not declare it (default 343980, htdemucs's 7.8 s) */
+	segment?: number
+	/** 'openunmix' | 'mask' | 'waveform': chunk length in seconds (default 30) */
+	chunk?: number
+	/** crossfade overlap in seconds between chunks (default 2); must be < chunk */
+	overlap?: number
+	/** resample to this rate for model inference, stems back to the input rate (default: the preset's rate, else the input's) */
+	targetRate?: number
+	/** STFT size for the spectral pipeline (default 4096, Open-Unmix's n_fft) */
+	n?: number
+	/** STFT hop for the spectral pipeline (default 1024, Open-Unmix's n_hop) */
+	hop?: number
+	/** passed through to @audio/neural-runtime's load() as `backend` */
+	device?: string
+	/** only 'float32' is implemented; anything else throws */
+	dtype?: 'float32'
+	progress?: (p: { chunk: number; totalChunks: number }) => void
+	/** overrides @audio/neural-runtime's load() — for tests, or a custom ORT setup */
+	session?: (model: ModelSpec, opts: SeparateOptions) => Session | Promise<Session>
+}
+
+export interface SeparateResult {
+	/** keys are the model's target names ('stem' for a single bare model) */
+	stems: Record<string, Float32Array[]>
+	sampleRate: number
+	/** mixture minus the sum of all stems, per channel, at the input rate */
+	residual: Float32Array[]
+}
+
+/** Separate a mixture into stems. See README for the ONNX I/O contract per modelType. */
+export default function separate(audio: AudioInput, opts: SeparateOptions): Promise<SeparateResult>
+
+// ------------------------------------------------------------- STFT / iSTFT
+
+export interface StftOptions {
+	/** FFT size, power of 2 (default 4096) */
+	n?: number
+	/** hop size (default 1024) */
+	hop?: number
+	/** analysis window, length n (default: periodic Hann) */
+	window?: Float64Array
+	/** torch.stft-compatible reflect-padding by n/2 on both ends (default true) */
+	center?: boolean
+}
+
+export interface ComplexStft {
+	/** re[frame] is a Float64Array(bins), bins = n/2+1 */
+	re: Float64Array[]
+	im: Float64Array[]
+	n: number
+	hop: number
+	bins: number
+	center: boolean
+}
+
+/** Mono complex STFT. Frame count = 1 + floor(x.length / hop) when center=true (torch.stft-compatible). */
+export function stft(x: Float32Array | Float64Array, opts?: StftOptions): ComplexStft
+
+export interface IstftOptions {
+	/** exact output length (torch.istft-compatible crop/pad); default: paddedLength − 2·(n/2) when center */
+	length?: number
+	/** synthesis window (default: the same periodic Hann as stft()) */
+	window?: Float64Array
+}
+
+/** Exact inverse of stft() via squared-window-normalized overlap-add (WOLA). */
+export function istft(frames: ComplexStft, opts?: IstftOptions): Float32Array
+
+// ------------------------------------------------------------- Wiener filter
+
+export interface WienerOptions {
+	/** EM refinement steps (default 1); 0 returns the initial softmask/phase-substituted estimate untouched */
+	iterations?: number
+	/** true = ratio mask (sums to the mixture exactly); false = magnitude with the mixture's phase (default, open-unmix's recommendation) */
+	softmask?: boolean
+	/** appends a 'residual' target = mixture − Σ(other targets), computed before EM */
+	residual?: boolean
+	/** regularization floor (default 1e-10) */
+	eps?: number
+	/** EM numerical-stability rescaling divisor (default 10, open-unmix-pytorch's scale_factor) */
+	scaleFactor?: number
+}
+
+/**
+ * Multichannel Wiener EM (Liutkus & Stöter, github.com/sigsep/norbert; Duong,
+ * Vincent, Gribonval, IEEE TASLP 2010) with open-unmix-pytorch's defaults.
+ * estimates: per-target magnitude, one Float64Array(bins) per frame per
+ * channel — a single-channel estimate broadcasts to all mixture channels.
+ */
+export function wienerFilter(
+	mixStft: ComplexStft[],
+	estimates: Record<string, Float64Array[][]>,
+	opts?: WienerOptions
+): Record<string, ComplexStft[]>

package/package.json

@@ -0,0 +1,61 @@
+{
+  "name": "@audio/neural-separate",
+  "version": "0.1.0",
+  "description": "Source separation (stems): Open-Unmix-class spectrogram masking + multichannel Wiener refinement, Hybrid Transformer Demucs, model-agnostic ONNX",
+  "type": "module",
+  "sideEffects": false,
+  "main": "separate.js",
+  "types": "./index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./index.d.ts",
+      "default": "./separate.js"
+    },
+    "./package.json": "./package.json"
+  },
+  "files": [
+    "separate.js",
+    "index.d.ts",
+    "scripts/export-openunmix.py",
+    "scripts/export-htdemucs.py",
+    "scripts/reference.py",
+    "NOTICE"
+  ],
+  "dependencies": {
+    "@audio/neural-runtime": "^0.1.0",
+    "@audio/resample-sinc": "^1.1.2",
+    "fourier-transform": "^2.4.1"
+  },
+  "devDependencies": {
+    "onnxruntime-node": "^1.29.0",
+    "tst": "^9.4.0"
+  },
+  "scripts": {
+    "test": "node test.js"
+  },
+  "keywords": [
+    "audio",
+    "ml",
+    "neural",
+    "separate",
+    "stems",
+    "open-unmix",
+    "demucs",
+    "wiener",
+    "onnx"
+  ],
+  "license": "MIT",
+  "author": "Dmitry Iv. <dfcreative@gmail.com>",
+  "engines": {
+    "node": ">=18"
+  },
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/audiojs/neural.git",
+    "directory": "packages/neural-separate"
+  },
+  "funding": "https://github.com/sponsors/audiojs",
+  "publishConfig": {
+    "access": "public"
+  }
+}

package/scripts/export-htdemucs.py

@@ -0,0 +1,230 @@
+#!/usr/bin/env python3
+"""Export Demucs v4 (Hybrid Transformer Demucs) to the ONNX contract
+@audio/neural-separate expects for modelType 'hybrid': the network without its
+STFT and iSTFT, which ONNX export cannot carry. The split is sevagh's
+demucs.onnx (github.com/sevagh/demucs.onnx, MIT); JS computes both transforms.
+
+    pip install torch demucs onnx onnxruntime
+    python3 export-htdemucs.py --model htdemucs --verify
+    python3 export-htdemucs.py --model htdemucs_ft --verify   # four source-specialist graphs
+
+One 7.8 s training segment (343980 samples at 44.1 kHz) per run:
+  inputs   mix       [1, 2, 343980]         waveform
+           mix_spec  [1, 4, 2048, 336]      STFT as complex-as-channels (L re, L im, R re, R im)
+  outputs  stems_spec  [1, S, 4, 2048, 336] frequency branch, same layout
+           stems_wave  [1, S, 2, 343980]    time branch
+  stem = iSTFT(stems_spec) + stems_wave
+
+Writes <out-dir>/<model>.onnx (htdemucs_ft: <out-dir>/<source>.onnx per
+specialist), weights embedded. Default out-dir: $AUDIO_NEURAL_CACHE/<model>, else
+~/.cache/audiojs/neural/<model>. Last run: torch 2.14, demucs 4.1.0,
+onnxruntime 1.30.
+
+How: HTDemucs.forward runs unmodified with its four transform methods swapped
+on the instance: _spec and _magnitude return the given spectrogram, _mask keeps
+the frequency-branch estimate, _ispec returns 0 so forward returns the time
+branch alone. sevagh edits a vendored copy of htdemucs.py to the same effect;
+patching keeps upstream's forward as the single source. --verify rebuilds the
+stems from the ONNX outputs with demucs's own _mask and _ispec and compares them
+with model(mix) on the same segment.
+
+Weights: Demucs's checkpoints "are not covered by the MIT license, and are
+provided only for scientific purposes" (the author,
+github.com/facebookresearch/demucs/issues/327). Memory: an export peaks near
+3.2 GB; htdemucs_ft exports and checks each specialist in its own process.
+"""
+import argparse
+import gc
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def build_core(model):
+    import torch
+    import torch.nn as nn
+
+    class HybridCore(nn.Module):
+        def __init__(self, m):
+            super().__init__()
+            self.m = m
+
+        def forward(self, mix, mix_spec):
+            m, box = self.m, {}
+            m._spec = lambda x: mix_spec
+            m._magnitude = lambda z: z
+            m._mask = lambda z, x: box.setdefault("spec", x)
+            m._ispec = lambda z, length=None, scale=0: torch.zeros((), dtype=mix.dtype)
+            try:
+                wave = m(mix)
+            finally:
+                for k in ("_spec", "_magnitude", "_mask", "_ispec"):
+                    del m.__dict__[k]
+            return box["spec"], wave
+
+    return HybridCore(model).eval()
+
+
+def cac_spec(model, mix):
+    """The mix's STFT as HTDemucs sees it (its own _spec and _magnitude)."""
+    return model._magnitude(model._spec(mix))
+
+
+def export_one(model, out_path, opset, verify):
+    import torch
+
+    L = int(model.segment * model.samplerate)
+    mix = torch.randn(1, model.audio_channels, L)
+    with torch.no_grad():
+        spec = cac_spec(model, mix)
+    # traced with grad enabled: under no_grad, nn.MultiheadAttention takes its fused
+    # fast path, aten::_native_multi_head_attention, which has no ONNX export
+    torch.onnx.export(
+        build_core(model),
+        (mix, spec),
+        str(out_path),
+        input_names=["mix", "mix_spec"],
+        output_names=["stems_spec", "stems_wave"],
+        opset_version=opset,
+        dynamo=False,
+    )
+    print(f"wrote {out_path} ({out_path.stat().st_size / 1e6:.1f} MB), segment {L} samples")
+    gc.collect()  # the traced graph: export peaks near 3 GB
+    if verify:
+        verify_export(model, out_path, L)
+
+
+def verify_export(model, onnx_path, L):
+    """Stems rebuilt from onnxruntime's outputs against model(mix): a random
+    segment and a harmonic one, relative max |diff| < 1e-3 (float32 through a
+    transformer; a transform mismatch would show at order 1)."""
+    import math
+    import numpy as np
+    import onnxruntime as ort
+    import torch
+
+    opts = ort.SessionOptions()
+    opts.enable_cpu_mem_arena = False  # give memory back between runs
+    opts.enable_mem_pattern = False
+    session = ort.InferenceSession(str(onnx_path), opts, providers=["CPUExecutionProvider"])
+    t = torch.arange(L) / model.samplerate
+    tone = sum(torch.sin(2 * math.pi * f * t) / k for k, f in enumerate((110, 220, 440, 880, 1760), 1))
+    for name, mix in (("noise", 0.1 * torch.randn(1, 2, L)), ("tones", 0.2 * torch.stack([tone, 0.8 * tone])[None])):
+        with torch.no_grad():
+            expected = model(mix).numpy()
+            spec = cac_spec(model, mix)
+            s, w = session.run(None, {"mix": mix.numpy(), "mix_spec": spec.numpy()})
+            actual = (model._ispec(model._mask(None, torch.from_numpy(s)), L) + torch.from_numpy(w)).numpy()
+        rel = np.abs(actual - expected).max() / max(np.abs(expected).max(), 1e-12)
+        ok = rel < 1e-3
+        print(f"  verify {onnx_path.name} {name}: max|diff|/max|y| = {rel:.2e} " + ("OK" if ok else "FAIL"))
+        if not ok:
+            raise SystemExit(f"{onnx_path}: verification failed ({rel:.2e} >= 1e-3)")
+    del session
+    gc.collect()
+
+
+def write_reference(bag, out_dir):
+    """demucs.api.Separator's path: whole-input normalization, apply_model with split=True,
+    overlap=0.25 and shifts=0 (the CLI's one random shift is not reproducible)."""
+    import torch
+    from demucs.apply import apply_model
+
+    sys.path.insert(0, str(Path(__file__).parent))
+    sys.dont_write_bytecode = True  # no __pycache__ beside the package's scripts
+    import reference
+
+    def separate(mix):
+        wav = torch.from_numpy(mix)
+        ref = wav.mean(0)
+        mean, std = ref.mean(), ref.std() + 1e-8
+        with torch.no_grad():
+            est = apply_model(bag, ((wav - mean) / std)[None], shifts=0, split=True, overlap=0.25)[0] * std + mean
+        return {s: est[k].numpy() for k, s in enumerate(bag.sources)}
+
+    reference.write(out_dir, separate)
+
+
+def bag_of(name):
+    """(signature, source) per model of a bag of source specialists (htdemucs_ft), read off the
+    bag's YAML in the demucs package without importing torch; None for a single model."""
+    import importlib.util
+    import yaml
+
+    spec = importlib.util.find_spec("demucs")
+    f = Path(spec.origin).parent / "remote" / f"{name}.yaml" if spec else None
+    bag = yaml.safe_load(f.read_text()) if f and f.exists() else {}
+    if len(bag.get("models", [])) < 2 or not bag.get("weights"):
+        return None
+    sources = ["drums", "bass", "other", "vocals"]  # demucs.pretrained.SOURCES, the weights' column order
+    return [(sig, sources[max(range(len(w)), key=lambda k: w[k])]) for sig, w in zip(bag["models"], bag["weights"])]
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--model", default="htdemucs", help="htdemucs | htdemucs_ft | htdemucs_6s")
+    ap.add_argument("--out-dir", default=None, help="default: $AUDIO_NEURAL_CACHE/<model> or ~/.cache/audiojs/neural/<model>")
+    ap.add_argument("--opset", type=int, default=17)
+    ap.add_argument("--verify", action="store_true", help="compare the rebuilt stems with model(mix), and write the "
+                    "reference separation test.js compares against (scripts/reference.py)")
+    ap.add_argument("--source", default=None, help=argparse.SUPPRESS)  # a bag's specialist, in its own process
+    ap.add_argument("--reference", action="store_true", help=argparse.SUPPRESS)  # a bag's reference, in its own process
+    ap.add_argument("--check", action="store_true", help=argparse.SUPPRESS)  # verify a specialist's export, in its own process
+    args = ap.parse_args()
+
+    cache = os.environ.get("AUDIO_NEURAL_CACHE") or Path.home() / ".cache" / "audiojs" / "neural"
+    out_dir = Path(args.out_dir) if args.out_dir else Path(cache) / args.model
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        bag = bag_of(args.model)
+    except ImportError:
+        raise SystemExit("torch and demucs are required: pip install torch demucs onnx onnxruntime")
+    if bag and not (args.reference or args.source):
+        # a bag of source specialists: one process each, since an export peaks near 3 GB and is
+        # not handed back; this one stays free of torch
+        run = lambda *extra: subprocess.run([sys.executable, __file__, "--model", args.model, "--out-dir", str(out_dir), "--opset", str(args.opset), *extra], check=True)
+        if args.verify:
+            run("--reference")
+        for _, src in bag:
+            run("--source", src)
+            if args.verify:
+                run("--source", src, "--check")
+        print(f"done: {args.model} in {out_dir}")
+        return
+
+    try:
+        from demucs.pretrained import get_model
+        from demucs.htdemucs import HTDemucs
+    except ImportError:
+        raise SystemExit("torch and demucs are required: pip install torch demucs onnx onnxruntime")
+
+    def load(name):
+        model = get_model(name)
+        for m in getattr(model, "models", [model]):
+            if not isinstance(m, HTDemucs):
+                raise SystemExit(f"{name}: not a Hybrid Transformer Demucs model")
+            m.eval()
+        return model
+
+    if args.reference:
+        write_reference(load(args.model), out_dir)
+    elif args.source:
+        # the one specialist, by its signature: a single HTDemucs; exported, or checked in a fresh process
+        model, path = load(dict((s, g) for g, s in bag)[args.source]), out_dir / f"{args.source}.onnx"
+        if args.check:
+            verify_export(model, path, int(model.segment * model.samplerate))
+        else:
+            export_one(model, path, args.opset, False)
+    else:
+        model = load(args.model)
+        if args.verify:
+            write_reference(model, out_dir)  # first, while memory is fresh
+            gc.collect()
+        export_one(getattr(model, "models", [model])[0], out_dir / f"{args.model}.onnx", args.opset, args.verify)
+        print(f"done: {args.model} in {out_dir}")
+
+
+if __name__ == "__main__":
+    sys.exit(main())