Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import math | |
| from torch.nn import functional as F | |
| from .attentions_mechanism import LayerNorm | |
| from .transforms import piecewise_rational_quadratic_transform | |
| from .residual_coupling_layers import Flip | |
| class Log(nn.Module): | |
| def forward(self, x, x_mask, reverse=False, **kwargs): | |
| if not reverse: | |
| y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask | |
| logdet = torch.sum(-y, [1, 2]) | |
| return y, logdet | |
| else: | |
| x = torch.exp(x) * x_mask | |
| return x | |
| class ElementwiseAffine(nn.Module): | |
| def __init__(self, channels): | |
| super().__init__() | |
| self.channels = channels | |
| self.m = nn.Parameter(torch.zeros(channels, 1)) | |
| self.logs = nn.Parameter(torch.zeros(channels, 1)) | |
| def forward(self, x, x_mask, reverse=False, **kwargs): | |
| if not reverse: | |
| y = self.m + torch.exp(self.logs) * x | |
| y = y * x_mask | |
| logdet = torch.sum(self.logs * x_mask, [1, 2]) | |
| return y, logdet | |
| else: | |
| x = (x - self.m) * torch.exp(-self.logs) * x_mask | |
| return x | |
| class DDSConv(nn.Module): | |
| """ | |
| Dialted and Depth-Separable Convolution | |
| """ | |
| def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0): | |
| super().__init__() | |
| self.channels = channels | |
| self.kernel_size = kernel_size | |
| self.n_layers = n_layers | |
| self.p_dropout = p_dropout | |
| self.drop = nn.Dropout(p_dropout) | |
| self.convs_sep = nn.ModuleList() | |
| self.convs_1x1 = nn.ModuleList() | |
| self.norms_1 = nn.ModuleList() | |
| self.norms_2 = nn.ModuleList() | |
| for i in range(n_layers): | |
| dilation = kernel_size**i | |
| padding = (kernel_size * dilation - dilation) // 2 | |
| self.convs_sep.append( | |
| nn.Conv1d( | |
| channels, | |
| channels, | |
| kernel_size, | |
| groups=channels, | |
| dilation=dilation, | |
| padding=padding, | |
| ) | |
| ) | |
| self.convs_1x1.append(nn.Conv1d(channels, channels, 1)) | |
| self.norms_1.append(LayerNorm(channels)) | |
| self.norms_2.append(LayerNorm(channels)) | |
| def forward(self, x, x_mask, g=None): | |
| if g is not None: | |
| x = x + g | |
| for i in range(self.n_layers): | |
| y = self.convs_sep[i](x * x_mask) | |
| y = self.norms_1[i](y) | |
| y = F.gelu(y) | |
| y = self.convs_1x1[i](y) | |
| y = self.norms_2[i](y) | |
| y = F.gelu(y) | |
| y = self.drop(y) | |
| x = x + y | |
| return x * x_mask | |
| class ConvFlow(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels, | |
| filter_channels, | |
| kernel_size, | |
| n_layers, | |
| num_bins=10, | |
| tail_bound=5.0, | |
| ): | |
| super().__init__() | |
| self.in_channels = in_channels | |
| self.filter_channels = filter_channels | |
| self.kernel_size = kernel_size | |
| self.n_layers = n_layers | |
| self.num_bins = num_bins | |
| self.tail_bound = tail_bound | |
| self.half_channels = in_channels // 2 | |
| self.pre = nn.Conv1d(self.half_channels, filter_channels, 1) | |
| self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0) | |
| self.proj = nn.Conv1d( | |
| filter_channels, self.half_channels * (num_bins * 3 - 1), 1 | |
| ) | |
| self.proj.weight.data.zero_() | |
| self.proj.bias.data.zero_() | |
| def forward(self, x, x_mask, g=None, reverse=False): | |
| x0, x1 = torch.split(x, [self.half_channels] * 2, 1) | |
| h = self.pre(x0) | |
| h = self.convs(h, x_mask, g=g) | |
| h = self.proj(h) * x_mask | |
| b, c, t = x0.shape | |
| h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?] | |
| unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels) | |
| unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt( | |
| self.filter_channels | |
| ) | |
| unnormalized_derivatives = h[..., 2 * self.num_bins :] | |
| x1, logabsdet = piecewise_rational_quadratic_transform( | |
| x1, | |
| unnormalized_widths, | |
| unnormalized_heights, | |
| unnormalized_derivatives, | |
| inverse=reverse, | |
| tails="linear", | |
| tail_bound=self.tail_bound, | |
| ) | |
| x = torch.cat([x0, x1], 1) * x_mask | |
| logdet = torch.sum(logabsdet * x_mask, [1, 2]) | |
| if not reverse: | |
| return x, logdet | |
| else: | |
| return x | |
| class StochasticDurationPredictor(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels, | |
| filter_channels, | |
| kernel_size, | |
| p_dropout, | |
| n_flows=4, | |
| gin_channels=0, | |
| ): | |
| super().__init__() | |
| filter_channels = in_channels # it needs to be removed from future version. | |
| self.in_channels = in_channels | |
| self.filter_channels = filter_channels | |
| self.kernel_size = kernel_size | |
| self.p_dropout = p_dropout | |
| self.n_flows = n_flows | |
| self.gin_channels = gin_channels | |
| self.log_flow = Log() | |
| self.flows = nn.ModuleList() | |
| self.flows.append(ElementwiseAffine(2)) | |
| for i in range(n_flows): | |
| self.flows.append( | |
| ConvFlow(2, filter_channels, kernel_size, n_layers=3) | |
| ) | |
| self.flows.append(Flip()) | |
| self.post_pre = nn.Conv1d(1, filter_channels, 1) | |
| self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1) | |
| self.post_convs = DDSConv( | |
| filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout | |
| ) | |
| self.post_flows = nn.ModuleList() | |
| self.post_flows.append(ElementwiseAffine(2)) | |
| for i in range(4): | |
| self.post_flows.append( | |
| ConvFlow(2, filter_channels, kernel_size, n_layers=3) | |
| ) | |
| self.post_flows.append(Flip()) | |
| self.pre = nn.Conv1d(in_channels, filter_channels, 1) | |
| self.proj = nn.Conv1d(filter_channels, filter_channels, 1) | |
| self.convs = DDSConv( | |
| filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout | |
| ) | |
| if gin_channels != 0: | |
| self.cond = nn.Conv1d(gin_channels, filter_channels, 1) | |
| def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0): | |
| x = torch.detach(x) | |
| x = self.pre(x) | |
| if g is not None: | |
| g = torch.detach(g) | |
| x = x + self.cond(g) | |
| x = self.convs(x, x_mask) | |
| x = self.proj(x) * x_mask | |
| if not reverse: | |
| flows = self.flows | |
| assert w is not None | |
| logdet_tot_q = 0 | |
| h_w = self.post_pre(w) | |
| h_w = self.post_convs(h_w, x_mask) | |
| h_w = self.post_proj(h_w) * x_mask | |
| e_q = ( | |
| torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) | |
| * x_mask | |
| ) | |
| z_q = e_q | |
| for flow in self.post_flows: | |
| z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w)) | |
| logdet_tot_q += logdet_q | |
| z_u, z1 = torch.split(z_q, [1, 1], 1) | |
| u = torch.sigmoid(z_u) * x_mask | |
| z0 = (w - u) * x_mask | |
| logdet_tot_q += torch.sum( | |
| (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2] | |
| ) | |
| logq = ( | |
| torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2]) | |
| - logdet_tot_q | |
| ) | |
| logdet_tot = 0 | |
| z0, logdet = self.log_flow(z0, x_mask) | |
| logdet_tot += logdet | |
| z = torch.cat([z0, z1], 1) | |
| for flow in flows: | |
| z, logdet = flow(z, x_mask, g=x, reverse=reverse) | |
| logdet_tot = logdet_tot + logdet | |
| nll = ( | |
| torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2]) | |
| - logdet_tot | |
| ) | |
| return nll + logq # [b] | |
| else: | |
| flows = list(reversed(self.flows)) | |
| flows = flows[:-2] + [flows[-1]] # remove a useless vflow | |
| z = ( | |
| torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) | |
| * noise_scale | |
| ) | |
| for flow in flows: | |
| z = flow(z, x_mask, g=x, reverse=reverse) | |
| z0, z1 = torch.split(z, [1, 1], 1) | |
| logw = z0 | |
| return logw |