# Copyright (c) 2023 Amphion. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. # This code is adopted from META's Encodec under MIT License # https://github.com/facebookresearch/encodec """MS-STFT discriminator, provided here for reference.""" import typing as tp import torchaudio import torch from torch import nn from einops import rearrange from modules.vocoder_blocks import * FeatureMapType = tp.List[torch.Tensor] LogitsType = torch.Tensor DiscriminatorOutput = tp.Tuple[tp.List[LogitsType], tp.List[FeatureMapType]] def get_2d_padding( kernel_size: tp.Tuple[int, int], dilation: tp.Tuple[int, int] = (1, 1) ): return ( ((kernel_size[0] - 1) * dilation[0]) // 2, ((kernel_size[1] - 1) * dilation[1]) // 2, ) class DiscriminatorSTFT(nn.Module): """STFT sub-discriminator. Args: filters (int): Number of filters in convolutions in_channels (int): Number of input channels. Default: 1 out_channels (int): Number of output channels. Default: 1 n_fft (int): Size of FFT for each scale. Default: 1024 hop_length (int): Length of hop between STFT windows for each scale. Default: 256 kernel_size (tuple of int): Inner Conv2d kernel sizes. Default: ``(3, 9)`` stride (tuple of int): Inner Conv2d strides. Default: ``(1, 2)`` dilations (list of int): Inner Conv2d dilation on the time dimension. Default: ``[1, 2, 4]`` win_length (int): Window size for each scale. Default: 1024 normalized (bool): Whether to normalize by magnitude after stft. Default: True norm (str): Normalization method. Default: `'weight_norm'` activation (str): Activation function. Default: `'LeakyReLU'` activation_params (dict): Parameters to provide to the activation function. growth (int): Growth factor for the filters. Default: 1 """ def __init__( self, filters: int, in_channels: int = 1, out_channels: int = 1, n_fft: int = 1024, hop_length: int = 256, win_length: int = 1024, max_filters: int = 1024, filters_scale: int = 1, kernel_size: tp.Tuple[int, int] = (3, 9), dilations: tp.List = [1, 2, 4], stride: tp.Tuple[int, int] = (1, 2), normalized: bool = True, norm: str = "weight_norm", activation: str = "LeakyReLU", activation_params: dict = {"negative_slope": 0.2}, ): super().__init__() assert len(kernel_size) == 2 assert len(stride) == 2 self.filters = filters self.in_channels = in_channels self.out_channels = out_channels self.n_fft = n_fft self.hop_length = hop_length self.win_length = win_length self.normalized = normalized self.activation = getattr(torch.nn, activation)(**activation_params) self.spec_transform = torchaudio.transforms.Spectrogram( n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window_fn=torch.hann_window, normalized=self.normalized, center=False, pad_mode=None, power=None, ) spec_channels = 2 * self.in_channels self.convs = nn.ModuleList() self.convs.append( NormConv2d( spec_channels, self.filters, kernel_size=kernel_size, padding=get_2d_padding(kernel_size), ) ) in_chs = min(filters_scale * self.filters, max_filters) for i, dilation in enumerate(dilations): out_chs = min((filters_scale ** (i + 1)) * self.filters, max_filters) self.convs.append( NormConv2d( in_chs, out_chs, kernel_size=kernel_size, stride=stride, dilation=(dilation, 1), padding=get_2d_padding(kernel_size, (dilation, 1)), norm=norm, ) ) in_chs = out_chs out_chs = min( (filters_scale ** (len(dilations) + 1)) * self.filters, max_filters ) self.convs.append( NormConv2d( in_chs, out_chs, kernel_size=(kernel_size[0], kernel_size[0]), padding=get_2d_padding((kernel_size[0], kernel_size[0])), norm=norm, ) ) self.conv_post = NormConv2d( out_chs, self.out_channels, kernel_size=(kernel_size[0], kernel_size[0]), padding=get_2d_padding((kernel_size[0], kernel_size[0])), norm=norm, ) def forward(self, x: torch.Tensor): """Discriminator STFT Module is the sub module of MultiScaleSTFTDiscriminator. Args: x (torch.Tensor): input tensor of shape [B, 1, Time] Returns: z: z is the output of the last convolutional layer of shape fmap: fmap is the list of feature maps of every convolutional layer of shape """ fmap = [] z = self.spec_transform(x) # [B, 2, Freq, Frames, 2] z = torch.cat([z.real, z.imag], dim=1) z = rearrange(z, "b c w t -> b c t w") for i, layer in enumerate(self.convs): z = layer(z) z = self.activation(z) fmap.append(z) z = self.conv_post(z) return z, fmap class MultiScaleSTFTDiscriminator(nn.Module): """Multi-Scale STFT (MS-STFT) discriminator. Args: filters (int): Number of filters in convolutions in_channels (int): Number of input channels. Default: 1 out_channels (int): Number of output channels. Default: 1 n_ffts (Sequence[int]): Size of FFT for each scale hop_lengths (Sequence[int]): Length of hop between STFT windows for each scale win_lengths (Sequence[int]): Window size for each scale **kwargs: additional args for STFTDiscriminator """ def __init__( self, cfg, in_channels: int = 1, out_channels: int = 1, n_ffts: tp.List[int] = [1024, 2048, 512], hop_lengths: tp.List[int] = [256, 512, 256], win_lengths: tp.List[int] = [1024, 2048, 512], **kwargs, ): self.cfg = cfg super().__init__() assert len(n_ffts) == len(hop_lengths) == len(win_lengths) self.discriminators = nn.ModuleList( [ DiscriminatorSTFT( filters=self.cfg.model.msstftd.filters, in_channels=in_channels, out_channels=out_channels, n_fft=n_ffts[i], win_length=win_lengths[i], hop_length=hop_lengths[i], **kwargs, ) for i in range(len(n_ffts)) ] ) self.num_discriminators = len(self.discriminators) def forward(self, y, y_hat) -> DiscriminatorOutput: """Multi-Scale STFT (MS-STFT) discriminator. Args: x (torch.Tensor): input waveform Returns: logits: list of every discriminator's output fmaps: list of every discriminator's feature maps, each feature maps is a list of Discriminator STFT's every layer """ y_d_rs = [] y_d_gs = [] fmap_rs = [] fmap_gs = [] for disc in self.discriminators: y_d_r, fmap_r = disc(y) y_d_g, fmap_g = disc(y_hat) y_d_rs.append(y_d_r) fmap_rs.append(fmap_r) y_d_gs.append(y_d_g) fmap_gs.append(fmap_g) return y_d_rs, y_d_gs, fmap_rs, fmap_gs