# Copyright (c) 2023 Amphion. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import numpy as np import torch from torch import nn, sin, pow from torch.nn import Parameter import torch.nn.functional as F from torch.nn.utils import weight_norm from .alias_free_torch import * from .quantize import * from einops import rearrange from einops.layers.torch import Rearrange from .transformer import TransformerEncoder from .gradient_reversal import GradientReversal from .melspec import MelSpectrogram def init_weights(m): if isinstance(m, nn.Conv1d): nn.init.trunc_normal_(m.weight, std=0.02) nn.init.constant_(m.bias, 0) def WNConv1d(*args, **kwargs): return weight_norm(nn.Conv1d(*args, **kwargs)) def WNConvTranspose1d(*args, **kwargs): return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) class CNNLSTM(nn.Module): def __init__(self, indim, outdim, head, global_pred=False): super().__init__() self.global_pred = global_pred self.model = nn.Sequential( ResidualUnit(indim, dilation=1), ResidualUnit(indim, dilation=2), ResidualUnit(indim, dilation=3), Activation1d(activation=SnakeBeta(indim, alpha_logscale=True)), Rearrange("b c t -> b t c"), ) self.heads = nn.ModuleList([nn.Linear(indim, outdim) for i in range(head)]) def forward(self, x): # x: [B, C, T] x = self.model(x) if self.global_pred: x = torch.mean(x, dim=1, keepdim=False) outs = [head(x) for head in self.heads] return outs class SnakeBeta(nn.Module): """ A modified Snake function which uses separate parameters for the magnitude of the periodic components Shape: - Input: (B, C, T) - Output: (B, C, T), same shape as the input Parameters: - alpha - trainable parameter that controls frequency - beta - trainable parameter that controls magnitude References: - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: https://arxiv.org/abs/2006.08195 Examples: >>> a1 = snakebeta(256) >>> x = torch.randn(256) >>> x = a1(x) """ def __init__( self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False ): """ Initialization. INPUT: - in_features: shape of the input - alpha - trainable parameter that controls frequency - beta - trainable parameter that controls magnitude alpha is initialized to 1 by default, higher values = higher-frequency. beta is initialized to 1 by default, higher values = higher-magnitude. alpha will be trained along with the rest of your model. """ super(SnakeBeta, self).__init__() self.in_features = in_features # initialize alpha self.alpha_logscale = alpha_logscale if self.alpha_logscale: # log scale alphas initialized to zeros self.alpha = Parameter(torch.zeros(in_features) * alpha) self.beta = Parameter(torch.zeros(in_features) * alpha) else: # linear scale alphas initialized to ones self.alpha = Parameter(torch.ones(in_features) * alpha) self.beta = Parameter(torch.ones(in_features) * alpha) self.alpha.requires_grad = alpha_trainable self.beta.requires_grad = alpha_trainable self.no_div_by_zero = 0.000000001 def forward(self, x): """ Forward pass of the function. Applies the function to the input elementwise. SnakeBeta := x + 1/b * sin^2 (xa) """ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] beta = self.beta.unsqueeze(0).unsqueeze(-1) if self.alpha_logscale: alpha = torch.exp(alpha) beta = torch.exp(beta) x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) return x class ResidualUnit(nn.Module): def __init__(self, dim: int = 16, dilation: int = 1): super().__init__() pad = ((7 - 1) * dilation) // 2 self.block = nn.Sequential( Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)), WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)), WNConv1d(dim, dim, kernel_size=1), ) def forward(self, x): return x + self.block(x) class EncoderBlock(nn.Module): def __init__(self, dim: int = 16, stride: int = 1): super().__init__() self.block = nn.Sequential( ResidualUnit(dim // 2, dilation=1), ResidualUnit(dim // 2, dilation=3), ResidualUnit(dim // 2, dilation=9), Activation1d(activation=SnakeBeta(dim // 2, alpha_logscale=True)), WNConv1d( dim // 2, dim, kernel_size=2 * stride, stride=stride, padding=stride // 2 + stride % 2, ), ) def forward(self, x): return self.block(x) class FACodecEncoder(nn.Module): def __init__( self, ngf=32, up_ratios=(2, 4, 5, 5), out_channels=1024, ): super().__init__() self.hop_length = np.prod(up_ratios) self.up_ratios = up_ratios # Create first convolution d_model = ngf self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)] # Create EncoderBlocks that double channels as they downsample by `stride` for stride in up_ratios: d_model *= 2 self.block += [EncoderBlock(d_model, stride=stride)] # Create last convolution self.block += [ Activation1d(activation=SnakeBeta(d_model, alpha_logscale=True)), WNConv1d(d_model, out_channels, kernel_size=3, padding=1), ] # Wrap black into nn.Sequential self.block = nn.Sequential(*self.block) self.enc_dim = d_model self.reset_parameters() def forward(self, x): out = self.block(x) return out def inference(self, x): return self.block(x) def remove_weight_norm(self): """Remove weight normalization module from all of the layers.""" def _remove_weight_norm(m): try: torch.nn.utils.remove_weight_norm(m) except ValueError: # this module didn't have weight norm return self.apply(_remove_weight_norm) def apply_weight_norm(self): """Apply weight normalization module from all of the layers.""" def _apply_weight_norm(m): if isinstance(m, nn.Conv1d): torch.nn.utils.weight_norm(m) self.apply(_apply_weight_norm) def reset_parameters(self): self.apply(init_weights) class DecoderBlock(nn.Module): def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1): super().__init__() self.block = nn.Sequential( Activation1d(activation=SnakeBeta(input_dim, alpha_logscale=True)), WNConvTranspose1d( input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=stride // 2 + stride % 2, output_padding=stride % 2, ), ResidualUnit(output_dim, dilation=1), ResidualUnit(output_dim, dilation=3), ResidualUnit(output_dim, dilation=9), ) def forward(self, x): return self.block(x) class FACodecDecoder(nn.Module): def __init__( self, in_channels=256, upsample_initial_channel=1536, ngf=32, up_ratios=(5, 5, 4, 2), vq_num_q_c=2, vq_num_q_p=1, vq_num_q_r=3, vq_dim=1024, vq_commit_weight=0.005, vq_weight_init=False, vq_full_commit_loss=False, codebook_dim=8, codebook_size_prosody=10, # true codebook size is equal to 2^codebook_size codebook_size_content=10, codebook_size_residual=10, quantizer_dropout=0.0, dropout_type="linear", use_gr_content_f0=False, use_gr_prosody_phone=False, use_gr_residual_f0=False, use_gr_residual_phone=False, use_gr_x_timbre=False, use_random_mask_residual=True, prob_random_mask_residual=0.75, ): super().__init__() self.hop_length = np.prod(up_ratios) self.ngf = ngf self.up_ratios = up_ratios self.use_random_mask_residual = use_random_mask_residual self.prob_random_mask_residual = prob_random_mask_residual self.vq_num_q_p = vq_num_q_p self.vq_num_q_c = vq_num_q_c self.vq_num_q_r = vq_num_q_r self.codebook_size_prosody = codebook_size_prosody self.codebook_size_content = codebook_size_content self.codebook_size_residual = codebook_size_residual quantizer_class = ResidualVQ self.quantizer = nn.ModuleList() # prosody quantizer = quantizer_class( num_quantizers=vq_num_q_p, dim=vq_dim, codebook_size=codebook_size_prosody, codebook_dim=codebook_dim, threshold_ema_dead_code=2, commitment=vq_commit_weight, weight_init=vq_weight_init, full_commit_loss=vq_full_commit_loss, quantizer_dropout=quantizer_dropout, dropout_type=dropout_type, ) self.quantizer.append(quantizer) # phone quantizer = quantizer_class( num_quantizers=vq_num_q_c, dim=vq_dim, codebook_size=codebook_size_content, codebook_dim=codebook_dim, threshold_ema_dead_code=2, commitment=vq_commit_weight, weight_init=vq_weight_init, full_commit_loss=vq_full_commit_loss, quantizer_dropout=quantizer_dropout, dropout_type=dropout_type, ) self.quantizer.append(quantizer) # residual if self.vq_num_q_r > 0: quantizer = quantizer_class( num_quantizers=vq_num_q_r, dim=vq_dim, codebook_size=codebook_size_residual, codebook_dim=codebook_dim, threshold_ema_dead_code=2, commitment=vq_commit_weight, weight_init=vq_weight_init, full_commit_loss=vq_full_commit_loss, quantizer_dropout=quantizer_dropout, dropout_type=dropout_type, ) self.quantizer.append(quantizer) # Add first conv layer channels = upsample_initial_channel layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)] # Add upsampling + MRF blocks for i, stride in enumerate(up_ratios): input_dim = channels // 2**i output_dim = channels // 2 ** (i + 1) layers += [DecoderBlock(input_dim, output_dim, stride)] # Add final conv layer layers += [ Activation1d(activation=SnakeBeta(output_dim, alpha_logscale=True)), WNConv1d(output_dim, 1, kernel_size=7, padding=3), nn.Tanh(), ] self.model = nn.Sequential(*layers) self.timbre_encoder = TransformerEncoder( enc_emb_tokens=None, encoder_layer=4, encoder_hidden=256, encoder_head=4, conv_filter_size=1024, conv_kernel_size=5, encoder_dropout=0.1, use_cln=False, ) self.timbre_linear = nn.Linear(in_channels, in_channels * 2) self.timbre_linear.bias.data[:in_channels] = 1 self.timbre_linear.bias.data[in_channels:] = 0 self.timbre_norm = nn.LayerNorm(in_channels, elementwise_affine=False) self.f0_predictor = CNNLSTM(in_channels, 1, 2) self.phone_predictor = CNNLSTM(in_channels, 5003, 1) self.use_gr_content_f0 = use_gr_content_f0 self.use_gr_prosody_phone = use_gr_prosody_phone self.use_gr_residual_f0 = use_gr_residual_f0 self.use_gr_residual_phone = use_gr_residual_phone self.use_gr_x_timbre = use_gr_x_timbre if self.vq_num_q_r > 0 and self.use_gr_residual_f0: self.res_f0_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) ) if self.vq_num_q_r > 0 and self.use_gr_residual_phone > 0: self.res_phone_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) ) if self.use_gr_content_f0: self.content_f0_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) ) if self.use_gr_prosody_phone: self.prosody_phone_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) ) if self.use_gr_x_timbre: self.x_timbre_predictor = nn.Sequential( GradientReversal(alpha=1), CNNLSTM(in_channels, 245200, 1, global_pred=True), ) self.reset_parameters() def quantize(self, x, n_quantizers=None): outs, qs, commit_loss, quantized_buf = 0, [], [], [] # prosody f0_input = x # (B, d, T) f0_quantizer = self.quantizer[0] out, q, commit, quantized = f0_quantizer(f0_input, n_quantizers=n_quantizers) outs += out qs.append(q) quantized_buf.append(quantized.sum(0)) commit_loss.append(commit) # phone phone_input = x phone_quantizer = self.quantizer[1] out, q, commit, quantized = phone_quantizer( phone_input, n_quantizers=n_quantizers ) outs += out qs.append(q) quantized_buf.append(quantized.sum(0)) commit_loss.append(commit) # residual if self.vq_num_q_r > 0: residual_quantizer = self.quantizer[2] residual_input = x - (quantized_buf[0] + quantized_buf[1]).detach() out, q, commit, quantized = residual_quantizer( residual_input, n_quantizers=n_quantizers ) outs += out qs.append(q) quantized_buf.append(quantized.sum(0)) # [L, B, C, T] -> [B, C, T] commit_loss.append(commit) qs = torch.cat(qs, dim=0) commit_loss = torch.cat(commit_loss, dim=0) return outs, qs, commit_loss, quantized_buf def forward( self, x, vq=True, get_vq=False, eval_vq=True, speaker_embedding=None, n_quantizers=None, quantized=None, ): if get_vq: return self.quantizer.get_emb() if vq is True: if eval_vq: self.quantizer.eval() x_timbre = x outs, qs, commit_loss, quantized_buf = self.quantize( x, n_quantizers=n_quantizers ) x_timbre = x_timbre.transpose(1, 2) x_timbre = self.timbre_encoder(x_timbre, None, None) x_timbre = x_timbre.transpose(1, 2) spk_embs = torch.mean(x_timbre, dim=2) return outs, qs, commit_loss, quantized_buf, spk_embs out = {} layer_0 = quantized[0] f0, uv = self.f0_predictor(layer_0) f0 = rearrange(f0, "... 1 -> ...") uv = rearrange(uv, "... 1 -> ...") layer_1 = quantized[1] (phone,) = self.phone_predictor(layer_1) out = {"f0": f0, "uv": uv, "phone": phone} if self.use_gr_prosody_phone: (prosody_phone,) = self.prosody_phone_predictor(layer_0) out["prosody_phone"] = prosody_phone if self.use_gr_content_f0: content_f0, content_uv = self.content_f0_predictor(layer_1) content_f0 = rearrange(content_f0, "... 1 -> ...") content_uv = rearrange(content_uv, "... 1 -> ...") out["content_f0"] = content_f0 out["content_uv"] = content_uv if self.vq_num_q_r > 0: layer_2 = quantized[2] if self.use_gr_residual_f0: res_f0, res_uv = self.res_f0_predictor(layer_2) res_f0 = rearrange(res_f0, "... 1 -> ...") res_uv = rearrange(res_uv, "... 1 -> ...") out["res_f0"] = res_f0 out["res_uv"] = res_uv if self.use_gr_residual_phone: (res_phone,) = self.res_phone_predictor(layer_2) out["res_phone"] = res_phone style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) gamma, beta = style.chunk(2, 1) # (B, d, 1) if self.vq_num_q_r > 0: if self.use_random_mask_residual: bsz = quantized[2].shape[0] res_mask = np.random.choice( [0, 1], size=bsz, p=[ self.prob_random_mask_residual, 1 - self.prob_random_mask_residual, ], ) res_mask = ( torch.from_numpy(res_mask).unsqueeze(1).unsqueeze(1) ) # (B, 1, 1) res_mask = res_mask.to( device=quantized[2].device, dtype=quantized[2].dtype ) x = ( quantized[0].detach() + quantized[1].detach() + quantized[2] * res_mask ) # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] * res_mask else: x = quantized[0].detach() + quantized[1].detach() + quantized[2] # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] else: x = quantized[0].detach() + quantized[1].detach() # x = quantized_perturbe[0].detach() + quantized[1].detach() if self.use_gr_x_timbre: (x_timbre,) = self.x_timbre_predictor(x) out["x_timbre"] = x_timbre x = x.transpose(1, 2) x = self.timbre_norm(x) x = x.transpose(1, 2) x = x * gamma + beta x = self.model(x) out["audio"] = x return out def vq2emb(self, vq, use_residual_code=True): # vq: [num_quantizer, B, T] self.quantizer = self.quantizer.eval() out = 0 out += self.quantizer[0].vq2emb(vq[0 : self.vq_num_q_p]) out += self.quantizer[1].vq2emb( vq[self.vq_num_q_p : self.vq_num_q_p + self.vq_num_q_c] ) if self.vq_num_q_r > 0 and use_residual_code: out += self.quantizer[2].vq2emb(vq[self.vq_num_q_p + self.vq_num_q_c :]) return out def inference(self, x, speaker_embedding): style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) gamma, beta = style.chunk(2, 1) # (B, d, 1) x = x.transpose(1, 2) x = self.timbre_norm(x) x = x.transpose(1, 2) x = x * gamma + beta x = self.model(x) return x def remove_weight_norm(self): """Remove weight normalization module from all of the layers.""" def _remove_weight_norm(m): try: torch.nn.utils.remove_weight_norm(m) except ValueError: # this module didn't have weight norm return self.apply(_remove_weight_norm) def apply_weight_norm(self): """Apply weight normalization module from all of the layers.""" def _apply_weight_norm(m): if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d): torch.nn.utils.weight_norm(m) self.apply(_apply_weight_norm) def reset_parameters(self): self.apply(init_weights) class FACodecRedecoder(nn.Module): def __init__( self, in_channels=256, upsample_initial_channel=1280, up_ratios=(5, 5, 4, 2), vq_num_q_c=2, vq_num_q_p=1, vq_num_q_r=3, vq_dim=256, codebook_size_prosody=10, codebook_size_content=10, codebook_size_residual=10, ): super().__init__() self.hop_length = np.prod(up_ratios) self.up_ratios = up_ratios self.vq_num_q_p = vq_num_q_p self.vq_num_q_c = vq_num_q_c self.vq_num_q_r = vq_num_q_r self.vq_dim = vq_dim self.codebook_size_prosody = codebook_size_prosody self.codebook_size_content = codebook_size_content self.codebook_size_residual = codebook_size_residual self.prosody_embs = nn.ModuleList() for i in range(self.vq_num_q_p): emb_tokens = nn.Embedding( num_embeddings=2**self.codebook_size_prosody, embedding_dim=self.vq_dim, ) emb_tokens.weight.data.normal_(mean=0.0, std=1e-5) self.prosody_embs.append(emb_tokens) self.content_embs = nn.ModuleList() for i in range(self.vq_num_q_c): emb_tokens = nn.Embedding( num_embeddings=2**self.codebook_size_content, embedding_dim=self.vq_dim, ) emb_tokens.weight.data.normal_(mean=0.0, std=1e-5) self.content_embs.append(emb_tokens) self.residual_embs = nn.ModuleList() for i in range(self.vq_num_q_r): emb_tokens = nn.Embedding( num_embeddings=2**self.codebook_size_residual, embedding_dim=self.vq_dim, ) emb_tokens.weight.data.normal_(mean=0.0, std=1e-5) self.residual_embs.append(emb_tokens) # Add first conv layer channels = upsample_initial_channel layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)] # Add upsampling + MRF blocks for i, stride in enumerate(up_ratios): input_dim = channels // 2**i output_dim = channels // 2 ** (i + 1) layers += [DecoderBlock(input_dim, output_dim, stride)] # Add final conv layer layers += [ Activation1d(activation=SnakeBeta(output_dim, alpha_logscale=True)), WNConv1d(output_dim, 1, kernel_size=7, padding=3), nn.Tanh(), ] self.model = nn.Sequential(*layers) self.timbre_linear = nn.Linear(in_channels, in_channels * 2) self.timbre_linear.bias.data[:in_channels] = 1 self.timbre_linear.bias.data[in_channels:] = 0 self.timbre_norm = nn.LayerNorm(in_channels, elementwise_affine=False) self.timbre_cond_prosody_enc = TransformerEncoder( enc_emb_tokens=None, encoder_layer=4, encoder_hidden=256, encoder_head=4, conv_filter_size=1024, conv_kernel_size=5, encoder_dropout=0.1, use_cln=True, cfg=None, ) def forward( self, vq, speaker_embedding, use_residual_code=False, ): x = 0 x_p = 0 for i in range(self.vq_num_q_p): x_p = x_p + self.prosody_embs[i](vq[i]) # (B, T, d) spk_cond = speaker_embedding.unsqueeze(1).expand(-1, x_p.shape[1], -1) x_p = self.timbre_cond_prosody_enc( x_p, key_padding_mask=None, condition=spk_cond ) x = x + x_p x_c = 0 for i in range(self.vq_num_q_c): x_c = x_c + self.content_embs[i](vq[self.vq_num_q_p + i]) x = x + x_c if use_residual_code: x_r = 0 for i in range(self.vq_num_q_r): x_r = x_r + self.residual_embs[i]( vq[self.vq_num_q_p + self.vq_num_q_c + i] ) x = x + x_r style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) gamma, beta = style.chunk(2, 1) # (B, d, 1) x = x.transpose(1, 2) x = self.timbre_norm(x) x = x.transpose(1, 2) x = x * gamma + beta x = self.model(x) return x def vq2emb(self, vq, speaker_embedding, use_residual=True): out = 0 x_t = 0 for i in range(self.vq_num_q_p): x_t += self.prosody_embs[i](vq[i]) # (B, T, d) spk_cond = speaker_embedding.unsqueeze(1).expand(-1, x_t.shape[1], -1) x_t = self.timbre_cond_prosody_enc( x_t, key_padding_mask=None, condition=spk_cond ) # prosody out += x_t # content for i in range(self.vq_num_q_c): out += self.content_embs[i](vq[self.vq_num_q_p + i]) # residual if use_residual: for i in range(self.vq_num_q_r): out += self.residual_embs[i](vq[self.vq_num_q_p + self.vq_num_q_c + i]) out = out.transpose(1, 2) # (B, T, d) -> (B, d, T) return out def inference(self, x, speaker_embedding): style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) gamma, beta = style.chunk(2, 1) # (B, d, 1) x = x.transpose(1, 2) x = self.timbre_norm(x) x = x.transpose(1, 2) x = x * gamma + beta x = self.model(x) return x class FACodecEncoderV2(nn.Module): def __init__( self, ngf=32, up_ratios=(2, 4, 5, 5), out_channels=1024, ): super().__init__() self.hop_length = np.prod(up_ratios) self.up_ratios = up_ratios # Create first convolution d_model = ngf self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)] # Create EncoderBlocks that double channels as they downsample by `stride` for stride in up_ratios: d_model *= 2 self.block += [EncoderBlock(d_model, stride=stride)] # Create last convolution self.block += [ Activation1d(activation=SnakeBeta(d_model, alpha_logscale=True)), WNConv1d(d_model, out_channels, kernel_size=3, padding=1), ] # Wrap black into nn.Sequential self.block = nn.Sequential(*self.block) self.enc_dim = d_model self.mel_transform = MelSpectrogram( n_fft=1024, num_mels=80, sampling_rate=16000, hop_size=200, win_size=800, fmin=0, fmax=8000, ) self.reset_parameters() def forward(self, x): out = self.block(x) return out def inference(self, x): return self.block(x) def get_prosody_feature(self, x): return self.mel_transform(x.squeeze(1))[:, :20, :] def remove_weight_norm(self): """Remove weight normalization module from all of the layers.""" def _remove_weight_norm(m): try: torch.nn.utils.remove_weight_norm(m) except ValueError: # this module didn't have weight norm return self.apply(_remove_weight_norm) def apply_weight_norm(self): """Apply weight normalization module from all of the layers.""" def _apply_weight_norm(m): if isinstance(m, nn.Conv1d): torch.nn.utils.weight_norm(m) self.apply(_apply_weight_norm) def reset_parameters(self): self.apply(init_weights) class FACodecDecoderV2(nn.Module): def __init__( self, in_channels=256, upsample_initial_channel=1536, ngf=32, up_ratios=(5, 5, 4, 2), vq_num_q_c=2, vq_num_q_p=1, vq_num_q_r=3, vq_dim=1024, vq_commit_weight=0.005, vq_weight_init=False, vq_full_commit_loss=False, codebook_dim=8, codebook_size_prosody=10, # true codebook size is equal to 2^codebook_size codebook_size_content=10, codebook_size_residual=10, quantizer_dropout=0.0, dropout_type="linear", use_gr_content_f0=False, use_gr_prosody_phone=False, use_gr_residual_f0=False, use_gr_residual_phone=False, use_gr_x_timbre=False, use_random_mask_residual=True, prob_random_mask_residual=0.75, ): super().__init__() self.hop_length = np.prod(up_ratios) self.ngf = ngf self.up_ratios = up_ratios self.use_random_mask_residual = use_random_mask_residual self.prob_random_mask_residual = prob_random_mask_residual self.vq_num_q_p = vq_num_q_p self.vq_num_q_c = vq_num_q_c self.vq_num_q_r = vq_num_q_r self.codebook_size_prosody = codebook_size_prosody self.codebook_size_content = codebook_size_content self.codebook_size_residual = codebook_size_residual quantizer_class = ResidualVQ self.quantizer = nn.ModuleList() # prosody quantizer = quantizer_class( num_quantizers=vq_num_q_p, dim=vq_dim, codebook_size=codebook_size_prosody, codebook_dim=codebook_dim, threshold_ema_dead_code=2, commitment=vq_commit_weight, weight_init=vq_weight_init, full_commit_loss=vq_full_commit_loss, quantizer_dropout=quantizer_dropout, dropout_type=dropout_type, ) self.quantizer.append(quantizer) # phone quantizer = quantizer_class( num_quantizers=vq_num_q_c, dim=vq_dim, codebook_size=codebook_size_content, codebook_dim=codebook_dim, threshold_ema_dead_code=2, commitment=vq_commit_weight, weight_init=vq_weight_init, full_commit_loss=vq_full_commit_loss, quantizer_dropout=quantizer_dropout, dropout_type=dropout_type, ) self.quantizer.append(quantizer) # residual if self.vq_num_q_r > 0: quantizer = quantizer_class( num_quantizers=vq_num_q_r, dim=vq_dim, codebook_size=codebook_size_residual, codebook_dim=codebook_dim, threshold_ema_dead_code=2, commitment=vq_commit_weight, weight_init=vq_weight_init, full_commit_loss=vq_full_commit_loss, quantizer_dropout=quantizer_dropout, dropout_type=dropout_type, ) self.quantizer.append(quantizer) # Add first conv layer channels = upsample_initial_channel layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)] # Add upsampling + MRF blocks for i, stride in enumerate(up_ratios): input_dim = channels // 2**i output_dim = channels // 2 ** (i + 1) layers += [DecoderBlock(input_dim, output_dim, stride)] # Add final conv layer layers += [ Activation1d(activation=SnakeBeta(output_dim, alpha_logscale=True)), WNConv1d(output_dim, 1, kernel_size=7, padding=3), nn.Tanh(), ] self.model = nn.Sequential(*layers) self.timbre_encoder = TransformerEncoder( enc_emb_tokens=None, encoder_layer=4, encoder_hidden=256, encoder_head=4, conv_filter_size=1024, conv_kernel_size=5, encoder_dropout=0.1, use_cln=False, ) self.timbre_linear = nn.Linear(in_channels, in_channels * 2) self.timbre_linear.bias.data[:in_channels] = 1 self.timbre_linear.bias.data[in_channels:] = 0 self.timbre_norm = nn.LayerNorm(in_channels, elementwise_affine=False) self.f0_predictor = CNNLSTM(in_channels, 1, 2) self.phone_predictor = CNNLSTM(in_channels, 5003, 1) self.use_gr_content_f0 = use_gr_content_f0 self.use_gr_prosody_phone = use_gr_prosody_phone self.use_gr_residual_f0 = use_gr_residual_f0 self.use_gr_residual_phone = use_gr_residual_phone self.use_gr_x_timbre = use_gr_x_timbre if self.vq_num_q_r > 0 and self.use_gr_residual_f0: self.res_f0_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) ) if self.vq_num_q_r > 0 and self.use_gr_residual_phone > 0: self.res_phone_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) ) if self.use_gr_content_f0: self.content_f0_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) ) if self.use_gr_prosody_phone: self.prosody_phone_predictor = nn.Sequential( GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) ) if self.use_gr_x_timbre: self.x_timbre_predictor = nn.Sequential( GradientReversal(alpha=1), CNNLSTM(in_channels, 245200, 1, global_pred=True), ) self.melspec_linear = nn.Linear(20, 256) self.melspec_encoder = TransformerEncoder( enc_emb_tokens=None, encoder_layer=4, encoder_hidden=256, encoder_head=4, conv_filter_size=1024, conv_kernel_size=5, encoder_dropout=0.1, use_cln=False, cfg=None, ) self.reset_parameters() def quantize(self, x, prosody_feature, n_quantizers=None): outs, qs, commit_loss, quantized_buf = 0, [], [], [] # prosody f0_input = prosody_feature.transpose(1, 2) # (B, T, 20) f0_input = self.melspec_linear(f0_input) f0_input = self.melspec_encoder(f0_input, None, None) f0_input = f0_input.transpose(1, 2) f0_quantizer = self.quantizer[0] out, q, commit, quantized = f0_quantizer(f0_input, n_quantizers=n_quantizers) outs += out qs.append(q) quantized_buf.append(quantized.sum(0)) commit_loss.append(commit) # phone phone_input = x phone_quantizer = self.quantizer[1] out, q, commit, quantized = phone_quantizer( phone_input, n_quantizers=n_quantizers ) outs += out qs.append(q) quantized_buf.append(quantized.sum(0)) commit_loss.append(commit) # residual if self.vq_num_q_r > 0: residual_quantizer = self.quantizer[2] residual_input = x - (quantized_buf[0] + quantized_buf[1]).detach() out, q, commit, quantized = residual_quantizer( residual_input, n_quantizers=n_quantizers ) outs += out qs.append(q) quantized_buf.append(quantized.sum(0)) # [L, B, C, T] -> [B, C, T] commit_loss.append(commit) qs = torch.cat(qs, dim=0) commit_loss = torch.cat(commit_loss, dim=0) return outs, qs, commit_loss, quantized_buf def forward( self, x, prosody_feature, vq=True, get_vq=False, eval_vq=True, speaker_embedding=None, n_quantizers=None, quantized=None, ): if get_vq: return self.quantizer.get_emb() if vq is True: if eval_vq: self.quantizer.eval() x_timbre = x outs, qs, commit_loss, quantized_buf = self.quantize( x, prosody_feature, n_quantizers=n_quantizers ) x_timbre = x_timbre.transpose(1, 2) x_timbre = self.timbre_encoder(x_timbre, None, None) x_timbre = x_timbre.transpose(1, 2) spk_embs = torch.mean(x_timbre, dim=2) return outs, qs, commit_loss, quantized_buf, spk_embs out = {} layer_0 = quantized[0] f0, uv = self.f0_predictor(layer_0) f0 = rearrange(f0, "... 1 -> ...") uv = rearrange(uv, "... 1 -> ...") layer_1 = quantized[1] (phone,) = self.phone_predictor(layer_1) out = {"f0": f0, "uv": uv, "phone": phone} if self.use_gr_prosody_phone: (prosody_phone,) = self.prosody_phone_predictor(layer_0) out["prosody_phone"] = prosody_phone if self.use_gr_content_f0: content_f0, content_uv = self.content_f0_predictor(layer_1) content_f0 = rearrange(content_f0, "... 1 -> ...") content_uv = rearrange(content_uv, "... 1 -> ...") out["content_f0"] = content_f0 out["content_uv"] = content_uv if self.vq_num_q_r > 0: layer_2 = quantized[2] if self.use_gr_residual_f0: res_f0, res_uv = self.res_f0_predictor(layer_2) res_f0 = rearrange(res_f0, "... 1 -> ...") res_uv = rearrange(res_uv, "... 1 -> ...") out["res_f0"] = res_f0 out["res_uv"] = res_uv if self.use_gr_residual_phone: (res_phone,) = self.res_phone_predictor(layer_2) out["res_phone"] = res_phone style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) gamma, beta = style.chunk(2, 1) # (B, d, 1) if self.vq_num_q_r > 0: if self.use_random_mask_residual: bsz = quantized[2].shape[0] res_mask = np.random.choice( [0, 1], size=bsz, p=[ self.prob_random_mask_residual, 1 - self.prob_random_mask_residual, ], ) res_mask = ( torch.from_numpy(res_mask).unsqueeze(1).unsqueeze(1) ) # (B, 1, 1) res_mask = res_mask.to( device=quantized[2].device, dtype=quantized[2].dtype ) x = ( quantized[0].detach() + quantized[1].detach() + quantized[2] * res_mask ) # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] * res_mask else: x = quantized[0].detach() + quantized[1].detach() + quantized[2] # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] else: x = quantized[0].detach() + quantized[1].detach() # x = quantized_perturbe[0].detach() + quantized[1].detach() if self.use_gr_x_timbre: (x_timbre,) = self.x_timbre_predictor(x) out["x_timbre"] = x_timbre x = x.transpose(1, 2) x = self.timbre_norm(x) x = x.transpose(1, 2) x = x * gamma + beta x = self.model(x) out["audio"] = x return out def vq2emb(self, vq, use_residual=True): # vq: [num_quantizer, B, T] self.quantizer = self.quantizer.eval() out = 0 out += self.quantizer[0].vq2emb(vq[0 : self.vq_num_q_p]) out += self.quantizer[1].vq2emb( vq[self.vq_num_q_p : self.vq_num_q_p + self.vq_num_q_c] ) if self.vq_num_q_r > 0 and use_residual: out += self.quantizer[2].vq2emb(vq[self.vq_num_q_p + self.vq_num_q_c :]) return out def inference(self, x, speaker_embedding): style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) gamma, beta = style.chunk(2, 1) # (B, d, 1) x = x.transpose(1, 2) x = self.timbre_norm(x) x = x.transpose(1, 2) x = x * gamma + beta x = self.model(x) return x def remove_weight_norm(self): """Remove weight normalization module from all of the layers.""" def _remove_weight_norm(m): try: torch.nn.utils.remove_weight_norm(m) except ValueError: # this module didn't have weight norm return self.apply(_remove_weight_norm) def apply_weight_norm(self): """Apply weight normalization module from all of the layers.""" def _apply_weight_norm(m): if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d): torch.nn.utils.weight_norm(m) self.apply(_apply_weight_norm) def reset_parameters(self): self.apply(init_weights)