surfmore/SimpleRVC

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 infer_pack import commons10from infer_pack.commons import init_weights, get_padding11from infer_pack.transforms import piecewise_rational_quadratic_transform12 13LRELU_SLOPE = 0.114 15 16class LayerNorm(nn.Module):17    def __init__(self, channels, eps=1e-5):18        super().__init__()19        self.channels = channels20        self.eps = eps21 22        self.gamma = nn.Parameter(torch.ones(channels))23        self.beta = nn.Parameter(torch.zeros(channels))24 25    def forward(self, x):26        x = x.transpose(1, -1)27        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)28        return x.transpose(1, -1)29 30 31class ConvReluNorm(nn.Module):32    def __init__(33        self,34        in_channels,35        hidden_channels,36        out_channels,37        kernel_size,38        n_layers,39        p_dropout,40    ):41        super().__init__()42        self.in_channels = in_channels43        self.hidden_channels = hidden_channels44        self.out_channels = out_channels45        self.kernel_size = kernel_size46        self.n_layers = n_layers47        self.p_dropout = p_dropout48        assert n_layers > 1, "Number of layers should be larger than 0."49 50        self.conv_layers = nn.ModuleList()51        self.norm_layers = nn.ModuleList()52        self.conv_layers.append(53            nn.Conv1d(54                in_channels, hidden_channels, kernel_size, padding=kernel_size // 255            )56        )57        self.norm_layers.append(LayerNorm(hidden_channels))58        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))59        for _ in range(n_layers - 1):60            self.conv_layers.append(61                nn.Conv1d(62                    hidden_channels,63                    hidden_channels,64                    kernel_size,65                    padding=kernel_size // 2,66                )67            )68            self.norm_layers.append(LayerNorm(hidden_channels))69        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)70        self.proj.weight.data.zero_()71        self.proj.bias.data.zero_()72 73    def forward(self, x, x_mask):74        x_org = x75        for i in range(self.n_layers):76            x = self.conv_layers[i](x * x_mask)77            x = self.norm_layers[i](x)78            x = self.relu_drop(x)79        x = x_org + self.proj(x)80        return x * x_mask81 82 83class DDSConv(nn.Module):84    """85    Dialted and Depth-Separable Convolution86    """87 88    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):89        super().__init__()90        self.channels = channels91        self.kernel_size = kernel_size92        self.n_layers = n_layers93        self.p_dropout = p_dropout94 95        self.drop = nn.Dropout(p_dropout)96        self.convs_sep = nn.ModuleList()97        self.convs_1x1 = nn.ModuleList()98        self.norms_1 = nn.ModuleList()99        self.norms_2 = nn.ModuleList()100        for i in range(n_layers):101            dilation = kernel_size**i102            padding = (kernel_size * dilation - dilation) // 2103            self.convs_sep.append(104                nn.Conv1d(105                    channels,106                    channels,107                    kernel_size,108                    groups=channels,109                    dilation=dilation,110                    padding=padding,111                )112            )113            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))114            self.norms_1.append(LayerNorm(channels))115            self.norms_2.append(LayerNorm(channels))116 117    def forward(self, x, x_mask, g=None):118        if g is not None:119            x = x + g120        for i in range(self.n_layers):121            y = self.convs_sep[i](x * x_mask)122            y = self.norms_1[i](y)123            y = F.gelu(y)124            y = self.convs_1x1[i](y)125            y = self.norms_2[i](y)126            y = F.gelu(y)127            y = self.drop(y)128            x = x + y129        return x * x_mask130 131 132class WN(torch.nn.Module):133    def __init__(134        self,135        hidden_channels,136        kernel_size,137        dilation_rate,138        n_layers,139        gin_channels=0,140        p_dropout=0,141    ):142        super(WN, self).__init__()143        assert kernel_size % 2 == 1144        self.hidden_channels = hidden_channels145        self.kernel_size = (kernel_size,)146        self.dilation_rate = dilation_rate147        self.n_layers = n_layers148        self.gin_channels = gin_channels149        self.p_dropout = p_dropout150 151        self.in_layers = torch.nn.ModuleList()152        self.res_skip_layers = torch.nn.ModuleList()153        self.drop = nn.Dropout(p_dropout)154 155        if gin_channels != 0:156            cond_layer = torch.nn.Conv1d(157                gin_channels, 2 * hidden_channels * n_layers, 1158            )159            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")160 161        for i in range(n_layers):162            dilation = dilation_rate**i163            padding = int((kernel_size * dilation - dilation) / 2)164            in_layer = torch.nn.Conv1d(165                hidden_channels,166                2 * hidden_channels,167                kernel_size,168                dilation=dilation,169                padding=padding,170            )171            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")172            self.in_layers.append(in_layer)173 174            # last one is not necessary175            if i < n_layers - 1:176                res_skip_channels = 2 * hidden_channels177            else:178                res_skip_channels = hidden_channels179 180            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)181            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")182            self.res_skip_layers.append(res_skip_layer)183 184    def forward(self, x, x_mask, g=None, **kwargs):185        output = torch.zeros_like(x)186        n_channels_tensor = torch.IntTensor([self.hidden_channels])187 188        if g is not None:189            g = self.cond_layer(g)190 191        for i in range(self.n_layers):192            x_in = self.in_layers[i](x)193            if g is not None:194                cond_offset = i * 2 * self.hidden_channels195                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]196            else:197                g_l = torch.zeros_like(x_in)198 199            acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)200            acts = self.drop(acts)201 202            res_skip_acts = self.res_skip_layers[i](acts)203            if i < self.n_layers - 1:204                res_acts = res_skip_acts[:, : self.hidden_channels, :]205                x = (x + res_acts) * x_mask206                output = output + res_skip_acts[:, self.hidden_channels :, :]207            else:208                output = output + res_skip_acts209        return output * x_mask210 211    def remove_weight_norm(self):212        if self.gin_channels != 0:213            torch.nn.utils.remove_weight_norm(self.cond_layer)214        for l in self.in_layers:215            torch.nn.utils.remove_weight_norm(l)216        for l in self.res_skip_layers:217            torch.nn.utils.remove_weight_norm(l)218 219 220class ResBlock1(torch.nn.Module):221    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):222        super(ResBlock1, self).__init__()223        self.convs1 = nn.ModuleList(224            [225                weight_norm(226                    Conv1d(227                        channels,228                        channels,229                        kernel_size,230                        1,231                        dilation=dilation[0],232                        padding=get_padding(kernel_size, dilation[0]),233                    )234                ),235                weight_norm(236                    Conv1d(237                        channels,238                        channels,239                        kernel_size,240                        1,241                        dilation=dilation[1],242                        padding=get_padding(kernel_size, dilation[1]),243                    )244                ),245                weight_norm(246                    Conv1d(247                        channels,248                        channels,249                        kernel_size,250                        1,251                        dilation=dilation[2],252                        padding=get_padding(kernel_size, dilation[2]),253                    )254                ),255            ]256        )257        self.convs1.apply(init_weights)258 259        self.convs2 = nn.ModuleList(260            [261                weight_norm(262                    Conv1d(263                        channels,264                        channels,265                        kernel_size,266                        1,267                        dilation=1,268                        padding=get_padding(kernel_size, 1),269                    )270                ),271                weight_norm(272                    Conv1d(273                        channels,274                        channels,275                        kernel_size,276                        1,277                        dilation=1,278                        padding=get_padding(kernel_size, 1),279                    )280                ),281                weight_norm(282                    Conv1d(283                        channels,284                        channels,285                        kernel_size,286                        1,287                        dilation=1,288                        padding=get_padding(kernel_size, 1),289                    )290                ),291            ]292        )293        self.convs2.apply(init_weights)294 295    def forward(self, x, x_mask=None):296        for c1, c2 in zip(self.convs1, self.convs2):297            xt = F.leaky_relu(x, LRELU_SLOPE)298            if x_mask is not None:299                xt = xt * x_mask300            xt = c1(xt)301            xt = F.leaky_relu(xt, LRELU_SLOPE)302            if x_mask is not None:303                xt = xt * x_mask304            xt = c2(xt)305            x = xt + x306        if x_mask is not None:307            x = x * x_mask308        return x309 310    def remove_weight_norm(self):311        for l in self.convs1:312            remove_weight_norm(l)313        for l in self.convs2:314            remove_weight_norm(l)315 316 317class ResBlock2(torch.nn.Module):318    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):319        super(ResBlock2, self).__init__()320        self.convs = nn.ModuleList(321            [322                weight_norm(323                    Conv1d(324                        channels,325                        channels,326                        kernel_size,327                        1,328                        dilation=dilation[0],329                        padding=get_padding(kernel_size, dilation[0]),330                    )331                ),332                weight_norm(333                    Conv1d(334                        channels,335                        channels,336                        kernel_size,337                        1,338                        dilation=dilation[1],339                        padding=get_padding(kernel_size, dilation[1]),340                    )341                ),342            ]343        )344        self.convs.apply(init_weights)345 346    def forward(self, x, x_mask=None):347        for c in self.convs:348            xt = F.leaky_relu(x, LRELU_SLOPE)349            if x_mask is not None:350                xt = xt * x_mask351            xt = c(xt)352            x = xt + x353        if x_mask is not None:354            x = x * x_mask355        return x356 357    def remove_weight_norm(self):358        for l in self.convs:359            remove_weight_norm(l)360 361 362class Log(nn.Module):363    def forward(self, x, x_mask, reverse=False, **kwargs):364        if not reverse:365            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask366            logdet = torch.sum(-y, [1, 2])367            return y, logdet368        else:369            x = torch.exp(x) * x_mask370            return x371 372 373class Flip(nn.Module):374    def forward(self, x, *args, reverse=False, **kwargs):375        x = torch.flip(x, [1])376        if not reverse:377            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)378            return x, logdet379        else:380            return x381 382 383class ElementwiseAffine(nn.Module):384    def __init__(self, channels):385        super().__init__()386        self.channels = channels387        self.m = nn.Parameter(torch.zeros(channels, 1))388        self.logs = nn.Parameter(torch.zeros(channels, 1))389 390    def forward(self, x, x_mask, reverse=False, **kwargs):391        if not reverse:392            y = self.m + torch.exp(self.logs) * x393            y = y * x_mask394            logdet = torch.sum(self.logs * x_mask, [1, 2])395            return y, logdet396        else:397            x = (x - self.m) * torch.exp(-self.logs) * x_mask398            return x399 400 401class ResidualCouplingLayer(nn.Module):402    def __init__(403        self,404        channels,405        hidden_channels,406        kernel_size,407        dilation_rate,408        n_layers,409        p_dropout=0,410        gin_channels=0,411        mean_only=False,412    ):413        assert channels % 2 == 0, "channels should be divisible by 2"414        super().__init__()415        self.channels = channels416        self.hidden_channels = hidden_channels417        self.kernel_size = kernel_size418        self.dilation_rate = dilation_rate419        self.n_layers = n_layers420        self.half_channels = channels // 2421        self.mean_only = mean_only422 423        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)424        self.enc = WN(425            hidden_channels,426            kernel_size,427            dilation_rate,428            n_layers,429            p_dropout=p_dropout,430            gin_channels=gin_channels,431        )432        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)433        self.post.weight.data.zero_()434        self.post.bias.data.zero_()435 436    def forward(self, x, x_mask, g=None, reverse=False):437        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)438        h = self.pre(x0) * x_mask439        h = self.enc(h, x_mask, g=g)440        stats = self.post(h) * x_mask441        if not self.mean_only:442            m, logs = torch.split(stats, [self.half_channels] * 2, 1)443        else:444            m = stats445            logs = torch.zeros_like(m)446 447        if not reverse:448            x1 = m + x1 * torch.exp(logs) * x_mask449            x = torch.cat([x0, x1], 1)450            logdet = torch.sum(logs, [1, 2])451            return x, logdet452        else:453            x1 = (x1 - m) * torch.exp(-logs) * x_mask454            x = torch.cat([x0, x1], 1)455            return x456 457    def remove_weight_norm(self):458        self.enc.remove_weight_norm()459 460 461class ConvFlow(nn.Module):462    def __init__(463        self,464        in_channels,465        filter_channels,466        kernel_size,467        n_layers,468        num_bins=10,469        tail_bound=5.0,470    ):471        super().__init__()472        self.in_channels = in_channels473        self.filter_channels = filter_channels474        self.kernel_size = kernel_size475        self.n_layers = n_layers476        self.num_bins = num_bins477        self.tail_bound = tail_bound478        self.half_channels = in_channels // 2479 480        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)481        self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)482        self.proj = nn.Conv1d(483            filter_channels, self.half_channels * (num_bins * 3 - 1), 1484        )485        self.proj.weight.data.zero_()486        self.proj.bias.data.zero_()487 488    def forward(self, x, x_mask, g=None, reverse=False):489        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)490        h = self.pre(x0)491        h = self.convs(h, x_mask, g=g)492        h = self.proj(h) * x_mask493 494        b, c, t = x0.shape495        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)  # [b, cx?, t] -> [b, c, t, ?]496 497        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)498        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(499            self.filter_channels500        )501        unnormalized_derivatives = h[..., 2 * self.num_bins :]502 503        x1, logabsdet = piecewise_rational_quadratic_transform(504            x1,505            unnormalized_widths,506            unnormalized_heights,507            unnormalized_derivatives,508            inverse=reverse,509            tails="linear",510            tail_bound=self.tail_bound,511        )512 513        x = torch.cat([x0, x1], 1) * x_mask514        logdet = torch.sum(logabsdet * x_mask, [1, 2])515        if not reverse:516            return x, logdet517        else:518            return x519