elicwhite/bs-roformer-sw-6stem-onnx

BS-Roformer-SW 6-stem — ONNX (browser-ready)

ONNX export of jarredou/BS-ROFO-SW-Fixed (BS-Roformer-SW 6-stem source separation: bass, drums, other, vocals, guitar, piano) re-packaged to run in the browser via onnxruntime-web with the WebGPU EP.

The companion repo with the build pipeline and a working browser demo is elicwhite/bs-roformer-web.

Files

File	Size	When to use
bs_roformer_sw_6stem_fp16.onnx	336 MB	Default. Weights stored fp16; the graph inserts Cast(fp16→fp32) after each weight initializer, so ORT folds back to fp32 compute on CPU/WASM (no precision loss). On WebGPU the EP uses fp16 accumulators internally regardless.
bs_roformer_sw_6stem_fp32.onnx	669 MB	Same graph, fp32 weights. Use if your runtime can't handle the Cast pattern or you want a single backend-agnostic source of truth.
reference.npz	12 MB	Deterministic input/output reference (PyTorch). Lets you validate your build with python validate_onnx.py --onnx ... --reference reference.npz or the browser's ?parity=1 mode.

I/O contract

The ONNX graph takes pre-computed STFT coefficients and returns post-masking STFT coefficients. STFT and iSTFT happen outside the graph (in JS in the browser, or in Python via wrapper.istft_from_spec).

Inputs (both shape [1, 2, 1025, T], float32):

• spec_real — real part of torch.stft(audio, n_fft=2048, hop_length=512, win_length=2048, window=hann, center=True, normalized=False)
• spec_imag — imaginary part of the same STFT

Outputs (both shape [1, 6, 2, 1025, T], float32):

• out_spec_real — masked spectrogram real part, per stem
• out_spec_imag — masked spectrogram imaginary part, per stem

Stem order (matches the BS-Roformer-SW training config):

0: bass    1: drums    2: other    3: vocals    4: guitar    5: piano

Trace details

Traced at a fixed chunk size: 176400 samples (4 s @ 44.1 kHz, T=345). The model's rotary embedding caches by sequence length, so the time dimension can't actually flex even if you mark it dynamic. Browser callers chunk audio into 4-s segments (typically with 25% overlap-add). Smaller T is faster per second of audio (O(T²) attention) but quality drops; larger T is slower per second of audio because attention cost grows faster than the per-chunk overhead saves.

WebGPU caveats baked into the graph

The fp32 / fp16 files have already been through two graph rewrites that the WebGPU EP needs to produce correct output:

1. Negative axes → positive. ORT-Web's WebGPU EP miscomputes Concat / Split with axis=-1 (silently produces zero-sized outputs).
2. 62-way Split / Concat → ≤7-way trees. Chrome enforces maxStorageBuffersPerShaderStage=8. BS-Roformer's 62-way band Split and each MaskEstimator's 62-way Concat would each need 63 storage buffers per shader stage, exceeding the limit and silently producing zeros. The fix splits those wide ops into trees of at-most-7-way variants, plus a Slice-based BandSplit.

Without these rewrites the model loads but produces all-zero output on WebGPU. The fp16 file additionally has all Conv/MatMul/Gemm weight initializers stored as fp16 with Cast(fp16→fp32) inserted right after each one, mirroring the demucs-next approach.

License

MIT — consistent with the upstream code lineage:

• Architecture: lucidrains/BS-RoFormer (MIT)
• Inference framework used to instantiate the model: ZFTurbo/Music-Source-Separation-Training (MIT)
• Export pipeline that produced these ONNX files: elicwhite/bs-roformer-web (MIT)

The pretrained weights themselves were rehosted by jarredou without a stated license. The original rehost owner has confirmed they were not involved in training and have no provenance information. Use accordingly.