ORI-Muchim/BlueArchiveTTS
58
1import copy2import math3import torch4from torch import nn5from torch.nn import functional as F6 7import commons8import modules9import attentions10import monotonic_align11 12from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d13from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm14from commons import init_weights, get_padding15 16from pqmf import PQMF17from stft import TorchSTFT, OnnxSTFT18 19AVAILABLE_FLOW_TYPES = ["pre_conv", "pre_conv2", "fft", "mono_layer_inter_residual", "mono_layer_post_residual"]20AVAILABLE_DURATION_DISCRIMINATOR_TYPES = {"dur_disc_1": "DurationDiscriminator", "dur_disc_2": "DurationDiscriminator2"}21 22 23class StochasticDurationPredictor(nn.Module):24 def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):25 super().__init__()26 filter_channels = in_channels # it needs to be removed from future version.27 self.in_channels = in_channels28 self.filter_channels = filter_channels29 self.kernel_size = kernel_size30 self.p_dropout = p_dropout31 self.n_flows = n_flows32 self.gin_channels = gin_channels33 34 self.log_flow = modules.Log()35 self.flows = nn.ModuleList()36 self.flows.append(modules.ElementwiseAffine(2))37 for i in range(n_flows):38 self.flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))39 self.flows.append(modules.Flip())40 41 self.post_pre = nn.Conv1d(1, filter_channels, 1)42 self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)43 self.post_convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)44 self.post_flows = nn.ModuleList()45 self.post_flows.append(modules.ElementwiseAffine(2))46 for i in range(4):47 self.post_flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))48 self.post_flows.append(modules.Flip())49 50 self.pre = nn.Conv1d(in_channels, filter_channels, 1)51 self.proj = nn.Conv1d(filter_channels, filter_channels, 1)52 self.convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)53 if gin_channels != 0:54 self.cond = nn.Conv1d(gin_channels, filter_channels, 1)55 56 def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):57 x = torch.detach(x)58 x = self.pre(x)59 if g is not None:60 g = torch.detach(g)61 x = x + self.cond(g)62 x = self.convs(x, x_mask)63 x = self.proj(x) * x_mask64 65 if not reverse:66 flows = self.flows67 assert w is not None68 69 logdet_tot_q = 070 h_w = self.post_pre(w)71 h_w = self.post_convs(h_w, x_mask)72 h_w = self.post_proj(h_w) * x_mask73 e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask74 z_q = e_q75 for flow in self.post_flows:76 z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))77 logdet_tot_q += logdet_q78 z_u, z1 = torch.split(z_q, [1, 1], 1)79 u = torch.sigmoid(z_u) * x_mask80 z0 = (w - u) * x_mask81 logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2])82 logq = torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q ** 2)) * x_mask, [1, 2]) - logdet_tot_q83 84 logdet_tot = 085 z0, logdet = self.log_flow(z0, x_mask)86 logdet_tot += logdet87 z = torch.cat([z0, z1], 1)88 for flow in flows:89 z, logdet = flow(z, x_mask, g=x, reverse=reverse)90 logdet_tot = logdet_tot + logdet91 nll = torch.sum(0.5 * (math.log(2 * math.pi) + (z ** 2)) * x_mask, [1, 2]) - logdet_tot92 return nll + logq # [b]93 else:94 flows = list(reversed(self.flows))95 flows = flows[:-2] + [flows[-1]] # remove a useless vflow96 z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale97 for flow in flows:98 z = flow(z, x_mask, g=x, reverse=reverse)99 z0, z1 = torch.split(z, [1, 1], 1)100 logw = z0101 return logw102 103 104class DurationPredictor(nn.Module):105 def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):106 super().__init__()107 108 self.in_channels = in_channels109 self.filter_channels = filter_channels110 self.kernel_size = kernel_size111 self.p_dropout = p_dropout112 self.gin_channels = gin_channels113 114 self.drop = nn.Dropout(p_dropout)115 self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)116 self.norm_1 = modules.LayerNorm(filter_channels)117 self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)118 self.norm_2 = modules.LayerNorm(filter_channels)119 self.proj = nn.Conv1d(filter_channels, 1, 1)120 121 if gin_channels != 0:122 self.cond = nn.Conv1d(gin_channels, in_channels, 1)123 124 def forward(self, x, x_mask, g=None):125 x = torch.detach(x)126 if g is not None:127 g = torch.detach(g)128 x = x + self.cond(g)129 x = self.conv_1(x * x_mask)130 x = torch.relu(x)131 x = self.norm_1(x)132 x = self.drop(x)133 x = self.conv_2(x * x_mask)134 x = torch.relu(x)135 x = self.norm_2(x)136 x = self.drop(x)137 x = self.proj(x * x_mask)138 return x * x_mask139 140 141class DurationDiscriminator(nn.Module): # vits2142 # TODO : not using "spk conditioning" for now according to the paper.143 # Can be a better discriminator if we use it.144 def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):145 super().__init__()146 147 self.in_channels = in_channels148 self.filter_channels = filter_channels149 self.kernel_size = kernel_size150 self.p_dropout = p_dropout151 self.gin_channels = gin_channels152 153 self.drop = nn.Dropout(p_dropout)154 self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)155 # self.norm_1 = modules.LayerNorm(filter_channels)156 self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)157 # self.norm_2 = modules.LayerNorm(filter_channels)158 self.dur_proj = nn.Conv1d(1, filter_channels, 1)159 160 self.pre_out_conv_1 = nn.Conv1d(2 * filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)161 self.pre_out_norm_1 = modules.LayerNorm(filter_channels)162 self.pre_out_conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)163 self.pre_out_norm_2 = modules.LayerNorm(filter_channels)164 165 # if gin_channels != 0:166 # self.cond = nn.Conv1d(gin_channels, in_channels, 1)167 168 self.output_layer = nn.Sequential(169 nn.Linear(filter_channels, 1),170 nn.Sigmoid()171 )172 173 def forward_probability(self, x, x_mask, dur, g=None):174 dur = self.dur_proj(dur)175 x = torch.cat([x, dur], dim=1)176 x = self.pre_out_conv_1(x * x_mask)177 # x = torch.relu(x)178 # x = self.pre_out_norm_1(x)179 # x = self.drop(x)180 x = self.pre_out_conv_2(x * x_mask)181 # x = torch.relu(x)182 # x = self.pre_out_norm_2(x)183 # x = self.drop(x)184 x = x * x_mask185 x = x.transpose(1, 2)186 output_prob = self.output_layer(x)187 return output_prob188 189 def forward(self, x, x_mask, dur_r, dur_hat, g=None):190 x = torch.detach(x)191 # if g is not None:192 # g = torch.detach(g)193 # x = x + self.cond(g)194 x = self.conv_1(x * x_mask)195 # x = torch.relu(x)196 # x = self.norm_1(x)197 # x = self.drop(x)198 x = self.conv_2(x * x_mask)199 # x = torch.relu(x)200 # x = self.norm_2(x)201 # x = self.drop(x)202 203 output_probs = []204 for dur in [dur_r, dur_hat]:205 output_prob = self.forward_probability(x, x_mask, dur, g)206 output_probs.append(output_prob)207 208 return output_probs209 210 211class DurationDiscriminator2(nn.Module): # vits2 - DurationDiscriminator2212 # TODO : not using "spk conditioning" for now according to the paper.213 # Can be a better discriminator if we use it.214 def __init__(215 self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0216 ):217 super().__init__()218 219 self.in_channels = in_channels220 self.filter_channels = filter_channels221 self.kernel_size = kernel_size222 self.p_dropout = p_dropout223 self.gin_channels = gin_channels224 225 self.conv_1 = nn.Conv1d(226 in_channels, filter_channels, kernel_size, padding=kernel_size // 2227 )228 self.norm_1 = modules.LayerNorm(filter_channels)229 self.conv_2 = nn.Conv1d(230 filter_channels, filter_channels, kernel_size, padding=kernel_size // 2231 )232 self.norm_2 = modules.LayerNorm(filter_channels)233 self.dur_proj = nn.Conv1d(1, filter_channels, 1)234 235 self.pre_out_conv_1 = nn.Conv1d(236 2 * filter_channels, filter_channels, kernel_size, padding=kernel_size // 2237 )238 self.pre_out_norm_1 = modules.LayerNorm(filter_channels)239 self.pre_out_conv_2 = nn.Conv1d(240 filter_channels, filter_channels, kernel_size, padding=kernel_size // 2241 )242 self.pre_out_norm_2 = modules.LayerNorm(filter_channels)243 244 # if gin_channels != 0:245 # self.cond = nn.Conv1d(gin_channels, in_channels, 1)246 247 self.output_layer = nn.Sequential(nn.Linear(filter_channels, 1), nn.Sigmoid())248 249 def forward_probability(self, x, x_mask, dur, g=None):250 dur = self.dur_proj(dur)251 x = torch.cat([x, dur], dim=1)252 x = self.pre_out_conv_1(x * x_mask)253 x = torch.relu(x)254 x = self.pre_out_norm_1(x)255 x = self.pre_out_conv_2(x * x_mask)256 x = torch.relu(x)257 x = self.pre_out_norm_2(x)258 x = x * x_mask259 x = x.transpose(1, 2)260 output_prob = self.output_layer(x)261 return output_prob262 263 def forward(self, x, x_mask, dur_r, dur_hat, g=None):264 x = torch.detach(x)265 # if g is not None:266 # g = torch.detach(g)267 # x = x + self.cond(g)268 x = self.conv_1(x * x_mask)269 x = torch.relu(x)270 x = self.norm_1(x)271 x = self.conv_2(x * x_mask)272 x = torch.relu(x)273 x = self.norm_2(x)274 275 output_probs = []276 for dur in [dur_r, dur_hat]:277 output_prob = self.forward_probability(x, x_mask, dur, g)278 output_probs.append([output_prob])279 280 return output_probs281 282 283class TextEncoder(nn.Module):284 def __init__(self,285 n_vocab,286 out_channels,287 hidden_channels,288 filter_channels,289 n_heads,290 n_layers,291 kernel_size,292 p_dropout,293 gin_channels=0):294 super().__init__()295 self.n_vocab = n_vocab296 self.out_channels = out_channels297 self.hidden_channels = hidden_channels298 self.filter_channels = filter_channels299 self.n_heads = n_heads300 self.n_layers = n_layers301 self.kernel_size = kernel_size302 self.p_dropout = p_dropout303 self.gin_channels = gin_channels304 self.emb = nn.Embedding(n_vocab, hidden_channels)305 nn.init.normal_(self.emb.weight, 0.0, hidden_channels ** -0.5)306 307 self.encoder = attentions.Encoder(308 hidden_channels,309 filter_channels,310 n_heads,311 n_layers,312 kernel_size,313 p_dropout,314 gin_channels=self.gin_channels)315 self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)316 317 def forward(self, x, x_lengths, g=None):318 x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h]319 x = torch.transpose(x, 1, -1) # [b, h, t]320 x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)321 322 x = self.encoder(x * x_mask, x_mask, g=g)323 stats = self.proj(x) * x_mask324 325 m, logs = torch.split(stats, self.out_channels, dim=1)326 return x, m, logs, x_mask327 328 329class ResidualCouplingTransformersLayer2(nn.Module): # vits2330 def __init__(331 self,332 channels,333 hidden_channels,334 kernel_size,335 dilation_rate,336 n_layers,337 p_dropout=0,338 gin_channels=0,339 mean_only=False,340 ):341 assert channels % 2 == 0, "channels should be divisible by 2"342 super().__init__()343 self.channels = channels344 self.hidden_channels = hidden_channels345 self.kernel_size = kernel_size346 self.dilation_rate = dilation_rate347 self.n_layers = n_layers348 self.half_channels = channels // 2349 self.mean_only = mean_only350 351 self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)352 self.pre_transformer = attentions.Encoder(353 hidden_channels,354 hidden_channels,355 n_heads=2,356 n_layers=1,357 kernel_size=kernel_size,358 p_dropout=p_dropout,359 # window_size=None,360 )361 self.enc = modules.WN(362 hidden_channels,363 kernel_size,364 dilation_rate,365 n_layers,366 p_dropout=p_dropout,367 gin_channels=gin_channels,368 )369 370 self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)371 self.post.weight.data.zero_()372 self.post.bias.data.zero_()373 374 def forward(self, x, x_mask, g=None, reverse=False):375 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)376 h = self.pre(x0) * x_mask377 h = h + self.pre_transformer(h * x_mask, x_mask) # vits2 residual connection378 h = self.enc(h, x_mask, g=g)379 stats = self.post(h) * x_mask380 if not self.mean_only:381 m, logs = torch.split(stats, [self.half_channels] * 2, 1)382 else:383 m = stats384 logs = torch.zeros_like(m)385 if not reverse:386 x1 = m + x1 * torch.exp(logs) * x_mask387 x = torch.cat([x0, x1], 1)388 logdet = torch.sum(logs, [1, 2])389 return x, logdet390 else:391 x1 = (x1 - m) * torch.exp(-logs) * x_mask392 x = torch.cat([x0, x1], 1)393 return x394 395 396class ResidualCouplingTransformersLayer(nn.Module): # vits2397 def __init__(398 self,399 channels,400 hidden_channels,401 kernel_size,402 dilation_rate,403 n_layers,404 p_dropout=0,405 gin_channels=0,406 mean_only=False,407 ):408 assert channels % 2 == 0, "channels should be divisible by 2"409 super().__init__()410 self.channels = channels411 self.hidden_channels = hidden_channels412 self.kernel_size = kernel_size413 self.dilation_rate = dilation_rate414 self.n_layers = n_layers415 self.half_channels = channels // 2416 self.mean_only = mean_only417 # vits2418 self.pre_transformer = attentions.Encoder(419 self.half_channels,420 self.half_channels,421 n_heads=2,422 n_layers=2,423 kernel_size=3,424 p_dropout=0.1,425 window_size=None426 )427 428 self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)429 self.enc = modules.WN(430 hidden_channels,431 kernel_size,432 dilation_rate,433 n_layers,434 p_dropout=p_dropout,435 gin_channels=gin_channels,436 )437 # vits2438 self.post_transformer = attentions.Encoder(439 self.hidden_channels,440 self.hidden_channels,441 n_heads=2,442 n_layers=2,443 kernel_size=3,444 p_dropout=0.1,445 window_size=None446 )447 448 self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)449 self.post.weight.data.zero_()450 self.post.bias.data.zero_()451 452 def forward(self, x, x_mask, g=None, reverse=False):453 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)454 x0_ = self.pre_transformer(x0 * x_mask, x_mask) # vits2455 x0_ = x0_ + x0 # vits2 residual connection456 h = self.pre(x0_) * x_mask # changed from x0 to x0_ to retain x0 for the flow457 h = self.enc(h, x_mask, g=g)458 459 # vits2 - (experimental;uncomment the following 2 line to use)460 # h_ = self.post_transformer(h, x_mask)461 # h = h + h_ #vits2 residual connection462 463 stats = self.post(h) * x_mask464 if not self.mean_only:465 m, logs = torch.split(stats, [self.half_channels] * 2, 1)466 else:467 m = stats468 logs = torch.zeros_like(m)469 if not reverse:470 x1 = m + x1 * torch.exp(logs) * x_mask471 x = torch.cat([x0, x1], 1)472 logdet = torch.sum(logs, [1, 2])473 return x, logdet474 else:475 x1 = (x1 - m) * torch.exp(-logs) * x_mask476 x = torch.cat([x0, x1], 1)477 return x478 479 def remove_weight_norm(self): # !480 self.enc.remove_weight_norm()481 482 483class FFTransformerCouplingLayer(nn.Module): # vits2484 def __init__(self,485 channels,486 hidden_channels,487 kernel_size,488 n_layers,489 n_heads,490 p_dropout=0,491 filter_channels=768,492 mean_only=False,493 gin_channels=0494 ):495 assert channels % 2 == 0, "channels should be divisible by 2"496 super().__init__()497 self.channels = channels498 self.hidden_channels = hidden_channels499 self.kernel_size = kernel_size500 self.n_layers = n_layers501 self.half_channels = channels // 2502 self.mean_only = mean_only503 504 self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)505 self.enc = attentions.FFT(506 hidden_channels,507 filter_channels,508 n_heads,509 n_layers,510 kernel_size,511 p_dropout,512 isflow=True,513 gin_channels=gin_channels514 )515 self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)516 self.post.weight.data.zero_()517 self.post.bias.data.zero_()518 519 def forward(self, x, x_mask, g=None, reverse=False):520 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)521 h = self.pre(x0) * x_mask522 h_ = self.enc(h, x_mask, g=g)523 h = h_ + h524 stats = self.post(h) * x_mask525 if not self.mean_only:526 m, logs = torch.split(stats, [self.half_channels] * 2, 1)527 else:528 m = stats529 logs = torch.zeros_like(m)530 531 if not reverse:532 x1 = m + x1 * torch.exp(logs) * x_mask533 x = torch.cat([x0, x1], 1)534 logdet = torch.sum(logs, [1, 2])535 return x, logdet536 else:537 x1 = (x1 - m) * torch.exp(-logs) * x_mask538 x = torch.cat([x0, x1], 1)539 return x540 541 542class MonoTransformerFlowLayer(nn.Module): # vits2543 def __init__(544 self,545 channels,546 hidden_channels,547 mean_only=False,548 residual_connection=False,549 # according to VITS-2 paper fig 1B set residual_connection=True550 ):551 assert channels % 2 == 0, "channels should be divisible by 2"552 super().__init__()553 self.channels = channels554 self.hidden_channels = hidden_channels555 self.half_channels = channels // 2556 self.mean_only = mean_only557 self.residual_connection = residual_connection558 # vits2559 self.pre_transformer = attentions.Encoder(560 self.half_channels,561 self.half_channels,562 n_heads=2,563 n_layers=2,564 kernel_size=3,565 p_dropout=0.1,566 window_size=None567 )568 569 self.post = nn.Conv1d(self.half_channels, self.half_channels * (2 - mean_only), 1)570 self.post.weight.data.zero_()571 self.post.bias.data.zero_()572 573 def forward(self, x, x_mask, g=None, reverse=False):574 if self.residual_connection:575 if not reverse:576 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)577 x0_ = x0 * x_mask578 x0_ = self.pre_transformer(x0, x_mask) # vits2579 stats = self.post(x0_) * x_mask580 if not self.mean_only:581 m, logs = torch.split(stats, [self.half_channels] * 2, 1)582 else:583 m = stats584 logs = torch.zeros_like(m)585 x1 = m + x1 * torch.exp(logs) * x_mask586 x_ = torch.cat([x0, x1], 1)587 x = x + x_588 logdet = torch.sum(torch.log(torch.exp(logs) + 1), [1, 2])589 logdet = logdet + torch.log(torch.tensor(2)) * (x0.shape[1] * x0.shape[2])590 return x, logdet591 592 else:593 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)594 x0 = x0 / 2595 x0_ = x0 * x_mask596 x0_ = self.pre_transformer(x0, x_mask) # vits2597 stats = self.post(x0_) * x_mask598 if not self.mean_only:599 m, logs = torch.split(stats, [self.half_channels] * 2, 1)600 else:601 m = stats602 logs = torch.zeros_like(m)603 x1_ = ((x1 - m) / (1 + torch.exp(-logs))) * x_mask604 x = torch.cat([x0, x1_], 1)605 return x606 else:607 x0, x1 = torch.split(x, [self.half_channels] * 2, 1)608 x0_ = self.pre_transformer(x0 * x_mask, x_mask) # vits2609 h = x0_ + x0 # vits2610 stats = self.post(h) * x_mask611 if not self.mean_only:612 m, logs = torch.split(stats, [self.half_channels] * 2, 1)613 else:614 m = stats615 logs = torch.zeros_like(m)616 if not reverse:617 x1 = m + x1 * torch.exp(logs) * x_mask618 x = torch.cat([x0, x1], 1)619 logdet = torch.sum(logs, [1, 2])620 return x, logdet621 else:622 x1 = (x1 - m) * torch.exp(-logs) * x_mask623 x = torch.cat([x0, x1], 1)624 return x625 626 627class ResidualCouplingTransformersBlock(nn.Module): # vits2628 def __init__(self,629 channels,630 hidden_channels,631 kernel_size,632 dilation_rate,633 n_layers,634 n_flows=4,635 gin_channels=0,636 use_transformer_flows=False,637 transformer_flow_type="pre_conv",638 ):639 super().__init__()640 self.channels = channels641 self.hidden_channels = hidden_channels642 self.kernel_size = kernel_size643 self.dilation_rate = dilation_rate644 self.n_layers = n_layers645 self.n_flows = n_flows646 self.gin_channels = gin_channels647 648 self.flows = nn.ModuleList()649 # TODO : clean up this mess650 if use_transformer_flows:651 if transformer_flow_type == "pre_conv":652 for i in range(n_flows):653 self.flows.append(654 ResidualCouplingTransformersLayer(655 channels,656 hidden_channels,657 kernel_size,658 dilation_rate,659 n_layers,660 gin_channels=gin_channels,661 mean_only=True662 )663 )664 self.flows.append(modules.Flip())665 elif transformer_flow_type == "pre_conv2":666 for i in range(n_flows):667 self.flows.append(668 ResidualCouplingTransformersLayer2(669 channels,670 hidden_channels,671 kernel_size,672 dilation_rate,673 n_layers,674 gin_channels=gin_channels,675 mean_only=True,676 )677 )678 self.flows.append(modules.Flip())679 elif transformer_flow_type == "fft":680 for i in range(n_flows):681 self.flows.append(682 FFTransformerCouplingLayer(683 channels,684 hidden_channels,685 kernel_size,686 dilation_rate,687 n_layers,688 gin_channels=gin_channels,689 mean_only=True690 )691 )692 self.flows.append(modules.Flip())693 elif transformer_flow_type == "mono_layer_inter_residual":694 for i in range(n_flows):695 self.flows.append(696 modules.ResidualCouplingLayer(697 channels,698 hidden_channels,699 kernel_size,700 dilation_rate,701 n_layers,702 gin_channels=gin_channels,703 mean_only=True704 )705 )706 self.flows.append(modules.Flip())707 self.flows.append(708 MonoTransformerFlowLayer(709 channels, hidden_channels, mean_only=True710 )711 )712 elif transformer_flow_type == "mono_layer_post_residual":713 for i in range(n_flows):714 self.flows.append(715 modules.ResidualCouplingLayer(716 channels,717 hidden_channels,718 kernel_size,719 dilation_rate,720 n_layers,721 gin_channels=gin_channels,722 mean_only=True,723 )724 )725 self.flows.append(modules.Flip())726 self.flows.append(727 MonoTransformerFlowLayer(728 channels, hidden_channels, mean_only=True,729 residual_connection=True730 )731 )732 else:733 for i in range(n_flows):734 self.flows.append(735 modules.ResidualCouplingLayer(736 channels,737 hidden_channels,738 kernel_size,739 dilation_rate,740 n_layers,741 gin_channels=gin_channels,742 mean_only=True743 )744 )745 self.flows.append(modules.Flip())746 747 def forward(self, x, x_mask, g=None, reverse=False):748 if not reverse:749 for flow in self.flows:750 x, _ = flow(x, x_mask, g=g, reverse=reverse)751 else:752 for flow in reversed(self.flows):753 x = flow(x, x_mask, g=g, reverse=reverse)754 return x755 756 def remove_weight_norm(self): # !757 for i, l in enumerate(self.flows):758 if i % 2 == 0:759 l.remove_weight_norm()760 761 762class ResidualCouplingBlock(nn.Module):763 def __init__(self,764 channels,765 hidden_channels,766 kernel_size,767 dilation_rate,768 n_layers,769 n_flows=4,770 gin_channels=0):771 super().__init__()772 self.channels = channels773 self.hidden_channels = hidden_channels774 self.kernel_size = kernel_size775 self.dilation_rate = dilation_rate776 self.n_layers = n_layers777 self.n_flows = n_flows778 self.gin_channels = gin_channels779 780 self.flows = nn.ModuleList()781 for i in range(n_flows):782 self.flows.append(783 modules.ResidualCouplingLayer(784 channels,785 hidden_channels,786 kernel_size,787 dilation_rate,788 n_layers,789 gin_channels=gin_channels,790 mean_only=True791 )792 )793 self.flows.append(modules.Flip())794 795 def forward(self, x, x_mask, g=None, reverse=False):796 if not reverse:797 for flow in self.flows:798 x, _ = flow(x, x_mask, g=g, reverse=reverse)799 else:800 for flow in reversed(self.flows):801 x = flow(x, x_mask, g=g, reverse=reverse)802 return x803 804 def remove_weight_norm(self): # !805 for i, l in enumerate(self.flows):806 if i % 2 == 0:807 l.remove_weight_norm()808 809 810class PosteriorEncoder(nn.Module):811 def __init__(self,812 in_channels,813 out_channels,814 hidden_channels,815 kernel_size,816 dilation_rate,817 n_layers,818 gin_channels=0):819 super().__init__()820 self.in_channels = in_channels821 self.out_channels = out_channels822 self.hidden_channels = hidden_channels823 self.kernel_size = kernel_size824 self.dilation_rate = dilation_rate825 self.n_layers = n_layers826 self.gin_channels = gin_channels827 828 self.pre = nn.Conv1d(in_channels, hidden_channels, 1)829 self.enc = modules.WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)830 self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)831 832 def forward(self, x, x_lengths, g=None):833 x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)834 x = self.pre(x) * x_mask835 x = self.enc(x, x_mask, g=g)836 stats = self.proj(x) * x_mask837 m, logs = torch.split(stats, self.out_channels, dim=1)838 z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask839 return z, m, logs, x_mask840 841 842class Generator(torch.nn.Module):843 def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates,844 upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):845 super(Generator, self).__init__()846 self.num_kernels = len(resblock_kernel_sizes)847 self.num_upsamples = len(upsample_rates)848 self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)849 resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2850 851 self.ups = nn.ModuleList()852 for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):853 self.ups.append(weight_norm(854 ConvTranspose1d(upsample_initial_channel // (2 ** i), upsample_initial_channel // (2 ** (i + 1)),855 k, u, padding=(k - u) // 2)))856 857 self.resblocks = nn.ModuleList()858 for i in range(len(self.ups)):859 ch = upsample_initial_channel // (2 ** (i + 1))860 for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):861 self.resblocks.append(resblock(ch, k, d))862 863 self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)864 self.ups.apply(init_weights)865 866 if gin_channels != 0:867 self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)868 869 def forward(self, x, g=None):870 x = self.conv_pre(x)871 if g is not None:872 x = x + self.cond(g)873 874 for i in range(self.num_upsamples):875 x = F.leaky_relu(x, modules.LRELU_SLOPE)876 x = self.ups[i](x)877 xs = None878 for j in range(self.num_kernels):879 if xs is None:880 xs = self.resblocks[i * self.num_kernels + j](x)881 else:882 xs += self.resblocks[i * self.num_kernels + j](x)883 x = xs / self.num_kernels884 x = F.leaky_relu(x)885 x = self.conv_post(x)886 x = torch.tanh(x)887 888 return x889 890 def remove_weight_norm(self):891 print('Removing weight norm...')892 for l in self.ups:893 remove_weight_norm(l)894 for l in self.resblocks:895 l.remove_weight_norm()896 897 898class iSTFT_Generator(torch.nn.Module):899 def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates,900 upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size,901 gin_channels=0, is_onnx=False):902 super(iSTFT_Generator, self).__init__()903 # self.h = h904 self.gen_istft_n_fft = gen_istft_n_fft905 self.gen_istft_hop_size = gen_istft_hop_size906 907 self.num_kernels = len(resblock_kernel_sizes)908 self.num_upsamples = len(upsample_rates)909 self.conv_pre = weight_norm(Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))910 resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2911 912 self.ups = nn.ModuleList()913 for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):914 self.ups.append(weight_norm(915 ConvTranspose1d(upsample_initial_channel // (2 ** i), upsample_initial_channel // (2 ** (i + 1)),916 k, u, padding=(k - u) // 2)))917 918 self.resblocks = nn.ModuleList()919 for i in range(len(self.ups)):920 ch = upsample_initial_channel // (2 ** (i + 1))921 for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):922 self.resblocks.append(resblock(ch, k, d))923 924 self.post_n_fft = self.gen_istft_n_fft925 self.conv_post = weight_norm(Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))926 self.ups.apply(init_weights)927 self.conv_post.apply(init_weights)928 self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))929 '''930 self.stft = TorchSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size,931 win_length=self.gen_istft_n_fft)932 '''933 # - for onnx934 if is_onnx == True:935 self.stft = OnnxSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size,936 win_length=self.gen_istft_n_fft)937 else:938 self.stft = TorchSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size,939 win_length=self.gen_istft_n_fft)940 941 def forward(self, x, g=None):942 x = self.conv_pre(x)943 for i in range(self.num_upsamples):944 x = F.leaky_relu(x, modules.LRELU_SLOPE)945 x = self.ups[i](x)946 xs = None947 for j in range(self.num_kernels):948 if xs is None:949 xs = self.resblocks[i * self.num_kernels + j](x)950 else:951 xs += self.resblocks[i * self.num_kernels + j](x)952 x = xs / self.num_kernels953 x = F.leaky_relu(x)954 x = self.reflection_pad(x)955 x = self.conv_post(x)956 spec = torch.exp(x[:, :self.post_n_fft // 2 + 1, :])957 phase = math.pi * torch.sin(x[:, self.post_n_fft // 2 + 1:, :])958 out = self.stft.inverse(spec, phase).to(x.device)959 return out, None960 961 def remove_weight_norm(self):962 print('Removing weight norm...')963 for l in self.ups:964 remove_weight_norm(l)965 for l in self.resblocks:966 l.remove_weight_norm()967 remove_weight_norm(self.conv_pre)968 remove_weight_norm(self.conv_post)969 970 971class Multiband_iSTFT_Generator(torch.nn.Module): # !972 def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates,973 upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, subbands,974 gin_channels=0, is_onnx=False):975 super(Multiband_iSTFT_Generator, self).__init__()976 # self.h = h977 self.subbands = subbands978 self.num_kernels = len(resblock_kernel_sizes)979 self.num_upsamples = len(upsample_rates)980 self.conv_pre = weight_norm(Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))981 resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2982 983 self.ups = nn.ModuleList()984 for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):985 self.ups.append(weight_norm(986 ConvTranspose1d(upsample_initial_channel // (2 ** i), upsample_initial_channel // (2 ** (i + 1)),987 k, u, padding=(k - u) // 2)))988 989 self.resblocks = nn.ModuleList()990 for i in range(len(self.ups)):991 ch = upsample_initial_channel // (2 ** (i + 1))992 for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):993 self.resblocks.append(resblock(ch, k, d))994 995 self.post_n_fft = gen_istft_n_fft996 self.ups.apply(init_weights)997 self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))998 self.reshape_pixelshuffle = []999 1000 self.subband_conv_post = weight_norm(Conv1d(ch, self.subbands * (self.post_n_fft + 2), 7, 1, padding=3))1001 1002 self.subband_conv_post.apply(init_weights)1003 1004 self.gen_istft_n_fft = gen_istft_n_fft1005 self.gen_istft_hop_size = gen_istft_hop_size1006 1007 #- for onnx1008 if is_onnx == True:1009 self.stft = OnnxSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size, win_length=self.gen_istft_n_fft)1010 else:1011 self.stft = TorchSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size, win_length=self.gen_istft_n_fft)1012 1013 def forward(self, x, g=None):1014 '''1015 stft = TorchSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size,1016 win_length=self.gen_istft_n_fft).to(x.device) # !1017 '''1018 stft = self.stft.to(x.device)1019 pqmf = PQMF(x.device)1020 1021 x = self.conv_pre(x) # [B, ch, length]1022 1023 for i in range(self.num_upsamples):1024 x = F.leaky_relu(x, modules.LRELU_SLOPE)1025 x = self.ups[i](x)1026 1027 xs = None1028 for j in range(self.num_kernels):1029 if xs is None:1030 xs = self.resblocks[i * self.num_kernels + j](x)1031 else:1032 xs += self.resblocks[i * self.num_kernels + j](x)1033 x = xs / self.num_kernels1034 1035 x = F.leaky_relu(x)1036 x = self.reflection_pad(x)1037 x = self.subband_conv_post(x)1038 x = torch.reshape(x, (x.shape[0], self.subbands, x.shape[1] // self.subbands, x.shape[-1]))1039 1040 spec = torch.exp(x[:, :, :self.post_n_fft // 2 + 1, :])1041 phase = math.pi * torch.sin(x[:, :, self.post_n_fft // 2 + 1:, :])1042 1043 y_mb_hat = stft.inverse(1044 torch.reshape(spec, (spec.shape[0] * self.subbands, self.gen_istft_n_fft // 2 + 1, spec.shape[-1])),1045 torch.reshape(phase, (phase.shape[0] * self.subbands, self.gen_istft_n_fft // 2 + 1, phase.shape[-1])))1046 y_mb_hat = torch.reshape(y_mb_hat, (x.shape[0], self.subbands, 1, y_mb_hat.shape[-1]))1047 y_mb_hat = y_mb_hat.squeeze(-2)1048 1049 y_g_hat = pqmf.synthesis(y_mb_hat)1050 1051 return y_g_hat, y_mb_hat1052 1053 def remove_weight_norm(self):1054 print('Removing weight norm...')1055 for l in self.ups:1056 remove_weight_norm(l)1057 for l in self.resblocks:1058 l.remove_weight_norm()1059 1060 1061class Multistream_iSTFT_Generator(torch.nn.Module):1062 def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates,1063 upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, subbands,1064 gin_channels=0, is_onnx=False):1065 super(Multistream_iSTFT_Generator, self).__init__()1066 # self.h = h1067 self.subbands = subbands1068 self.num_kernels = len(resblock_kernel_sizes)1069 self.num_upsamples = len(upsample_rates)1070 self.conv_pre = weight_norm(Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))1071 resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock21072 1073 self.ups = nn.ModuleList()1074 for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):1075 self.ups.append(weight_norm(1076 ConvTranspose1d(upsample_initial_channel // (2 ** i), upsample_initial_channel // (2 ** (i + 1)),1077 k, u, padding=(k - u) // 2)))1078 1079 self.resblocks = nn.ModuleList()1080 for i in range(len(self.ups)):1081 ch = upsample_initial_channel // (2 ** (i + 1))1082 for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):1083 self.resblocks.append(resblock(ch, k, d))1084 1085 self.post_n_fft = gen_istft_n_fft1086 self.ups.apply(init_weights)1087 self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))1088 self.reshape_pixelshuffle = []1089 1090 self.subband_conv_post = weight_norm(Conv1d(ch, self.subbands * (self.post_n_fft + 2), 7, 1, padding=3))1091 1092 self.subband_conv_post.apply(init_weights)1093 1094 self.gen_istft_n_fft = gen_istft_n_fft1095 self.gen_istft_hop_size = gen_istft_hop_size1096 1097 updown_filter = torch.zeros((self.subbands, self.subbands, self.subbands)).float()1098 for k in range(self.subbands):1099 updown_filter[k, k, 0] = 1.01100 self.register_buffer("updown_filter", updown_filter)1101 #self.multistream_conv_post = weight_norm(Conv1d(4, 1, kernel_size=63, bias=False, padding=get_padding(63, 1)))1102 self.multistream_conv_post = weight_norm(Conv1d(self.subbands, 1, kernel_size=63, bias=False, padding=get_padding(63, 1))) # from MB-iSTFT-VITS-44100-Ja1103 self.multistream_conv_post.apply(init_weights)1104 1105 #- for onnx1106 if is_onnx == True:1107 self.stft = OnnxSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size, win_length=self.gen_istft_n_fft)1108 else:1109 self.stft = TorchSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size, win_length=self.gen_istft_n_fft)1110 1111 def forward(self, x, g=None):1112 '''1113 stft = TorchSTFT(filter_length=self.gen_istft_n_fft, hop_length=self.gen_istft_hop_size,1114 win_length=self.gen_istft_n_fft).to(x.device) # !1115 '''1116 stft = self.stft.to(x.device)1117 1118 # pqmf = PQMF(x.device)1119 1120 x = self.conv_pre(x) # [B, ch, length]1121 1122 for i in range(self.num_upsamples):1123 1124 x = F.leaky_relu(x, modules.LRELU_SLOPE)1125 x = self.ups[i](x)1126 1127 xs = None1128 for j in range(self.num_kernels):1129 if xs is None:1130 xs = self.resblocks[i * self.num_kernels + j](x)1131 else:1132 xs += self.resblocks[i * self.num_kernels + j](x)1133 x = xs / self.num_kernels1134 1135 x = F.leaky_relu(x)1136 x = self.reflection_pad(x)1137 x = self.subband_conv_post(x)1138 x = torch.reshape(x, (x.shape[0], self.subbands, x.shape[1] // self.subbands, x.shape[-1]))1139 1140 spec = torch.exp(x[:, :, :self.post_n_fft // 2 + 1, :])1141 phase = math.pi * torch.sin(x[:, :, self.post_n_fft // 2 + 1:, :])1142 1143 y_mb_hat = stft.inverse(1144 torch.reshape(spec, (spec.shape[0] * self.subbands, self.gen_istft_n_fft // 2 + 1, spec.shape[-1])),1145 torch.reshape(phase, (phase.shape[0] * self.subbands, self.gen_istft_n_fft // 2 + 1, phase.shape[-1])))1146 y_mb_hat = torch.reshape(y_mb_hat, (x.shape[0], self.subbands, 1, y_mb_hat.shape[-1]))1147 y_mb_hat = y_mb_hat.squeeze(-2)1148 1149 #y_mb_hat = F.conv_transpose1d(y_mb_hat, self.updown_filter.cuda(x.device) * self.subbands, stride=self.subbands)1150 y_mb_hat = F.conv_transpose1d(y_mb_hat, self.updown_filter.to(x.device) * self.subbands, stride=self.subbands)1151 1152 y_g_hat = self.multistream_conv_post(y_mb_hat)1153 1154 return y_g_hat, y_mb_hat1155 1156 def remove_weight_norm(self):1157 print('Removing weight norm...')1158 for l in self.ups:1159 remove_weight_norm(l)1160 for l in self.resblocks:1161 l.remove_weight_norm()1162 1163 1164class DiscriminatorP(torch.nn.Module):1165 def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):1166 super(DiscriminatorP, self).__init__()1167 self.period = period1168 self.use_spectral_norm = use_spectral_norm1169 norm_f = weight_norm if use_spectral_norm == False else spectral_norm1170 self.convs = nn.ModuleList([1171 norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),1172 norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),1173 norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),1174 norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),1175 norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(get_padding(kernel_size, 1), 0))),1176 ])1177 self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))1178 1179 def forward(self, x):1180 fmap = []1181 1182 # 1d to 2d1183 b, c, t = x.shape1184 if t % self.period != 0: # pad first1185 n_pad = self.period - (t % self.period)1186 x = F.pad(x, (0, n_pad), "reflect")1187 t = t + n_pad1188 x = x.view(b, c, t // self.period, self.period)1189 1190 for l in self.convs:1191 x = l(x)1192 x = F.leaky_relu(x, modules.LRELU_SLOPE)1193 fmap.append(x)1194 x = self.conv_post(x)1195 fmap.append(x)1196 x = torch.flatten(x, 1, -1)1197 1198 return x, fmap1199 1200