| 1 | # Modified from https://github.com/echocatzh/conv-stft/blob/master/conv_stft/conv_stft.py |
| 2 | |
| 3 | # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. |
| 4 | # |
| 5 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | # you may not use this file except in compliance with the License. |
| 7 | # You may obtain a copy of the License at |
| 8 | # |
| 9 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | # |
| 11 | # Unless required by applicable law or agreed to in writing, software |
| 12 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. |
| 16 | |
| 17 | # MIT License |
| 18 | |
| 19 | # Copyright (c) 2020 Shimin Zhang |
| 20 | |
| 21 | # Permission is hereby granted, free of charge, to any person obtaining a copy |
| 22 | # of this software and associated documentation files (the "Software"), to deal |
| 23 | # in the Software without restriction, including without limitation the rights |
| 24 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 25 | # copies of the Software, and to permit persons to whom the Software is |
| 26 | # furnished to do so, subject to the following conditions: |
| 27 | |
| 28 | # The above copyright notice and this permission notice shall be included in all |
| 29 | # copies or substantial portions of the Software. |
| 30 | |
| 31 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 32 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 33 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 34 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 35 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 36 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 37 | # SOFTWARE. |
| 38 | |
| 39 | import torch as th |
| 40 | import torch.nn.functional as F |
| 41 | from scipy.signal import check_COLA, get_window |
| 42 | |
| 43 | |
| 44 | support_clp_op = None |
| 45 | if th.__version__ >= "1.7.0": |
| 46 | from torch.fft import rfft as fft |
| 47 | |
| 48 | support_clp_op = True |
| 49 | else: |
| 50 | from torch import rfft as fft |
| 51 | |
| 52 | |
| 53 | class STFT(th.nn.Module): |
| 54 | def __init__( |
| 55 | self, |
| 56 | win_len=1024, |
| 57 | win_hop=512, |
| 58 | fft_len=1024, |
| 59 | enframe_mode="continue", |
| 60 | win_type="hann", |
| 61 | win_sqrt=False, |
| 62 | pad_center=True, |
| 63 | ): |
| 64 | """ |
| 65 | Implement of STFT using 1D convolution and 1D transpose convolutions. |
| 66 | Implement of framing the signal in 2 ways, `break` and `continue`. |
| 67 | `break` method is a kaldi-like framing. |
| 68 | `continue` method is a librosa-like framing. |
| 69 | |
| 70 | More information about `perfect reconstruction`: |
| 71 | 1. https://ww2.mathworks.cn/help/signal/ref/stft.html |
| 72 | 2. https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html |
| 73 | |
| 74 | Args: |
| 75 | win_len (int): Number of points in one frame. Defaults to 1024. |
| 76 | win_hop (int): Number of framing stride. Defaults to 512. |
| 77 | fft_len (int): Number of DFT points. Defaults to 1024. |
| 78 | enframe_mode (str, optional): `break` and `continue`. Defaults to 'continue'. |
| 79 | win_type (str, optional): The type of window to create. Defaults to 'hann'. |
| 80 | win_sqrt (bool, optional): using square root window. Defaults to True. |
| 81 | pad_center (bool, optional): `perfect reconstruction` opts. Defaults to True. |
| 82 | """ |
| 83 | super(STFT, self).__init__() |
| 84 | assert enframe_mode in ["break", "continue"] |
| 85 | assert fft_len >= win_len |
| 86 | self.win_len = win_len |
| 87 | self.win_hop = win_hop |
| 88 | self.fft_len = fft_len |
| 89 | self.mode = enframe_mode |
| 90 | self.win_type = win_type |
| 91 | self.win_sqrt = win_sqrt |
| 92 | self.pad_center = pad_center |
| 93 | self.pad_amount = self.fft_len // 2 |
| 94 | |
| 95 | en_k, fft_k, ifft_k, ola_k = self.__init_kernel__() |
| 96 | self.register_buffer("en_k", en_k) |
| 97 | self.register_buffer("fft_k", fft_k) |
| 98 | self.register_buffer("ifft_k", ifft_k) |
| 99 | self.register_buffer("ola_k", ola_k) |
| 100 | |
| 101 | def __init_kernel__(self): |
| 102 | """ |
| 103 | Generate enframe_kernel, fft_kernel, ifft_kernel and overlap-add kernel. |
| 104 | ** enframe_kernel: Using conv1d layer and identity matrix. |
| 105 | ** fft_kernel: Using linear layer for matrix multiplication. In fact, |
| 106 | enframe_kernel and fft_kernel can be combined, But for the sake of |
| 107 | readability, I took the two apart. |
| 108 | ** ifft_kernel, pinv of fft_kernel. |
| 109 | ** overlap-add kernel, just like enframe_kernel, but transposed. |
| 110 | |
| 111 | Returns: |
| 112 | tuple: four kernels. |
| 113 | """ |
| 114 | enframed_kernel = th.eye(self.fft_len)[:, None, :] |
| 115 | if support_clp_op: |
| 116 | tmp = fft(th.eye(self.fft_len)) |
| 117 | fft_kernel = th.stack([tmp.real, tmp.imag], dim=2) |
| 118 | else: |
| 119 | fft_kernel = fft(th.eye(self.fft_len), 1) |
| 120 | if self.mode == "break": |
| 121 | enframed_kernel = th.eye(self.win_len)[:, None, :] |
| 122 | fft_kernel = fft_kernel[: self.win_len] |
| 123 | fft_kernel = th.cat((fft_kernel[:, :, 0], fft_kernel[:, :, 1]), dim=1) |
| 124 | ifft_kernel = th.pinverse(fft_kernel)[:, None, :] |
| 125 | window = get_window(self.win_type, self.win_len) |
| 126 | |
| 127 | self.perfect_reconstruct = check_COLA(window, self.win_len, self.win_len - self.win_hop) |
| 128 | window = th.FloatTensor(window) |
| 129 | if self.mode == "continue": |
| 130 | left_pad = (self.fft_len - self.win_len) // 2 |
| 131 | right_pad = left_pad + (self.fft_len - self.win_len) % 2 |
| 132 | window = F.pad(window, (left_pad, right_pad)) |
| 133 | if self.win_sqrt: |
| 134 | self.padded_window = window |
| 135 | window = th.sqrt(window) |
| 136 | else: |
| 137 | self.padded_window = window**2 |
| 138 | |
| 139 | fft_kernel = fft_kernel.T * window |
| 140 | ifft_kernel = ifft_kernel * window |
| 141 | ola_kernel = th.eye(self.fft_len)[: self.win_len, None, :] |
| 142 | if self.mode == "continue": |
| 143 | ola_kernel = th.eye(self.fft_len)[:, None, : self.fft_len] |
| 144 | return enframed_kernel, fft_kernel, ifft_kernel, ola_kernel |
| 145 | |
| 146 | def is_perfect(self): |
| 147 | """ |
| 148 | Whether the parameters win_len, win_hop and win_sqrt |
| 149 | obey constants overlap-add(COLA) |
| 150 | |
| 151 | Returns: |
| 152 | bool: Return true if parameters obey COLA. |
| 153 | """ |
| 154 | return self.perfect_reconstruct and self.pad_center |
| 155 | |
| 156 | def transform(self, inputs, return_type="complex"): |
| 157 | """Take input data (audio) to STFT domain. |
| 158 | |
| 159 | Args: |
| 160 | inputs (tensor): Tensor of floats, with shape (num_batch, num_samples) |
| 161 | return_type (str, optional): return (mag, phase) when `magphase`, |
| 162 | return (real, imag) when `realimag` and complex(real, imag) when `complex`. |
| 163 | Defaults to 'complex'. |
| 164 | |
| 165 | Returns: |
| 166 | tuple: (mag, phase) when `magphase`, return (real, imag) when |
| 167 | `realimag`. Defaults to 'complex', each elements with shape |
| 168 | [num_batch, num_frequencies, num_frames] |
| 169 | """ |
| 170 | assert return_type in ["magphase", "realimag", "complex"] |
| 171 | if inputs.dim() == 2: |
| 172 | inputs = th.unsqueeze(inputs, 1) |
| 173 | self.num_samples = inputs.size(-1) |
| 174 | if self.pad_center: |
| 175 | inputs = F.pad(inputs, (self.pad_amount, self.pad_amount), mode="reflect") |
| 176 | enframe_inputs = F.conv1d(inputs, self.en_k, stride=self.win_hop) |
| 177 | outputs = th.transpose(enframe_inputs, 1, 2) |
| 178 | outputs = F.linear(outputs, self.fft_k) |
| 179 | outputs = th.transpose(outputs, 1, 2) |
| 180 | dim = self.fft_len // 2 + 1 |
| 181 | real = outputs[:, :dim, :] |
| 182 | imag = outputs[:, dim:, :] |
| 183 | if return_type == "realimag": |
| 184 | return real, imag |
| 185 | elif return_type == "complex": |
| 186 | assert support_clp_op |
| 187 | return th.complex(real, imag) |
| 188 | else: |
| 189 | mags = th.sqrt(real**2 + imag**2) |
| 190 | phase = th.atan2(imag, real) |
| 191 | return mags, phase |
| 192 | |
| 193 | def inverse(self, input1, input2=None, input_type="magphase"): |
| 194 | """Call the inverse STFT (iSTFT), given tensors produced |
| 195 | by the `transform` function. |
| 196 | |
| 197 | Args: |
| 198 | input1 (tensors): Magnitude/Real-part of STFT with shape |
| 199 | [num_batch, num_frequencies, num_frames] |
| 200 | input2 (tensors): Phase/Imag-part of STFT with shape |
| 201 | [num_batch, num_frequencies, num_frames] |
| 202 | input_type (str, optional): Mathematical meaning of input tensor's. |
| 203 | Defaults to 'magphase'. |
| 204 | |
| 205 | Returns: |
| 206 | tensors: Reconstructed audio given magnitude and phase. Of |
| 207 | shape [num_batch, num_samples] |
| 208 | """ |
| 209 | assert input_type in ["magphase", "realimag"] |
| 210 | if input_type == "realimag": |
| 211 | real, imag = None, None |
| 212 | if support_clp_op and th.is_complex(input1): |
| 213 | real, imag = input1.real, input1.imag |
| 214 | else: |
| 215 | real, imag = input1, input2 |
| 216 | else: |
| 217 | real = input1 * th.cos(input2) |
| 218 | imag = input1 * th.sin(input2) |
| 219 | inputs = th.cat([real, imag], dim=1) |
| 220 | outputs = F.conv_transpose1d(inputs, self.ifft_k, stride=self.win_hop) |
| 221 | t = (self.padded_window[None, :, None]).repeat(1, 1, inputs.size(-1)) |
| 222 | t = t.to(inputs.device) |
| 223 | coff = F.conv_transpose1d(t, self.ola_k, stride=self.win_hop) |
| 224 | |
| 225 | num_frames = input1.size(-1) |
| 226 | num_samples = num_frames * self.win_hop |
| 227 | |
| 228 | rm_start, rm_end = self.pad_amount, self.pad_amount + num_samples |
| 229 | |
| 230 | outputs = outputs[..., rm_start:rm_end] |
| 231 | coff = coff[..., rm_start:rm_end] |
| 232 | coffidx = th.where(coff > 1e-8) |
| 233 | outputs[coffidx] = outputs[coffidx] / (coff[coffidx]) |
| 234 | return outputs.squeeze(dim=1) |
| 235 | |
| 236 | def forward(self, inputs): |
| 237 | """Take input data (audio) to STFT domain and then back to audio. |
| 238 | |
| 239 | Args: |
| 240 | inputs (tensor): Tensor of floats, with shape [num_batch, num_samples] |
| 241 | |
| 242 | Returns: |
| 243 | tensor: Reconstructed audio given magnitude and phase. |
| 244 | Of shape [num_batch, num_samples] |
| 245 | """ |
| 246 | mag, phase = self.transform(inputs) |
| 247 | rec_wav = self.inverse(mag, phase) |
| 248 | return rec_wav |
| 249 |