from typing import List import torch from torch import nn from torch.nn import Module from torch.nn import functional as F from torch.nn.utils import parametrize from .kernel_predictor import KernelPredictor class LVCBlock(Module): r"""The location-variable convolutions block. To efficiently capture the local information of the condition, location-variable convolution (LVC) obtained better sound quality and speed while maintaining the model size. The kernels of the LVC layers are predicted using a kernel predictor that takes the log-mel-spectrogram as the input. The kernel predictor is connected to a residual stack. One kernel predictor simultaneously predicts the kernels of all LVC layers in one residual stack. Args: in_channels (int): The number of input channels. cond_channels (int): The number of conditioning channels. stride (int): The stride of the convolutional layers. dilations (List[int]): A list of dilation values for the convolutional layers. lReLU_slope (float): The slope of the LeakyReLU activation function. conv_kernel_size (int): The kernel size of the convolutional layers. cond_hop_length (int): The hop length of the conditioning sequence. kpnet_hidden_channels (int): The number of hidden channels in the kernel predictor network. kpnet_conv_size (int): The kernel size of the convolutional layers in the kernel predictor network. kpnet_dropout (float): The dropout rate for the kernel predictor network. Attributes: cond_hop_length (int): The hop length of the conditioning sequence. conv_layers (int): The number of convolutional layers. conv_kernel_size (int): The kernel size of the convolutional layers. kernel_predictor (KernelPredictor): The kernel predictor network. convt_pre (nn.Sequential): The convolutional transpose layer. conv_blocks (nn.ModuleList): The list of convolutional blocks. """ def __init__( self, in_channels: int, cond_channels: int, stride: int, dilations: List[int] = [1, 3, 9, 27], lReLU_slope: float = 0.2, conv_kernel_size: int = 3, cond_hop_length: int = 256, kpnet_hidden_channels: int = 64, kpnet_conv_size: int = 3, kpnet_dropout: float = 0.0, ): super().__init__() self.cond_hop_length = cond_hop_length self.conv_layers = len(dilations) self.conv_kernel_size = conv_kernel_size self.kernel_predictor = KernelPredictor( cond_channels=cond_channels, conv_in_channels=in_channels, conv_out_channels=2 * in_channels, conv_layers=len(dilations), conv_kernel_size=conv_kernel_size, kpnet_hidden_channels=kpnet_hidden_channels, kpnet_conv_size=kpnet_conv_size, kpnet_dropout=kpnet_dropout, lReLU_slope=lReLU_slope, ) self.convt_pre = nn.Sequential( nn.LeakyReLU(lReLU_slope), nn.utils.parametrizations.weight_norm( nn.ConvTranspose1d( in_channels, in_channels, 2 * stride, stride=stride, padding=stride // 2 + stride % 2, output_padding=stride % 2, ), ), ) self.conv_blocks = nn.ModuleList( [ nn.Sequential( nn.LeakyReLU(lReLU_slope), nn.utils.parametrizations.weight_norm( nn.Conv1d( in_channels, in_channels, conv_kernel_size, padding=dilation * (conv_kernel_size - 1) // 2, dilation=dilation, ), ), nn.LeakyReLU(lReLU_slope), ) for dilation in dilations ], ) def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor: r"""Forward propagation of the location-variable convolutions. Args: x (Tensor): The input sequence (batch, in_channels, in_length). c (Tensor): The conditioning sequence (batch, cond_channels, cond_length). Returns: Tensor: The output sequence (batch, in_channels, in_length). """ _, in_channels, _ = x.shape # (B, c_g, L') x = self.convt_pre(x) # (B, c_g, stride * L') kernels, bias = self.kernel_predictor(c) for i, conv in enumerate(self.conv_blocks): output = conv(x) # (B, c_g, stride * L') k = kernels[:, i, :, :, :, :] # (B, 2 * c_g, c_g, kernel_size, cond_length) b = bias[:, i, :, :] # (B, 2 * c_g, cond_length) output = self.location_variable_convolution( output, k, b, hop_size=self.cond_hop_length, ) # (B, 2 * c_g, stride * L'): LVC x = x + torch.sigmoid(output[:, :in_channels, :]) * torch.tanh( output[:, in_channels:, :], ) # (B, c_g, stride * L'): GAU return x def location_variable_convolution( self, x: torch.Tensor, kernel: torch.Tensor, bias: torch.Tensor, dilation: int = 1, hop_size: int = 256, ) -> torch.Tensor: r"""Perform location-variable convolution operation on the input sequence (x) using the local convolution kernel. Time: 414 μs ± 309 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each), test on NVIDIA V100. Args: x (Tensor): The input sequence (batch, in_channels, in_length). kernel (Tensor): The local convolution kernel (batch, in_channel, out_channels, kernel_size, kernel_length). bias (Tensor): The bias for the local convolution (batch, out_channels, kernel_length). dilation (int): The dilation of convolution. hop_size (int): The hop_size of the conditioning sequence. Returns: (Tensor): The output sequence after performing local convolution. (batch, out_channels, in_length). """ batch, _, in_length = x.shape batch, _, out_channels, kernel_size, kernel_length = kernel.shape assert in_length == ( kernel_length * hop_size ), "length of (x, kernel) is not matched" padding = dilation * int((kernel_size - 1) / 2) x = F.pad( x, (padding, padding), "constant", 0, ) # (batch, in_channels, in_length + 2*padding) x = x.unfold( 2, hop_size + 2 * padding, hop_size, ) # (batch, in_channels, kernel_length, hop_size + 2*padding) if hop_size < dilation: x = F.pad(x, (0, dilation), "constant", 0) x = x.unfold( 3, dilation, dilation, ) # (batch, in_channels, kernel_length, (hop_size + 2*padding)/dilation, dilation) x = x[:, :, :, :, :hop_size] x = x.transpose( 3, 4, ) # (batch, in_channels, kernel_length, dilation, (hop_size + 2*padding)/dilation) x = x.unfold( 4, kernel_size, 1, ) # (batch, in_channels, kernel_length, dilation, _, kernel_size) o = torch.einsum("bildsk,biokl->bolsd", x, kernel) o = o.contiguous(memory_format=torch.channels_last_3d) bias = ( bias.unsqueeze(-1) .unsqueeze(-1) .contiguous(memory_format=torch.channels_last_3d) ) o = o + bias return o.contiguous().view(batch, out_channels, -1) def remove_weight_norm(self) -> None: r"""Remove weight normalization from the convolutional layers in the LVCBlock. This method removes weight normalization from the kernel predictor and all convolutional layers in the LVCBlock. """ self.kernel_predictor.remove_weight_norm() parametrize.remove_parametrizations(self.convt_pre[1], "weight") for block in self.conv_blocks: parametrize.remove_parametrizations(block[1], "weight") # type: ignore