ORI-Muchim/BlueArchiveTTS
58
1"""2BSD 3-Clause License3Copyright (c) 2017, Prem Seetharaman4All rights reserved.5* Redistribution and use in source and binary forms, with or without6 modification, are permitted provided that the following conditions are met:7* Redistributions of source code must retain the above copyright notice,8 this list of conditions and the following disclaimer.9* Redistributions in binary form must reproduce the above copyright notice, this10 list of conditions and the following disclaimer in the11 documentation and/or other materials provided with the distribution.12* Neither the name of the copyright holder nor the names of its13 contributors may be used to endorse or promote products derived from this14 software without specific prior written permission.15THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND16ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED17WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE18DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR19ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES20(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;21LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON22ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT23(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS24SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.25"""26 27import torch28import numpy as np29import torch.nn.functional as F30from torch.autograd import Variable31from scipy.signal import get_window32from librosa.util import pad_center, tiny33import librosa.util as librosa_util34 35def window_sumsquare(window, n_frames, hop_length=200, win_length=800,36 n_fft=800, dtype=np.float32, norm=None):37 """38 # from librosa 0.639 Compute the sum-square envelope of a window function at a given hop length.40 This is used to estimate modulation effects induced by windowing41 observations in short-time fourier transforms.42 Parameters43 ----------44 window : string, tuple, number, callable, or list-like45 Window specification, as in `get_window`46 n_frames : int > 047 The number of analysis frames48 hop_length : int > 049 The number of samples to advance between frames50 win_length : [optional]51 The length of the window function. By default, this matches `n_fft`.52 n_fft : int > 053 The length of each analysis frame.54 dtype : np.dtype55 The data type of the output56 Returns57 -------58 wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`59 The sum-squared envelope of the window function60 """61 if win_length is None:62 win_length = n_fft63 64 n = n_fft + hop_length * (n_frames - 1)65 x = np.zeros(n, dtype=dtype)66 67 # Compute the squared window at the desired length68 win_sq = get_window(window, win_length, fftbins=True)69 win_sq = librosa_util.normalize(win_sq, norm=norm)**270 win_sq = librosa_util.pad_center(win_sq, n_fft)71 72 # Fill the envelope73 for i in range(n_frames):74 sample = i * hop_length75 x[sample:min(n, sample + n_fft)] += win_sq[:max(0, min(n_fft, n - sample))]76 return x77 78 79class STFT(torch.nn.Module):80 """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""81 def __init__(self, filter_length=800, hop_length=200, win_length=800,82 window='hann'):83 super(STFT, self).__init__()84 self.filter_length = filter_length85 self.hop_length = hop_length86 self.win_length = win_length87 self.window = window88 self.forward_transform = None89 scale = self.filter_length / self.hop_length90 fourier_basis = np.fft.fft(np.eye(self.filter_length))91 92 cutoff = int((self.filter_length / 2 + 1))93 fourier_basis = np.vstack([np.real(fourier_basis[:cutoff, :]),94 np.imag(fourier_basis[:cutoff, :])])95 96 forward_basis = torch.FloatTensor(fourier_basis[:, None, :])97 inverse_basis = torch.FloatTensor(98 np.linalg.pinv(scale * fourier_basis).T[:, None, :])99 100 if window is not None:101 assert(filter_length >= win_length)102 # get window and zero center pad it to filter_length103 fft_window = get_window(window, win_length, fftbins=True)104 fft_window = pad_center(fft_window, filter_length)105 fft_window = torch.from_numpy(fft_window).float()106 107 # window the bases108 forward_basis *= fft_window109 inverse_basis *= fft_window110 111 self.register_buffer('forward_basis', forward_basis.float())112 self.register_buffer('inverse_basis', inverse_basis.float())113 114 def transform(self, input_data):115 num_batches = input_data.size(0)116 num_samples = input_data.size(1)117 118 self.num_samples = num_samples119 120 # similar to librosa, reflect-pad the input121 input_data = input_data.view(num_batches, 1, num_samples)122 input_data = F.pad(123 input_data.unsqueeze(1),124 (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0),125 mode='reflect')126 input_data = input_data.squeeze(1)127 128 forward_transform = F.conv1d(129 input_data,130 Variable(self.forward_basis, requires_grad=False),131 stride=self.hop_length,132 padding=0)133 134 cutoff = int((self.filter_length / 2) + 1)135 real_part = forward_transform[:, :cutoff, :]136 imag_part = forward_transform[:, cutoff:, :]137 138 magnitude = torch.sqrt(real_part**2 + imag_part**2)139 phase = torch.autograd.Variable(140 torch.atan2(imag_part.data, real_part.data))141 142 return magnitude, phase143 144 def inverse(self, magnitude, phase):145 recombine_magnitude_phase = torch.cat(146 [magnitude*torch.cos(phase), magnitude*torch.sin(phase)], dim=1)147 148 inverse_transform = F.conv_transpose1d(149 recombine_magnitude_phase,150 Variable(self.inverse_basis, requires_grad=False),151 stride=self.hop_length,152 padding=0)153 154 if self.window is not None:155 window_sum = window_sumsquare(156 self.window, magnitude.size(-1), hop_length=self.hop_length,157 win_length=self.win_length, n_fft=self.filter_length,158 dtype=np.float32)159 # remove modulation effects160 approx_nonzero_indices = torch.from_numpy(161 np.where(window_sum > tiny(window_sum))[0])162 window_sum = torch.autograd.Variable(163 torch.from_numpy(window_sum), requires_grad=False)164 window_sum = window_sum.to(inverse_transform.device()) if magnitude.is_cuda else window_sum165 inverse_transform[:, :, approx_nonzero_indices] /= window_sum[approx_nonzero_indices]166 167 # scale by hop ratio168 inverse_transform *= float(self.filter_length) / self.hop_length169 170 inverse_transform = inverse_transform[:, :, int(self.filter_length/2):]171 inverse_transform = inverse_transform[:, :, :-int(self.filter_length/2):]172 173 return inverse_transform174 175 def forward(self, input_data):176 self.magnitude, self.phase = self.transform(input_data)177 reconstruction = self.inverse(self.magnitude, self.phase)178 return reconstruction179 180 181class OnnxSTFT(torch.nn.Module):182 """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""183 def __init__(self, filter_length=800, hop_length=200, win_length=800,184 window='hann'):185 super(OnnxSTFT, self).__init__()186 self.filter_length = filter_length187 self.hop_length = hop_length188 self.win_length = win_length189 self.window = window190 self.forward_transform = None191 scale = self.filter_length / self.hop_length192 fourier_basis = np.fft.fft(np.eye(self.filter_length))193 194 cutoff = int((self.filter_length / 2 + 1))195 fourier_basis = np.vstack([np.real(fourier_basis[:cutoff, :]),196 np.imag(fourier_basis[:cutoff, :])])197 198 forward_basis = torch.FloatTensor(fourier_basis[:, None, :])199 inverse_basis = torch.FloatTensor(200 np.linalg.pinv(scale * fourier_basis).T[:, None, :])201 202 if window is not None:203 assert(filter_length >= win_length)204 # get window and zero center pad it to filter_length205 fft_window = get_window(window, win_length, fftbins=True)206 fft_window = pad_center(fft_window, filter_length)207 fft_window = torch.from_numpy(fft_window).float()208 209 # window the bases210 forward_basis *= fft_window211 inverse_basis *= fft_window212 213 self.register_buffer('forward_basis', forward_basis.float())214 self.register_buffer('inverse_basis', inverse_basis.float())215 216 def transform(self, input_data):217 num_batches = input_data.size(0)218 num_samples = input_data.size(1)219 220 self.num_samples = num_samples221 222 # similar to librosa, reflect-pad the input223 input_data = input_data.view(num_batches, 1, num_samples)224 input_data = F.pad(225 input_data.unsqueeze(1),226 (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0),227 mode='reflect')228 input_data = input_data.squeeze(1)229 230 forward_transform = F.conv1d(231 input_data,232 Variable(self.forward_basis, requires_grad=False),233 stride=self.hop_length,234 padding=0)235 236 cutoff = int((self.filter_length / 2) + 1)237 real_part = forward_transform[:, :cutoff, :]238 imag_part = forward_transform[:, cutoff:, :]239 240 magnitude = torch.sqrt(real_part**2 + imag_part**2)241 phase = torch.autograd.Variable(242 torch.atan2(imag_part.data, real_part.data))243 244 return magnitude, phase245 246 def inverse(self, magnitude, phase):247 recombine_magnitude_phase = torch.cat(248 [magnitude*torch.cos(phase), magnitude*torch.sin(phase)], dim=1)249 250 inverse_transform = F.conv_transpose1d(251 recombine_magnitude_phase,252 Variable(self.inverse_basis, requires_grad=False),253 stride=self.hop_length,254 padding=0)255 256 inverse_transform = inverse_transform[:, :, int(self.filter_length/2):]257 inverse_transform = inverse_transform[:, :, :-int(self.filter_length/2):]258 259 return inverse_transform260 261 def forward(self, input_data):262 self.magnitude, self.phase = self.transform(input_data)263 reconstruction = self.inverse(self.magnitude, self.phase)264 return reconstruction265 266 267class TorchSTFT(torch.nn.Module):268 def __init__(self, filter_length=800, hop_length=200, win_length=800, window='hann'):269 super().__init__()270 self.filter_length = filter_length271 self.hop_length = hop_length272 self.win_length = win_length273 self.window = torch.from_numpy(get_window(window, win_length, fftbins=True).astype(np.float32))274 275 def transform(self, input_data):276 forward_transform = torch.stft(277 input_data,278 self.filter_length, self.hop_length, self.win_length, window=self.window,279 return_complex=True)280 281 return torch.abs(forward_transform), torch.angle(forward_transform)282 283 def inverse(self, magnitude, phase):284 inverse_transform = torch.istft(285 magnitude * torch.exp(phase * 1j),286 self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device))287 288 return inverse_transform.unsqueeze(-2) # unsqueeze to stay consistent with conv_transpose1d implementation289 290 def forward(self, input_data):291 self.magnitude, self.phase = self.transform(input_data)292 reconstruction = self.inverse(self.magnitude, self.phase)293 return reconstruction294 295 296