Spaces:
Sleeping
Sleeping
| from typing import List, Tuple | |
| import torch | |
| from torch import Tensor, nn | |
| from torch.nn import AvgPool1d, Conv1d, Module | |
| import torch.nn.functional as F | |
| from torch.nn.utils import spectral_norm, weight_norm | |
| from models.config import HifiGanPretrainingConfig | |
| # Leaky ReLU slope | |
| LRELU_SLOPE = HifiGanPretrainingConfig.lReLU_slope | |
| class DiscriminatorS(Module): | |
| def __init__(self, use_spectral_norm: bool = False): | |
| r"""Initialize the DiscriminatorS module. | |
| Args: | |
| use_spectral_norm (bool, optional): Whether to use spectral normalization. Defaults to False. | |
| """ | |
| super().__init__() | |
| norm_f = weight_norm if not use_spectral_norm else spectral_norm | |
| self.convs = nn.ModuleList( | |
| [ | |
| norm_f(Conv1d(1, 128, 15, 1, padding=7)), | |
| norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)), | |
| norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)), | |
| norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)), | |
| norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)), | |
| norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)), | |
| norm_f(Conv1d(1024, 1024, 5, 1, padding=2)), | |
| ], | |
| ) | |
| self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1)) | |
| def forward(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]: | |
| r"""Forward pass of the DiscriminatorS module. | |
| Args: | |
| x (Tensor): The input tensor. | |
| Returns: | |
| Tuple[Tensor, List[Tensor]]: The output tensor and a list of feature maps. | |
| """ | |
| fmap = [] | |
| for layer in self.convs: | |
| x = layer(x) | |
| x = F.leaky_relu(x, LRELU_SLOPE) | |
| fmap.append(x) | |
| x = self.conv_post(x) | |
| fmap.append(x) | |
| x = torch.flatten(x, 1, -1) | |
| return x, fmap | |
| class MultiScaleDiscriminator(Module): | |
| def __init__(self): | |
| r"""Initialize the MultiScaleDiscriminator module.""" | |
| super().__init__() | |
| self.discriminators = nn.ModuleList( | |
| [ | |
| DiscriminatorS(use_spectral_norm=True), | |
| DiscriminatorS(), | |
| DiscriminatorS(), | |
| ], | |
| ) | |
| self.meanpools = nn.ModuleList( | |
| [ | |
| AvgPool1d(4, 2, padding=2), | |
| AvgPool1d(4, 2, padding=2), | |
| ], | |
| ) | |
| def forward( | |
| self, | |
| y: Tensor, | |
| y_hat: Tensor, | |
| ) -> Tuple[ | |
| List[Tensor], | |
| List[Tensor], | |
| List[Tensor], | |
| List[Tensor], | |
| ]: | |
| r"""Forward pass of the MultiScaleDiscriminator module. | |
| Args: | |
| y (Tensor): The real audio tensor. | |
| y_hat (Tensor): The generated audio tensor. | |
| Returns: | |
| Tuple[List[Tensor], List[Tensor], List[Tensor], List[Tensor]]: | |
| A tuple containing lists of discriminator outputs and feature maps for real and generated audio. | |
| """ | |
| y_d_rs = [] | |
| y_d_gs = [] | |
| fmap_rs = [] | |
| fmap_gs = [] | |
| for i, discriminator in enumerate(self.discriminators): | |
| if i != 0: | |
| y = self.meanpools[i - 1](y) | |
| y_hat = self.meanpools[i - 1](y_hat) | |
| y_d_r, fmap_r = discriminator(y) | |
| y_d_g, fmap_g = discriminator(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 | |