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.
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:
Negative axes → positive. ORT-Web's WebGPU EP miscomputes Concat /
Split with axis=-1 (silently produces zero-sized outputs).
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.
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.
