ISTFT Head module for predicting STFT complex coefficients. Args: dim (int): Hidden dimension of the model. n_fft (int): Size of Fourier transform. hop_length (int): The distance between neighboring sliding window frames, which should align with
| 469 | |
| 470 | |
| 471 | class ISTFTHead(nn.Module): |
| 472 | """ |
| 473 | ISTFT Head module for predicting STFT complex coefficients. |
| 474 | |
| 475 | Args: |
| 476 | dim (int): Hidden dimension of the model. |
| 477 | n_fft (int): Size of Fourier transform. |
| 478 | hop_length (int): The distance between neighboring sliding window frames, which should align with |
| 479 | the resolution of the input features. |
| 480 | padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
| 481 | """ |
| 482 | |
| 483 | def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): |
| 484 | super().__init__() |
| 485 | self.hop_length = hop_length |
| 486 | out_dim = n_fft + 2 |
| 487 | self.out = torch.nn.Linear(dim, out_dim) |
| 488 | self.istft = ISTFT( |
| 489 | n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding |
| 490 | ) |
| 491 | |
| 492 | def forward(self, x: torch.Tensor, x_len: torch.Tensor) -> torch.Tensor: |
| 493 | """ |
| 494 | Forward pass of the ISTFTHead module. |
| 495 | |
| 496 | Args: |
| 497 | x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, |
| 498 | L is the sequence length, and H denotes the model dimension. |
| 499 | |
| 500 | Returns: |
| 501 | Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. |
| 502 | """ |
| 503 | x_pred = self.out(x) |
| 504 | x_pred = x_pred.transpose(1, 2) |
| 505 | mag, p = x_pred.chunk(2, dim=1) |
| 506 | mag = torch.exp(mag) |
| 507 | mag = torch.clip( |
| 508 | mag, max=1e2 |
| 509 | ) # safeguard to prevent excessively large magnitudes |
| 510 | # wrapping happens here. These two lines produce real and imaginary value |
| 511 | x = torch.cos(p) |
| 512 | y = torch.sin(p) |
| 513 | # recalculating phase here does not produce anything new |
| 514 | # only costs time |
| 515 | # phase = torch.atan2(y, x) |
| 516 | # S = mag * torch.exp(phase * 1j) |
| 517 | # better directly produce the complex value |
| 518 | S = mag * (x + 1j * y) |
| 519 | audio = self.istft(S) |
| 520 | audio_length = x_len * self.hop_length |
| 521 | return audio, audio_length |
| 522 | |
| 523 | def forward_chunk( |
| 524 | self, x: torch.Tensor, cache: torch.Tensor = None, last_chunk: bool = False |
| 525 | ): |
| 526 | """ISTFTHead can be adapted in streaming inference without retraining. |
| 527 | |
| 528 | Args: |