Hilley/ChatTTS-OpenVoice
61
1import math2import torch3from torch import nn4from torch.nn import functional as F5 6from . import commons7from . import modules8from . import attentions9 10from torch.nn import Conv1d, ConvTranspose1d, Conv2d11from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm12 13from .commons import init_weights, get_padding14 15 16class TextEncoder(nn.Module):17 def __init__(self,18 n_vocab,19 out_channels,20 hidden_channels,21 filter_channels,22 n_heads,23 n_layers,24 kernel_size,25 p_dropout):26 super().__init__()27 self.n_vocab = n_vocab28 self.out_channels = out_channels29 self.hidden_channels = hidden_channels30 self.filter_channels = filter_channels31 self.n_heads = n_heads32 self.n_layers = n_layers33 self.kernel_size = kernel_size34 self.p_dropout = p_dropout35 36 self.emb = nn.Embedding(n_vocab, hidden_channels)37 nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)38 39 self.encoder = attentions.Encoder(40 hidden_channels,41 filter_channels,42 n_heads,43 n_layers,44 kernel_size,45 p_dropout)46 self.proj= nn.Conv1d(hidden_channels, out_channels * 2, 1)47 48 def forward(self, x, x_lengths):49 x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h]50 x = torch.transpose(x, 1, -1) # [b, h, t]51 x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)52 53 x = self.encoder(x * x_mask, x_mask)54 stats = self.proj(x) * x_mask55 56 m, logs = torch.split(stats, self.out_channels, dim=1)57 return x, m, logs, x_mask58 59 60class DurationPredictor(nn.Module):61 def __init__(62 self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=063 ):64 super().__init__()65 66 self.in_channels = in_channels67 self.filter_channels = filter_channels68 self.kernel_size = kernel_size69 self.p_dropout = p_dropout70 self.gin_channels = gin_channels71 72 self.drop = nn.Dropout(p_dropout)73 self.conv_1 = nn.Conv1d(74 in_channels, filter_channels, kernel_size, padding=kernel_size // 275 )76 self.norm_1 = modules.LayerNorm(filter_channels)77 self.conv_2 = nn.Conv1d(78 filter_channels, filter_channels, kernel_size, padding=kernel_size // 279 )80 self.norm_2 = modules.LayerNorm(filter_channels)81 self.proj = nn.Conv1d(filter_channels, 1, 1)82 83 if gin_channels != 0:84 self.cond = nn.Conv1d(gin_channels, in_channels, 1)85 86 def forward(self, x, x_mask, g=None):87 x = torch.detach(x)88 if g is not None:89 g = torch.detach(g)90 x = x + self.cond(g)91 x = self.conv_1(x * x_mask)92 x = torch.relu(x)93 x = self.norm_1(x)94 x = self.drop(x)95 x = self.conv_2(x * x_mask)96 x = torch.relu(x)97 x = self.norm_2(x)98 x = self.drop(x)99 x = self.proj(x * x_mask)100 return x * x_mask101 102class StochasticDurationPredictor(nn.Module):103 def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):104 super().__init__()105 filter_channels = in_channels # it needs to be removed from future version.106 self.in_channels = in_channels107 self.filter_channels = filter_channels108 self.kernel_size = kernel_size109 self.p_dropout = p_dropout110 self.n_flows = n_flows111 self.gin_channels = gin_channels112 113 self.log_flow = modules.Log()114 self.flows = nn.ModuleList()115 self.flows.append(modules.ElementwiseAffine(2))116 for i in range(n_flows):117 self.flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))118 self.flows.append(modules.Flip())119 120 self.post_pre = nn.Conv1d(1, filter_channels, 1)121 self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)122 self.post_convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)123 self.post_flows = nn.ModuleList()124 self.post_flows.append(modules.ElementwiseAffine(2))125 for i in range(4):126 self.post_flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))127 self.post_flows.append(modules.Flip())128 129 self.pre = nn.Conv1d(in_channels, filter_channels, 1)130 self.proj = nn.Conv1d(filter_channels, filter_channels, 1)131 self.convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)132 if gin_channels != 0:133 self.cond = nn.Conv1d(gin_channels, filter_channels, 1)134 135 def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):136 x = torch.detach(x)137 x = self.pre(x)138 if g is not None:139 g = torch.detach(g)140 x = x + self.cond(g)141 x = self.convs(x, x_mask)142 x = self.proj(x) * x_mask143 144 if not reverse:145 flows = self.flows146 assert w is not None147 148 logdet_tot_q = 0149 h_w = self.post_pre(w)150 h_w = self.post_convs(h_w, x_mask)151 h_w = self.post_proj(h_w) * x_mask152 e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask153 z_q = e_q154 for flow in self.post_flows:155 z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))156 logdet_tot_q += logdet_q157 z_u, z1 = torch.split(z_q, [1, 1], 1)158 u = torch.sigmoid(z_u) * x_mask159 z0 = (w - u) * x_mask160 logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1,2])161 logq = torch.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - logdet_tot_q162 163 logdet_tot = 0164 z0, logdet = self.log_flow(z0, x_mask)165 logdet_tot += logdet166 z = torch.cat([z0, z1], 1)167 for flow in flows:168 z, logdet = flow(z, x_mask, g=x, reverse=reverse)169 logdet_tot = logdet_tot + logdet170 nll = torch.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - logdet_tot171 return nll + logq # [b]172 else:173 flows = list(reversed(self.flows))174 flows = flows[:-2] + [flows[-1]] # remove a useless vflow175 z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale176 for flow in flows:177 z = flow(z, x_mask, g=x, reverse=reverse)178 z0, z1 = torch.split(z, [1, 1], 1)179 logw = z0180 return logw181 182class PosteriorEncoder(nn.Module):183 def __init__(184 self,185 in_channels,186 out_channels,187 hidden_channels,188 kernel_size,189 dilation_rate,190 n_layers,191 gin_channels=0,192 ):193 super().__init__()194 self.in_channels = in_channels195 self.out_channels = out_channels196 self.hidden_channels = hidden_channels197 self.kernel_size = kernel_size198 self.dilation_rate = dilation_rate199 self.n_layers = n_layers200 self.gin_channels = gin_channels201 202 self.pre = nn.Conv1d(in_channels, hidden_channels, 1)203 self.enc = modules.WN(204 hidden_channels,205 kernel_size,206 dilation_rate,207 n_layers,208 gin_channels=gin_channels,209 )210 self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)211 212 def forward(self, x, x_lengths, g=None, tau=1.0):213 x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(214 x.dtype215 )216 x = self.pre(x) * x_mask217 x = self.enc(x, x_mask, g=g)218 stats = self.proj(x) * x_mask219 m, logs = torch.split(stats, self.out_channels, dim=1)220 z = (m + torch.randn_like(m) * tau * torch.exp(logs)) * x_mask221 return z, m, logs, x_mask222 223 224class Generator(torch.nn.Module):225 def __init__(226 self,227 initial_channel,228 resblock,229 resblock_kernel_sizes,230 resblock_dilation_sizes,231 upsample_rates,232 upsample_initial_channel,233 upsample_kernel_sizes,234 gin_channels=0,235 ):236 super(Generator, self).__init__()237 self.num_kernels = len(resblock_kernel_sizes)238 self.num_upsamples = len(upsample_rates)239 self.conv_pre = Conv1d(240 initial_channel, upsample_initial_channel, 7, 1, padding=3241 )242 resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2243 244 self.ups = nn.ModuleList()245 for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):246 self.ups.append(247 weight_norm(248 ConvTranspose1d(249 upsample_initial_channel // (2**i),250 upsample_initial_channel // (2 ** (i + 1)),251 k,252 u,253 padding=(k - u) // 2,254 )255 )256 )257 258 self.resblocks = nn.ModuleList()259 for i in range(len(self.ups)):260 ch = upsample_initial_channel // (2 ** (i + 1))261 for j, (k, d) in enumerate(262 zip(resblock_kernel_sizes, resblock_dilation_sizes)263 ):264 self.resblocks.append(resblock(ch, k, d))265 266 self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)267 self.ups.apply(init_weights)268 269 if gin_channels != 0:270 self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)271 272 def forward(self, x, g=None):273 x = self.conv_pre(x)274 if g is not None:275 x = x + self.cond(g)276 277 for i in range(self.num_upsamples):278 x = F.leaky_relu(x, modules.LRELU_SLOPE)279 x = self.ups[i](x)280 xs = None281 for j in range(self.num_kernels):282 if xs is None:283 xs = self.resblocks[i * self.num_kernels + j](x)284 else:285 xs += self.resblocks[i * self.num_kernels + j](x)286 x = xs / self.num_kernels287 x = F.leaky_relu(x)288 x = self.conv_post(x)289 x = torch.tanh(x)290 291 return x292 293 def remove_weight_norm(self):294 print("Removing weight norm...")295 for layer in self.ups:296 remove_weight_norm(layer)297 for layer in self.resblocks:298 layer.remove_weight_norm()299 300 301class ReferenceEncoder(nn.Module):302 """303 inputs --- [N, Ty/r, n_mels*r] mels304 outputs --- [N, ref_enc_gru_size]305 """306 307 def __init__(self, spec_channels, gin_channels=0, layernorm=True):308 super().__init__()309 self.spec_channels = spec_channels310 ref_enc_filters = [32, 32, 64, 64, 128, 128]311 K = len(ref_enc_filters)312 filters = [1] + ref_enc_filters313 convs = [314 weight_norm(315 nn.Conv2d(316 in_channels=filters[i],317 out_channels=filters[i + 1],318 kernel_size=(3, 3),319 stride=(2, 2),320 padding=(1, 1),321 )322 )323 for i in range(K)324 ]325 self.convs = nn.ModuleList(convs)326 327 out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)328 self.gru = nn.GRU(329 input_size=ref_enc_filters[-1] * out_channels,330 hidden_size=256 // 2,331 batch_first=True,332 )333 self.proj = nn.Linear(128, gin_channels)334 if layernorm:335 self.layernorm = nn.LayerNorm(self.spec_channels)336 else:337 self.layernorm = None338 339 def forward(self, inputs, mask=None):340 N = inputs.size(0)341 342 out = inputs.view(N, 1, -1, self.spec_channels) # [N, 1, Ty, n_freqs]343 if self.layernorm is not None:344 out = self.layernorm(out)345 346 for conv in self.convs:347 out = conv(out)348 # out = wn(out)349 out = F.relu(out) # [N, 128, Ty//2^K, n_mels//2^K]350 351 out = out.transpose(1, 2) # [N, Ty//2^K, 128, n_mels//2^K]352 T = out.size(1)353 N = out.size(0)354 out = out.contiguous().view(N, T, -1) # [N, Ty//2^K, 128*n_mels//2^K]355 356 self.gru.flatten_parameters()357 memory, out = self.gru(out) # out --- [1, N, 128]358 359 return self.proj(out.squeeze(0))360 361 def calculate_channels(self, L, kernel_size, stride, pad, n_convs):362 for i in range(n_convs):363 L = (L - kernel_size + 2 * pad) // stride + 1364 return L365 366 367class ResidualCouplingBlock(nn.Module):368 def __init__(self,369 channels,370 hidden_channels,371 kernel_size,372 dilation_rate,373 n_layers,374 n_flows=4,375 gin_channels=0):376 super().__init__()377 self.channels = channels378 self.hidden_channels = hidden_channels379 self.kernel_size = kernel_size380 self.dilation_rate = dilation_rate381 self.n_layers = n_layers382 self.n_flows = n_flows383 self.gin_channels = gin_channels384 385 self.flows = nn.ModuleList()386 for i in range(n_flows):387 self.flows.append(modules.ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))388 self.flows.append(modules.Flip())389 390 def forward(self, x, x_mask, g=None, reverse=False):391 if not reverse:392 for flow in self.flows:393 x, _ = flow(x, x_mask, g=g, reverse=reverse)394 else:395 for flow in reversed(self.flows):396 x = flow(x, x_mask, g=g, reverse=reverse)397 return x398 399class SynthesizerTrn(nn.Module):400 """401 Synthesizer for Training402 """403 404 def __init__(405 self,406 n_vocab,407 spec_channels,408 inter_channels,409 hidden_channels,410 filter_channels,411 n_heads,412 n_layers,413 kernel_size,414 p_dropout,415 resblock,416 resblock_kernel_sizes,417 resblock_dilation_sizes,418 upsample_rates,419 upsample_initial_channel,420 upsample_kernel_sizes,421 n_speakers=256,422 gin_channels=256,423 **kwargs424 ):425 super().__init__()426 427 self.dec = Generator(428 inter_channels,429 resblock,430 resblock_kernel_sizes,431 resblock_dilation_sizes,432 upsample_rates,433 upsample_initial_channel,434 upsample_kernel_sizes,435 gin_channels=gin_channels,436 )437 self.enc_q = PosteriorEncoder(438 spec_channels,439 inter_channels,440 hidden_channels,441 5,442 1,443 16,444 gin_channels=gin_channels,445 )446 447 self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)448 449 self.n_speakers = n_speakers450 if n_speakers == 0:451 self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)452 else:453 self.enc_p = TextEncoder(n_vocab,454 inter_channels,455 hidden_channels,456 filter_channels,457 n_heads,458 n_layers,459 kernel_size,460 p_dropout)461 self.sdp = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels)462 self.dp = DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)463 self.emb_g = nn.Embedding(n_speakers, gin_channels)464 465 def infer(self, x, x_lengths, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., sdp_ratio=0.2, max_len=None):466 x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths)467 if self.n_speakers > 0:468 g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]469 else:470 g = None471 472 logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * sdp_ratio \473 + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)474 475 w = torch.exp(logw) * x_mask * length_scale476 w_ceil = torch.ceil(w)477 y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()478 y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(x_mask.dtype)479 attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)480 attn = commons.generate_path(w_ceil, attn_mask)481 482 m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']483 logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']484 485 z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale486 z = self.flow(z_p, y_mask, g=g, reverse=True)487 o = self.dec((z * y_mask)[:,:,:max_len], g=g)488 return o, attn, y_mask, (z, z_p, m_p, logs_p)489 490 def voice_conversion(self, y, y_lengths, sid_src, sid_tgt, tau=1.0):491 g_src = sid_src492 g_tgt = sid_tgt493 z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src, tau=tau)494 z_p = self.flow(z, y_mask, g=g_src)495 z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True)496 o_hat = self.dec(z_hat * y_mask, g=g_tgt)497 return o_hat, y_mask, (z, z_p, z_hat)498