chilge/nemo

1import copy2import math3import numpy as np4import scipy5import torch6from torch import nn7from torch.nn import functional as F8 9from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d10from torch.nn.utils import weight_norm, remove_weight_norm11 12import commons13from commons import init_weights, get_padding14 15 16LRELU_SLOPE = 0.117 18 19class LayerNorm(nn.Module):20  def __init__(self, channels, eps=1e-5):21    super().__init__()22    self.channels = channels23    self.eps = eps24 25    self.gamma = nn.Parameter(torch.ones(channels))26    self.beta = nn.Parameter(torch.zeros(channels))27 28  def forward(self, x):29    x = x.transpose(1, -1)30    x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)31    return x.transpose(1, -1)32 33 34class ConvReluNorm(nn.Module):35  def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):36    super().__init__()37    self.in_channels = in_channels38    self.hidden_channels = hidden_channels39    self.out_channels = out_channels40    self.kernel_size = kernel_size41    self.n_layers = n_layers42    self.p_dropout = p_dropout43    assert n_layers > 1, "Number of layers should be larger than 0."44 45    self.conv_layers = nn.ModuleList()46    self.norm_layers = nn.ModuleList()47    self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))48    self.norm_layers.append(LayerNorm(hidden_channels))49    self.relu_drop = nn.Sequential(50        nn.ReLU(),51        nn.Dropout(p_dropout))52    for _ in range(n_layers-1):53      self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))54      self.norm_layers.append(LayerNorm(hidden_channels))55    self.proj = nn.Conv1d(hidden_channels, out_channels, 1)56    self.proj.weight.data.zero_()57    self.proj.bias.data.zero_()58 59  def forward(self, x, x_mask):60    x_org = x61    for i in range(self.n_layers):62      x = self.conv_layers[i](x * x_mask)63      x = self.norm_layers[i](x)64      x = self.relu_drop(x)65    x = x_org + self.proj(x)66    return x * x_mask67 68 69class DDSConv(nn.Module):70  """71  Dialted and Depth-Separable Convolution72  """73  def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):74    super().__init__()75    self.channels = channels76    self.kernel_size = kernel_size77    self.n_layers = n_layers78    self.p_dropout = p_dropout79 80    self.drop = nn.Dropout(p_dropout)81    self.convs_sep = nn.ModuleList()82    self.convs_1x1 = nn.ModuleList()83    self.norms_1 = nn.ModuleList()84    self.norms_2 = nn.ModuleList()85    for i in range(n_layers):86      dilation = kernel_size ** i87      padding = (kernel_size * dilation - dilation) // 288      self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, 89          groups=channels, dilation=dilation, padding=padding90      ))91      self.convs_1x1.append(nn.Conv1d(channels, channels, 1))92      self.norms_1.append(LayerNorm(channels))93      self.norms_2.append(LayerNorm(channels))94 95  def forward(self, x, x_mask, g=None):96    if g is not None:97      x = x + g98    for i in range(self.n_layers):99      y = self.convs_sep[i](x * x_mask)100      y = self.norms_1[i](y)101      y = F.gelu(y)102      y = self.convs_1x1[i](y)103      y = self.norms_2[i](y)104      y = F.gelu(y)105      y = self.drop(y)106      x = x + y107    return x * x_mask108 109 110class WN(torch.nn.Module):111  def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):112    super(WN, self).__init__()113    assert(kernel_size % 2 == 1)114    self.hidden_channels =hidden_channels115    self.kernel_size = kernel_size,116    self.dilation_rate = dilation_rate117    self.n_layers = n_layers118    self.gin_channels = gin_channels119    self.p_dropout = p_dropout120 121    self.in_layers = torch.nn.ModuleList()122    self.res_skip_layers = torch.nn.ModuleList()123    self.drop = nn.Dropout(p_dropout)124 125    if gin_channels != 0:126      cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)127      self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')128 129    for i in range(n_layers):130      dilation = dilation_rate ** i131      padding = int((kernel_size * dilation - dilation) / 2)132      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,133                                 dilation=dilation, padding=padding)134      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')135      self.in_layers.append(in_layer)136 137      # last one is not necessary138      if i < n_layers - 1:139        res_skip_channels = 2 * hidden_channels140      else:141        res_skip_channels = hidden_channels142 143      res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)144      res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')145      self.res_skip_layers.append(res_skip_layer)146 147  def forward(self, x, x_mask, g=None, **kwargs):148    output = torch.zeros_like(x)149    n_channels_tensor = torch.IntTensor([self.hidden_channels])150 151    if g is not None:152      g = self.cond_layer(g)153 154    for i in range(self.n_layers):155      x_in = self.in_layers[i](x)156      if g is not None:157        cond_offset = i * 2 * self.hidden_channels158        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]159      else:160        g_l = torch.zeros_like(x_in)161 162      acts = commons.fused_add_tanh_sigmoid_multiply(163          x_in,164          g_l,165          n_channels_tensor)166      acts = self.drop(acts)167 168      res_skip_acts = self.res_skip_layers[i](acts)169      if i < self.n_layers - 1:170        res_acts = res_skip_acts[:,:self.hidden_channels,:]171        x = (x + res_acts) * x_mask172        output = output + res_skip_acts[:,self.hidden_channels:,:]173      else:174        output = output + res_skip_acts175    return output * x_mask176 177  def remove_weight_norm(self):178    if self.gin_channels != 0:179      torch.nn.utils.remove_weight_norm(self.cond_layer)180    for l in self.in_layers:181      torch.nn.utils.remove_weight_norm(l)182    for l in self.res_skip_layers:183     torch.nn.utils.remove_weight_norm(l)184 185 186class ResBlock1(torch.nn.Module):187    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):188        super(ResBlock1, self).__init__()189        self.convs1 = nn.ModuleList([190            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],191                               padding=get_padding(kernel_size, dilation[0]))),192            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],193                               padding=get_padding(kernel_size, dilation[1]))),194            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],195                               padding=get_padding(kernel_size, dilation[2])))196        ])197        self.convs1.apply(init_weights)198 199        self.convs2 = nn.ModuleList([200            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,201                               padding=get_padding(kernel_size, 1))),202            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,203                               padding=get_padding(kernel_size, 1))),204            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,205                               padding=get_padding(kernel_size, 1)))206        ])207        self.convs2.apply(init_weights)208 209    def forward(self, x, x_mask=None):210        for c1, c2 in zip(self.convs1, self.convs2):211            xt = F.leaky_relu(x, LRELU_SLOPE)212            if x_mask is not None:213                xt = xt * x_mask214            xt = c1(xt)215            xt = F.leaky_relu(xt, LRELU_SLOPE)216            if x_mask is not None:217                xt = xt * x_mask218            xt = c2(xt)219            x = xt + x220        if x_mask is not None:221            x = x * x_mask222        return x223 224    def remove_weight_norm(self):225        for l in self.convs1:226            remove_weight_norm(l)227        for l in self.convs2:228            remove_weight_norm(l)229 230 231class ResBlock2(torch.nn.Module):232    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):233        super(ResBlock2, self).__init__()234        self.convs = nn.ModuleList([235            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],236                               padding=get_padding(kernel_size, dilation[0]))),237            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],238                               padding=get_padding(kernel_size, dilation[1])))239        ])240        self.convs.apply(init_weights)241 242    def forward(self, x, x_mask=None):243        for c in self.convs:244            xt = F.leaky_relu(x, LRELU_SLOPE)245            if x_mask is not None:246                xt = xt * x_mask247            xt = c(xt)248            x = xt + x249        if x_mask is not None:250            x = x * x_mask251        return x252 253    def remove_weight_norm(self):254        for l in self.convs:255            remove_weight_norm(l)256 257 258class Log(nn.Module):259  def forward(self, x, x_mask, reverse=False, **kwargs):260    if not reverse:261      y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask262      logdet = torch.sum(-y, [1, 2])263      return y, logdet264    else:265      x = torch.exp(x) * x_mask266      return x267    268 269class Flip(nn.Module):270  def forward(self, x, *args, reverse=False, **kwargs):271    x = torch.flip(x, [1])272    if not reverse:273      logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)274      return x, logdet275    else:276      return x277 278 279class ElementwiseAffine(nn.Module):280  def __init__(self, channels):281    super().__init__()282    self.channels = channels283    self.m = nn.Parameter(torch.zeros(channels,1))284    self.logs = nn.Parameter(torch.zeros(channels,1))285 286  def forward(self, x, x_mask, reverse=False, **kwargs):287    if not reverse:288      y = self.m + torch.exp(self.logs) * x289      y = y * x_mask290      logdet = torch.sum(self.logs * x_mask, [1,2])291      return y, logdet292    else:293      x = (x - self.m) * torch.exp(-self.logs) * x_mask294      return x295 296 297class ResidualCouplingLayer(nn.Module):298  def __init__(self,299      channels,300      hidden_channels,301      kernel_size,302      dilation_rate,303      n_layers,304      p_dropout=0,305      gin_channels=0,306      mean_only=False):307    assert channels % 2 == 0, "channels should be divisible by 2"308    super().__init__()309    self.channels = channels310    self.hidden_channels = hidden_channels311    self.kernel_size = kernel_size312    self.dilation_rate = dilation_rate313    self.n_layers = n_layers314    self.half_channels = channels // 2315    self.mean_only = mean_only316 317    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)318    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)319    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)320    self.post.weight.data.zero_()321    self.post.bias.data.zero_()322 323  def forward(self, x, x_mask, g=None, reverse=False):324    x0, x1 = torch.split(x, [self.half_channels]*2, 1)325    h = self.pre(x0) * x_mask326    h = self.enc(h, x_mask, g=g)327    stats = self.post(h) * x_mask328    if not self.mean_only:329      m, logs = torch.split(stats, [self.half_channels]*2, 1)330    else:331      m = stats332      logs = torch.zeros_like(m)333 334    if not reverse:335      x1 = m + x1 * torch.exp(logs) * x_mask336      x = torch.cat([x0, x1], 1)337      logdet = torch.sum(logs, [1,2])338      return x, logdet339    else:340      x1 = (x1 - m) * torch.exp(-logs) * x_mask341      x = torch.cat([x0, x1], 1)342      return x343