@audio/neural-separate 0.4.0 → 0.5.0

package/README.md

@@ -28,6 +28,7 @@ In the browser `@audio/neural-runtime` runs `onnxruntime-web` instead (`device:
 | `'htdemucs_ft'` | four fine-tuned HTDemucs, one per source | 4 × 174 MB | research only | 8.67 · 9.45 · 9.68 · 5.58 | vocals only: 1.1 · not measured · not measured |
 | `'scnet-large'` | SCNet-large: band-split convolutions around dual-path LSTMs, on the complex spectrogram | 45 MB (export 169 MB) | MIT | 11.00 · 10.27 · 8.21 · 6.87 | 2.1 · 7.8 · 5.0 |
 | `'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 |
+| `'mel-roformer'` | Mel-Band RoFormer (Kimberley Jensen's vocals model): axial transformers over time and mel bands, on the complex spectrogram; written for the GPU | 464 MB, float16 weights (export 928 MB) | MIT | vocals only: 12.08 | 1.4 · not measured · 0.56 (Node's WebGPU; the export 0.44) |
 
 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.
 
@@ -67,8 +68,9 @@ the float32 graph's.
 ## Weights
 
 Nothing ships in this package. `scnet-large`, `scnet`, `mrx` and `tiger` read a compact file each, `<model>.int8.onnx`
-(Size, below), hosted on Hugging Face (`audiojs/scnet-large`, `audiojs/scnet`, `audiojs/mrx`, `audiojs/tiger-dnr`, each
-with its model card and licence): without `weights`, `https://huggingface.co/<repo>/resolve/<revision>/<file>` at the
+(Size, below), and `mel-roformer` its float16 weights, `mel-roformer.fp16.onnx` (or `roformer.js`'s `mel-roformer.engine.bin`), hosted on Hugging Face
+(`audiojs/scnet-large`, `audiojs/scnet`, `audiojs/mrx`, `audiojs/tiger-dnr`, `audiojs/mel-roformer`, each with its model
+card and licence): without `weights`, `https://huggingface.co/<repo>/resolve/<revision>/<file>` at the
 revision `REVISIONS` pins, fetched once and checked against its SHA-256 (the browser keeps it in the Cache API, Node in
 the neural cache, below). Every model can also be exported from its original checkpoint:
 
@@ -88,6 +90,10 @@ python3 node_modules/@audio/neural-separate/scripts/export-mrx.py --verify
 pip install safetensors
 python3 node_modules/@audio/neural-separate/scripts/export-tiger.py --verify
 
+pip install einops beartype rotary-embedding-torch librosa
+python3 node_modules/@audio/neural-separate/scripts/export-roformer.py --verify --fp16 --engine
+node node_modules/@audio/neural-separate/scripts/engine.mjs       # roformer.js in Chromium against onnxruntime (playwright)
+
 python3 node_modules/@audio/neural-separate/scripts/compact.py --model tiger --verify   # tiger.int8.onnx from tiger.onnx
 ```
 
@@ -97,6 +103,8 @@ They land in `$AUDIO_NEURAL_CACHE/<model>/`, default `~/.cache/audiojs/neural/<m
 - **`export-scnet.py`** fetches SCNet's code from ZFTurbo/Music-Source-Separation-Training at a pinned commit and the author's MUSDB18-HQ checkpoint from that repository's releases, and exports the network between its STFT and iSTFT (the same split as demucs.onnx's): `mix_spec [1, 4, 2049, 476]` in, `stems_spec [1, 16, 2049, 476]` out, one 11 s segment. Two changes to the graph, neither to its function: the rFFT and irFFT over time in its separation network become products with cosine and sine matrices (torch 2.14 exports no `fft_rfft`), and each `GroupNorm(1, C)` reduces its mean and variance one axis at a time: exported whole, onnxruntime summed a 10-million-value group in float32 and the stems came out 52 to 55 dB from PyTorch's, against 110 dB now (PyTorch's own float32 against float64: 110 dB). `--fp16` also writes float16 weights.
 - **`export-mrx.py`** fetches MRX's code and MERL's checkpoint (`default_mrx_pre_trained_weights.pth`, trained with the SNR loss so its stems keep the mixture's level; the three others are scale-invariant or trained on loudness- or EQ-adapted stems, and their stems come out at another level: on Divide and Remaster v3 their SNR is negative where their SI-SDR is up to 1 dB higher) from merlresearch/cocktail-fork-separation at a pinned commit, and exports the network between its STFTs and iSTFTs: each channel's three magnitude spectrograms `mag_1024 [B, 513, T]`, `mag_2048 [B, 1025, T]`, `mag_8192 [B, 4097, T]` in (scaled 1/√n, torch.stft's `normalized=True`), the masks `mask_<n> [B, 3, F, T]` out, music, speech and sfx; B and T free.
 - **`export-tiger.py`** fetches TIGER's model code (JusperLee/TIGER, `look2hear/models/tiger_dnr.py` and the two layer modules it reads) at a pinned commit and its Divide and Remaster weights (huggingface.co/JusperLee/TIGER-DnR) at a pinned revision, and exports TIGERDNR's three models between their STFT and iSTFT as one graph, each its kept source (the first's dialogue, the second's effects, the third's music): `mix_spec [1, 2, 1025, 1034]` (one channel's STFT, n 2048, hop 512, Hann, unnormalized) in, `stems_spec [1, 6, 1025, 1034]` (dialogue, music, effects, re and im) out, one 12 s segment. Three changes to the graph, none to its function: `F.adaptive_avg_pool1d`, which the exporter takes only for output sizes dividing the input's, becomes a cumulative sum read at torch's windows; UConvBlock's global sum starts at its first term, not at a zeros tensor exported as a constant the size of the features; each `GroupNorm(1, C)` reduces axis by axis, as SCNet's. ONNX Runtime runs the graph in 12 s segments at about twice real time on a CPU: three models of 15,000 small operators each, every sample in three segments.
+- **`export-roformer.py`** fetches Mel-Band RoFormer's code and Kim's config from ZFTurbo/Music-Source-Separation-Training at a pinned commit and Kim's checkpoint (huggingface.co/KimberleyJSN/melbandroformer, MIT) at a pinned revision, checked by its SHA-256, and exports the network between its STFT and iSTFT: `mix_spec [1, 4, 1025, 801]` in, `stems_spec [1, 4, 1025, 801]` (the vocals) out, one 8 s segment every 4 s. It is written for the GPU, its function unchanged (`--verify`: 1e-5 to 2e-4 of the peak from MelBandRoformer.forward, whose own float32 stands 5e-3 from its float64 on tones): the 60 bands' projections and mask MLPs as batched products, bands of near widths grouped and zero-padded; each RMSNorm's scale folded into the weights it feeds, the attention's scale into the queries, its gates' projection joined to q, k, v; the rotary embedding's interleaved pairs turned to halves (the same permutation of q's and k's features), so it is two products with cos and sin tables and one with a fixed matrix; the overlapping mel bands' masks averaged by two gathers, not a scatter. 826 operators where an export of BS-RoFormer as written has about 3,000, and SCNet's LSTMs none: on Node's WebGPU (Metal, M4 Max) 2.4 s a segment against 9 s on its CPU; SCNet-large's LSTMs take 19 of its 26 s a segment there. And no tensor passes 100 MB, each transformer run on slices of its sequences: onnxruntime 1.30's WebGPU Softmax over a tensor past WebGPU's default 128 MiB storage binding came out wrong one run in two (36 dB from the CPU's at 307 MB, every run right at 123 MB), and the stems with it. `--fp16`: float16 weights, float32 compute, half the download, 3.0 s a segment and 48 dB from the float32 file's stems.
+- **`roformer.js`**: in a browser whose GPU offers subgroup matrices (`chromium-experimental-subgroup-matrix`: Chromium on Metal), `mel-roformer` runs on kernels of its own instead of onnxruntime, from `mel-roformer.engine.bin` and `.json` (`export-roformer.py --engine`: the same tensors, hosted beside the ONNX file; a browser downloads one or the other). Matrix products and attention on the GPU's matrix units, float32: band split, norms, the q, k, v projection with the rotation and gates, attention in two passes over the keys (the rows' maxima, then the output in registers), feed-forward with GELU, the mask MLP with tanh and GLU, the masks averaged and applied; 163 dispatches a segment, recorded once; kept on the GPU a minute after a separation, so the next starts at once (its weights and 2 GB of buffers made in about a second the first time). Against onnxruntime's CPU on the same weights 101.5 dB (`scripts/engine.mjs`); in Chromium on an M4 Max under the same load as onnxruntime-web's WebGPU there, 1.5 to 1.8 s a segment against 2.6 to 3.3; GPU time by kernel: the time attention a third, the matrix products half (about 3 TFLOPS each, against onnxruntime's 1.8 for its float32 products and 0.5 to 0.75 for its attention). Elsewhere (Node, Firefox, Safari, a GPU without them) onnxruntime runs the ONNX file as before.
 - **`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.
 
 ## Size

package/index.d.ts

@@ -12,7 +12,7 @@ export type AudioInput = Float32Array[] | { channelData: Float32Array[]; sampleR
 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' | 'scnet-large' | 'scnet' | 'mrx' | 'tiger'
+export type ModelName = 'umxhq' | 'htdemucs' | 'htdemucs_ft' | 'scnet-large' | 'scnet' | 'mrx' | 'tiger' | 'mel-roformer'
 
 /** One graph whose output stacks several sources on its S axis */
 export type MultiGraph = { url: ModelSpec; targets: string[] }
@@ -74,7 +74,7 @@ export interface ModelPreset {
 export const models: Record<ModelName, ModelPreset>
 
 /** 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. */
-export const REVISIONS: Record<'scnet-large' | 'scnet' | 'mrx' | 'tiger', string>
+export const REVISIONS: Record<'scnet-large' | 'scnet' | 'mrx' | 'tiger' | 'mel-roformer', string>
 
 /** A neural-runtime-shaped session — enough of it to drive separate() with a test double. */
 export interface Session {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@audio/neural-separate",
-  "version": "0.4.0",
+  "version": "0.5.0",
   "description": "Source separation (stems): a song's by SCNet (MIT weights), Hybrid Transformer Demucs, Open-Unmix with multichannel Wiener refinement; a soundtrack's dialogue, music and effects by MRX (MIT weights) and TIGER (Apache-2.0); model-agnostic ONNX",
   "type": "module",
   "sideEffects": false,
@@ -15,18 +15,21 @@
   },
   "files": [
     "separate.js",
+    "roformer.js",
     "index.d.ts",
     "scripts/export-openunmix.py",
     "scripts/export-htdemucs.py",
     "scripts/export-scnet.py",
     "scripts/export-mrx.py",
     "scripts/export-tiger.py",
+    "scripts/export-roformer.py",
+    "scripts/engine.mjs",
     "scripts/reference.py",
     "scripts/compact.py",
     "NOTICE"
   ],
   "dependencies": {
-    "@audio/neural-runtime": "^0.1.0",
+    "@audio/neural-runtime": "^0.1.1",
     "@audio/resample-sinc": "^1.1.2",
     "fourier-transform": "^2.4.1"
   },

package/roformer.js

@@ -0,0 +1,396 @@
+// Mel-Band RoFormer on WebGPU, kernels of its own: the GPU form scripts/export-roformer.py writes (--engine), one 8 s
+// segment a run, what its ONNX graph computes (47.8 dB from onnxruntime's CPU on the float32 graph: the float16 weights'
+// rounding, as the float16 ONNX file's). Matrix products and attention on the GPU's matrix units through
+// chromium_experimental_subgroup_matrix (Apple's simdgroup matrices: Chromium on Metal), float32 throughout: under the same
+// load, 1.6 s a segment where onnxruntime-web's WebGPU takes 2.6 to 3.3 (an M4 Max). Attention in two passes over the
+// keys (the rows' maxima, then the output in registers), never past 128 MiB a dispatch needs no binding beyond it.
+//   let rf = await engine(meta, bytes)                  // null where the GPU lacks subgroup matrices
+//   let { stems_spec } = await rf.run({ mix_spec })     // as an onnxruntime session: [1, 4, F, T] in and out
+const TILE = 32
+
+/** A device with what the kernels need, or null */
+export async function gpu() {
+  let adapter = await globalThis.navigator?.gpu?.requestAdapter({ powerPreference: 'high-performance' }).catch(() => null)
+  if (!adapter || adapter.info?.isFallbackAdapter || !['subgroups', 'chromium-experimental-subgroup-matrix'].every(f => adapter.features.has(f))) return null
+  let cfg = [...(adapter.info.subgroupMatrixConfigs ?? [])]
+  if (!cfg.some(c => c.componentType === 'f32' && c.M === 8 && c.N === 8 && c.K === 8)) return null
+  return adapter.requestDevice({
+    requiredFeatures: ['subgroups', 'chromium-experimental-subgroup-matrix'],
+    requiredLimits: { maxComputeWorkgroupStorageSize: 32768, maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize, maxBufferSize: adapter.limits.maxBufferSize },
+  }).catch(() => null)
+}
+
+/** The engine as an onnxruntime-shaped session ({ inputs, outputs, run(feeds), free() }), or null without such a GPU */
+export async function engine(meta, bytes, device) {
+  device ??= await gpu()
+  if (!device) return null
+  let rf = await roformer(device, meta, bytes), { freqs: F, frames: T } = meta
+  return {
+    inputs: [{ name: 'mix_spec', dims: [1, 4, F, T], type: 'float32' }], outputs: [{ name: 'stems_spec', dims: [1, 4, F, T], type: 'float32' }],
+    backend: 'roformer',
+    async run(feeds) { return { stems_spec: { data: await rf.run(feeds.mix_spec.data), dims: [1, 4, F, T], type: 'float32' } } },
+    free() { device.destroy() },
+  }
+}
+
+export default async function roformer(device, meta, bytes) {
+  const { frames: T, freqs: F, channels: S, dim: D, heads: H, depth, bands: B, groups } = meta
+  const DH = 64, FF = 1536, MH = 1536, Tp = Math.ceil(T / TILE) * TILE, QKV = 3 * H * DH, QKVG = Math.ceil((QKV + H) / TILE) * TILE
+  const COLS = F * S * 2, LT = Math.ceil(T / 16) * 16, LF = Math.ceil(B / 16) * 16
+  const pad16 = w => Math.ceil(w / 16) * 16
+  const G = groups.map(([b0, b1, W]) => ({ b0, b1, n: b1 - b0, W, Wp: pad16(W) }))
+
+  // ---------------------------------------------------------------- tensors
+  const view = n => {
+    const t = meta.tensors[n], len = t.shape.reduce((a, b) => a * b, 1)
+    return t.type === 'float16' ? new Float16Array(bytes.buffer, bytes.byteOffset + t.offset, len)
+      : t.type === 'int32' ? new Int32Array(bytes.buffer, bytes.byteOffset + t.offset, len)
+      : new Float32Array(bytes.buffer, bytes.byteOffset + t.offset, len)
+  }
+  const usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC
+  const buf = (n, type = 'f32') => device.createBuffer({ size: Math.max(16, n * 4), usage })
+  const upload = data => { const b = device.createBuffer({ size: Math.max(16, data.byteLength), usage, mappedAtCreation: true }); new data.constructor(b.getMappedRange()).set(data); b.unmap(); return b }
+  const f32 = a => a instanceof Float32Array ? a : Float32Array.from(a)
+  // a weight [batch, K, N] to [batch, Kp, Np], zeros past (rows, cols)
+  const padded = (a, batch, K, N, Kp, Np, place = (k, n) => k * Np + n) => {
+    const o = new Float32Array(batch * Kp * Np)
+    for (let b = 0; b < batch; b++) for (let k = 0; k < K; k++) for (let n = 0; n < N; n++) o[b * Kp * Np + place(k, n)] = a[(b * K + k) * N + n]
+    return o
+  }
+
+  const W = { split: [], splitB: [], splitIdx: [], m2: [], m2B: [] }
+  for (const [g, gr] of G.entries()) {
+    W.split.push(upload(padded(view(`split_w.${g}`), gr.n, gr.W, D, gr.Wp, D)))
+    W.splitB.push(upload(f32(view(`split_b.${g}`))))
+    const idx = view(`si${g}`), pi = new Int32Array(gr.n * gr.Wp).fill(COLS)
+    for (let j = 0; j < gr.n; j++) for (let e = 0; e < gr.W; e++) pi[j * gr.Wp + e] = idx[j * gr.W + e]
+    W.splitIdx.push(upload(pi))
+    // the last mask layer: its value half to [0, W), its gate half to [Wp, Wp + W)
+    const w = view(`m2_w.${g}`), b = view(`m2_b.${g}`), Wp = gr.Wp, wo = new Float32Array(gr.n * MH * 2 * Wp), bo = new Float32Array(gr.n * 2 * Wp)
+    for (let j = 0; j < gr.n; j++) {
+      for (let k = 0; k < MH; k++) for (let e = 0; e < gr.W; e++) {
+        wo[(j * MH + k) * 2 * Wp + e] = w[(j * MH + k) * 2 * gr.W + e]
+        wo[(j * MH + k) * 2 * Wp + Wp + e] = w[(j * MH + k) * 2 * gr.W + gr.W + e]
+      }
+      for (let e = 0; e < gr.W; e++) { bo[j * 2 * Wp + e] = b[j * 2 * gr.W + e]; bo[j * 2 * Wp + Wp + e] = b[j * 2 * gr.W + gr.W + e] }
+    }
+    W.m2.push(upload(wo)); W.m2B.push(upload(bo))
+  }
+  W.mw = [0, 1].map(i => upload(f32(view(`mw.${i}`)))); W.mb = [0, 1].map(i => upload(f32(view(`mb.${i}`))))
+  W.blocks = []
+  for (let d = 0; d < depth; d++) for (let k = 0; k < 2; k++) {
+    const p = `blocks.${d}.${k}.`, a = p + 'layers.0.0.', f = p + 'layers.0.1.'
+    W.blocks.push({
+      out: upload(f32(view(p + 'out'))),
+      qkvg: upload(padded(view(a + 'qkvg'), 1, D, QKV + H, D, QKVG)), gb: upload(f32(view(a + 'gb'))), wo: upload(f32(view(a + 'wo'))),
+      w1: upload(f32(view(f + 'w1'))), b1: upload(f32(view(f + 'b1'))), w2: upload(f32(view(f + 'w2'))), b2: upload(f32(view(f + 'b2'))),
+    })
+  }
+  const rope = k => { const c = view(`blocks.0.${k}.layers.0.0.cos`), s = view(`blocks.0.${k}.layers.0.0.sin`); return [upload(f32(c)), upload(f32(s))] }
+  const [cosT, sinT] = rope(0), [cosF, sinF] = rope(1)
+  // the masks' sources: per output column, two (base, stride: the group's Wp) into the GLU outputs, and its weight
+  const i1 = view('i1'), i2 = view('i2'), wt = view('wt'), gofs = [], cols = []
+  let mOff = 0, cOff = 0
+  for (const gr of G) { gofs.push(mOff); cols.push(cOff); mOff += gr.n * Tp * gr.Wp; cOff += gr.n * gr.W }
+  const src = c => {
+    if (c >= cOff) return [0, 0, 0]
+    let g = cols.findLastIndex(o => o <= c), gr = G[g], r = c - cols[g], j = Math.floor(r / gr.W), e = r % gr.W
+    return [gofs[g] + j * Tp * gr.Wp + e, gr.Wp, 1]
+  }
+  const map = new Uint32Array(COLS * 8)                   // per column: base1, stride1, on1, base2, stride2, on2, -, -
+  for (let o = 0; o < COLS; o++) map.set([...src(i1[o]), ...src(i2[o])], o * 8)
+  const maskMap = upload(map)
+  const maskWt = upload(Float32Array.from({ length: COLS }, (_, o) => wt[o]))
+
+  // ---------------------------------------------------------------- activations
+  const E = buf(Tp * (COLS + 1)), X = buf(B * Tp * D), XN = buf(B * Tp * D), Y = buf(B * Tp * QKVG)
+  const att = Math.max(B * H * LT, T * H * LF) * DH
+  const Qb = buf(att), Kb = buf(att), Vb = buf(att), Ob = buf(att), Gt = buf(B * Tp * H), Z = buf(B * Tp * H * DH)
+  const FFH = buf(B * Tp * FF), H2 = buf(B * Tp * MH), Mg = buf(mOff), OUT = buf(4 * F * T)
+  const splitA = G.map(gr => buf(gr.n * Tp * gr.Wp)), glu = G.map(gr => buf(gr.n * Tp * 2 * gr.Wp)), scratch = buf(4)
+
+  // ---------------------------------------------------------------- kernels
+  const pipes = new Map()
+  const pipe = (key, code) => pipes.get(key) ?? pipes.set(key, device.createComputePipeline({ layout: 'auto', compute: { module: device.createShaderModule({ code }), entryPoint: 'main' } })).get(key)
+  const uni = (...v) => { const b = device.createBuffer({ size: Math.ceil(v.length / 4) * 16, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST }); device.queue.writeBuffer(b, 0, new Uint32Array(v.map(x => typeof x === 'number' && !Number.isInteger(x) ? new Uint32Array(new Float32Array([x]).buffer)[0] : x))); return b }
+  const steps = []                                          // [pipeline, bindings, [x, y, z]]
+  const step = (p, bindings, grid) => steps.push([p, device.createBindGroup({ layout: p.getBindGroupLayout(0), entries: bindings.map((b, i) => ({ binding: i, resource: { buffer: b } })) }), grid])
+
+  const HEAD = `enable subgroups;\nenable chromium_experimental_subgroup_matrix;\n`
+  // C[b] (M×N, row stride ldc, at c0 + b·sc) = A[b] (M×K, at a0 + b·sa) · Bm[b] (K×N, at b0 + b·sb), then the epilogue:
+  // ep: '' as is; 'b' + bias (at bias0 + b·sbias); then 'g' gelu or 't' tanh; 'r' C += it (a residual)
+  const MM = ep => pipe('mm' + ep, HEAD + `
+@group(0) @binding(0) var<storage, read> A: array<f32>;
+@group(0) @binding(1) var<storage, read> Bm: array<f32>;
+@group(0) @binding(2) var<storage, read_write> C: array<f32>;
+@group(0) @binding(3) var<storage, read> bias: array<f32>;
+struct P { M: u32, N: u32, K: u32, a0: u32, sa: u32, b0: u32, sb: u32, c0: u32, sc: u32, ldc: u32, bias0: u32, sbias: u32 }
+@group(0) @binding(4) var<uniform> p: P;
+var<workgroup> tile: array<f32, 1024>;
+@compute @workgroup_size(32)
+fn main(@builtin(workgroup_id) wg: vec3u, @builtin(subgroup_invocation_id) lane: u32) {
+  _ = bias[0];   // bound whether read or not: the layout is the module's own
+  let r0 = wg.y * 32u; let n0 = wg.x * 32u; let b = wg.z;
+  let ab = p.a0 + b * p.sa + r0 * p.K; let bb = p.b0 + b * p.sb + n0;
+  ${[0, 1, 2, 3].map(i => [0, 1, 2, 3].map(j => `var c${i}${j}: subgroup_matrix_result<f32, 8, 8>;`).join(' ')).join('\n  ')}
+  for (var k = 0u; k < p.K; k += 8u) {
+    ${[0, 1, 2, 3].map(i => `let a${i} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(&A, ab + ${i * 8}u * p.K + k, false, p.K);`).join('\n    ')}
+    ${[0, 1, 2, 3].map(j => `let b${j} = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>>(&Bm, bb + k * p.N + ${j * 8}u, false, p.N);`).join('\n    ')}
+    ${[0, 1, 2, 3].map(i => [0, 1, 2, 3].map(j => `c${i}${j} = subgroupMatrixMultiplyAccumulate(a${i}, b${j}, c${i}${j});`).join(' ')).join('\n    ')}
+  }
+  ${[0, 1, 2, 3].map(i => [0, 1, 2, 3].map(j => `subgroupMatrixStore(&tile, ${i * 8 * 32 + j * 8}u, c${i}${j}, false, 32u);`).join(' ')).join('\n  ')}
+  workgroupBarrier();
+  for (var e = lane; e < 1024u; e += 32u) {
+    let r = r0 + e / 32u; let n = n0 + e % 32u;
+    if (r < p.M && n < p.N) {
+      var v = tile[e];
+      ${ep.includes('b') ? 'v += bias[p.bias0 + b * p.sbias + n];' : ''}
+      ${ep.includes('g') ? 'v = 0.5 * v * (1.0 + erf_(v * 0.7071067811865476));' : ep.includes('t') ? 'v = tanh(v);' : ''}
+      let o = p.c0 + b * p.sc + r * p.ldc + n;
+      ${ep.includes('r') ? 'C[o] += v;' : 'C[o] = v;'}
+    }
+  }
+}
+fn erf_(x: f32) -> f32 {   // Abramowitz & Stegun 7.1.26 is 1.5e-7; GELU's own float32 erf is as close
+  let t = 1.0 / (1.0 + 0.3275911 * abs(x));
+  let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-x * x);
+  return select(-y, y, x >= 0.0);
+}`)
+  const mm = (ep, A, Bm, C, bias, { M, N, K, batch = 1, a0 = 0, sa = 0, b0 = 0, sb = 0, c0 = 0, sc = 0, ldc = N, bias0 = 0, sbias = 0 }) =>
+    step(MM(ep), [A, Bm, C, bias ?? maskWt, uni(M, N, K, a0, sa, b0, sb, c0, sc, ldc, bias0, sbias)], [N / 32, M / 32, batch])
+
+  // rows of n: y = x / max(|x|, 1e-12), or in place, scaled by g (a vector): x = x / max(|x|, 1e-12) · g
+  const NORM = scaled => pipe('norm' + scaled, HEAD + `
+@group(0) @binding(0) var<storage, read_write> x: array<f32>;
+@group(0) @binding(1) var<storage, read_write> y: array<f32>;
+@group(0) @binding(2) var<storage, read> g: array<f32>;
+struct P { rows: u32, n: u32 }
+@group(0) @binding(3) var<uniform> p: P;
+@compute @workgroup_size(32)
+fn main(@builtin(workgroup_id) wg: vec3u, @builtin(subgroup_invocation_id) lane: u32) {
+  _ = g[0]; _ = y[0];
+  let r = wg.x + wg.y * 65535u;
+  if (r >= p.rows) { return; }
+  var s = 0.0;
+  for (var i = lane; i < p.n; i += 32u) { let v = x[r * p.n + i]; s += v * v; }
+  let k = 1.0 / max(sqrt(subgroupAdd(s)), 1e-12);
+  for (var i = lane; i < p.n; i += 32u) { ${scaled ? 'x[r * p.n + i] = x[r * p.n + i] * k * g[i];' : 'y[r * p.n + i] = x[r * p.n + i] * k;'} }
+}`)
+  const grid1 = n => [Math.min(n, 65535), Math.ceil(n / 65535), 1]
+  const norm = (x, y, rows, n, g) => step(NORM(!!g), [x, g ? scratch : y, g ?? maskWt, uni(rows, n)], grid1(rows))
+
+  // a group's band inputs: band j's columns of E at frame t (none past T), normalized: A[j, t, e]
+  const SPLIT = pipe('split', HEAD + `
+@group(0) @binding(0) var<storage, read> E: array<f32>;
+@group(0) @binding(1) var<storage, read> idx: array<i32>;
+@group(0) @binding(2) var<storage, read_write> A: array<f32>;
+struct P { Wp: u32, Tp: u32, T: u32, cols: u32 }
+@group(0) @binding(3) var<uniform> p: P;
+@compute @workgroup_size(32)
+fn main(@builtin(workgroup_id) wg: vec3u, @builtin(subgroup_invocation_id) lane: u32) {
+  let t = wg.x; let j = wg.y;
+  var s = 0.0;
+  for (var e = lane; e < p.Wp; e += 32u) { let v = select(0.0, E[t * (p.cols + 1u) + u32(idx[j * p.Wp + e])], t < p.T); s += v * v; }
+  let k = 1.0 / max(sqrt(subgroupAdd(s)), 1e-12);
+  for (var e = lane; e < p.Wp; e += 32u) {
+    let v = select(0.0, E[t * (p.cols + 1u) + u32(idx[j * p.Wp + e])], t < p.T);
+    A[(j * p.Tp + t) * p.Wp + e] = v * k;
+  }
+}`)
+
+  // q, k, v of each head to sequences, the rotation on q and k, the gates: token r = band·Tp + t; over time a sequence is
+  // a band, its positions frames (n T, padded to L); over frequency a frame, its positions bands (n B, padded to L)
+  const ROPE = pipe('rope', HEAD + `
+@group(0) @binding(0) var<storage, read> Yq: array<f32>;
+@group(0) @binding(1) var<storage, read_write> Q: array<f32>;
+@group(0) @binding(2) var<storage, read_write> K: array<f32>;
+@group(0) @binding(3) var<storage, read_write> V: array<f32>;
+@group(0) @binding(4) var<storage, read_write> Gt: array<f32>;
+@group(0) @binding(5) var<storage, read> cs: array<f32>;
+@group(0) @binding(6) var<storage, read> sn: array<f32>;
+@group(0) @binding(7) var<storage, read> gb: array<f32>;
+struct P { time: u32, Tp: u32, T: u32, L: u32, ld: u32 }
+@group(0) @binding(8) var<uniform> p: P;
+@compute @workgroup_size(64)
+fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) li: vec3u) {
+  let t = wg.x; let band = wg.y; let h = wg.z; let d = li.x;
+  if (t >= p.T) { return; }
+  let r = band * p.Tp + t;
+  var seq = band; var pos = t;
+  if (p.time == 0u) { seq = t; pos = band; }
+  let o = ((seq * ${H}u + h) * p.L + pos) * 64u + d;
+  let y = r * p.ld + h * 64u;
+  let rot = select(Yq[y + d - 32u], -Yq[y + d + 32u], d < 32u);
+  let rk = select(Yq[y + ${H * DH}u + d - 32u], -Yq[y + ${H * DH}u + d + 32u], d < 32u);
+  Q[o] = Yq[y + d] * cs[pos * 64u + d] + rot * sn[pos * 64u + d];
+  K[o] = Yq[y + ${H * DH}u + d] * cs[pos * 64u + d] + rk * sn[pos * 64u + d];
+  V[o] = Yq[y + ${2 * H * DH}u + d];
+  if (d == 0u) { Gt[r * ${H}u + h] = 1.0 / (1.0 + exp(-(Yq[r * p.ld + ${QKV}u + h] + gb[h]))); }
+}`)
+  // and back: Z[r, h·64 + d] = O · gate
+  const MERGE = pipe('merge', HEAD + `
+@group(0) @binding(0) var<storage, read> O: array<f32>;
+@group(0) @binding(1) var<storage, read> Gt: array<f32>;
+@group(0) @binding(2) var<storage, read_write> Zb: array<f32>;
+struct P { time: u32, Tp: u32, T: u32, L: u32 }
+@group(0) @binding(3) var<uniform> p: P;
+@compute @workgroup_size(64)
+fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) li: vec3u) {
+  let t = wg.x; let band = wg.y; let h = wg.z; let d = li.x;
+  let r = band * p.Tp + t;
+  if (t >= p.T) { Zb[r * ${H * DH}u + h * 64u + d] = 0.0; return; }
+  var seq = band; var pos = t;
+  if (p.time == 0u) { seq = t; pos = band; }
+  Zb[r * ${H * DH}u + h * 64u + d] = O[((seq * ${H}u + h) * p.L + pos) * 64u + d] * Gt[r * ${H}u + h];
+}`)
+  // attention in two passes over the keys, 32 a block, 16 query rows a subgroup: the rows' maxima first (Q Kᵀ alone),
+  // then exp(S − max) and the output accumulated in registers, no rescaling; keys past n masked
+  const BR = 16, BC = 32, SG = 4, I2 = [0, 1], J4 = [0, 1, 2, 3], D8 = [...Array(8).keys()]
+  const scores = `${I2.map(i => J4.map(j => `var s${i}_${j}: subgroup_matrix_result<f32, 8, 8>;`).join(' ')).join('\n    ')}
+    for (var k = 0u; k < 8u; k++) {
+      ${I2.map(i => `let q${i} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(&Q, base + (r0 + ${i * 8}u) * 64u + k * 8u, false, 64u);`).join('\n      ')}
+      ${J4.map(j => `let k${j} = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>>(&K, base + (c0 + ${j * 8}u) * 64u + k * 8u, true, 64u);`).join('\n      ')}
+      ${I2.map(i => J4.map(j => `s${i}_${j} = subgroupMatrixMultiplyAccumulate(q${i}, k${j}, s${i}_${j});`).join(' ')).join('\n      ')}
+    }
+    ${I2.map(i => J4.map(j => `subgroupMatrixStore(&S[sg], ${i * 8 * BC + j * 8}u, s${i}_${j}, false, ${BC}u);`).join(' ')).join('\n    ')}
+    workgroupBarrier();`
+  const ATT = pipe('att', HEAD + `
+@group(0) @binding(0) var<storage, read> Q: array<f32>;
+@group(0) @binding(1) var<storage, read> K: array<f32>;
+@group(0) @binding(2) var<storage, read> V: array<f32>;
+@group(0) @binding(3) var<storage, read_write> O: array<f32>;
+struct P { n: u32, np: u32 }
+@group(0) @binding(4) var<uniform> p: P;
+var<workgroup> S: array<array<f32, ${BR * BC}>, ${SG}>;
+var<workgroup> Ow: array<array<f32, ${BR * 64}>, ${SG}>;
+@compute @workgroup_size(${32 * SG})
+fn main(@builtin(workgroup_id) wg: vec3u, @builtin(subgroup_id) sg: u32, @builtin(subgroup_invocation_id) lane: u32) {
+  let b = wg.y;
+  let r0 = (wg.x * ${SG}u + sg) * ${BR}u;
+  let base = b * p.np * 64u;
+  let row = lane / 2u; let half = (lane % 2u) * ${BC / 2}u;    // two lanes a row, 16 columns each
+  var mx = -1e30;
+  for (var c0 = 0u; c0 < p.n; c0 += ${BC}u) {
+    ${scores}
+    for (var j = 0u; j < ${BC / 2}u; j++) { if (c0 + half + j < p.n) { mx = max(mx, S[sg][row * ${BC}u + half + j]); } }
+    workgroupBarrier();
+  }
+  mx = max(mx, subgroupShuffleXor(mx, 1u));
+  var l = 0.0;
+  ${I2.map(i => D8.map(t => `var o${i}_${t}: subgroup_matrix_result<f32, 8, 8>;`).join(' ')).join('\n  ')}
+  for (var c0 = 0u; c0 < p.n; c0 += ${BC}u) {
+    ${scores}
+    for (var j = 0u; j < ${BC / 2}u; j++) {
+      let at = row * ${BC}u + half + j;
+      let e = select(0.0, exp(S[sg][at] - mx), c0 + half + j < p.n);
+      S[sg][at] = e; l += e;
+    }
+    workgroupBarrier();
+    ${I2.map(i => J4.map(j => `let p${i}_${j} = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(&S[sg], ${i * 8 * BC + j * 8}u, false, ${BC}u);`).join(' ')).join('\n    ')}
+    ${D8.map(t => J4.map(j => `let v${j}_${t} = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>>(&V, base + (c0 + ${j * 8}u) * 64u + ${t * 8}u, false, 64u);`).join(' ') + '\n    ' + I2.map(i => J4.map(j => `o${i}_${t} = subgroupMatrixMultiplyAccumulate(p${i}_${j}, v${j}_${t}, o${i}_${t});`).join(' ')).join(' ')).join('\n    ')}
+    workgroupBarrier();
+  }
+  l += subgroupShuffleXor(l, 1u);
+  ${I2.map(i => D8.map(t => `subgroupMatrixStore(&Ow[sg], ${i * 8 * 64 + t * 8}u, o${i}_${t}, false, 64u);`).join(' ')).join('\n  ')}
+  workgroupBarrier();
+  // lane pairs hold their row's sum; each lane writes half its row
+  for (var d = 0u; d < 32u; d++) {
+    let r = r0 + row;
+    if (r < p.n) { O[base + r * 64u + half * 2u + d] = Ow[sg][row * 64u + half * 2u + d] / l; }
+  }
+}`)
+  // GLU: value half by the sigmoid of the gate half
+  const GLU = pipe('glu', HEAD + `
+@group(0) @binding(0) var<storage, read> Yg: array<f32>;
+@group(0) @binding(1) var<storage, read_write> Mo: array<f32>;
+struct P { rows: u32, Wp: u32, at: u32 }
+@group(0) @binding(2) var<uniform> p: P;
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gi: vec3u) {
+  let i = gi.x + gi.y * 65535u * 64u;
+  if (i >= p.rows * p.Wp) { return; }
+  let r = i / p.Wp; let e = i % p.Wp;
+  Mo[p.at + i] = Yg[r * 2u * p.Wp + e] / (1.0 + exp(-Yg[r * 2u * p.Wp + p.Wp + e]));
+}`)
+  // the masks averaged over the bands that cover a column, times the spectrum: out planes [s·2 + c][f][t]
+  const MASK = pipe('mask', HEAD + `
+@group(0) @binding(0) var<storage, read> Mo: array<f32>;
+@group(0) @binding(1) var<storage, read> map: array<u32>;
+@group(0) @binding(2) var<storage, read> wt: array<f32>;
+@group(0) @binding(3) var<storage, read> E: array<f32>;
+@group(0) @binding(4) var<storage, read_write> Out: array<f32>;
+struct P { T: u32, F: u32, cols: u32 }
+@group(0) @binding(5) var<uniform> p: P;
+fn mask(t: u32, o: u32) -> f32 {
+  let a = map[o * 8u]; let sa = map[o * 8u + 1u]; let on1 = map[o * 8u + 2u];
+  let b = map[o * 8u + 3u]; let sb = map[o * 8u + 4u]; let on2 = map[o * 8u + 5u];
+  var v = 0.0;
+  if (on1 == 1u) { v += Mo[a + t * sa]; }
+  if (on2 == 1u) { v += Mo[b + t * sb]; }
+  return v * wt[o];
+}
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gi: vec3u) {
+  let t = gi.x; let fs = gi.y;          // fs = f·S + s
+  if (t >= p.T) { return; }
+  let f = fs / ${S}u; let s = fs % ${S}u;
+  let o = fs * 2u;
+  let er = E[t * (p.cols + 1u) + o]; let ei = E[t * (p.cols + 1u) + o + 1u];
+  let mr = mask(t, o); let mi = mask(t, o + 1u);
+  Out[((s * 2u) * p.F + f) * p.T + t] = er * mr - ei * mi;
+  Out[((s * 2u + 1u) * p.F + f) * p.T + t] = er * mi + ei * mr;
+}`)
+
+  // ---------------------------------------------------------------- the run, recorded once
+  for (const [g, gr] of G.entries()) {
+    step(SPLIT, [E, W.splitIdx[g], splitA[g], uni(gr.Wp, Tp, T, COLS)], [Tp, gr.n, 1])
+    mm('b', splitA[g], W.split[g], X, W.splitB[g], { M: Tp, N: D, K: gr.Wp, batch: gr.n, sa: Tp * gr.Wp, sb: gr.Wp * D, c0: gr.b0 * Tp * D, sc: Tp * D, sbias: D })
+  }
+  for (const [i, w] of W.blocks.entries()) {
+    const time = i % 2 === 0, L = time ? LT : LF, n = time ? T : B, seqs = time ? B : T
+    norm(X, XN, B * Tp, D)
+    mm('', XN, w.qkvg, Y, null, { M: B * Tp, N: QKVG, K: D })
+    step(ROPE, [Y, Qb, Kb, Vb, Gt, time ? cosT : cosF, time ? sinT : sinF, w.gb, uni(time ? 1 : 0, Tp, T, L, QKVG)], [Tp, B, H])
+    step(ATT, [Qb, Kb, Vb, Ob, uni(n, L)], [Math.ceil(L / (BR * SG)), seqs * H, 1])
+    step(MERGE, [Ob, Gt, Z, uni(time ? 1 : 0, Tp, T, L)], [Tp, B, H])
+    mm('r', Z, w.wo, X, null, { M: B * Tp, N: D, K: H * DH })
+    norm(X, XN, B * Tp, D)
+    mm('bg', XN, w.w1, FFH, w.b1, { M: B * Tp, N: FF, K: D })
+    mm('br', FFH, w.w2, X, w.b2, { M: B * Tp, N: D, K: FF })
+    norm(X, X, B * Tp, D, w.out)
+  }
+  mm('bt', X, W.mw[0], FFH, W.mb[0], { M: Tp, N: MH, K: D, batch: B, sa: Tp * D, sb: D * MH, sc: Tp * MH, sbias: MH })
+  mm('bt', FFH, W.mw[1], H2, W.mb[1], { M: Tp, N: MH, K: MH, batch: B, sa: Tp * MH, sb: MH * MH, sc: Tp * MH, sbias: MH })
+  for (const [g, gr] of G.entries()) {
+    mm('b', H2, W.m2[g], glu[g], W.m2B[g], { M: Tp, N: 2 * gr.Wp, K: MH, batch: gr.n, a0: gr.b0 * Tp * MH, sa: Tp * MH, sb: MH * 2 * gr.Wp, sc: Tp * 2 * gr.Wp, sbias: 2 * gr.Wp })
+    const n = gr.n * Tp * gr.Wp
+    step(GLU, [glu[g], Mg, uni(gr.n * Tp, gr.Wp, gofs[g])], [Math.min(Math.ceil(n / 64), 65535), Math.ceil(n / 64 / 65535), 1])
+  }
+  step(MASK, [Mg, maskMap, maskWt, E, OUT, uni(T, F, COLS)], [Math.ceil(T / 64), F * S, 1])
+  const zero = device.createCommandEncoder()
+  for (const b of [Qb, Kb, Vb]) zero.clearBuffer(b)
+  device.queue.submit([zero.finish()])
+
+  const stage = new Float32Array(Tp * (COLS + 1)), read = device.createBuffer({ size: 4 * F * T * 4, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ })
+  return {
+    steps: steps.length,
+    async run(spec) {
+      // planes [s·2 + c][f][t] → E[t][(f·S + s)·2 + c], a zero column past them, rows past T zero
+      stage.fill(0)
+      for (let s = 0; s < S; s++) for (let c = 0; c < 2; c++) for (let f = 0; f < F; f++) {
+        const plane = ((s * 2 + c) * F + f) * T, col = (f * S + s) * 2 + c
+        for (let t = 0; t < T; t++) stage[t * (COLS + 1) + col] = spec[plane + t]
+      }
+      device.queue.writeBuffer(E, 0, stage)
+      const e = device.createCommandEncoder(), pass = e.beginComputePass()
+      for (const [p, bg, grid] of steps) { pass.setPipeline(p); pass.setBindGroup(0, bg); pass.dispatchWorkgroups(...grid) }
+      pass.end()
+      e.copyBufferToBuffer(OUT, 0, read, 0, 4 * F * T * 4)
+      device.queue.submit([e.finish()])
+      await read.mapAsync(GPUMapMode.READ)
+      const out = new Float32Array(read.getMappedRange().slice(0))
+      read.unmap()
+      return out
+    },
+  }
+}

package/scripts/engine.mjs

@@ -0,0 +1,58 @@
+#!/usr/bin/env node
+// roformer.js, the package's own WebGPU engine for Mel-Band RoFormer, against onnxruntime on the same weights: Chromium on
+// the GPU (playwright's 'chromium' channel: Metal here) runs the engine on one segment of scripts/reference.py's mix,
+// onnxruntime-node's CPU the float16 ONNX file on the same spectrum; prints how far apart they are and the engine's time.
+//
+//   node scripts/engine.mjs            (needs playwright, onnxruntime-node, and in the neural cache, mel-roformer/:
+//                                       export-roformer.py --verify --engine, compact.py --model mel-roformer --as fp16)
+import { createServer } from 'node:http'
+import { readFileSync } from 'node:fs'
+import path from 'node:path'
+import os from 'node:os'
+import { fileURLToPath } from 'node:url'
+import { chromium } from 'playwright'
+import * as ort from 'onnxruntime-node'
+import { stft } from '../separate.js'
+
+const HERE = path.dirname(fileURLToPath(import.meta.url)), CACHE = path.join(process.env.AUDIO_NEURAL_CACHE || path.join(os.homedir(), '.cache', 'audiojs', 'neural'), 'mel-roformer')
+const F = 1025, T = 801, L = 352800, f32 = f => { let b = readFileSync(f); return new Float32Array(b.buffer, b.byteOffset, b.byteLength / 4) }
+
+// one segment's spectrum, as separate.js's complexSegment makes it: [L re, L im, R re, R im][f][t]
+const mix = f32(path.join(CACHE, 'test.f32')), N = mix.length / 2, spec = new Float32Array(4 * F * T)
+for (let c = 0; c < 2; c++) {
+  let { re, im } = stft(Float64Array.from(mix.subarray(c * N, c * N + L)), { n: 2048, hop: 441 })
+  for (let f = 0; f < F; f++) for (let t = 0; t < T; t++) { spec[(2 * c * F + f) * T + t] = re[t][f]; spec[((2 * c + 1) * F + f) * T + t] = im[t][f] }
+}
+let cpu = await ort.InferenceSession.create(path.join(CACHE, 'mel-roformer.fp16.onnx'))
+let want = (await cpu.run({ mix_spec: new ort.Tensor('float32', spec, [1, 4, F, T]) })).stems_spec.data
+
+const files = { '/roformer.js': path.join(HERE, '..', 'roformer.js'), '/meta': path.join(CACHE, 'mel-roformer.engine.json'), '/bin': path.join(CACHE, 'mel-roformer.engine.bin') }
+const server = createServer((q, r) => {
+  if (q.url === '/') return r.writeHead(200, { 'content-type': 'text/html' }).end('<!doctype html><title>engine</title>')
+  if (q.url === '/spec') return r.writeHead(200).end(Buffer.from(spec.buffer))
+  if (!files[q.url]) return r.writeHead(404).end()
+  r.writeHead(200, { 'content-type': q.url.endsWith('.js') ? 'text/javascript' : 'application/octet-stream' }).end(readFileSync(files[q.url]))
+})
+await new Promise(r => server.listen(0, '127.0.0.1', r))
+const browser = await chromium.launch({ channel: 'chromium', args: ['--enable-unsafe-webgpu'] }), page = await browser.newPage()
+await page.goto(`http://127.0.0.1:${server.address().port}/`)
+const got = await page.evaluate(async () => {
+  const { engine } = await import('/roformer.js')
+  const meta = await (await fetch('/meta')).json(), bytes = new Uint8Array(await (await fetch('/bin')).arrayBuffer())
+  let t = performance.now()
+  const s = await engine(meta, bytes)
+  if (!s) return { none: true }
+  const ready = performance.now() - t, spec = new Float32Array(await (await fetch('/spec')).arrayBuffer()), times = []
+  let out
+  for (let k = 0; k < 4; k++) { t = performance.now(); out = (await s.run({ mix_spec: { data: spec } })).stems_spec.data; times.push(performance.now() - t) }
+  return { ready, times, out: [...out] }
+})
+await browser.close(); server.close()
+if (got.none) console.log('no GPU with subgroup matrices here: the engine stands aside for onnxruntime')
+else {
+  let e = 0, p = 0
+  for (let i = 0; i < want.length; i++) { e += (got.out[i] - want[i]) ** 2; p += want[i] ** 2 }
+  console.log(`engine ready in ${(got.ready / 1000).toFixed(1)} s; a segment ${got.times.map(t => t.toFixed(0)).join(', ')} ms; ` +
+    `${(10 * Math.log10(p / e)).toFixed(1)} dB from onnxruntime's CPU on the same weights`)
+}
+process.exit(0)

