A vocal-separation model cannot eat a whole song at once. The network in Wunjo Make takes a fixed-size block of audio, runs an STFT, predicts a mask, runs an inverse STFT, and hands you back a block of the same size. So a four-minute track has to be cut into blocks, processed one at a time, and glued back together.
The naive way to glue blocks back is to lay them end to end. That clicks. Every seam becomes an audible tick, and across a song you get a rhythmic stutter that no listener will forgive.
This article is about the fix, which lives entirely inside one method, demix(), in portable/src/sound_processing/separator/mdx/model.py. The fix is a sliding window with overlap and a weighted average. I want to walk through it line by line with the real constants from the code, because the trick is small and the payoff is large.
A long take has to be cut into pieces before a model can touch it. Photo: Unsplash
Why cutting at hard edges clicks
Picture cutting a long ribbon into sections, processing each section, then taping them back in order. If the cuts are clean and the sections line up perfectly, fine. But a neural net does not return perfectly continuous audio at a block boundary. The sample at the very end of block 1 and the sample at the very start of block 2 were predicted by two separate forward passes that never saw each other. The waveform value can jump.
A jump in a waveform is a step. A step is a click. Your ear is very good at hearing steps, because a sudden discontinuity contains energy across the whole frequency spectrum. That is the tick you hear at every seam.
So you have two problems at once. The model output near a block edge is less reliable (the STFT has less context there), and butting two blocks together creates a discontinuity. Overlap-add solves both with the same move.
The constants, with real numbers
Before the loop, the model settings get computed in initialize_model_settings():
self.n_bins = self.n_fft // 2 + 1
self.trim = self.n_fft // 2
self.chunk_size = self.hop_length * (self.segment_size - 1)
self.gen_size = self.chunk_size - 2 * self.trim
The fixed knobs in the MDX separator are these:
self.segment_size = 256
self.overlap = 0.25
self.hop_length = 1024
self.batch_size = 1
Plug in the two that do not depend on the model file. The chunk size is hop_length * (segment_size - 1), which is 1024 * 255 = 261120 samples. At a 44100 Hz sample rate that is about 5.9 seconds of audio per block. So the model chews through the song roughly six seconds at a time.
trim is n_fft // 2, where n_fft comes from the model file (mdx_n_fft_scale_set). It is the padding the STFT needs on each side so the transform has context. gen_size is the useful interior of a chunk, chunk_size - 2 * trim, the part you actually keep after throwing away the padded edges.
Hold on to one number: chunk size is 261120 samples. Everything else is relative to it.
The window: a fade-in and fade-out in one array
Here is the plain version. Instead of taping the ribbon sections edge to edge, you let each section fade in at its start and fade out at its end, and you let neighboring sections overlap. Where two sections overlap, you add them. The fades are shaped so that the overlapping fades sum back to one. No seam, because there is no single point where one block stops and the next starts cold.
The shape that does this is a Hann window, np.hanning. It is a smooth bell that is zero at both ends and one in the middle.
window = None
if overlap != 0:
window = np.hanning(chunk_size_actual)
window = np.tile(window[None, None, :], (1, 2, 1))
chunk_size_actual is the length of the current chunk, which equals chunk_size for every chunk except the last one (more on that below). np.tile just copies the one-dimensional bell across both stereo channels so it lines up with the (1, 2, samples) shape of the audio.
The window is created fresh per chunk, keyed on the actual length, which matters for the short final chunk.
The step: how far the window slides
step = int((1 - overlap) * chunk_size)
With overlap = 0.25 and chunk_size = 261120, the step is int(0.75 * 261120) = 195840 samples. So each chunk starts 195840 samples after the previous one, but each chunk is 261120 samples long. The difference, 261120 - 195840 = 65280 samples, is the overlap region. At 44100 Hz that is about 1.5 seconds where two consecutive windowed chunks both contribute.
The loop walks the whole padded mixture in steps:
for i in tqdm(range(0, mixture.shape[-1], step)):
start = i
end = min(i + chunk_size, mixture.shape[-1])
start advances by step, end is start + chunk_size clamped to the array length. That clamp is what makes the last chunk shorter, and it is why the window is sized from chunk_size_actual = end - start rather than a constant.
The accumulators: result and divider
This is the heart of it. Two arrays, both zeroed and both the full length of the padded mixture:
result = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)
divider = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)
Think of result as the running sum of all the windowed, predicted audio, and divider as the running sum of all the window weights at each sample position. They grow together, chunk by chunk.
Inside the loop, after the model runs:
if window is not None:
tar_waves[..., :chunk_size_actual] *= window
divider[..., start:end] += window
else:
divider[..., start:end] += 1
result[..., start:end] += tar_waves[..., : end - start]
So for each chunk: multiply the model output by the bell (fade in, fade out), add that into result at this position, and add the same bell into divider at this position. The fade makes the chunk contribute almost nothing at its edges and full weight in its middle. Where two chunks overlap, both their bells land in divider, so divider records exactly how much total weight piled up at every sample.
The else branch (divider += 1) is the no-overlap fallback. If overlap were zero, every sample gets weight one and you are back to plain concatenation. The code keeps that path, but the default overlap = 0.25 never takes it.
The normalization: divide by the weight you actually applied
After the loop:
epsilon = 1e-8 # A small constant to avoid division by zero
tar_waves = result / (divider + epsilon)
This one line is the whole point. You summed up weighted predictions in result. You summed up the weights in divider. Dividing gives you a weighted average at every single sample. A sample that only one chunk covered gets that chunk's value. A sample in the overlap region gets a blend of two chunks, weighted by how confident each window was there (the fading-out tail of the left chunk, the fading-in head of the right chunk).
Because the blend is continuous, the waveform is continuous. No step, no click.
The epsilon is a guard. At the very ends of the padded array the window can taper to zero, so divider can be zero there, and 0 / 0 is a NaN that would poison the audio. Adding 1e-8 first turns 0 / 0 into 0 / 1e-8 = 0, which is silence, which is exactly what you want at a sample no chunk really covered.
This pattern, sum weighted signal over sum of weights, is the classic weighted overlap-add (WOLA). It is the same family of idea that makes the STFT itself invertible. Here it is applied one level up, across model-sized blocks instead of FFT frames.
The padding and the trim, so the edges line up
Two more pieces make the bookkeeping exact. First, padding before the loop:
pad = gen_size + self.trim - ((mix.shape[-1]) % gen_size)
mixture = np.concatenate(
(np.zeros((2, self.trim), dtype="float32"), mix, np.zeros((2, pad), dtype="float32")), 1)
Silence of length trim is prepended, and enough silence is appended to round the length up so the chunking covers everything. The leading trim of silence gives the very first real sample the same STFT context every interior sample gets, so the model is not guessing at the start.
Second, after normalization, the trim comes back off:
tar_waves_ = np.vstack(tar_waves_)[:, :, self.trim: -self.trim]
tar_waves = np.concatenate(tar_waves_, axis=-1)[:, : mix.shape[-1]]
[:, :, self.trim: -self.trim] drops the padded edges, and [:, : mix.shape[-1]] cuts the result back to the exact length of the input you handed in. You get out exactly as many samples as you put in, click-free.
Gotchas you will actually hit
A few things that bit me or would have:
The last chunk is shorter, and that is fine. Because end is clamped with min(...), the final chunk's length is whatever is left. The window is sized from chunk_size_actual = end - start, so the bell still fits. There is also a zero-pad just before the tensor conversion to make the model see a full-size block:
mix_part_ = mixture[:, start:end]
if end != i + chunk_size:
pad_size = (i + chunk_size) - end
mix_part_ = np.concatenate((mix_part_, np.zeros((2, pad_size), dtype="float32")), axis=-1)
but only the first end - start samples of the model output are accumulated (tar_waves[..., : end - start]), so the padding never leaks into the result.
Do not forget the epsilon. If you reimplement this and divide by divider directly, you will get NaNs at the tapered ends and your exported WAV will be full of garbage or silence-with-pops. The + 1e-8 is not decoration.
Overlap is a quality-versus-time dial, not free. With overlap = 0.25 the step is 75 percent of a chunk, so consecutive chunks share 25 percent. Push overlap to 0.5 and you process roughly twice as many chunks for smoother seams. The match-mix pass in the same method deliberately uses a tiny overlap = 0.02, because that pass is reconstructing the mixture, not the delicate vocal, and does not need the expensive smoothing:
if is_match_mix:
chunk_size = self.hop_length * (self.segment_size - 1)
overlap = 0.02
The window must match the chunk slice exactly. tar_waves[..., :chunk_size_actual] *= window and divider[..., start:end] += window both use the same length. If those two ever disagree, the divide stops being a true average and the seams come back. Keeping both keyed on chunk_size_actual is what keeps it honest.
 |
