Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges. See issue: https://github.com/pytorch/pytorch/issues/62323 Specifically, in the context o
| 321 | |
| 322 | |
| 323 | class ISTFT(nn.Module): |
| 324 | """ |
| 325 | Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with |
| 326 | windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges. |
| 327 | See issue: https://github.com/pytorch/pytorch/issues/62323 |
| 328 | Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs. |
| 329 | The NOLA constraint is met as we trim padded samples anyway. |
| 330 | |
| 331 | Args: |
| 332 | n_fft (int): Size of Fourier transform. |
| 333 | hop_length (int): The distance between neighboring sliding window frames. |
| 334 | win_length (int): The size of window frame and STFT filter. |
| 335 | padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
| 336 | """ |
| 337 | |
| 338 | def __init__( |
| 339 | self, n_fft: int, hop_length: int, win_length: int, padding: str = "same" |
| 340 | ): |
| 341 | super().__init__() |
| 342 | assert padding in ["center", "same"], "Padding must be 'center' or 'same'." |
| 343 | self.padding = padding |
| 344 | self.n_fft = n_fft |
| 345 | self.hop_length = hop_length |
| 346 | self.win_length = win_length |
| 347 | window = torch.hann_window(win_length) |
| 348 | self.register_buffer("window", window) |
| 349 | |
| 350 | def forward(self, spec: torch.Tensor) -> torch.Tensor: |
| 351 | """ |
| 352 | Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. |
| 353 | |
| 354 | Args: |
| 355 | spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size, |
| 356 | N is the number of frequency bins, and T is the number of time frames. |
| 357 | |
| 358 | Returns: |
| 359 | Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal. |
| 360 | """ |
| 361 | if self.padding == "center": |
| 362 | # Fallback to pytorch native implementation |
| 363 | return torch.istft( |
| 364 | spec, |
| 365 | self.n_fft, |
| 366 | self.hop_length, |
| 367 | self.win_length, |
| 368 | self.window, |
| 369 | center=True, |
| 370 | ) |
| 371 | elif self.padding == "same": |
| 372 | pad = (self.win_length - self.hop_length) // 2 |
| 373 | else: |
| 374 | raise ValueError("Padding must be 'center' or 'same'.") |
| 375 | |
| 376 | assert spec.dim() == 3, "Expected a 3D tensor as input" |
| 377 | B, N, T = spec.shape |
| 378 | |
| 379 | # Inverse FFT |
| 380 | ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") |