import torch import torch.nn.functional as F import torch.nn as nn from torch.nn import Conv1d, ConvTranspose1d from torch.nn.utils.parametrizations import weight_norm import math import numpy as np def get_padding(kernel_size, dilation=1): return int((kernel_size*dilation - dilation)/2) def _tile(x, length=None): x = x.repeat(1, 1, int(length / x.shape[2]) + 1)[:, :, :length] return x class AdaIN1d(nn.Module): # used by HiFiGan & ProsodyPredictor def __init__(self, style_dim, num_features): super().__init__() self.norm = nn.InstanceNorm1d(num_features, affine=False) self.fc = nn.Linear(style_dim, num_features*2) def forward(self, x, s): # x = torch.Size([1, 512, 248]) same as output # s = torch.Size([1, 7, 1, 128]) s = self.fc(s.transpose(1, 2)).transpose(1, 2) s = _tile(s, length=x.shape[2]) gamma, beta = torch.chunk(s, chunks=2, dim=1) return (1+gamma) * self.norm(x) + beta class AdaINResBlock1(torch.nn.Module): def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64): super(AdaINResBlock1, self).__init__() self.convs1 = nn.ModuleList([ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding=get_padding(kernel_size, dilation[0]))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding=get_padding(kernel_size, dilation[1]))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding=get_padding(kernel_size, dilation[2]))) ]) # self.convs1.apply(init_weights) self.convs2 = nn.ModuleList([ weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))), weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))) ]) # self.convs2.apply(init_weights) self.adain1 = nn.ModuleList([ AdaIN1d(style_dim, channels), AdaIN1d(style_dim, channels), AdaIN1d(style_dim, channels), ]) self.adain2 = nn.ModuleList([ AdaIN1d(style_dim, channels), AdaIN1d(style_dim, channels), AdaIN1d(style_dim, channels), ]) self.alpha1 = nn.ParameterList( [nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))]) self.alpha2 = nn.ParameterList( [nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))]) def forward(self, x, s): for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2): xt = n1(x, s) # THIS IS ADAIN - EXPECTS conv1d dims xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2) # Snake1D xt = c1(xt) xt = n2(xt, s) # THIS IS ADAIN - EXPECTS conv1d dims xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2) # Snake1D xt = c2(xt) x = xt + x return x class SourceModuleHnNSF(torch.nn.Module): def __init__(self): super().__init__() self.harmonic_num = 8 self.l_linear = torch.nn.Linear(self.harmonic_num + 1, 1) self.upsample_scale = 300 def forward(self, x): # -- x = torch.multiply(x, torch.FloatTensor( [[range(1, self.harmonic_num + 2)]]).to(x.device)) # [1, 145200, 9] # modulo of negative f0_values => -21 % 10 = 9 as -3*10 + 9 = 21 NOTICE THAT f0_values IS SIGNED rad_values = x / 25647 #).clamp(0, 1) # rad_values = torch.where(torch.logical_or(rad_values < 0, rad_values > 1), 0.5, rad_values) rad_values = rad_values % 1 # % of neg values rad_values = F.interpolate(rad_values.transpose(1, 2), scale_factor=1/self.upsample_scale, mode='linear').transpose(1, 2) # 1.89 sounds also nice has woofer at punctuation phase = torch.cumsum(rad_values, dim=1) * 1.84 * np.pi phase = F.interpolate(phase.transpose(1, 2) * self.upsample_scale, scale_factor=self.upsample_scale, mode='linear').transpose(1, 2) x = .009 * phase.sin() # -- x = self.l_linear(x).tanh() return x class Generator(torch.nn.Module): def __init__(self, style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes): super(Generator, self).__init__() self.num_kernels = len(resblock_kernel_sizes) self.num_upsamples = len(upsample_rates) self.m_source = SourceModuleHnNSF() self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates)) self.noise_convs = nn.ModuleList() self.ups = nn.ModuleList() self.noise_res = nn.ModuleList() for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): c_cur = upsample_initial_channel // (2 ** (i + 1)) self.ups.append(weight_norm(ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//( 2**(i+1)), k, u, padding=(u//2 + u % 2), output_padding=u % 2))) if i + 1 < len(upsample_rates): stride_f0 = np.prod(upsample_rates[i + 1:]) self.noise_convs.append(Conv1d( 1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2)) self.noise_res.append(AdaINResBlock1( c_cur, 7, [1, 3, 5], style_dim)) else: self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1)) self.noise_res.append(AdaINResBlock1( c_cur, 11, [1, 3, 5], style_dim)) self.resblocks = nn.ModuleList() self.alphas = nn.ParameterList() self.alphas.append(nn.Parameter( torch.ones(1, upsample_initial_channel, 1))) for i in range(len(self.ups)): ch = upsample_initial_channel//(2**(i+1)) self.alphas.append(nn.Parameter(torch.ones(1, ch, 1))) for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): self.resblocks.append(AdaINResBlock1(ch, k, d, style_dim)) self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3)) def forward(self, x, s, f0): # x.shape=torch.Size([1, 512, 484]) s.shape=torch.Size([1, 1, 1, 128]) f0.shape=torch.Size([1, 484]) GENERAT 249 f0 = self.f0_upsamp(f0).transpose(1, 2) # x.shape=torch.Size([1, 512, 484]) s.shape=torch.Size([1, 1, 1, 128]) f0.shape=torch.Size([1, 145200, 1]) GENERAT 253 # [1, 145400, 1] f0 enters already upsampled to full wav 24kHz length har_source = self.m_source(f0) har_source = har_source.transpose(1, 2) for i in range(self.num_upsamples): x = x + (1 / self.alphas[i]) * (torch.sin(self.alphas[i] * x) ** 2) x_source = self.noise_convs[i](har_source) x_source = self.noise_res[i](x_source, s) x = self.ups[i](x) x = x + x_source xs = None for j in range(self.num_kernels): if xs is None: xs = self.resblocks[i*self.num_kernels+j](x, s) else: xs += self.resblocks[i*self.num_kernels+j](x, s) x = xs / self.num_kernels # x = x + (1 / self.alphas[i+1]) * (torch.sin(self.alphas[i+1] * x) ** 2) # noisy x = self.conv_post(x) x = torch.tanh(x) return x class AdainResBlk1d(nn.Module): # also used in ProsodyPredictor() def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2), upsample='none', dropout_p=0.0): super().__init__() self.actv = actv self.upsample_type = upsample self.upsample = UpSample1d(upsample) self.learned_sc = dim_in != dim_out self._build_weights(dim_in, dim_out, style_dim) if upsample == 'none': self.pool = nn.Identity() else: self.pool = weight_norm(nn.ConvTranspose1d( dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1)) def _build_weights(self, dim_in, dim_out, style_dim): self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1)) self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1)) self.norm1 = AdaIN1d(style_dim, dim_in) self.norm2 = AdaIN1d(style_dim, dim_out) if self.learned_sc: self.conv1x1 = weight_norm( nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False)) def _shortcut(self, x): x = self.upsample(x) if self.learned_sc: x = self.conv1x1(x) return x def _residual(self, x, s): x = self.norm1(x, s) x = self.actv(x) x = self.pool(x) x = self.conv1(x) x = self.norm2(x, s) x = self.actv(x) x = self.conv2(x) return x def forward(self, x, s): out = self._residual(x, s) out = (out + self._shortcut(x)) / math.sqrt(2) return out class UpSample1d(nn.Module): def __init__(self, layer_type): super().__init__() self.layer_type = layer_type def forward(self, x): if self.layer_type == 'none': return x else: return F.interpolate(x, scale_factor=2, mode='nearest-exact') class Decoder(nn.Module): def __init__(self, dim_in=512, F0_channel=512, style_dim=64, dim_out=80, resblock_kernel_sizes=[3, 7, 11], upsample_rates=[10, 5, 3, 2], upsample_initial_channel=512, resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], upsample_kernel_sizes=[20, 10, 6, 4]): super().__init__() self.decode = nn.ModuleList() self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim) self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim)) self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim)) self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim)) self.decode.append(AdainResBlk1d( 1024 + 2 + 64, 512, style_dim, upsample=True)) self.F0_conv = weight_norm( nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1)) # smooth self.N_conv = weight_norm( nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1)) self.asr_res = nn.Sequential( weight_norm(nn.Conv1d(512, 64, kernel_size=1)), ) self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes) def forward(self, asr=None, F0_curve=None, N=None, s=None): F0 = self.F0_conv(F0_curve) N = self.N_conv(N) x = torch.cat([asr, F0, N], axis=1) x = self.encode(x, s) asr_res = self.asr_res(asr) res = True for block in self.decode: if res: x = torch.cat([x, asr_res, F0, N], axis=1) x = block(x, s) if block.upsample_type != "none": res = False x = self.generator(x, s, F0_curve) return x