@audio/neural-separate 0.4.0 → 0.6.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.
 
@@ -62,13 +63,22 @@ together (0.9 GB of it the browser idle): wasm `scnet` 3.0 GB, `scnet-large` 4.6
 on wasm 9 and 0.3 minutes. Float16 compute does not pay: `tiger`'s on WebGPU is 13 % faster and its music 4.7 dB from
 the float32 graph's.
 
+A voice's room (de-reverb):
+
+| `model` | Architecture | Files | Weights | Targets | Real-time factor: Node CPU · WebGPU |
+|---|---|---|---|---|---|
+| `'dereverb-room'` | anvuew's room de-reverb: a mono BS-RoFormer (axial transformers over time and 63 contiguous bands of the complex spectrogram, dim 128, depth 16, 29.4 M parameters), trained on voices in rooms simulated with pyroomacoustics | 60 MB, float16 weights (export 122 MB) | GPL-3.0 | `dry`: the voice with its room off, early reflections and tail; the residual is the room | 2.3 · 0.65 (Node's WebGPU) |
+
+Each channel heard apart, 8.7 s segments every 4.35 s as `demix()` runs it (`num_overlap: 2`), nothing normalized. Measured through [`audio`](https://github.com/audiojs/audio)'s `dereverb()` (`bench/rx/dereverb.mjs`, test): VoiceBank test utterances in 19 MIT IR Survey rooms (Traer & McDermott, PNAS 2016; T60 0.38–1.85 s), each response split at direct + 50 ms, PESQ against the dry take 1.39 → 2.65, STOI 0.77 → 0.94, DNSMOS OVRL 2.19 → 2.98, the late tail 15.6 dB down; iZotope RX 12 De-reverb tuned 1.45, its Dialogue Isolate with the reverb off 1.79. It takes the early reflections too: the output stands 5.2 dB under the direct sound and its first 50 ms, where the voice is 10 dB over the tail. Music is room to it (MUSDB18 previews, GuitarSet, sung tones and four pieces lose 6 dB in their worst octave as they are, and 9 dB of themselves in a room); `audio` passes music by a speech/music guard. The real-time factors are of 20 s of stereo on a busy 14-core M4 Max (load averages over 100), so upper bounds; the float16 file stands 74 dB from the PyTorch forward through this package's pipeline, the float32 export 123 dB.
+
 `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. `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 +98,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 +111,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
@@ -279,6 +295,7 @@ Open-source stem separation has these lineages; this package runs SCNet, Open-Un
 | **SCNet** (`scnet`, `scnet-large`; starrytong) | MIT | MIT: "I confirm that the released SCNet and SCNet-large pretrained weights are distributed under the MIT License, consistent with the source code. You are welcome to redistribute the original checkpoints and format-converted versions, including ONNX exports, as part of your MIT-licensed tool, with appropriate attribution." (the author in [starrytong/SCNet#35](https://github.com/starrytong/SCNet/issues/35), 2026) | trained on MUSDB18-HQ; hosted with its configs by [ZFTurbo/Music-Source-Separation-Training](https://github.com/ZFTurbo/Music-Source-Separation-Training) (MIT), whose own MUSDB18-HQ checkpoints (SCNet XL, BS-RoFormer) carry its author's MIT statement only for some files ([#254](https://github.com/ZFTurbo/Music-Source-Separation-Training/issues/254)) |
 | **MRX** (MERL; `mrx`) | MIT | MIT: the repository's [`.reuse/dep5`](https://github.com/merlresearch/cocktail-fork-separation/blob/main/.reuse/dep5) names the four checkpoints, "Copyright: 2023 Mitsubishi Electric Research Laboratories (MERL) / License: MIT" | trained on Divide and Remaster v2, whose music and effects come from FMA and FSD50K clips each under its own licence, some non-commercial (the DnR v3 paper's account) |
 | **TIGER** (Tsinghua; `tiger`) | MIT (repository LICENSE; its README badge says Apache 2.0) | Apache 2.0: the [model card](https://huggingface.co/JusperLee/TIGER-DnR) states `license: apache-2.0` | trained on Divide and Remaster v1 (652 test mixtures in its paper), sourced as v2 is |
+| **anvuew's de-reverb** (`dereverb-room`) | the architecture: BS-RoFormer (Phil Wang's [BS-RoFormer](https://github.com/lucidrains/BS-RoFormer), MIT, in ZFTurbo's Music-Source-Separation-Training) | **GPL-3.0** ([model card](https://huggingface.co/anvuew/dereverb_room)): the converted copy is hosted under the same licence, fetched at run time, never bundled | trained on convolutional reverbs from pyroomacoustics, the dry voices not named |
 | **Bandit v2** (DnR v3) | Apache 2.0 | **CC BY-SA 4.0** ([Zenodo](https://zenodo.org/records/12701995)): a converted copy is an adaptation under the same terms | not shipped; BandIt v1's DnR checkpoints are CC BY-NC 4.0 ([Zenodo](https://zenodo.org/records/10160698)) |
 | **Spleeter** (Deezer) | MIT | **undocumented** | the README licenses only "the code of Spleeter"; the pretrained weights' license is an open, unresolved question ([deezer/spleeter#898](https://github.com/deezer/spleeter/issues/898)) — do not assume MIT |
 

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.6.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,519 @@
+#!/usr/bin/env python3
+"""Export Mel-Band RoFormer (Wang, Lu, Won, Chen, Song, "Mel-Band RoFormer for Music Source Separation", ISMIR LBD
+2023) and Band-Split RoFormer (Lu, Wang, Kong, Hung, "Music Source Separation with Band-Split RoPE Transformer", ICASSP
+2024) 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
+    python3 export-roformer.py --model dereverb-room --verify --fp16
+
+The models, each run by ZFTurbo/Music-Source-Separation-Training's (MIT) models/bs_roformer code at a pinned commit:
+  mel-roformer    Kimberley Jensen's vocals Mel-Band RoFormer (huggingface.co/KimberleyJSN/melbandroformer, MIT), with
+                  the config that repository runs it by (configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml)
+  dereverb-room   anvuew's room de-reverb, a mono BS-RoFormer (huggingface.co/anvuew/dereverb_room, GPL-3.0): a voice
+                  in a room in, the voice dry out; trained on voices in rooms pyroomacoustics simulated (its card); its
+                  own config, dim 128, depth 16, 63 bands, 29.4 M parameters
+Each checkpoint and config at a pinned revision, the checkpoint checked by its SHA-256. What follows is told of
+mel-roformer; dereverb-room is one channel, 8.7 s segments (384000 samples, hop 512, 751 frames), its bands contiguous
+(each frequency in one band, so its masks are taken as they come) and one RMSNorm after its transformers, not one each.
+
+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 (a path in REPO or a URL), class, segment (samples), the stem's name
+    "mel-roformer": ("https://huggingface.co/KimberleyJSN/melbandroformer/resolve/ac9b0614ab3cd7f77219e18ba494dfd93956c348/MelBandRoformer.ckpt",
+                     "87201f4d31afb5bc79993230fc49446918425574db48c01c405e44f365c7559e",
+                     "configs/KimberleyJensen/config_vocals_mel_band_roformer_kj.yaml", "MelBandRoformer", 352800, "vocals"),
+    "dereverb-room": ("https://huggingface.co/anvuew/dereverb_room/resolve/0b85f5b80b7f779b2dfe80f33a1b35b38af9376d/dereverb_room_anvuew_sdr_13.7432.ckpt",
+                      "2edec521f09e26341c1923dc82c8c52dbc86478b42b9999f679535743c970cb3",
+                      "https://huggingface.co/anvuew/dereverb_room/resolve/0b85f5b80b7f779b2dfe80f33a1b35b38af9376d/dereverb_room_anvuew.yaml",
+                      "BSRoformer", 384000, "dry"),
+}
+SEGMENT = 352800  # the model's, set by main()
+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"):
+    """The model 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", "bs_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
+    from models.bs_roformer.bs_roformer import BSRoformer
+
+    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, cls, _, _ = MODELS[name]
+    cfg = yaml.load(fetch(config if "://" in config else 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": MelBandRoformer, "BSRoformer": BSRoformer}[cls](**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 bands(model):
+    """(frequencies, (f·S + s) of each band's pairs, bands in order, frequencies per band): Mel-Band RoFormer's
+    overlapping mel bands, BS-RoFormer's contiguous ones"""
+    S = model.audio_channels
+    if hasattr(model, "freq_indices"):
+        return model.num_bands_per_freq.numel(), model.freq_indices.tolist(), model.num_freqs_per_band.tolist()
+    nf = [d // (2 * S) for d in model.band_split.dim_inputs]
+    return sum(nf), list(range(sum(nf) * S)), nf
+
+
+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:
+    """The model's 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
+        S = model.audio_channels
+        Fq, fi, nf = bands(model)
+        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))]))
+                # BS-RoFormer's one norm after the transformers (Mel-Band RoFormer's end each transformer)
+                fn = getattr(model, "final_norm", None)
+                s.final = nn.Parameter(rms(fn).float(), requires_grad=False) if fn is not None else None
+                # 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
+                if s.final is not None:
+                    x = x / x.norm(dim=-1, keepdim=True).clamp(min=1e-12) * s.final
+                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 = None if isinstance(tr.norm, nn.Identity) else 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 if s.out is None else 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 bands(model)[2]])}")
+    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 = bands(model)[2]
+    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(C):
+    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, C, SEGMENT)), ("tones", 0.2 * torch.stack([tone, 0.8 * tone][:C])[None]))
+
+
+def verify(model, run, label, tol):
+    """The stem rebuilt from run(spec) against the model's forward: noise and tones."""
+    import numpy as np
+    import torch
+
+    for name, mix in mixes(model.audio_channels):
+        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, stem):
+    """The reference separation: the model's forward on segments as scripts/reference.py's chunked() cuts them
+    (step 0.5, no standardization: Kim's config normalizes nothing, nor anvuew's), over a 20 s mix with a noise floor;
+    a mono model's channel by channel, as separate.js runs a mono preset."""
+    import torch
+
+    import numpy as np
+
+    sys.path.insert(0, str(Path(__file__).parent))
+    import reference as ref
+
+    def run(seg):
+        with torch.no_grad():
+            y = model(torch.from_numpy(seg)[None])                     # (1, C, L), BS-RoFormer's (1, 1 stem, C, L)
+            return (y[0] if y.ndim == 4 else y[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
+    if model.audio_channels == 1:
+        each = lambda mix: np.concatenate([ref.chunked(mix[c:c + 1], run, SEGMENT, STEP)[0] for c in range(len(mix))])
+    else:
+        each = lambda mix: ref.chunked(mix, run, SEGMENT, STEP)[0]
+    ref.write(out_dir, lambda mix: {stem: each(mix)}, 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 the model's 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()
+    if args.engine and MODELS[args.model][3] != "MelBandRoformer":
+        raise SystemExit("--engine: roformer.js runs Mel-Band RoFormer alone")
+    global SEGMENT
+    SEGMENT = MODELS[args.model][4]
+
+    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, MODELS[args.model][5])
+    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,24 @@ 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' },
+}
+
+// anvuew's room de-reverb (a mono BS-RoFormer, GPL-3.0 weights) as scripts/export-roformer.py exports it: its STFT
+// (Hann of 2048, hop 512, unnormalized), 8.7 s segments of 751 frames every 4.35 s, as Music-Source-Separation-Training's
+// demix() runs it (num_overlap 2); each channel apart, nothing normalized. Its one stem is the voice dry; the residual
+// is the room
+const ROOM = {
+	modelType: 'complex', sampleRate: 44100, targets: ['dry'], sources: ['dry'], n: 2048, hop: 512, window: 'hann',
+	normalized: false, segment: 384000, frames: 751, step: 0.5, standardize: false, mono: true, gpu: true, script: 'export-roformer.py',
+}
+
 // 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 +396,10 @@ 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' },
+	// float16 weights, float32 compute: 60 MB, its output 49 dB from the export's
+	'dereverb-room': { file: 'dereverb-room.fp16.onnx', repo: 'audiojs/dereverb-room', sha256: '99797e086a71d1cb57d50f3b5be93dcb7c89c49d1537c60487a488a24aaee9ee' },
 }
 
 // The hosted revision of each compact file's repository (a commit), its URL https://huggingface.co/<repo>/resolve/<revision>/
@@ -388,6 +410,8 @@ export const REVISIONS = {
 	scnet: 'e71bcdb42bce1d62b7fae25ba90fcbcf90f73f46',
 	mrx: 'dedc823a0bea2ed085d19e8b7f32589d35169ee8',
 	tiger: '1c959ab7ff8794dc5663b104ba301687ab95d2ab',
+	'mel-roformer': 'c1aba02c00a0bee852d6754fe6ede8164f90a4b3',
+	'dereverb-room': 'b4c2d14da9f2199099656195b06fd8253fd16e1d',
 }
 
 const hosted = name => models[name].repo && REVISIONS[name] ? `https://huggingface.co/${models[name].repo}/resolve/${REVISIONS[name]}/${models[name].file}` : null
@@ -406,6 +430,11 @@ 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'] },
+	// anvuew's room de-reverb (BS-RoFormer, GPL-3.0 weights) as scripts/export-roformer.py exports it: a voice in a room
+	// in, the voice dry out
+	'dereverb-room': { ...ROOM, ...FILES['dereverb-room'] },
 }
 
 async function cacheDir() {
@@ -452,6 +481,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 +564,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 +1026,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))