import torch import torch.nn as nn import numpy as np from torch.nn import Conv1d from torch.nn import ConvTranspose1d from torch.nn.utils import weight_norm from torch.nn.utils import remove_weight_norm from .nsf import SourceModuleHnNSF from .bigv import init_weights, SnakeBeta, AMPBlock from .alias import Activation1d class SpeakerAdapter(nn.Module): def __init__(self, speaker_dim, adapter_dim, epsilon=1e-5 ): super(SpeakerAdapter, self).__init__() self.speaker_dim = speaker_dim self.adapter_dim = adapter_dim self.epsilon = epsilon self.W_scale = nn.Linear(self.speaker_dim, self.adapter_dim) self.W_bias = nn.Linear(self.speaker_dim, self.adapter_dim) self.reset_parameters() def reset_parameters(self): torch.nn.init.constant_(self.W_scale.weight, 0.0) torch.nn.init.constant_(self.W_scale.bias, 1.0) torch.nn.init.constant_(self.W_bias.weight, 0.0) torch.nn.init.constant_(self.W_bias.bias, 0.0) def forward(self, x, speaker_embedding): x = x.transpose(1, -1) mean = x.mean(dim=-1, keepdim=True) var = ((x - mean) ** 2).mean(dim=-1, keepdim=True) std = (var + self.epsilon).sqrt() y = (x - mean) / std scale = self.W_scale(speaker_embedding) bias = self.W_bias(speaker_embedding) y *= scale.unsqueeze(1) y += bias.unsqueeze(1) y = y.transpose(1, -1) return y class Generator(torch.nn.Module): # this is our main BigVGAN model. Applies anti-aliased periodic activation for resblocks. def __init__(self, hp): super(Generator, self).__init__() self.hp = hp self.num_kernels = len(hp.gen.resblock_kernel_sizes) self.num_upsamples = len(hp.gen.upsample_rates) # speaker adaper, 256 should change by what speaker encoder you use self.adapter = SpeakerAdapter(hp.vits.spk_dim, hp.gen.upsample_input) # pre conv self.conv_pre = nn.utils.weight_norm( Conv1d(hp.gen.upsample_input, hp.gen.upsample_initial_channel, 7, 1, padding=3)) # nsf self.f0_upsamp = torch.nn.Upsample( scale_factor=np.prod(hp.gen.upsample_rates)) self.m_source = SourceModuleHnNSF() self.noise_convs = nn.ModuleList() # transposed conv-based upsamplers. does not apply anti-aliasing self.ups = nn.ModuleList() for i, (u, k) in enumerate(zip(hp.gen.upsample_rates, hp.gen.upsample_kernel_sizes)): # print(f'ups: {i} {k}, {u}, {(k - u) // 2}') # base self.ups.append(nn.ModuleList([ weight_norm(ConvTranspose1d(hp.gen.upsample_initial_channel // (2 ** i), hp.gen.upsample_initial_channel // ( 2 ** (i + 1)), k, u, padding=(k - u) // 2)) ])) # nsf if i + 1 < len(hp.gen.upsample_rates): stride_f0 = np.prod(hp.gen.upsample_rates[i + 1:]) stride_f0 = int(stride_f0) self.noise_convs.append( Conv1d( 1, hp.gen.upsample_initial_channel // (2 ** (i + 1)), kernel_size=stride_f0 * 2, stride=stride_f0, padding=stride_f0 // 2, ) ) else: self.noise_convs.append( Conv1d(1, hp.gen.upsample_initial_channel // (2 ** (i + 1)), kernel_size=1) ) # residual blocks using anti-aliased multi-periodicity composition modules (AMP) self.resblocks = nn.ModuleList() for i in range(len(self.ups)): ch = hp.gen.upsample_initial_channel // (2 ** (i + 1)) for k, d in zip(hp.gen.resblock_kernel_sizes, hp.gen.resblock_dilation_sizes): self.resblocks.append(AMPBlock(hp, ch, k, d)) # post conv activation_post = SnakeBeta(ch, alpha_logscale=True) self.activation_post = Activation1d(activation=activation_post) self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3)) # weight initialization for i in range(len(self.ups)): self.ups[i].apply(init_weights) self.conv_post.apply(init_weights) def forward(self, spk, x, f0): # adapter x = self.adapter(x, spk) # nsf f0 = f0[:, None] f0 = self.f0_upsamp(f0).transpose(1, 2) har_source = self.m_source(f0) har_source = har_source.transpose(1, 2) x = self.conv_pre(x) for i in range(self.num_upsamples): # upsampling for i_up in range(len(self.ups[i])): x = self.ups[i][i_up](x) # nsf x_source = self.noise_convs[i](har_source) x = x + x_source # AMP blocks xs = None for j in range(self.num_kernels): if xs is None: xs = self.resblocks[i * self.num_kernels + j](x) else: xs += self.resblocks[i * self.num_kernels + j](x) x = xs / self.num_kernels # post conv x = self.activation_post(x) x = self.conv_post(x) x = torch.tanh(x) return x def remove_weight_norm(self): for l in self.ups: for l_i in l: remove_weight_norm(l_i) for l in self.resblocks: l.remove_weight_norm() remove_weight_norm(self.conv_pre) remove_weight_norm(self.conv_post) def eval(self, inference=False): super(Generator, self).eval() # don't remove weight norm while validation in training loop if inference: self.remove_weight_norm() def pitch2source(self, f0): f0 = f0[:, None] f0 = self.f0_upsamp(f0).transpose(1, 2) # [1,len,1] har_source = self.m_source(f0) har_source = har_source.transpose(1, 2) # [1,1,len] return har_source def source2wav(self, audio): MAX_WAV_VALUE = 32768.0 audio = audio.squeeze() audio = MAX_WAV_VALUE * audio audio = audio.clamp(min=-MAX_WAV_VALUE, max=MAX_WAV_VALUE-1) audio = audio.short() return audio.cpu().detach().numpy() def inference(self, spk, x, har_source): # adapter x = self.adapter(x, spk) x = self.conv_pre(x) for i in range(self.num_upsamples): # upsampling for i_up in range(len(self.ups[i])): x = self.ups[i][i_up](x) # nsf x_source = self.noise_convs[i](har_source) x = x + x_source # AMP blocks xs = None for j in range(self.num_kernels): if xs is None: xs = self.resblocks[i * self.num_kernels + j](x) else: xs += self.resblocks[i * self.num_kernels + j](x) x = xs / self.num_kernels # post conv x = self.activation_post(x) x = self.conv_post(x) x = torch.tanh(x) return x