vocal-separation-roformer-onnx

Vocal Separation Core (Mel-Band RoFormer) — ONNX / fp16 / WebGPU

Not used by Musetric any more, and its licence is not confirmed. These weights are an ONNX export of SYHFT Big V1. By its weights that model is a fine-tune of a model published without a licence, so the MIT label on its own card cannot grant the rights of its parent, and this repository makes no licence claim of its own. Musetric separates vocals with musetric/aname-mel-band-roformer-duality-onnx (Apache-2.0) instead (R-28). The files stay for the builds that pinned them.

ONNX export of a Mel-Band RoFormer vocal source-separation core, packaged for the musetric packages/ai runtime (onnxruntime-web on WebGPU).

The graph is the neural network core: it takes a precomputed STFT representation and returns per-bin complex masks. The mel-band gather/average tables are baked into the ONNX graph, so the host does not need sidecar /tables/* assets. STFT, iSTFT, chunking and complex packing run host-side (browser WGSL + FFT). This is not a drop-in PyTorch checkpoint.

Intended uses & limitations

Intended:

• Vocals / instrumental separation as the first stage of an audio pipeline.
• Client/edge inference via WebGPU through onnxruntime-web, on desktop and on mobile GPUs alike.

Out of scope:

• Standalone use without a host that computes the STFT input, applies the per-bin masks, and runs iSTFT (see musetric packages/ai).
• Use in other training frameworks — this is an inference-only export.

Limitations:

• One static time window, T = 1100 (~11 s), the model's full reference context, rounded down to a multiple of four. See the window note below.
• Training-data provenance of the upstream weights is undocumented.

How to use

The session runs the core; the host supplies stft_repr and consumes masks.

import * as ort from 'onnxruntime-web/webgpu';

// .onnx and .onnx.data must sit in the same directory; .data loads automatically.
const session = await ort.InferenceSession.create('syhft_core_t1100.onnx', {
  executionProviders: [{ name: 'webgpu', storageBufferCacheMode: 'simple' }],
});

// stftRepr: Float32Array of shape [1, 2050, 1100, 2], produced host-side from one
// ~11 s audio chunk (n_fft=2048, hop=441, 44.1 kHz, stereo).
const input = new ort.Tensor('float32', stftRepr, [1, 2050, 1100, 2]);
const { masks } = await session.run({ stft_repr: input });
// masks: float32 [1, 2050, 1100, 2] -> apply to STFT, then iSTFT host-side.

See the musetric packages/ai host code for the full STFT/iSTFT and chunk-recombination pipeline.

Files

One graph. Keep the .onnx next to its .onnx.data.

Signature — opset ai.onnx 23, IR 10, no com.microsoft import. fp16 weights, fp32 graph I/O, T = 1100:

How these graphs are written

Attention is split along the query axis into 64-row blocks. Written in one shot, each layer's score tensor is [60, 8, T, T] fp16 — 1110 MiB at this window — and a mobile storage-buffer binding stops working past 256 MiB whatever limit the adapter declares, silently returning zeros. Softmax normalizes each query row over the full key axis on its own, so splitting the queries is exact: same values, same work, no context given up. The peak score tensor becomes [60, 8, 64, T], 64.5 MiB.

RMSNorm is the fused ai.onnx::RMSNormalization, in fp16, with epsilon = 1e-9. Its WebGPU kernel accumulates the sum of squares in f32 inside the shader, which is what makes fp16 safe here; an fp32 island instead costs 387 MiB of activations at this window. The epsilon is not the 1e-12 an fp32 RMSNorm would carry: the kernel normalizes by the reciprocal 1/sqrt(mean(x^2) + epsilon) and casts it to the tensor dtype, so on an fp16 node any row whose mean(x^2) + epsilon falls below 1/65504^2 = 2.33e-10 gets +inf, then 0 * inf = NaN, and one NaN row takes the whole output out through the next attention softmax. Rows that small are ordinary — padded STFT edges, silent frames, dead band-split features, 0.089 % of rows on a real chunk.

Feed-forward blocks run on row chunks, and the fused qkv projection is three projections. Written whole, one feed-forward is a single [66000, 384] x [384, 1536] matmul, and a mobile GPU runs it as one uninterruptible dispatch: 546 ms on an Adreno 660. Twenty-four of those plus twenty-four more at 361 ms are 43 % of a window inside 48 of its 1938 dispatches, and nothing at the runtime level reaches inside one dispatch — a flush threshold only decides how many ride in a submit. That is the floor a compositor freeze sits on. Chunking the block over eight row groups is exact in fp16 because rows of a matmul are independent; splitting the reduction axis instead reorders the sums and costs roughly 66 dB a layer. The wide [rows, 1536] intermediate is never materialized at full height either, so peak GPU memory falls rather than rises. Unfusing qkv is the same move on the second-longest matmul, and free, because the rearrange that follows split it apart anyway. The q, k, v and output projections that remain then run on four row groups each, for the same reason.

Wide Concat/Split are re-treed to <=8-wide so every shader stays at <=9 storage buffers, under the strictest shipping cap (Dawn/Metal on macOS reports maxStorageBuffersPerShaderStage = 10).

No normalization or softmax sees more than 65 535 rows. The onnxruntime WebGPU kernels for RMSNormalization and Softmax run one workgroup per row and reduce that row in workgroup-shared memory, so they carry no bounds check. Past maxComputeWorkgroupsPerDimension (65 535, the WebGPU default) the dispatch becomes a ceil(sqrt(rows))^2 square, and its spare workgroups write past the end of the output. At this window the transformer norms see [60, 1100, 384] and [1100, 60, 384], 66 000 rows each, and the band-attention softmax sees [1100, 8, 60, 60], 528 000 rows. Where the backend clamps the stray index instead of dropping the write, as Chrome does on Metal, an exact-size output buffer (storageBufferCacheMode: 'simple', as in the snippet above) gets its last row overwritten: RMSNormalization returns +inf there, and attention turns every mask into NaN. The graph splits each such tensor along its batch axis first, two chunks for the norms and nine for the softmax scores, and concatenates the results. Rows never interact in either kernel, so the split is exact.

All five rewrites are compatibility, footprint and latency measures, and all five are exact. None changes a weight value; the qkv split only stores one projection matrix as the three slices the graph took from it anyway.

Why the window is 1100 and not 1101

On an Adreno 750, onnxruntime-web computes an fp16 MatMul with visibly less precision whenever the reduction length K or the output width N is not a multiple of four. Measured on that adapter against the CPU execution provider, M = 4096, both operands dynamic:

A multiple of two is not enough. Both attention matmuls carry the window in exactly that position — q @ kᵀ has N = T and attn @ v has K = T — so a window of 1101 put every attention in every layer on the bad path, while the projections (K = 384, N = 512 or 1536) stayed exact. End to end that cost 15 dB on that adapter and nothing anywhere else.

Rounding the window down to 1100 removes it: 26.9 dB becomes 41.7 dB on the Adreno 750, desktop is unchanged, and the price is one frame of context, 11.00 s instead of 11.01 s. Any future window should stay a multiple of four.

Validation

Against the PyTorch first-stage reference at the same window, which isolates conversion and execution-provider error:

The 1e-9 epsilon costs nothing measurable against 1e-12: 45.68 dB versus 45.70 dB on a real chunk, and 45.07 dB with a fifth of the frames zeroed, where 1e-12 returns NaN.

Against the earlier non-blocked export of the same weights, on a desktop NVIDIA GPU through the shipping host, one 30 s track: 52.7 dB on the vocals stem and 55.8 dB on the instrumental — the two graphs sit closer to each other than either sits to torch, which is what an exact rewrite plus fp16 rounding looks like. Two A/B pairs put the cost at 3 % and 7 % more wall clock at the same peak GPU memory, so the blocking is not a desktop optimization; it is what lets one graph serve every target.

One real 11 s chunk, identical input bytes, Chrome 151 on both phones, measured against this graph on the CPU execution provider:

All three are ordinary fp16 execution error, and the two phones agree to the digit.

What the row split changed

The tables above were measured on an earlier revision of this core, the query-blocked export at commit 98064f6. The rewrite that separates revision 4f72f5b from it is exact, so they carry over; what follows was measured on 4f72f5b directly.

One OnePlus 9RT (Adreno 660), one 11 s window through the shipping host, hot run, host submitting one dispatch at a time:

The 546 ms family is gone from the profile. What sets the floor now is a 239 ms dispatch over the input, twelve output projections at 192 ms and thirty-six q/k/v at 185 ms; all three yield to the same treatment if it is ever needed.

On one real track both revisions score 57.9 dB against the torch reference at the same window and agree to the last digit, with the same maximum deviation of 3.7e-02. Against each other they are 82.9 dB apart on that window and 88.2 dB on synthetic input. Choose the window by energy: on a quiet intro both score 3.6 dB and also agree to the digit, which is what validating on near-silence looks like rather than parity.

These four numbers were measured on the published bytes of 4f72f5b, on the CPU execution provider.

What the row cap changed

The row cap is exact as well, so everything above carries over. Measured on the bytes of this revision against 4f72f5b: one real 11 s window, identical input bytes, storageBufferCacheMode: 'simple', sessions alternated A-B-B-A on each device, second run of each session:

On these three backends the stray writes leave the output intact, which is why 4f72f5b already computed correct masks there. The overflow shows where the index is clamped: on Chrome/Metal with 'simple', 4f72f5b returns NaN masks. These bytes were not run on Metal. The cost is mostly memory: under 'simple' the buffer cache keeps a buffer of every new chunk size for the whole session, about 270 MB, while run time moves by 2–3 %.

What the projection split changed

The projection split is exact as well. Measured against 06f07ae on the bytes of this revision: a deterministic [1, 2050, 1100, 2] input, sessions paced one submission at a time as the host does, one first run and five runs per load, three loads of each revision per device in A-B-B-A-A-B order with 120 s idle between loads. The foreign job wait is how long a one-dispatch job submitted every 100 ms on another GPU device waits for its turn:

Both phones throttle after minutes of this: single runs stretch to 1.3-1.7x on either revision, and the worst wait of a load follows them rather than the graph. The weights file does not change between the two revisions; only the .onnx does.

Re-run the parity gate on the exact published bytes before relying on it.

Source & lineage

Code license and weight license are separate; ONNX conversion does not change the weight license. Documented only as far as it is verifiable.

• Architecture: Mel-Band RoFormer (arXiv:2310.01809).
• Reference implementation: lucidrains/BS-RoFormer.
• Training framework / config: ZFTurbo Music-Source-Separation-Training.
• Direct weight source: SYH99999/MelBandRoformerBigSYHFTV1Fast @ 96f4ae8e3f690e51ef26b3bef84531c944f5341b, MIT.
• Export tooling: scripts/onnx/roformer in musetric/lab. This revision is build_full_onnx.py with --core-only --fuse-rmsnorm --attn-block 64 --all-fp16 --frames 1100 --split-rows 8 --split-projections 4, followed by split_concat_webgpu.py; it passes fp16_epsilon_audit.py and dispatch_rows_audit.py, and reuse_external_data.py points it at the weights file of revision 06f07ae; the recipe is converting.md there.

The base checkpoint the upstream fine-tuned from is not documented upstream; we do not assert a chain we cannot verify. This export preserves the upstream MIT license; we do not claim authorship of the original weights.

License & citation

MIT, inherited from the upstream weights.

@article{wang2023melbandroformer,
  title={Mel-Band RoFormer for Music Source Separation},
  author={Wang, Ju-Chiang and Lu, Wei-Tsung and Won, Minz},
  journal={arXiv preprint arXiv:2310.01809},
  year={2023}
}