Avoid STFT padding failure on short terminal chunks

inference_chunk.py

@@ -79,8 +79,24 @@ def inference(args, device):
             window_sum = torch.zeros_like(Noisy_wav).to(device)
             for c in range(Noisy_wav.shape[0]): # for multi-channel speech
                 noisy_wav = Noisy_wav[c:c+1,:]
+                minimum_safe_input = n_fft_scaled // 2 + 1
                 for i in range(max(1, math.ceil((noisy_wav.shape[1]-chunk_size)/hop_length)+1)):
-                    noisy_wav_chunk = noisy_wav[:, i*hop_length : i*hop_length+chunk_size]
+                    emit_start = i * hop_length
+                    emit_end = min(emit_start + chunk_size, noisy_wav.shape[1])
+                    model_start = emit_start
+
+                    # torch.stft reflection padding fails when the input is no
+                    # longer than n_fft // 2. For a short final piece, reuse
+                    # preceding audio as context and emit only that piece.
+                    if emit_end - emit_start < minimum_safe_input:
+                        model_start = max(0, emit_end - chunk_size)
+
+                    noisy_wav_chunk = noisy_wav[:, model_start:emit_end]
+                    if noisy_wav_chunk.shape[1] < minimum_safe_input:
+                        noisy_wav_chunk = nn.functional.pad(
+                            noisy_wav_chunk,
+                            (0, minimum_safe_input - noisy_wav_chunk.shape[1]),
+                        )
 
                     noisy_mag, noisy_pha, noisy_com = mag_phase_stft(
                         noisy_wav_chunk,
@@ -100,8 +116,10 @@ def inference(args, device):
                     audio_g = mag_phase_istft(amp_g, pha_g, n_fft_scaled, hop_size_scaled, win_size_scaled, compress_factor)
                     audio_g = pad_or_trim_to_match(noisy_wav_chunk.detach(), audio_g, pad_value=1e-8)  # Align lengths using epsilon padding
 
-                    audio_enhanced[c:c+1,i*hop_length:i*hop_length+chunk_size] += audio_g*window[0:audio_g.shape[1]] 
-                    window_sum[c:c+1,i*hop_length:i*hop_length+chunk_size] += window[0:audio_g.shape[1]]
+                    local_start = emit_start - model_start
+                    local_end = local_start + (emit_end - emit_start)
+                    audio_enhanced[c:c+1,emit_start:emit_end] += audio_g[:,local_start:local_end]*window[local_start:local_end]
+                    window_sum[c:c+1,emit_start:emit_end] += window[local_start:local_end]
                     #norm[c:c+1,i*hop_length:i*hop_length+chunk_size] += 1.0
             nonzero_indices = (window_sum > 1e-8)
             audio_enhanced[:,nonzero_indices[0]] = audio_enhanced[:,nonzero_indices[0]]/window_sum[:,nonzero_indices[0]]