yt-precount-scnet-tran-onnx

SCNet Tran — ONNX core for in-browser 4-stem separation

ONNX export of the core network of SCNet Tran (SCNet with the separation network's bi-LSTM replaced by transformer layers), prepared for real-time-ish music source separation inside a web browser via onnxruntime-web.

Stems: drums, bass, other, vocals.

This repository exists to host the weight file for a browser extension. It claims no new rights over the weights — see Provenance and license.

Why "core"?

Two things had to change before this model could run under ONNX Runtime Web:

1. STFT / iSTFT are not in the graph. torch.stft / torch.istft have no stable ONNX equivalent — ONNX DFT cannot set inverse=True and onesided=True at the same time. So the exported graph takes a spectrogram and returns a spectrogram; the consumer performs STFT/iSTFT itself.

2. The internal FFT was replaced by matrix multiplication. FeatureConversion, in the middle of the separation network, calls torch.fft.rfft / irfft. That is part of the model, not I/O, so it cannot be moved out. It was replaced with multiplication by precomputed DFT basis matrices:

   rfft(norm="ortho"):  real = x @ C,   C[t,k] =  cos(2πkt/N)/√N
                        imag = x @ S,   S[t,k] = -sin(2πkt/N)/√N
   
   irfft(norm="ortho"): x = real @ Ci + imag @ Si
                        Ci[k,t] =  w_k·cos(2πkt/N)/√N
                        Si[k,t] = -w_k·sin(2πkt/N)/√N,   w_0 = w_{N/2} = 1, otherwise 2
   

   Substitution error against the original PyTorch model: max abs 1.04e-06.

Files

SHA-256: e2c6e2807e1deb937150c2c2d21db57b597388a67460706242e6f23a2d8f9c56

Interface

input   spec  (1, 4, 2049, 120)  float32   — view_as_real(stft) permuted/reshaped
output  out   (8, 2049, 120, 2)  float32   — pre-view_as_complex; 4 stems × 2 channels

Stem order along the output's first axis: ['drums', 'bass', 'other', 'vocals'], each followed by its two channels.

STFT convention (must match exactly)

Taken from the original scnet_tran.yaml plus torch.stft defaults. Getting any of these wrong still produces audio — it just produces wrong audio, so verify against reference tensors rather than by ear.

iSTFT is the exact inverse in torch.istft's order: undo normalization (×√n_fft) → irfft per frame → overlap-add → divide by the accumulated window-square envelope (with a rectangular window this is 4 in the interior and 2–3 at the edges — dividing by a constant 4 corrupts the boundaries) → strip the 2048-sample center padding.

Chunking

chunk   = 121275 samples (2.75 s @ 44100 Hz)
overlap = 2  →  step = 60637
fade    = 12127  (chunk // 10, linear in/out for overlap-add)

Per chunk, pad to a multiple of hop_length such that the frame count is odd (121275 → 121856, T = 120), run the model, then drop the 581 padded samples. Reflect-pad the whole input by chunk - step at both ends and crop afterwards. First chunk skips the fade-in; last chunk skips the fade-out.

Why 2.75 s and not the original 11 s: the 11 s chunk runs out of memory under onnxruntime-web's WASM backend, and even on WebGPU the small chunk was faster in our measurements.

Measured performance

Chrome, Apple Silicon, onnxruntime-web, 2.75 s chunk:

A 5-minute song at 2× overlap (221 chunks) takes roughly 50 s on WebGPU.

Agreement with the original PyTorch model: correlation 1.000000, max abs difference 6.43e-04.

Note that ONNX Runtime Web's WebGPU backend has no LSTM kernel, which is why the transformer variant (SCNet Tran) is used here rather than plain SCNet — the LSTM variant silently falls back to CPU.

Conversion notes

• Exported with torch.onnx.export, opset_version=17, dynamo=False.
• FeatureConversion modules are swapped for matmul equivalents after observing the actual sequence length each one receives via forward pre-hooks — the layers see different T because of differing downsampling, so guessing the lengths does not work.
• fp16 conversion was attempted and abandoned: intermediate activations exceed fp16 range and the converted graph produces NaNs.

Provenance and license

This file is a format conversion. The rights below belong to their respective holders, and nothing here grants rights that the upstream sources do not already grant.

Because of the MUSDB18-HQ terms, treat this model as non-commercial. The extension it was prepared for is distributed free of charge with no advertising and no revenue.

If you are the rights holder of any layer above and want the hosting or the wording changed, please open a discussion on this repository.