| |
| |
| |
|
|
| import torch |
| from einops import rearrange |
| from torch import nn |
| from torch.nn.utils.parametrizations import weight_norm |
|
|
|
|
| def get_2d_padding(kernel_size: tuple[int, int], dilation: tuple[int, int] = (1, 1)): |
| return (((kernel_size[0] - 1) * dilation[0]) // 2, ((kernel_size[1] - 1) * dilation[1]) // 2) |
|
|
|
|
| class SpectrogramDiscriminator(nn.Module): |
| def __init__( |
| self, |
| frequency_bins: int, |
| channels: int = 32, |
| kernel_size: tuple[int, int] = (3, 3), |
| dilation: list[int] = [1, 2, 4], |
| bands: tuple[tuple[float, float], ...] = ((0.0, 0.2), (0.2, 0.4), (0.4, 0.6), (0.6, 0.8), (0.8, 1.0)), |
| use_downsample: bool = True, |
| ): |
| super().__init__() |
| self.bands = [(int(b[0] * frequency_bins), int(b[1] * frequency_bins)) for b in bands] |
|
|
| self.stacks = nn.ModuleList() |
| for _ in self.bands: |
| stack = nn.ModuleList( |
| [weight_norm(nn.Conv2d(1, channels, kernel_size, padding=get_2d_padding(kernel_size)))] |
| ) |
|
|
| for d in dilation: |
| |
| pad = get_2d_padding(kernel_size, (d, 1)) |
| stack.append(weight_norm(nn.Conv2d(channels, channels, kernel_size, dilation=(d, 1), padding=pad))) |
|
|
| stack.append(weight_norm(nn.Conv2d(channels, channels, kernel_size, padding=get_2d_padding(kernel_size)))) |
|
|
| self.stacks.append(stack) |
|
|
| self.conv_post = weight_norm(nn.Conv2d(channels, 1, kernel_size, padding=get_2d_padding(kernel_size))) |
| if use_downsample: |
| self.downsample = nn.AvgPool2d(4, stride=2, padding=1, count_include_pad=False) |
| else: |
| self.downsample = nn.Identity() |
|
|
| def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: |
| """ |
| Args: |
| x (Tensor): Input spectrogram (B, C, F, T). |
| Returns: |
| output (Tensor): Discriminator output. |
| intermediates (list[Tensor]): List of intermediate feature maps. |
| """ |
| if x.dim() == 3: |
| x = x.unsqueeze(1) |
| assert x.dim() == 4, f"Expected 4D input, got {x.dim()}D" |
|
|
| |
| x = rearrange(x, "b c f t -> b c t f") |
| x_bands = [x[..., b[0] : b[1]] for b in self.bands] |
|
|
| x = [] |
| intermediates = [] |
| for x_band, stack in zip(x_bands, self.stacks): |
| for layer in stack: |
| x_band = layer(x_band) |
| x_band = torch.nn.functional.leaky_relu(x_band, 0.1) |
| intermediates.append(x_band) |
| x.append(x_band) |
|
|
| |
| x = torch.cat(x, dim=-1) |
| x = self.conv_post(x) |
| x = self.downsample(x) |
| return x, intermediates |
|
|
