| import torch |
| import torch.nn.functional as F |
| from torch.utils.checkpoint import checkpoint |
| from packaging import version |
| is_pytorch2_1 = version.parse(torch.__version__) >= version.parse("2.1.0") |
| if is_pytorch2_1: |
| from torch.nn.utils.parametrizations import spectral_norm, weight_norm |
| else: |
| from torch.nn.utils.parametrizations import spectral_norm |
| from torch.nn.utils import weight_norm |
|
|
| from .commons import get_padding |
| from .residuals import LRELU_SLOPE |
|
|
|
|
| class MultiPeriodDiscriminator(torch.nn.Module): |
|
|
| def __init__( |
| self, |
| use_spectral_norm: bool = False, |
| checkpointing: bool = False, |
| version: str = "v2", |
| ): |
| super().__init__() |
|
|
| if version == "v1": |
| periods = [2, 3, 5, 7, 11, 17] |
| resolutions = [] |
| elif version == "v2": |
| periods = [2, 3, 5, 7, 11, 17, 23, 37] |
| resolutions = [] |
| elif version == "v3": |
| periods = [2, 3, 5, 7, 11] |
| resolutions = [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]] |
|
|
| self.checkpointing = checkpointing |
| self.discriminators = torch.nn.ModuleList( |
| [DiscriminatorS(use_spectral_norm=use_spectral_norm)] |
| + [DiscriminatorP(p, use_spectral_norm=use_spectral_norm) for p in periods] |
| + [ |
| DiscriminatorR(r, use_spectral_norm=use_spectral_norm) |
| for r in resolutions |
| ] |
| ) |
|
|
| def forward(self, y, y_hat): |
| y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], [] |
| for d in self.discriminators: |
| if self.training and self.checkpointing: |
| y_d_r, fmap_r = checkpoint(d, y, use_reentrant=False) |
| y_d_g, fmap_g = checkpoint(d, y_hat, use_reentrant=False) |
| else: |
| y_d_r, fmap_r = d(y) |
| y_d_g, fmap_g = d(y_hat) |
| y_d_rs.append(y_d_r) |
| y_d_gs.append(y_d_g) |
| fmap_rs.append(fmap_r) |
| fmap_gs.append(fmap_g) |
|
|
| return y_d_rs, y_d_gs, fmap_rs, fmap_gs |
|
|
|
|
| class DiscriminatorS(torch.nn.Module): |
|
|
| def __init__(self, use_spectral_norm: bool = False): |
| super().__init__() |
|
|
| norm_f = spectral_norm if use_spectral_norm else weight_norm |
| self.convs = torch.nn.ModuleList( |
| [ |
| norm_f(torch.nn.Conv1d(1, 16, 15, 1, padding=7)), |
| norm_f(torch.nn.Conv1d(16, 64, 41, 4, groups=4, padding=20)), |
| norm_f(torch.nn.Conv1d(64, 256, 41, 4, groups=16, padding=20)), |
| norm_f(torch.nn.Conv1d(256, 1024, 41, 4, groups=64, padding=20)), |
| norm_f(torch.nn.Conv1d(1024, 1024, 41, 4, groups=256, padding=20)), |
| norm_f(torch.nn.Conv1d(1024, 1024, 5, 1, padding=2)), |
| ] |
| ) |
| self.conv_post = norm_f(torch.nn.Conv1d(1024, 1, 3, 1, padding=1)) |
| self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE) |
|
|
| def forward(self, x): |
| fmap = [] |
| for conv in self.convs: |
| x = self.lrelu(conv(x)) |
| fmap.append(x) |
| x = self.conv_post(x) |
| fmap.append(x) |
| x = torch.flatten(x, 1, -1) |
| return x, fmap |
|
|
|
|
| class DiscriminatorP(torch.nn.Module): |
|
|
| def __init__( |
| self, |
| period: int, |
| kernel_size: int = 5, |
| stride: int = 3, |
| use_spectral_norm: bool = False, |
| ): |
| super().__init__() |
| self.period = period |
| norm_f = spectral_norm if use_spectral_norm else weight_norm |
|
|
| in_channels = [1, 32, 128, 512, 1024] |
| out_channels = [32, 128, 512, 1024, 1024] |
| strides = [3, 3, 3, 3, 1] |
|
|
| self.convs = torch.nn.ModuleList( |
| [ |
| norm_f( |
| torch.nn.Conv2d( |
| in_ch, |
| out_ch, |
| (kernel_size, 1), |
| (s, 1), |
| padding=(get_padding(kernel_size, 1), 0), |
| ) |
| ) |
| for in_ch, out_ch, s in zip(in_channels, out_channels, strides) |
| ] |
| ) |
|
|
| self.conv_post = norm_f(torch.nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))) |
| self.lrelu = torch.nn.LeakyReLU(LRELU_SLOPE) |
|
|
| def forward(self, x): |
| fmap = [] |
| b, c, t = x.shape |
| if t % self.period != 0: |
| n_pad = self.period - (t % self.period) |
| x = torch.nn.functional.pad(x, (0, n_pad), "reflect") |
| x = x.view(b, c, -1, self.period) |
|
|
| for conv in self.convs: |
| x = self.lrelu(conv(x)) |
| fmap.append(x) |
| x = self.conv_post(x) |
| fmap.append(x) |
| x = torch.flatten(x, 1, -1) |
| return x, fmap |
|
|
|
|
| class DiscriminatorR(torch.nn.Module): |
| def __init__(self, resolution, use_spectral_norm=False): |
| super().__init__() |
|
|
| self.resolution = resolution |
| self.lrelu_slope = 0.1 |
| norm_f = spectral_norm if use_spectral_norm else weight_norm |
|
|
| self.convs = torch.nn.ModuleList( |
| [ |
| norm_f( |
| torch.nn.Conv2d( |
| 1, |
| 32, |
| (3, 9), |
| padding=(1, 4), |
| ) |
| ), |
| norm_f( |
| torch.nn.Conv2d( |
| 32, |
| 32, |
| (3, 9), |
| stride=(1, 2), |
| padding=(1, 4), |
| ) |
| ), |
| norm_f( |
| torch.nn.Conv2d( |
| 32, |
| 32, |
| (3, 9), |
| stride=(1, 2), |
| padding=(1, 4), |
| ) |
| ), |
| norm_f( |
| torch.nn.Conv2d( |
| 32, |
| 32, |
| (3, 9), |
| stride=(1, 2), |
| padding=(1, 4), |
| ) |
| ), |
| norm_f( |
| torch.nn.Conv2d( |
| 32, |
| 32, |
| (3, 3), |
| padding=(1, 1), |
| ) |
| ), |
| ] |
| ) |
| self.conv_post = norm_f(torch.nn.Conv2d(32, 1, (3, 3), padding=(1, 1))) |
|
|
| def forward(self, x): |
| fmap = [] |
|
|
| x = self.spectrogram(x).unsqueeze(1) |
|
|
| for layer in self.convs: |
| x = F.leaky_relu(layer(x), self.lrelu_slope) |
| fmap.append(x) |
| x = self.conv_post(x) |
| fmap.append(x) |
|
|
| return torch.flatten(x, 1, -1), fmap |
|
|
| def spectrogram(self, x): |
| n_fft, hop_length, win_length = self.resolution |
| pad = int((n_fft - hop_length) / 2) |
| x = F.pad( |
| x, |
| (pad, pad), |
| mode="reflect", |
| ).squeeze(1) |
| x = torch.stft( |
| x, |
| n_fft=n_fft, |
| hop_length=hop_length, |
| win_length=win_length, |
| window=torch.ones(win_length, device=x.device), |
| center=False, |
| return_complex=True, |
| ) |
|
|
| mag = torch.norm(torch.view_as_real(x), p=2, dim=-1) |
|
|
| return mag |
|
|
