Views
No views yet
2-stem model at
present." This is the missing half.| file | size | notes |
|---|---|---|
{vocals,drums,bass,other}.fp16.onnx | 19.7 MB each | use these — self-contained |
{vocals,drums,bass,other}.onnx | 39.4 MB each | fp32, self-contained |
| TF vs PyTorch port, per stem | vocals 2.6e-04 · drums 1.3e-02 · bass 1.1e-03 · other 2.6e-03 |
sum(4 stems) − mix | −153.1 dB |
| Inference, all four stems | 1.69 s for 30 s of 44.1 kHz stereo (~18× realtime, CPU) |
1// configs/2stems/base_config.json
2"model": { "type": "unet.unet", "params": {} }
3// ^^ defaults: LeakyReLU(0.2) + ReLU
4
5// configs/4stems/base_config.json
6"model": { "type": "unet.unet", "params": {
7 "conv_activation": "ELU", "deconv_activation": "ELU" } }unet.py hardcodes LeakyReLU/ReLU, and nothing about the shapes
tells you.conv_activation,
deconv_activation) so the next person doesn't have to find this the hard way.input x : float32 [2, num_splits, 512, 1024] # [channels, splits, frames, bins]
output y : float32 [2, num_splits, 512, 1024] # that stem's magnitude estimatenum_splits is dynamic. The graph takes magnitudes and returns magnitudes —
no complex numbers cross the boundary, which is why this ports cleanly where
time-domain models don't. The STFT, the mask and the iSTFT are yours to do.np.hanning is
symmetric and is not this — it differs by one sample, and that sample is
the difference between matching training-time spectrograms and merely
resembling them.abs(stft(x))[..., :1024] (1024 of 2049 bins), padded and partitioned to
512-frame splits.1total = sum(e ** 2 for e in estimates.values()) + 1e-10
2mask = (estimate ** 2 + 1e-10 / 4) / totalaverage, not zeroszeros discards it, which is a −23 dB hole in the
reconstruction. average carries the per-frame mean up and reconstructs
exactly (the −153.1 dB above is with average; with zeros it is −23.0 dB).1import numpy as np, onnxruntime as ort, soundfile as sf
2
3N_FFT, HOP, T, F, BINS = 4096, 1024, 512, 1024, 2049
4PAD = N_FFT - HOP
5W = np.hanning(N_FFT + 1)[:-1] # PERIODIC. np.hanning(N_FFT) is symmetric
6STEMS = ("vocals", "drums", "bass", "other")
7
8def stft(x):
9 n = int(np.ceil((PAD + len(x)) / HOP))
10 p = np.zeros((n - 1) * HOP + N_FFT)
11 p[PAD:PAD + len(x)] = x # front pad: see note below
12 idx = np.arange(N_FFT)[None, :] + HOP * np.arange(n)[:, None]
13 return np.fft.rfft(p[idx] * W, axis=-1)
14
15def istft(spec, length):
16 frames = np.fft.irfft(spec, n=N_FFT, axis=-1)
17 total = (len(spec) - 1) * HOP + N_FFT
18 out, wsum = np.zeros(total), np.zeros(total)
19 for i in range(len(spec)):
20 at = i * HOP
21 out[at:at + N_FFT] += frames[i] * W
22 wsum[at:at + N_FFT] += W ** 2
23 out = np.divide(out, wsum, out=np.zeros_like(out), where=wsum > 1e-8)
24 return out[PAD:PAD + length]
25
26wave, sr = sf.read("song.wav", dtype="float64") # 44.1 kHz stereo
27assert sr == 44100 and wave.shape[1] == 2
28n = len(wave)
29
30spec = np.stack([stft(wave[:, c]) for c in range(2)])
31frames = spec.shape[1]
32splits = int(np.ceil(frames / T))
33mag = np.zeros((2, splits * T, F), dtype=np.float32)
34mag[:, :frames] = np.abs(spec[:, :, :F])
35net_in = mag.reshape(2, splits, T, F)
36
37est = {}
38for s in STEMS: # all four: the mask needs every one
39 sess = ort.InferenceSession(f"{s}.fp16.onnx", providers=["CPUExecutionProvider"])
40 out = sess.run(["y"], {"x": net_in})[0]
41 est[s] = out.reshape(2, -1, F)[:, :frames]
42
43stems = {}
44denom = sum(e ** 2 for e in est.values()) + 1e-10
45for s, e in est.items():
46 mask = (e ** 2 + 1e-10 / len(est)) / denom
47 # extend 1024 -> 2049 bins with the per-frame mean ("average", not "zeros")
48 tail = np.repeat(mask.mean(axis=-1, keepdims=True), BINS - F, axis=-1)
49 full = np.concatenate([mask, tail], axis=-1)
50 # applied to the ORIGINAL complex spectrum: the phase is already correct
51 stems[s] = np.stack([istft(spec[c] * full[c], n) for c in range(2)], axis=-1)
52
53# The invariant. Check this, not your ears -- see below.
54res = 10 * np.log10(np.mean((sum(stems.values()) - wave) ** 2) / np.mean(wave ** 2))
55print(f"sum(stems) - mix = {res:.1f} dB") # -153 dB
56
57for s, y in stems.items():
58 sf.write(f"out_{s}.wav", y, sr, subtype="FLOAT")sum(stems) must equal the mix.
It does, to −153.1 dB. If yours doesn't, your window or your hop is wrong,
and no amount of listening will tell you which — every one of those mistakes
produces audio that sounds approximately right.sf.write defaults to PCM_16, and
16-bit quantisation of four stems costs ~75 dB on its own — enough to turn
−153 into −77 and send you hunting a bug that isn't there. Hence
subtype="FLOAT" above.zeros is not a bug. With Spleeter's default mask_extension, −23 dB is
the correct answer: that is the >11 kHz band being discarded, exactly as
asked.stft above prepends N_FFT-HOP zeros, which Spleeter does not. Alignment is
not part of the contract — the U-Net is convolutional and translation-equivariant
in time — and without the pad, reconstruction is exact in theory and broken in
practice: on the ramp-in a periodic Hann is ~1e-7, so W² is ~1e-13, and
dividing by it turns float noise into the one stretch of signal with no
redundancy to spare. Unpadded, the round-trip error is 2.3. Padded, it is
1e-15.WINDOW_COMPENSATION_FACTOR here. Spleeter's 2/3
constant and its inverse_stft_window_fn exist to undo TensorFlow-specific
normalisation; plain weighted overlap-add inverts this forward with no constants
at all.1@article{spleeter2020,
2 doi = {10.21105/joss.02154},
3 author = {Romain Hennequin and Anis Khlif and Felix Voituret and Manuel Moussallam},
4 title = {Spleeter: a fast and efficient music source separation tool with pre-trained models},
5 journal = {Journal of Open Source Software},
6 volume = {5}, number = {50}, pages = {2154}, year = {2020}
7}