About the author. I'm Wlad Radchenko, a software engineer. The code in this article comes from Wunjo Make (open source), local software for video makers, and Wunjo Design, an offline PWA for designers. Get in touch to find more on GitHub and LinkedIn. |
Try it on your own audio
The whole mechanism is one loop and one divide. If you want to see it work, separate a track in Wunjo Make and listen to a long sustained note that crosses a chunk boundary. With overlap on it is smooth. Set overlap = 0 in your own copy and you will hear the ticks come back roughly every 4.4 seconds (the step length), which is the clearest proof that the window is doing the work.
The code is in portable/src/sound_processing/separator/mdx/model.py in the Wunjo Make repo. The MDX separation approach this builds on comes from KUIELab-MDX-Net.
References
- Kim, Lee, Lee, Lee. "KUIELab-MDX-Net: A Two-Stream Neural Network for Music Demixing." MDX Workshop / ISMIR 2021. arXiv:2111.12203. https://arxiv.org/abs/2111.12203
- Mitsufuji et al. "Music Demixing Challenge 2021." Frontiers in Signal Processing, 2022. arXiv:2108.13559. https://arxiv.org/abs/2108.13559
- Défossez. "Hybrid Spectrogram and Waveform Source Separation." (Demucs v3.) 2021. arXiv:2111.03600. https://arxiv.org/abs/2111.03600
Top comments (0)