Hilley/ChatTTS-OpenVoice
61
1import math2import torch3from torch import nn4from torch.nn import functional as F5 6from torch.nn import Conv1d7from torch.nn.utils import weight_norm, remove_weight_norm8 9from . import commons10from .commons import init_weights, get_padding11from .transforms import piecewise_rational_quadratic_transform12from .attentions import Encoder13 14LRELU_SLOPE = 0.115 16 17class LayerNorm(nn.Module):18 def __init__(self, channels, eps=1e-5):19 super().__init__()20 self.channels = channels21 self.eps = eps22 23 self.gamma = nn.Parameter(torch.ones(channels))24 self.beta = nn.Parameter(torch.zeros(channels))25 26 def forward(self, x):27 x = x.transpose(1, -1)28 x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)29 return x.transpose(1, -1)30 31 32class ConvReluNorm(nn.Module):33 def __init__(34 self,35 in_channels,36 hidden_channels,37 out_channels,38 kernel_size,39 n_layers,40 p_dropout,41 ):42 super().__init__()43 self.in_channels = in_channels44 self.hidden_channels = hidden_channels45 self.out_channels = out_channels46 self.kernel_size = kernel_size47 self.n_layers = n_layers48 self.p_dropout = p_dropout49 assert n_layers > 1, "Number of layers should be larger than 0."50 51 self.conv_layers = nn.ModuleList()52 self.norm_layers = nn.ModuleList()53 self.conv_layers.append(54 nn.Conv1d(55 in_channels, hidden_channels, kernel_size, padding=kernel_size // 256 )57 )58 self.norm_layers.append(LayerNorm(hidden_channels))59 self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))60 for _ in range(n_layers - 1):61 self.conv_layers.append(62 nn.Conv1d(63 hidden_channels,64 hidden_channels,65 kernel_size,66 padding=kernel_size // 2,67 )68 )69 self.norm_layers.append(LayerNorm(hidden_channels))70 self.proj = nn.Conv1d(hidden_channels, out_channels, 1)71 self.proj.weight.data.zero_()72 self.proj.bias.data.zero_()73 74 def forward(self, x, x_mask):75 x_org = x76 for i in range(self.n_layers):77 x = self.conv_layers[i](x * x_mask)78 x = self.norm_layers[i](x)79 x = self.relu_drop(x)80 x = x_org + self.proj(x)81 return x * x_mask82 83 84class DDSConv(nn.Module):85 """86 Dilated and Depth-Separable Convolution87 """88 89 def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):90 super().__init__()91 self.channels = channels92 self.kernel_size = kernel_size93 self.n_layers = n_layers94 self.p_dropout = p_dropout95 96 self.drop = nn.Dropout(p_dropout)97 self.convs_sep = nn.ModuleList()98 self.convs_1x1 = nn.ModuleList()99 self.norms_1 = nn.ModuleList()100 self.norms_2 = nn.ModuleList()101 for i in range(n_layers):102 dilation = kernel_size**i103 padding = (kernel_size * dilation - dilation) // 2104 self.convs_sep.append(105 nn.Conv1d(106 channels,107 channels,108 kernel_size,109 groups=channels,110 dilation=dilation,111 padding=padding,112 )113 )114 self.convs_1x1.append(nn.Conv1d(channels, channels, 1))115 self.norms_1.append(LayerNorm(channels))116 self.norms_2.append(LayerNorm(channels))117 118 def forward(self, x, x_mask, g=None):119 if g is not None:120 x = x + g121 for i in range(self.n_layers):122 y = self.convs_sep[i](x * x_mask)123 y = self.norms_1[i](y)124 y = F.gelu(y)125 y = self.convs_1x1[i](y)126 y = self.norms_2[i](y)127 y = F.gelu(y)128 y = self.drop(y)129 x = x + y130 return x * x_mask131 132 133class WN(torch.nn.Module):134 def __init__(135 self,136 hidden_channels,137 kernel_size,138 dilation_rate,139 n_layers,140 gin_channels=0,141 p_dropout=0,142 ):143 super(WN, self).__init__()144 assert kernel_size % 2 == 1145 self.hidden_channels = hidden_channels146 self.kernel_size = (kernel_size,)147 self.dilation_rate = dilation_rate148 self.n_layers = n_layers149 self.gin_channels = gin_channels150 self.p_dropout = p_dropout151 152 self.in_layers = torch.nn.ModuleList()153 self.res_skip_layers = torch.nn.ModuleList()154 self.drop = nn.Dropout(p_dropout)155 156 if gin_channels != 0:157 cond_layer = torch.nn.Conv1d(158 gin_channels, 2 * hidden_channels * n_layers, 1159 )160 self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")161 162 for i in range(n_layers):163 dilation = dilation_rate**i164 padding = int((kernel_size * dilation - dilation) / 2)165 in_layer = torch.nn.Conv1d(166 hidden_channels,167 2 * hidden_channels,168 kernel_size,169 dilation=dilation,170 padding=padding,171 )172 in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")173 self.in_layers.append(in_layer)174 175 # last one is not necessary176 if i < n_layers - 1:177 res_skip_channels = 2 * hidden_channels178 else:179 res_skip_channels = hidden_channels180 181 res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)182 res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")183 self.res_skip_layers.append(res_skip_layer)184 185 def forward(self, x, x_mask, g=None, **kwargs):186 output = torch.zeros_like(x)187 n_channels_tensor = torch.IntTensor([self.hidden_channels])188 189 if g is not None:190 g = self.cond_layer(g)191 192 for i in range(self.n_layers):193 x_in = self.in_layers[i](x)194 if g is not None:195 cond_offset = i * 2 * self.hidden_channels196 g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]197 else:198 g_l = torch.zeros_like(x_in)199 200 acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)201 acts = self.drop(acts)202 203 res_skip_acts = self.res_skip_layers[i](acts)204 if i < self.n_layers - 1:205 res_acts = res_skip_acts[:, : self.hidden_channels, :]206 x = (x + res_acts) * x_mask207 output = output + res_skip_acts[:, self.hidden_channels :, :]208 else:209 output = output + res_skip_acts210 return output * x_mask211 212 def remove_weight_norm(self):213 if self.gin_channels != 0:214 torch.nn.utils.remove_weight_norm(self.cond_layer)215 for l in self.in_layers:216 torch.nn.utils.remove_weight_norm(l)217 for l in self.res_skip_layers:218 torch.nn.utils.remove_weight_norm(l)219 220 221class ResBlock1(torch.nn.Module):222 def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):223 super(ResBlock1, self).__init__()224 self.convs1 = nn.ModuleList(225 [226 weight_norm(227 Conv1d(228 channels,229 channels,230 kernel_size,231 1,232 dilation=dilation[0],233 padding=get_padding(kernel_size, dilation[0]),234 )235 ),236 weight_norm(237 Conv1d(238 channels,239 channels,240 kernel_size,241 1,242 dilation=dilation[1],243 padding=get_padding(kernel_size, dilation[1]),244 )245 ),246 weight_norm(247 Conv1d(248 channels,249 channels,250 kernel_size,251 1,252 dilation=dilation[2],253 padding=get_padding(kernel_size, dilation[2]),254 )255 ),256 ]257 )258 self.convs1.apply(init_weights)259 260 self.convs2 = nn.ModuleList(261 [262 weight_norm(263 Conv1d(264 channels,265 channels,266 kernel_size,267 1,268 dilation=1,269 padding=get_padding(kernel_size, 1),270 )271 ),272 weight_norm(273 Conv1d(274 channels,275 channels,276 kernel_size,277 1,278 dilation=1,279 padding=get_padding(kernel_size, 1),280 )281 ),282 weight_norm(283 Conv1d(284 channels,285 channels,286 kernel_size,287 1,288 dilation=1,289 padding=get_padding(kernel_size, 1),290 )291 ),292 ]293 )294 self.convs2.apply(init_weights)295 296 def forward(self, x, x_mask=None):297 for c1, c2 in zip(self.convs1, self.convs2):298 xt = F.leaky_relu(x, LRELU_SLOPE)299 if x_mask is not None:300 xt = xt * x_mask301 xt = c1(xt)302 xt = F.leaky_relu(xt, LRELU_SLOPE)303 if x_mask is not None:304 xt = xt * x_mask305 xt = c2(xt)306 x = xt + x307 if x_mask is not None:308 x = x * x_mask309 return x310 311 def remove_weight_norm(self):312 for l in self.convs1:313 remove_weight_norm(l)314 for l in self.convs2:315 remove_weight_norm(l)316 317 318class ResBlock2(torch.nn.Module):319 def __init__(self, channels, kernel_size=3, dilation=(1, 3)):320 super(ResBlock2, self).__init__()321 self.convs = nn.ModuleList(322 [323 weight_norm(324 Conv1d(325 channels,326 channels,327 kernel_size,328 1,329 dilation=dilation[0],330 padding=get_padding(kernel_size, dilation[0]),331 )332 ),333 weight_norm(334 Conv1d(335 channels,336 channels,337 kernel_size,338 1,339 dilation=dilation[1],340 padding=get_padding(kernel_size, dilation[1]),341 )342 ),343 ]344 )345 self.convs.apply(init_weights)346 347 def forward(self, x, x_mask=None):348 for c in self.convs:349 xt = F.leaky_relu(x, LRELU_SLOPE)350 if x_mask is not None:351 xt = xt * x_mask352 xt = c(xt)353 x = xt + x354 if x_mask is not None:355 x = x * x_mask356 return x357 358 def remove_weight_norm(self):359 for l in self.convs:360 remove_weight_norm(l)361 362 363class Log(nn.Module):364 def forward(self, x, x_mask, reverse=False, **kwargs):365 if not reverse:366 y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask367 logdet = torch.sum(-y, [1, 2])368 return y, logdet369 else:370 x = torch.exp(x) * x_mask371 return x372 373 374class Flip(nn.Module):375 def forward(self, x, *args, reverse=False, **kwargs):376 x = torch.flip(x, [1])377 if not reverse:378 logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)379 return x, logdet380 else:381 return x382 383 384class ElementwiseAffine(nn.Module):385 def __init__(self, channels):386 super().__init__()387 self.channels = channels388 self.m = nn.Parameter(torch.zeros(channels, 1))389 self.logs = nn.Parameter(torch.zeros(channels, 1))390 391 def forward(self, x, x_mask, reverse=False, **kwargs):392 if not reverse:393 y = self.m + torch.exp(self.logs) * x394 y = y * x_mask395 logdet = torch.sum(self.logs * x_mask, [1, 2])396 return y, logdet397 else:398 x = (x - self.m) * torch.exp(-self.logs) * x_mask399 return x400 401 402class ResidualCouplingLayer(nn.Module):403 def __init__(404 self,405 channels,406 hidden_channels,407 kernel_size,408 dilation_rate,409 n_layers,410 p_dropout=0,411 gin_channels=0,412 mean_only=False,413 ):414 assert channels % 2 == 0, "channels should be divisible by 2"415 super().__init__()416 self.channels = channels417 self.hidden_channels = hidden_channels418 self.kernel_size = kernel_size419 self.dilation_rate = dilation_rate420 self.n_layers = n_layers421 self.half_channels = channels // 2422 self.mean_only = mean_only423 424 self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)425 self.enc = WN(426 hidden_channels,427 kernel_size,428 dilation_rate,429 n_layers,430 p_dropout=p_dropout,431 gin_channels=gin_channels,432 )433 self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)434 self.post.weight.data.zero_()435 self.post.bias.data.zero_()436 437 def forward(self, x, x_mask, g=None, reverse=False):438 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)439 h = self.pre(x0) * x_mask440 h = self.enc(h, x_mask, g=g)441 stats = self.post(h) * x_mask442 if not self.mean_only:443 m, logs = torch.split(stats, [self.half_channels] * 2, 1)444 else:445 m = stats446 logs = torch.zeros_like(m)447 448 if not reverse:449 x1 = m + x1 * torch.exp(logs) * x_mask450 x = torch.cat([x0, x1], 1)451 logdet = torch.sum(logs, [1, 2])452 return x, logdet453 else:454 x1 = (x1 - m) * torch.exp(-logs) * x_mask455 x = torch.cat([x0, x1], 1)456 return x457 458 459class ConvFlow(nn.Module):460 def __init__(461 self,462 in_channels,463 filter_channels,464 kernel_size,465 n_layers,466 num_bins=10,467 tail_bound=5.0,468 ):469 super().__init__()470 self.in_channels = in_channels471 self.filter_channels = filter_channels472 self.kernel_size = kernel_size473 self.n_layers = n_layers474 self.num_bins = num_bins475 self.tail_bound = tail_bound476 self.half_channels = in_channels // 2477 478 self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)479 self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)480 self.proj = nn.Conv1d(481 filter_channels, self.half_channels * (num_bins * 3 - 1), 1482 )483 self.proj.weight.data.zero_()484 self.proj.bias.data.zero_()485 486 def forward(self, x, x_mask, g=None, reverse=False):487 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)488 h = self.pre(x0)489 h = self.convs(h, x_mask, g=g)490 h = self.proj(h) * x_mask491 492 b, c, t = x0.shape493 h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]494 495 unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)496 unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(497 self.filter_channels498 )499 unnormalized_derivatives = h[..., 2 * self.num_bins :]500 501 x1, logabsdet = piecewise_rational_quadratic_transform(502 x1,503 unnormalized_widths,504 unnormalized_heights,505 unnormalized_derivatives,506 inverse=reverse,507 tails="linear",508 tail_bound=self.tail_bound,509 )510 511 x = torch.cat([x0, x1], 1) * x_mask512 logdet = torch.sum(logabsdet * x_mask, [1, 2])513 if not reverse:514 return x, logdet515 else:516 return x517 518 519class TransformerCouplingLayer(nn.Module):520 def __init__(521 self,522 channels,523 hidden_channels,524 kernel_size,525 n_layers,526 n_heads,527 p_dropout=0,528 filter_channels=0,529 mean_only=False,530 wn_sharing_parameter=None,531 gin_channels=0,532 ):533 assert n_layers == 3, n_layers534 assert channels % 2 == 0, "channels should be divisible by 2"535 super().__init__()536 self.channels = channels537 self.hidden_channels = hidden_channels538 self.kernel_size = kernel_size539 self.n_layers = n_layers540 self.half_channels = channels // 2541 self.mean_only = mean_only542 543 self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)544 self.enc = (545 Encoder(546 hidden_channels,547 filter_channels,548 n_heads,549 n_layers,550 kernel_size,551 p_dropout,552 isflow=True,553 gin_channels=gin_channels,554 )555 if wn_sharing_parameter is None556 else wn_sharing_parameter557 )558 self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)559 self.post.weight.data.zero_()560 self.post.bias.data.zero_()561 562 def forward(self, x, x_mask, g=None, reverse=False):563 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)564 h = self.pre(x0) * x_mask565 h = self.enc(h, x_mask, g=g)566 stats = self.post(h) * x_mask567 if not self.mean_only:568 m, logs = torch.split(stats, [self.half_channels] * 2, 1)569 else:570 m = stats571 logs = torch.zeros_like(m)572 573 if not reverse:574 x1 = m + x1 * torch.exp(logs) * x_mask575 x = torch.cat([x0, x1], 1)576 logdet = torch.sum(logs, [1, 2])577 return x, logdet578 else:579 x1 = (x1 - m) * torch.exp(-logs) * x_mask580 x = torch.cat([x0, x1], 1)581 return x582 583 x1, logabsdet = piecewise_rational_quadratic_transform(584 x1,585 unnormalized_widths,586 unnormalized_heights,587 unnormalized_derivatives,588 inverse=reverse,589 tails="linear",590 tail_bound=self.tail_bound,591 )592 593 x = torch.cat([x0, x1], 1) * x_mask594 logdet = torch.sum(logabsdet * x_mask, [1, 2])595 if not reverse:596 return x, logdet597 else:598 return x599