# -*- coding: utf-8 -*- # Copyright 2019 Tomoki Hayashi # MIT License (https://opensource.org/licenses/MIT) """Parallel WaveGAN Modules.""" import logging import math import torch from torch import nn from modules.parallel_wavegan.layers import Conv1d from modules.parallel_wavegan.layers import Conv1d1x1 from modules.parallel_wavegan.layers import ResidualBlock from modules.parallel_wavegan.layers import upsample from modules.parallel_wavegan import models class ParallelWaveGANGenerator(torch.nn.Module): """Parallel WaveGAN Generator module.""" def __init__(self, in_channels=1, out_channels=1, kernel_size=3, layers=30, stacks=3, residual_channels=64, gate_channels=128, skip_channels=64, aux_channels=80, aux_context_window=2, dropout=0.0, bias=True, use_weight_norm=True, use_causal_conv=False, upsample_conditional_features=True, upsample_net="ConvInUpsampleNetwork", upsample_params={"upsample_scales": [4, 4, 4, 4]}, use_pitch_embed=False, ): """Initialize Parallel WaveGAN Generator module. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. kernel_size (int): Kernel size of dilated convolution. layers (int): Number of residual block layers. stacks (int): Number of stacks i.e., dilation cycles. residual_channels (int): Number of channels in residual conv. gate_channels (int): Number of channels in gated conv. skip_channels (int): Number of channels in skip conv. aux_channels (int): Number of channels for auxiliary feature conv. aux_context_window (int): Context window size for auxiliary feature. dropout (float): Dropout rate. 0.0 means no dropout applied. bias (bool): Whether to use bias parameter in conv layer. use_weight_norm (bool): Whether to use weight norm. If set to true, it will be applied to all of the conv layers. use_causal_conv (bool): Whether to use causal structure. upsample_conditional_features (bool): Whether to use upsampling network. upsample_net (str): Upsampling network architecture. upsample_params (dict): Upsampling network parameters. """ super(ParallelWaveGANGenerator, self).__init__() self.in_channels = in_channels self.out_channels = out_channels self.aux_channels = aux_channels self.layers = layers self.stacks = stacks self.kernel_size = kernel_size # check the number of layers and stacks assert layers % stacks == 0 layers_per_stack = layers // stacks # define first convolution self.first_conv = Conv1d1x1(in_channels, residual_channels, bias=True) # define conv + upsampling network if upsample_conditional_features: upsample_params.update({ "use_causal_conv": use_causal_conv, }) if upsample_net == "MelGANGenerator": assert aux_context_window == 0 upsample_params.update({ "use_weight_norm": False, # not to apply twice "use_final_nonlinear_activation": False, }) self.upsample_net = getattr(models, upsample_net)(**upsample_params) else: if upsample_net == "ConvInUpsampleNetwork": upsample_params.update({ "aux_channels": aux_channels, "aux_context_window": aux_context_window, }) self.upsample_net = getattr(upsample, upsample_net)(**upsample_params) else: self.upsample_net = None # define residual blocks self.conv_layers = torch.nn.ModuleList() for layer in range(layers): dilation = 2 ** (layer % layers_per_stack) conv = ResidualBlock( kernel_size=kernel_size, residual_channels=residual_channels, gate_channels=gate_channels, skip_channels=skip_channels, aux_channels=aux_channels, dilation=dilation, dropout=dropout, bias=bias, use_causal_conv=use_causal_conv, ) self.conv_layers += [conv] # define output layers self.last_conv_layers = torch.nn.ModuleList([ torch.nn.ReLU(inplace=True), Conv1d1x1(skip_channels, skip_channels, bias=True), torch.nn.ReLU(inplace=True), Conv1d1x1(skip_channels, out_channels, bias=True), ]) self.use_pitch_embed = use_pitch_embed if use_pitch_embed: self.pitch_embed = nn.Embedding(300, aux_channels, 0) self.c_proj = nn.Linear(2 * aux_channels, aux_channels) # apply weight norm if use_weight_norm: self.apply_weight_norm() def forward(self, x, c=None, pitch=None, **kwargs): """Calculate forward propagation. Args: x (Tensor): Input noise signal (B, C_in, T). c (Tensor): Local conditioning auxiliary features (B, C ,T'). pitch (Tensor): Local conditioning pitch (B, T'). Returns: Tensor: Output tensor (B, C_out, T) """ # perform upsampling if c is not None and self.upsample_net is not None: if self.use_pitch_embed: p = self.pitch_embed(pitch) c = self.c_proj(torch.cat([c.transpose(1, 2), p], -1)).transpose(1, 2) c = self.upsample_net(c) assert c.size(-1) == x.size(-1), (c.size(-1), x.size(-1)) # encode to hidden representation x = self.first_conv(x) skips = 0 for f in self.conv_layers: x, h = f(x, c) skips += h skips *= math.sqrt(1.0 / len(self.conv_layers)) # apply final layers x = skips for f in self.last_conv_layers: x = f(x) return x def remove_weight_norm(self): """Remove weight normalization module from all of the layers.""" def _remove_weight_norm(m): try: logging.debug(f"Weight norm is removed from {m}.") torch.nn.utils.remove_weight_norm(m) except ValueError: # this module didn't have weight norm return self.apply(_remove_weight_norm) def apply_weight_norm(self): """Apply weight normalization module from all of the layers.""" def _apply_weight_norm(m): if isinstance(m, torch.nn.Conv1d) or isinstance(m, torch.nn.Conv2d): torch.nn.utils.weight_norm(m) logging.debug(f"Weight norm is applied to {m}.") self.apply(_apply_weight_norm) @staticmethod def _get_receptive_field_size(layers, stacks, kernel_size, dilation=lambda x: 2 ** x): assert layers % stacks == 0 layers_per_cycle = layers // stacks dilations = [dilation(i % layers_per_cycle) for i in range(layers)] return (kernel_size - 1) * sum(dilations) + 1 @property def receptive_field_size(self): """Return receptive field size.""" return self._get_receptive_field_size(self.layers, self.stacks, self.kernel_size) class ParallelWaveGANDiscriminator(torch.nn.Module): """Parallel WaveGAN Discriminator module.""" def __init__(self, in_channels=1, out_channels=1, kernel_size=3, layers=10, conv_channels=64, dilation_factor=1, nonlinear_activation="LeakyReLU", nonlinear_activation_params={"negative_slope": 0.2}, bias=True, use_weight_norm=True, ): """Initialize Parallel WaveGAN Discriminator module. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. kernel_size (int): Number of output channels. layers (int): Number of conv layers. conv_channels (int): Number of chnn layers. dilation_factor (int): Dilation factor. For example, if dilation_factor = 2, the dilation will be 2, 4, 8, ..., and so on. nonlinear_activation (str): Nonlinear function after each conv. nonlinear_activation_params (dict): Nonlinear function parameters bias (bool): Whether to use bias parameter in conv. use_weight_norm (bool) Whether to use weight norm. If set to true, it will be applied to all of the conv layers. """ super(ParallelWaveGANDiscriminator, self).__init__() assert (kernel_size - 1) % 2 == 0, "Not support even number kernel size." assert dilation_factor > 0, "Dilation factor must be > 0." self.conv_layers = torch.nn.ModuleList() conv_in_channels = in_channels for i in range(layers - 1): if i == 0: dilation = 1 else: dilation = i if dilation_factor == 1 else dilation_factor ** i conv_in_channels = conv_channels padding = (kernel_size - 1) // 2 * dilation conv_layer = [ Conv1d(conv_in_channels, conv_channels, kernel_size=kernel_size, padding=padding, dilation=dilation, bias=bias), getattr(torch.nn, nonlinear_activation)(inplace=True, **nonlinear_activation_params) ] self.conv_layers += conv_layer padding = (kernel_size - 1) // 2 last_conv_layer = Conv1d( conv_in_channels, out_channels, kernel_size=kernel_size, padding=padding, bias=bias) self.conv_layers += [last_conv_layer] # apply weight norm if use_weight_norm: self.apply_weight_norm() def forward(self, x): """Calculate forward propagation. Args: x (Tensor): Input noise signal (B, 1, T). Returns: Tensor: Output tensor (B, 1, T) """ for f in self.conv_layers: x = f(x) return x def apply_weight_norm(self): """Apply weight normalization module from all of the layers.""" def _apply_weight_norm(m): if isinstance(m, torch.nn.Conv1d) or isinstance(m, torch.nn.Conv2d): torch.nn.utils.weight_norm(m) logging.debug(f"Weight norm is applied to {m}.") self.apply(_apply_weight_norm) def remove_weight_norm(self): """Remove weight normalization module from all of the layers.""" def _remove_weight_norm(m): try: logging.debug(f"Weight norm is removed from {m}.") torch.nn.utils.remove_weight_norm(m) except ValueError: # this module didn't have weight norm return self.apply(_remove_weight_norm) class ResidualParallelWaveGANDiscriminator(torch.nn.Module): """Parallel WaveGAN Discriminator module.""" def __init__(self, in_channels=1, out_channels=1, kernel_size=3, layers=30, stacks=3, residual_channels=64, gate_channels=128, skip_channels=64, dropout=0.0, bias=True, use_weight_norm=True, use_causal_conv=False, nonlinear_activation="LeakyReLU", nonlinear_activation_params={"negative_slope": 0.2}, ): """Initialize Parallel WaveGAN Discriminator module. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. kernel_size (int): Kernel size of dilated convolution. layers (int): Number of residual block layers. stacks (int): Number of stacks i.e., dilation cycles. residual_channels (int): Number of channels in residual conv. gate_channels (int): Number of channels in gated conv. skip_channels (int): Number of channels in skip conv. dropout (float): Dropout rate. 0.0 means no dropout applied. bias (bool): Whether to use bias parameter in conv. use_weight_norm (bool): Whether to use weight norm. If set to true, it will be applied to all of the conv layers. use_causal_conv (bool): Whether to use causal structure. nonlinear_activation_params (dict): Nonlinear function parameters """ super(ResidualParallelWaveGANDiscriminator, self).__init__() assert (kernel_size - 1) % 2 == 0, "Not support even number kernel size." self.in_channels = in_channels self.out_channels = out_channels self.layers = layers self.stacks = stacks self.kernel_size = kernel_size # check the number of layers and stacks assert layers % stacks == 0 layers_per_stack = layers // stacks # define first convolution self.first_conv = torch.nn.Sequential( Conv1d1x1(in_channels, residual_channels, bias=True), getattr(torch.nn, nonlinear_activation)( inplace=True, **nonlinear_activation_params), ) # define residual blocks self.conv_layers = torch.nn.ModuleList() for layer in range(layers): dilation = 2 ** (layer % layers_per_stack) conv = ResidualBlock( kernel_size=kernel_size, residual_channels=residual_channels, gate_channels=gate_channels, skip_channels=skip_channels, aux_channels=-1, dilation=dilation, dropout=dropout, bias=bias, use_causal_conv=use_causal_conv, ) self.conv_layers += [conv] # define output layers self.last_conv_layers = torch.nn.ModuleList([ getattr(torch.nn, nonlinear_activation)( inplace=True, **nonlinear_activation_params), Conv1d1x1(skip_channels, skip_channels, bias=True), getattr(torch.nn, nonlinear_activation)( inplace=True, **nonlinear_activation_params), Conv1d1x1(skip_channels, out_channels, bias=True), ]) # apply weight norm if use_weight_norm: self.apply_weight_norm() def forward(self, x): """Calculate forward propagation. Args: x (Tensor): Input noise signal (B, 1, T). Returns: Tensor: Output tensor (B, 1, T) """ x = self.first_conv(x) skips = 0 for f in self.conv_layers: x, h = f(x, None) skips += h skips *= math.sqrt(1.0 / len(self.conv_layers)) # apply final layers x = skips for f in self.last_conv_layers: x = f(x) return x def apply_weight_norm(self): """Apply weight normalization module from all of the layers.""" def _apply_weight_norm(m): if isinstance(m, torch.nn.Conv1d) or isinstance(m, torch.nn.Conv2d): torch.nn.utils.weight_norm(m) logging.debug(f"Weight norm is applied to {m}.") self.apply(_apply_weight_norm) def remove_weight_norm(self): """Remove weight normalization module from all of the layers.""" def _remove_weight_norm(m): try: logging.debug(f"Weight norm is removed from {m}.") torch.nn.utils.remove_weight_norm(m) except ValueError: # this module didn't have weight norm return self.apply(_remove_weight_norm)