Avoid STFT padding failure on short terminal chunks

#2
by faraday - opened

Summary

With a 44.1 kHz sample input (which I had created with:
ffmpeg -i <input wav> -t 00:02:00 -c copy <output wav> to cut first 2-minutes.

I've encountered padding problem with overlap=0, chunk-size 5s, 10s, 30s cases:
"Padding size should be less than the corresponding input dimension, but got:
padding (882, 882) at dimension 2 of input [1, 1, 32]"

This PR fixes the cases where short pieces left cause this problem with Pytorch reflective padding.

Traceback

Traceback (most recent call last):
File "/content/RE-USE/./inference_chunk.py", line 133, in
main()
~~~~^^
File "/content/RE-USE/./inference_chunk.py", line 130, in main
inference(args, device)
~~~~~~~~~^^^^^^^^^^^^^^
File "/content/RE-USE/./inference_chunk.py", line 85, in inference
noisy_mag, noisy_pha, noisy_com = mag_phase_stft(
~~~~~~~~~~~~~~^
noisy_wav_chunk,
^^^^^^^^^^^^^^^^
...<5 lines>...
addeps=False
^^^^^^^^^^^^
)
^
File "/content/RE-USE/models/stfts.py", line 35, in mag_phase_stft
stft_spec = torch.stft(
y, n_fft,
...<5 lines>...
normalized=False,
return_complex=True)
File "/usr/local/lib/python3.13/dist-packages/torch/functional.py", line 679, in stft
input = F.pad(input.view(extended_shape), [pad, pad], pad_mode)
File "/usr/local/lib/python3.13/dist-packages/torch/nn/functional.py", line 5430, in pad
return torch._C._nn.pad(input, pad, mode, value)
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Argument #4: Padding size should be less than the corresponding input dimension, but got: padding (882, 882) at dimension 2 of input [1, 1, 32]

Change

When a chunk is too short for STFT reflection padding, the inference code:

  • reuses the preceding chunk-sized audio as model context;
  • emits only the samples belonging to the short terminal piece;
  • zero-pads recordings that are themselves shorter than the safe STFT input
    length;
  • preserves the existing chunk schedule and output length.
    Normal-sized chunks continue to use the existing path.

Validation

The fix was reproduced and tested on an NVIDIA L4 GPU using 16 kHz and
44.1 kHz mono inputs with 5-, 10-, and 30-second chunks and zero overlap.
All configurations completed successfully preserving input sample rate and frame count.
Manually listened for sanity check.

The error can be produced synthetically but if needed, I can share my sample input which is based on a public standard dataset licensed CC BY SA 4.0

faraday changed pull request status to open

Sample to trigger the issue:

Start with this 44.1 kHz sample from CORAAL dataset
https://huggingface.co/datasets/zsayers/CORAAL/blob/main/wav/DCA_se2_ag1_f_07_1.wav

(this is a mirror, original page for CORAAL is https://oraal.github.io/coraal)

Then trim to 2-minutes by executing:
ffmpeg -i DCA_se2_ag1_f_07_1.wav -t 00:02:00 -c copy DCA_se2_ag1_f_07_1_trim_2m.wav

Now running NVIDIA RE-USE inference on DCA_se2_ag1_f_07_1_trim_2m.wav file should trigger the padding size error.

License for the sample: CC-BY-NC-SA
Kendall, Tyler and Charlie Farrington. 2023. The Corpus of Regional African American Language. Version 2023.06. Eugene, OR: The Online Resources for African American Language Project. [https://doi.org/10.7264/1ad5-6t35]

NVIDIA org

Thanks for pointing this out!
I’ve already updated the code so that the last chunk always has the full chunk_size by looking back from the end of the input (which should have better output quality).

Cannot merge
This branch has merge conflicts in the following files:
  • inference_chunk.py

Sign up or log in to comment