package/scripts/export-roformer.py

@@ -0,0 +1,481 @@
+#!/usr/bin/env python3
+"""Export Mel-Band RoFormer (Wang, Lu, Won, Chen, Song, "Mel-Band RoFormer for Music Source Separation", ISMIR LBD
+2023) to the ONNX contract @audio/neural-separate expects for modelType 'complex': the network between its STFT and
+iSTFT, which JS computes, written for the GPU.
+
+    pip install torch onnx onnxruntime pyyaml einops beartype rotary-embedding-torch librosa
+    python3 export-roformer.py --model mel-roformer --verify
+
+The model is Kimberley Jensen's vocals Mel-Band RoFormer (huggingface.co/KimberleyJSN/melbandroformer, MIT), with the
+config ZFTurbo/Music-Source-Separation-Training (MIT) runs it by (configs/KimberleyJensen/
+config_vocals_mel_band_roformer_kj.yaml) and that repository's models/bs_roformer code, both at a pinned commit.
+
+One 8 s segment (352800 samples at 44.1 kHz, 801 frames) per run:
+  input    mix_spec    [1, 4, 1025, 801]     STFT, complex as channels (L re, L im, R re, R im)
+  output   stems_spec  [1, 4, 1025, 801]     the vocals, per channel, re and im
+The STFT is the model's own: n_fft 2048, hop 441, Hann window of 2048, unnormalized, centered with reflect padding.
+
+The graph computes what MelBandRoformer.forward computes between its STFT and iSTFT, rearranged so a GPU runs it in few
+and large operations; none of it changes the function (--verify: against the PyTorch forward):
+  - the 60 bands' input projections and mask MLPs run as batched matrix products, bands of near widths grouped and
+    zero-padded to the widest of their group (10 groups, none wider than 1.3 times its narrowest), not 60
+    small products each;
+  - each RMSNorm's scale (√d·γ) is folded into the weights of the linear layer it feeds; the attention's query scale
+    into the query weights; its gates' projection joins the q, k, v projection;
+  - the rotary embedding's interleaved pairs become halves (the same permutation of q's and k's head features leaves
+    q·k as it was), so the rotation is two products with tables of cos and sin and one with a fixed 64×64 matrix;
+  - the masks of overlapping mel bands are averaged by two gathers and a weight, not a scatter, the DC bin's weight 0
+    (zero_dc);
+  - no tensor passes 100 MB: each transformer runs on slices of its sequences, the mask MLP by group of bands
+    (onnxruntime's WebGPU Softmax went wrong past WebGPU's default 128 MiB storage binding).
+
+Writes <out-dir>/mel-roformer.onnx (928 MB), and with --fp16 mel-roformer.fp16.onnx (float16 weights, float32 compute,
+half the download). Products computed in float16 measured slower on WebGPU (3.2 against 2.6 s a segment): not offered.
+Default out-dir: $AUDIO_NEURAL_CACHE/<model>, else ~/.cache/audiojs/neural/<model>. --verify compares the stems rebuilt
+from onnxruntime's output with MelBandRoformer.forward on the same segment and writes the reference separation test.js
+compares the JS pipeline against (scripts/reference.py, chunked as separate.js chunks: step 0.5, Kim's num_overlap 2).
+Last run: torch 2.14, onnxruntime 1.31.
+"""
+import argparse
+import gc
+import math
+import os
+import sys
+import urllib.request
+from pathlib import Path
+
+REPO = "https://raw.githubusercontent.com/ZFTurbo/Music-Source-Separation-Training/84b1eac0887756b4f1a9d7a1ff49105939749ed2/"
+MODELS = {  # checkpoint URL, its SHA-256, config path in REPO
+    "mel-roformer": ("https://huggingface.co/KimberleyJSN/melbandroformer/resolve/ac9b0614ab3cd7f77219e18ba494dfd93956c348/MelBandRoformer.ckpt",
+                     "87201f4d31afb5bc79993230fc49446918425574db48c01c405e44f365c7559e",
+                     "configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml"),
+}
+SEGMENT = 352800
+STEP = 0.5
+
+
+def fetch(url, path):
+    if not path.exists():
+        print(f"fetching {url}")
+        urllib.request.urlretrieve(url, str(path) + ".part")
+        os.replace(str(path) + ".part", path)
+    return path
+
+
+def load(src, name="mel-roformer"):
+    """MelBandRoformer with its weights: the model code, config and checkpoint fetched into `src` once."""
+    import hashlib
+    import torch
+    import yaml
+
+    pkg = src / "roformer_code" / "models" / "bs_roformer"
+    pkg.mkdir(parents=True, exist_ok=True)
+    (pkg.parent / "__init__.py").touch()
+    (pkg / "__init__.py").touch()
+    for f in ("mel_band_roformer.py", "attend.py"):
+        fetch(REPO + "models/bs_roformer/" + f, pkg / f)
+    sys.path.insert(0, str(src / "roformer_code"))
+    sys.dont_write_bytecode = True
+    from models.bs_roformer.mel_band_roformer import MelBandRoformer
+
+    class Loader(yaml.SafeLoader): pass  # the config's resolutions are a !!python/tuple
+    Loader.add_constructor("tag:yaml.org,2002:python/tuple", lambda l, n: tuple(l.construct_sequence(n)))
+    url, sha, config = MODELS[name]
+    cfg = yaml.load(fetch(REPO + config, src / Path(config).name).read_text(), Loader=Loader)
+    ckpt = fetch(url, src / Path(url).name)
+    h = hashlib.sha256()
+    with open(ckpt, "rb") as f:
+        for b in iter(lambda: f.read(1 << 24), b""):
+            h.update(b)
+    if h.hexdigest() != sha:
+        raise SystemExit(f"{ckpt}: SHA-256 {h.hexdigest()}, expected {sha}")
+    model = MelBandRoformer(**cfg["model"])
+    state = torch.load(ckpt, map_location="cpu", weights_only=True)
+    model.load_state_dict(state.get("state_dict", state) if isinstance(state, dict) else state)
+    return model.eval(), cfg
+
+
+def stft(model, x):
+    """The model's STFT: (B, C, L) -> (B, 2C, F, T), complex as channels."""
+    import torch
+
+    B, C, L = x.shape
+    z = torch.view_as_real(torch.stft(x.reshape(-1, L), **model.stft_kwargs, window=model.stft_window_fn(), return_complex=True))
+    return z.permute(0, 3, 1, 2).reshape(B, 2 * C, z.shape[1], z.shape[2])
+
+
+def istft(model, y, length):
+    """(B, 2·S·C, F, T) -> (B, S, C, L)"""
+    import torch
+
+    B, _, Fr, T = y.shape
+    y = torch.view_as_complex(y.reshape(-1, 2, Fr, T).permute(0, 2, 3, 1).contiguous())
+    x = torch.istft(y, **model.stft_kwargs, window=model.stft_window_fn(), length=length)
+    return x.reshape(B, -1, model.audio_channels, length)
+
+
+def groups(widths, slack=1.3):
+    """Contiguous runs of bands whose widest is at most `slack` times their narrowest: [(first, last + 1, widest)]"""
+    out, b0 = [], 0
+    for b in range(1, len(widths) + 1):
+        if b == len(widths) or max(widths[b0:b + 1]) > slack * min(widths[b0:b + 1]):
+            out.append((b0, b, max(widths[b0:b])))
+            b0 = b
+    return out
+
+
+class Core:
+    """MelBandRoformer.forward between its STFT and its iSTFT, for the GPU (the module docstring says how)."""
+
+    @staticmethod
+    def build(model, T):
+        import torch
+        import torch.nn as nn
+        import torch.nn.functional as F
+
+        f64 = torch.float64
+        Fq = model.num_bands_per_freq.numel()                       # 1025
+        S = model.audio_channels
+        fi = model.freq_indices.tolist()                            # (f·S + s) per band's pair, bands in order
+        nf = model.num_freqs_per_band.tolist()
+        widths = [2 * n * S for n in nf]                            # a band's input: its (f, s) pairs, re and im
+        G = groups(widths)
+        lin = torch.matmul
+        # every tensor under WebGPU's default storage binding (128 MiB, maxStorageBufferBindingSize): onnxruntime
+        # 1.30's WebGPU Softmax over a larger one came out wrong one run in two (36 dB from the CPU's at 307 MB, right
+        # at 123 MB), so a transformer runs on as many sequences at a time as keep its largest, the attention's
+        # scores (heads × n × n) or its feed-forward's hidden layer (n × 4·dim), under 100 MB
+        BUDGET = 100e6
+
+        def sliced(f, x):
+            per = max(f.biggest(x.shape[1]), 1)
+            size = max(1, int(BUDGET // per))
+            return torch.cat([f(x[i:i + size]) for i in range(0, x.shape[0], size)]) if size < x.shape[0] else f(x)
+
+        rot = lambda n, emb: emb(emb.get_seq_pos(n, device="cpu", dtype=torch.float32))  # (n, dh) interleaved
+
+        def rms(m):
+            return (m.gamma.detach().double() * m.scale)
+
+        class Mod(nn.Module):
+            def __init__(s):
+                super().__init__()
+                # the STFT as columns: E[t, (f·S + s)·2 + c], and a zero column past them (index Fq·S·2)
+                pos = []                                            # each band's input columns, in its order
+                o = 0
+                for n in nf:
+                    pos.append([fi[k] * 2 + c for k in range(o, o + n * S) for c in range(2)])
+                    o += n * S
+                zero = Fq * S * 2
+                bs = model.band_split.to_features
+                s.split_idx, s.split_w, s.split_b = [], nn.ParameterList(), nn.ParameterList()
+                for b0, b1, W in G:
+                    idx = torch.full((b1 - b0, W), zero, dtype=torch.long)
+                    w = torch.zeros(b1 - b0, W, bs[0][1].out_features, dtype=f64)
+                    bias = torch.zeros(b1 - b0, 1, w.shape[2], dtype=f64)
+                    for j, b in enumerate(range(b0, b1)):
+                        norm, lin = bs[b][0], bs[b][1]
+                        idx[j, :widths[b]] = torch.tensor(pos[b])
+                        w[j, :widths[b]] = (lin.weight.detach().double() * rms(norm)[None]).T
+                        bias[j, 0] = lin.bias.detach().double()
+                    s.register_buffer(f"si{len(s.split_idx)}", idx)
+                    s.split_idx.append(f"si{len(s.split_idx)}")
+                    s.split_w.append(nn.Parameter(w.float(), requires_grad=False))
+                    s.split_b.append(nn.Parameter(bias.float(), requires_grad=False))
+                # transformers: time then frequency, per depth
+                s.blocks = nn.ModuleList()
+                for time, freq in model.layers:
+                    s.blocks.append(nn.ModuleList([Block(time, T), Block(freq, len(nf))]))
+                # mask estimator (one stem): its MLP's linear layers before the last batched over bands (their shapes
+                # alike), the last by group, then GLU
+                me = model.mask_estimators[0].to_freqs
+                lin = [[m for m in me[b][0] if isinstance(m, nn.Linear)] for b in range(len(nf))]
+                s.mw = nn.ParameterList([nn.Parameter(torch.stack([l[i].weight.detach().T for l in lin]), requires_grad=False) for i in range(len(lin[0]) - 1)])
+                s.mb = nn.ParameterList([nn.Parameter(torch.stack([l[i].bias.detach() for l in lin])[:, None], requires_grad=False) for i in range(len(lin[0]) - 1)])
+                s.m2_w, s.m2_b = nn.ParameterList(), nn.ParameterList()
+                col, cols = 0, {}
+                for b0, b1, W in G:
+                    H = lin[0][-1].in_features
+                    w = torch.zeros(b1 - b0, H, 2 * W)
+                    bias = torch.zeros(b1 - b0, 1, 2 * W)
+                    for j, b in enumerate(range(b0, b1)):
+                        l2 = lin[b][-1]
+                        n = widths[b]
+                        w[j, :, :n], w[j, :, W:W + n] = l2.weight.detach()[:n].T, l2.weight.detach()[n:].T
+                        bias[j, 0, :n], bias[j, 0, W:W + n] = l2.bias.detach()[:n], l2.bias.detach()[n:]
+                        for e in range(n):
+                            cols.setdefault(pos[b][e], []).append(col + j * W + e)
+                    s.m2_w.append(nn.Parameter(w, requires_grad=False))
+                    s.m2_b.append(nn.Parameter(bias, requires_grad=False))
+                    col += (b1 - b0) * W
+                # each output column's sources among the masks (at most two bands overlap), the zero column for none,
+                # and its weight: 1 over the bands that cover it, 0 at DC
+                i1 = torch.full((zero,), col, dtype=torch.long)
+                i2 = torch.full((zero,), col, dtype=torch.long)
+                wt = torch.zeros(zero)
+                for o, src in cols.items():
+                    assert len(src) <= 2, "three bands over one frequency"
+                    i1[o] = src[0]
+                    i2[o] = src[-1] if len(src) == 2 else col
+                    wt[o] = 0 if o // (S * 2) == 0 and model.zero_dc else 1 / len(src)
+                s.register_buffer("i1", i1)
+                s.register_buffer("i2", i2)
+                s.register_buffer("wt", wt)
+                s.Fq, s.S = Fq, S
+
+            def forward(s, mix_spec):
+                Fq, S = s.Fq, s.S
+                T_ = mix_spec.shape[3]
+                # (1, 2S, F, T) -> E (T, F·S·2): [t, (f·S + s)·2 + c]
+                E = mix_spec.reshape(S, 2, Fq, T_).permute(3, 2, 0, 1).reshape(T_, Fq * S * 2)
+                Ez = torch.cat([E, torch.zeros(T_, 1, dtype=E.dtype)], 1)
+                xs = []
+                for k, name in enumerate(s.split_idx):
+                    idx = getattr(s, name)                          # (nb, W)
+                    x = Ez[:, idx].permute(1, 0, 2)                 # (nb, T, W)
+                    x = x / x.norm(dim=-1, keepdim=True).clamp(min=1e-12)
+                    xs.append(lin(x, s.split_w[k]) + s.split_b[k])
+                x = torch.cat(xs, 0)                                # (B, T, D)
+                for time, freq in s.blocks:
+                    x = sliced(time, x)                             # sequences over time, one per band
+                    x = sliced(freq, x.transpose(0, 1)).transpose(0, 1)  # over bands, one per frame
+                ms = []
+                for k, (lo, hi, W) in enumerate(G):
+                    h = x[lo:hi]
+                    for w, b in zip(s.mw, s.mb):
+                        h = torch.tanh(lin(h, w[lo:hi]) + b[lo:hi])  # (nb, T, H)
+                    y = lin(h, s.m2_w[k]) + s.m2_b[k]               # (nb, T, 2W)
+                    m = y[..., :W] * torch.sigmoid(y[..., W:])      # GLU
+                    ms.append(m.permute(1, 0, 2).reshape(T_, (hi - lo) * W))
+                M = torch.cat(ms + [torch.zeros(T_, 1)], 1)         # (T, columns + 1)
+                mask = (M[:, s.i1] + M[:, s.i2]) * s.wt             # (T, F·S·2)
+                e = E.reshape(T_, Fq, S, 2)
+                mk = mask.reshape(T_, Fq, S, 2)
+                er, ei, mr, mi = e[..., 0], e[..., 1], mk[..., 0], mk[..., 1]
+                out = torch.stack([er * mr - ei * mi, er * mi + ei * mr], -1)   # (T, F, S, 2)
+                return out.permute(2, 3, 1, 0).reshape(1, 2 * S, Fq, T_)
+
+        class Block(nn.Module):
+            """One Transformer of the model (its layers and output norm) over sequences of n: x (batch, n, D)."""
+
+            def __init__(s, tr, n):
+                super().__init__()
+                s.layers = nn.ModuleList()
+                for attn, ff in tr.layers:
+                    s.layers.append(nn.ModuleList([Attn(attn, n), FF(ff)]))
+                s.out = nn.Parameter(rms(tr.norm).float(), requires_grad=False)
+
+            def biggest(s, n):
+                """bytes of the largest tensor one sequence of n makes"""
+                a, ff = s.layers[0]
+                return 4 * max(a.h * n * n, n * ff.w1.shape[1], n * a.qkvg.shape[1])
+
+            def forward(s, x):
+                for attn, ff in s.layers:
+                    x = attn(x) + x
+                    x = ff(x) + x
+                return x / x.norm(dim=-1, keepdim=True).clamp(min=1e-12) * s.out
+
+        class Attn(nn.Module):
+            def __init__(s, a, n):
+                super().__init__()
+                h, dh = a.heads, a.to_qkv.out_features // (3 * a.heads)
+                g = rms(a.norm)[None]
+                wq, wk, wv = (a.to_qkv.weight.detach().double() * g).reshape(3, h, dh, -1)
+                # interleaved pairs -> halves, the same permutation of q's and k's features
+                perm = torch.cat([torch.arange(0, dh, 2), torch.arange(1, dh, 2)])
+                wq, wk = wq[:, perm] * a.scale, wk[:, perm]
+                wg = a.to_gates.weight.detach().double() * g
+                s.qkvg = nn.Parameter(torch.cat([wq.reshape(h * dh, -1), wk.reshape(h * dh, -1), wv.reshape(h * dh, -1), wg]).T.float(), requires_grad=False)
+                s.gb = nn.Parameter(a.to_gates.bias.detach().float(), requires_grad=False)
+                s.wo = nn.Parameter(a.to_out[0].weight.detach().T.float(), requires_grad=False)
+                freqs = rot(n, a.rotary_embed).double()[:, perm]  # (n, dh)
+                s.register_buffer("cos", freqs.cos().float())
+                s.register_buffer("sin", freqs.sin().float())
+                R = torch.zeros(dh, dh)                             # x @ R: (-x[half:], x[:half])
+                R[torch.arange(dh // 2) + dh // 2, torch.arange(dh // 2)] = -1
+                R[torch.arange(dh // 2), torch.arange(dh // 2) + dh // 2] = 1
+                s.register_buffer("R", R)
+                s.h, s.dh = h, dh
+
+            def forward(s, x):
+                Bt, n, _ = x.shape
+                h, dh = s.h, s.dh
+                x = x / x.norm(dim=-1, keepdim=True).clamp(min=1e-12)
+                y = lin(x, s.qkvg)
+                q, k, v = (y[..., i * h * dh:(i + 1) * h * dh].reshape(Bt, n, h, dh).transpose(1, 2) for i in range(3))
+                gates = torch.sigmoid(y[..., 3 * h * dh:] + s.gb)  # (Bt, n, h)
+                q = q * s.cos + torch.matmul(q, s.R) * s.sin
+                k = k * s.cos + torch.matmul(k, s.R) * s.sin
+                o = torch.matmul(torch.softmax(torch.matmul(q, k.transpose(-1, -2)), -1), v)  # (Bt, h, n, dh)
+                o = o * gates.transpose(1, 2)[..., None]
+                return lin(o.transpose(1, 2).reshape(Bt, n, h * dh), s.wo)
+
+        class FF(nn.Module):
+            def __init__(s, ff):
+                super().__init__()
+                norm, l1, _, _, l2, _ = ff.net
+                s.w1 = nn.Parameter((l1.weight.detach().double() * rms(norm)[None]).T.float(), requires_grad=False)
+                s.b1 = nn.Parameter(l1.bias.detach().float(), requires_grad=False)
+                s.w2 = nn.Parameter(l2.weight.detach().T.float(), requires_grad=False)
+                s.b2 = nn.Parameter(l2.bias.detach().float(), requires_grad=False)
+
+            def forward(s, x):
+                x = x / x.norm(dim=-1, keepdim=True).clamp(min=1e-12)
+                return lin(F.gelu(lin(x, s.w1) + s.b1), s.w2) + s.b2
+
+        return Mod().eval()
+
+
+def frames(model):
+    import torch
+
+    return stft(model, torch.zeros(1, model.audio_channels, SEGMENT)).shape[3]
+
+
+def export(model, path, opset):
+    import torch
+
+    T = frames(model)
+    core = Core.build(model, T)
+    spec = stft(model, torch.randn(1, model.audio_channels, SEGMENT))
+    with torch.no_grad():
+        torch.onnx.export(core, (spec,), str(path), input_names=["mix_spec"], output_names=["stems_spec"],
+                          opset_version=opset, dynamo=False)
+    print(f"wrote {path} ({path.stat().st_size / 1e6:.1f} MB), segment {SEGMENT} samples, spec {list(spec.shape)}, groups {groups([2 * n * model.audio_channels for n in model.num_freqs_per_band.tolist()])}")
+    return core
+
+
+def engine(model, out_dir, name):
+    """The GPU form's tensors for the package's own WebGPU engine (roformer.js): <name>.engine.bin, each tensor in turn
+    (weights float16, tables float32, indices int32), and <name>.engine.json naming them, their type, shape and offset,
+    with the band groups and the sizes the engine reads."""
+    import json
+    import numpy as np
+    import torch
+
+    T = frames(model)
+    core = Core.build(model, T)
+    tensors, blob, off = {}, bytearray(), 0
+    for n, t in list(core.named_parameters()) + list(core.named_buffers()):
+        a = t.detach().numpy()
+        kind = "int32" if a.dtype.kind in "iu" else "float16" if n.rsplit(".", 1)[-1] in ("qkvg", "wo", "w1", "w2") or n.startswith(("split_w.", "mw.", "m2_w.")) else "float32"
+        b = a.astype(kind).tobytes()
+        pad = -len(blob) % 16
+        blob += bytes(pad)
+        tensors[n] = {"type": kind, "shape": list(a.shape), "offset": len(blob)}
+        blob += b
+    nf = model.num_freqs_per_band.tolist()
+    meta = {
+        "model": name, "frames": T, "freqs": core.Fq, "channels": core.S, "dim": model.layers[0][0].layers[0][0].to_qkv.in_features,
+        "heads": model.layers[0][0].layers[0][0].heads, "depth": len(model.layers), "bands": len(nf),
+        "groups": [list(g) for g in groups([2 * n * core.S for n in nf])], "widths": [2 * n * core.S for n in nf],
+        "tensors": tensors,
+    }
+    (out_dir / f"{name}.engine.bin").write_bytes(bytes(blob))
+    (out_dir / f"{name}.engine.json").write_text(json.dumps(meta))
+    print(f"wrote {out_dir / (name + '.engine.bin')} ({len(blob) / 1e6:.1f} MB) and its .json, {len(tensors)} tensors")
+
+
+def fp16(path, out):
+    """float16 weights, float32 compute (scripts/compact.py --as fp16)."""
+    sys.path.insert(0, str(Path(__file__).parent))
+    import compact
+
+    compact.compact(path, out, "fp16")
+    print(f"wrote {out} ({out.stat().st_size / 1e6:.1f} MB)")
+
+
+def session(path):
+    import onnxruntime as ort
+
+    o = ort.SessionOptions()
+    o.enable_cpu_mem_arena = False
+    o.enable_mem_pattern = False
+    return ort.InferenceSession(str(path), o, providers=["CPUExecutionProvider"])
+
+
+def mixes():
+    import torch
+
+    t = torch.arange(SEGMENT) / 44100
+    tone = sum(torch.sin(2 * math.pi * f * t) / k for k, f in enumerate((110, 220, 440, 880, 1760), 1))
+    return (("noise", 0.1 * torch.randn(1, 2, SEGMENT)), ("tones", 0.2 * torch.stack([tone, 0.8 * tone])[None]))
+
+
+def verify(model, run, label, tol):
+    """The vocals rebuilt from run(spec) against MelBandRoformer.forward: noise and tones."""
+    import numpy as np
+    import torch
+
+    for name, mix in mixes():
+        with torch.no_grad():
+            want = model(mix).numpy()
+            got = istft(model, torch.as_tensor(run(stft(model, mix))), SEGMENT)[:, 0].numpy()
+        want = want[..., :got.shape[-1]]
+        rel = np.abs(got - want).max() / max(np.abs(want).max(), 1e-12)
+        print(f"  verify {label} {name}: max|diff|/max|y| = {rel:.2e} " + ("OK" if rel < tol else "FAIL"))
+        if rel >= tol:
+            raise SystemExit(f"{label}: verification failed ({rel:.2e} >= {tol})")
+
+
+def reference(model, out_dir):
+    """The reference separation: MelBandRoformer.forward on segments as scripts/reference.py's chunked() cuts them
+    (step 0.5, no standardization: Kim's config normalizes nothing), over a 20 s mix with a noise floor."""
+    import torch
+
+    sys.path.insert(0, str(Path(__file__).parent))
+    import reference as ref
+
+    def run(seg):
+        with torch.no_grad():
+            return model(torch.from_numpy(seg)[None])[0][None].numpy()   # (1 stem, C, L)
+
+    # under the mix a noise floor 60 dB down: each band normalizes its input, and a silent band's is its rounding
+    ref.write(out_dir, lambda mix: {"vocals": ref.chunked(mix, run, SEGMENT, STEP)[0]}, 20.0, 1e-3)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--model", default="mel-roformer", choices=list(MODELS))
+    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("--fp16", action="store_true", help="also write <model>.fp16.onnx (float16 weights, float32 compute)")
+    ap.add_argument("--verify", action="store_true", help="compare with MelBandRoformer.forward; write the reference separation")
+    ap.add_argument("--engine", action="store_true", help="also write <model>.engine.bin/.json: the tensors for roformer.js, the package's WebGPU engine")
+    args = ap.parse_args()
+
+    cache = Path(os.environ.get("AUDIO_NEURAL_CACHE") or Path.home() / ".cache" / "audiojs" / "neural")
+    out_dir = Path(args.out_dir) if args.out_dir else cache / args.model
+    out_dir.mkdir(parents=True, exist_ok=True)
+    src = out_dir / "src"
+    src.mkdir(exist_ok=True)
+    try:
+        model, _ = load(src, args.model)
+    except ImportError as e:
+        raise SystemExit(f"required: pip install torch onnx onnxruntime pyyaml einops beartype rotary-embedding-torch librosa ({e})")
+    path = out_dir / f"{args.model}.onnx"
+    core = export(model, path, args.opset)
+    if args.verify:
+        import torch
+
+        with torch.no_grad():
+            verify(model, lambda spec: core(spec), "the GPU form in torch", 1e-4)
+        # onnxruntime against torch: on the tones 1.7e-4, where torch's own float32 stands 4.6e-3 from its float64
+        s = session(path)
+        verify(model, lambda spec: s.run(None, {"mix_spec": spec.numpy()})[0], path.name, 1e-3)
+        del s
+        gc.collect()
+        reference(model, out_dir)
+    if args.fp16:
+        fp16(path, out_dir / f"{args.model}.fp16.onnx")
+        if args.verify:
+            s = session(out_dir / f"{args.model}.fp16.onnx")
+            verify(model, lambda spec: s.run(None, {"mix_spec": spec.numpy()})[0], f"{args.model}.fp16.onnx", 2e-2)
+    if args.engine:
+        engine(model, out_dir, args.model)
+    print(f"done: {args.model} in {out_dir}")
+
+
+if __name__ == "__main__":
+    sys.exit(main())

package/scripts/reference.py

@@ -64,9 +64,12 @@ def chunked(mix, run, L, step=0.25):
     return out[..., border:border + N] if wide else out
 
 
-def write(out_dir, separate, seconds=9.0):
-    """separate(mix (2, N) float32) -> { target: (2, N) array }"""
+def write(out_dir, separate, seconds=9.0, floor=0.0):
+    """separate(mix (2, N) float32) -> { target: (2, N) array }. `floor`: white noise of that RMS under the mix, so no
+    band of it is silent (a model that normalizes each band's input turns a silent band's rounding into its direction)"""
     mix = test_mix(seconds)
+    if floor:
+        mix = (mix + floor * np.random.RandomState(1).randn(*mix.shape)).astype(np.float32)
     mix.tofile(out_dir / "test.f32")
     for name, x in separate(mix).items():
         np.asarray(x, dtype=np.float32).tofile(out_dir / f"test.{name}.f32")

package/separate.js

@@ -27,7 +27,7 @@
 
 import { fft, ifft } from 'fourier-transform'
 import resampleSinc from '@audio/resample-sinc'
-import { load as neuralLoad, tensor, fetchModel } from '@audio/neural-runtime'
+import { load as neuralLoad, tensor, fetchModel, fetchJson, backends } from '@audio/neural-runtime'
 
 const PI2 = Math.PI * 2
 const isNode = typeof process !== 'undefined' && !!process.versions?.node
@@ -371,6 +371,15 @@ const TIGER = {
 	standardize: false, mono: true, script: 'export-tiger.py',
 }
 
+// Mel-Band RoFormer's STFT (Hann of 2048, hop 441, unnormalized), 8 s segments of 801 frames every 4 s, as
+// Music-Source-Separation-Training's demix() runs Kim's vocals model (num_overlap 2); nothing normalized
+const MEL = {
+	modelType: 'complex', sampleRate: 44100, targets: ['vocals'], sources: ['vocals'], n: 2048, hop: 441, window: 'hann',
+	normalized: false, segment: 352800, frames: 801, step: 0.5, standardize: false, gpu: true, script: 'export-roformer.py',
+	// roformer.js's tensors, for a browser whose GPU has subgroup matrices: the engine runs in onnxruntime's place
+	engine: { meta: 'mel-roformer.engine.json', bin: 'mel-roformer.engine.bin', sha256: '140fedcd9d03a8f918b235688c69f4dfdd86b08f9d430277659c095688a026ff' },
+}
+
 // The compact files (scripts/compact.py: int8 weights, a few float16, float32 compute), each within 0.05 dB of its
 // export's median SDR or SNR per stem (README, Size); their Hugging Face repositories, and SHA-256 checked when fetched
 const FILES = {
@@ -378,6 +387,8 @@ const FILES = {
 	scnet: { file: 'scnet.int8.onnx', repo: 'audiojs/scnet', sha256: '98228931494151762a1c4ab1ec7899a894b1f81fd4a509921a9d2f9bacc50845' },
 	mrx: { file: 'mrx.int8.onnx', repo: 'audiojs/mrx', sha256: 'd876c92d224f5d8cd2207d528b23369da895a6de71fb24132416fcb16f9f8424' },
 	tiger: { file: 'tiger.int8.onnx', repo: 'audiojs/tiger-dnr', sha256: '2189859ac7e89cf4d9d284c6f393e737aaa2b3233ed7ee424ad622012bce7bca' },
+	// float16 weights, float32 compute: half the download; int8 (259 MB) ran slower on WebGPU and 10 dB further from float32
+	'mel-roformer': { file: 'mel-roformer.fp16.onnx', repo: 'audiojs/mel-roformer', sha256: 'adcbbb826f12453fbb86652b503242d565a9c81cdef1bca2fe0e54752eba91ef' },
 }
 
 // The hosted revision of each compact file's repository (a commit), its URL https://huggingface.co/<repo>/resolve/<revision>/
@@ -388,6 +399,7 @@ export const REVISIONS = {
 	scnet: 'e71bcdb42bce1d62b7fae25ba90fcbcf90f73f46',
 	mrx: 'dedc823a0bea2ed085d19e8b7f32589d35169ee8',
 	tiger: '1c959ab7ff8794dc5663b104ba301687ab95d2ab',
+	'mel-roformer': 'c1aba02c00a0bee852d6754fe6ede8164f90a4b3',
 }
 
 const hosted = name => models[name].repo && REVISIONS[name] ? `https://huggingface.co/${models[name].repo}/resolve/${REVISIONS[name]}/${models[name].file}` : null
@@ -406,6 +418,8 @@ export const models = {
 	mrx: { ...MRX, ...FILES.mrx },
 	// TIGER (JusperLee/TIGER, Apache-2.0 weights): the same three stems, three band-split models, about 50 times MRX's time
 	tiger: { ...TIGER, ...FILES.tiger },
+	// Mel-Band RoFormer (Kimberley Jensen's vocals model, MIT weights) as scripts/export-roformer.py exports it for the GPU
+	'mel-roformer': { ...MEL, ...FILES['mel-roformer'] },
 }
 
 async function cacheDir() {
@@ -452,6 +466,38 @@ async function local(graphs, name, opts) {
 		`, or pass opts.weights: a URL or directory holding ${name}/<file>.onnx`)
 }
 
+// A preset's own WebGPU engine (roformer.js) as a session, its tensors from `weights` or the hosted revision, or null
+// where the GPU has no subgroup matrices (onnxruntime then): in the browser. Kept after a separation for KEEP ms, so the
+// next one starts at once (its 466 MB of weights and 2 GB of buffers on the GPU, made in about a second), then let go
+const KEEP = 60_000
+let kept = null
+function engineSession(opts, p) {
+	let name = opts.model, key = `${name}:${opts.weights ?? ''}:${p.engine.sha256}`
+	if (kept?.key === key) { clearTimeout(kept.timer); return kept.session }
+	kept?.session.then(s => s?.drop())
+	let mine = { key, timer: null }
+	// free() after a separation lets go KEEP ms later, unless the next one has taken it
+	mine.session = makeEngine(opts, p, name).then(s => s && {
+		...s, drop: s.free,
+		free() { clearTimeout(mine.timer); mine.timer = setTimeout(() => { if (kept === mine) kept = null; s.free() }, KEEP) },
+	})
+	return (kept = mine).session
+}
+async function makeEngine(opts, p, name) {
+	let { engine, gpu } = await import('./roformer.js'), device = await gpu()
+	if (!device) return null
+	try {
+		let base = opts.weights != null ? `${await weightsBase(opts, name)}${name}/` : `https://huggingface.co/${p.repo}/resolve/${REVISIONS[name]}/`
+		let [meta, bytes] = await Promise.all([fetchJson(base + p.engine.meta), checked(base + p.engine.bin, p.engine.sha256)])
+		return await engine(meta, bytes, device)
+	} catch (e) {
+		// its files not there (weights without them), or the GPU refusing its buffers: onnxruntime, as without it
+		device.destroy()
+		console.warn(`neural-separate: ${name}'s own engine stands aside (${e.message})`)
+		return null
+	}
+}
+
 // The hosted compact file into the cache, checked against its SHA-256, written whole or not at all → its file URL
 async function fetched(name, dest) {
 	let [fs, path, { pathToFileURL }] = await Promise.all([import('node:fs'), import('node:path'), import('node:url')])
@@ -503,12 +549,18 @@ async function resolveModel(opts) {
 // No pooled arena: one htdemucs segment peaks at 2.7 GB RSS in onnxruntime-node instead of 3.2 GB.
 const sessionOptions = { enableCpuMemArena: false, enableMemPattern: false }
 
-// A graph's session; a fetched file with a SHA-256 checked against it first
-function loader(opts) {
+// A graph's session; a fetched file with a SHA-256 checked against it first. A preset written for the GPU (`gpu`) runs
+// on WebGPU where there is one (onnxruntime-web's in a browser, onnxruntime-node's in Node) unless `device` says
+// otherwise: SCNet's LSTMs and the others' many small operators run no faster there
+function loader(opts, preset) {
+	let device = opts.device ?? (preset?.gpu ? backends().then(b => b.includes('webgpu') ? 'webgpu' : undefined) : undefined)
+	let own = !opts.device && !opts.session && preset?.engine && !isNode ? engineSession(opts, preset) : null
 	return async g => {
 		if (opts.session) return opts.session(g.spec, opts)
+		let s = own && await own
+		if (s) return s
 		let spec = g.sha256 && !g.spec.startsWith('file:') ? await checked(g.spec, g.sha256) : g.spec
-		return neuralLoad(spec, { backend: opts.device, sessionOptions })
+		return neuralLoad(spec, { backend: await device, sessionOptions })
 	}
 }
 
@@ -959,7 +1011,7 @@ export default async function separate(audio, opts = {}) {
 	if (!channelData[0].length) return { stems: Object.fromEntries(targets.map(t => [t, empty()])), sampleRate: rate, residual: empty() }
 	let targetRate = opts.targetRate ?? modelRate ?? rate
 	let proc = targetRate !== rate ? channelData.map(c => resampleSinc(c, { from: rate, to: targetRate })) : channelData
-	let load = loader(opts)
+	let load = loader(opts, preset)
 	// the input at the loudness the model trained at, its stems scaled back (a silent input as it is)
 	let level = preset?.loudness == null ? -Infinity : lufs(proc, targetRate), gain = Number.isFinite(level) ? 10 ** ((preset.loudness - level) / 20) : 1
 	if (gain !== 1) proc = proc.map(c => c.map(v => v * gain))