Spaces:
Running on Zero
Running on Zero
| from distutils.version import LooseVersion | |
| import torch | |
| from torch_complex.tensor import ComplexTensor | |
| from espnet2.enh.encoder.abs_encoder import AbsEncoder | |
| from espnet2.layers.stft import Stft | |
| is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0") | |
| class STFTEncoder(AbsEncoder): | |
| """STFT encoder for speech enhancement and separation""" | |
| def __init__( | |
| self, | |
| n_fft: int = 512, | |
| win_length: int = None, | |
| hop_length: int = 128, | |
| window="hann", | |
| center: bool = True, | |
| normalized: bool = False, | |
| onesided: bool = True, | |
| use_builtin_complex: bool = True, | |
| ): | |
| super().__init__() | |
| self.stft = Stft( | |
| n_fft=n_fft, | |
| win_length=win_length, | |
| hop_length=hop_length, | |
| window=window, | |
| center=center, | |
| normalized=normalized, | |
| onesided=onesided, | |
| ) | |
| self._output_dim = n_fft // 2 + 1 if onesided else n_fft | |
| self.use_builtin_complex = use_builtin_complex | |
| def output_dim(self) -> int: | |
| return self._output_dim | |
| def forward(self, input: torch.Tensor, ilens: torch.Tensor): | |
| """Forward. | |
| Args: | |
| input (torch.Tensor): mixed speech [Batch, sample] | |
| ilens (torch.Tensor): input lengths [Batch] | |
| """ | |
| spectrum, flens = self.stft(input, ilens) | |
| if is_torch_1_9_plus and self.use_builtin_complex: | |
| spectrum = torch.complex(spectrum[..., 0], spectrum[..., 1]) | |
| else: | |
| spectrum = ComplexTensor(spectrum[..., 0], spectrum[..., 1]) | |
| return spectrum, flens | |