Spaces:
Sleeping
Sleeping
File size: 4,799 Bytes
9d61c9b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
from typing import Any, Tuple
import torch
from torch import nn
from torch.nn import Module
import torch.nn.functional as F
from torch.nn.utils.parametrizations import spectral_norm, weight_norm
from models.config import VocoderModelConfig
class DiscriminatorR(Module):
r"""A class representing the Residual Discriminator network for a UnivNet vocoder.
Args:
resolution (Tuple): A tuple containing the number of FFT points, hop length, and window length.
model_config (VocoderModelConfig): A configuration object for the UnivNet model.
"""
def __init__(
self,
resolution: Tuple[int, int, int],
model_config: VocoderModelConfig,
):
super().__init__()
self.resolution = resolution
self.LRELU_SLOPE = model_config.mrd.lReLU_slope
# Use spectral normalization or weight normalization based on the configuration
norm_f: Any = (
spectral_norm if model_config.mrd.use_spectral_norm else weight_norm
)
# Define the convolutional layers
self.convs = nn.ModuleList(
[
norm_f(
nn.Conv2d(
1,
32,
(3, 9),
padding=(1, 4),
),
),
norm_f(
nn.Conv2d(
32,
32,
(3, 9),
stride=(1, 2),
padding=(1, 4),
),
),
norm_f(
nn.Conv2d(
32,
32,
(3, 9),
stride=(1, 2),
padding=(1, 4),
),
),
norm_f(
nn.Conv2d(
32,
32,
(3, 9),
stride=(1, 2),
padding=(1, 4),
),
),
norm_f(
nn.Conv2d(
32,
32,
(3, 3),
padding=(1, 1),
),
),
],
)
self.conv_post = norm_f(
nn.Conv2d(
32,
1,
(3, 3),
padding=(1, 1),
),
)
def forward(self, x: torch.Tensor) -> tuple[list[torch.Tensor], torch.Tensor]:
r"""Forward pass of the DiscriminatorR class.
Args:
x (torch.Tensor): The input tensor.
Returns:
tuple: A tuple containing the intermediate feature maps and the output tensor.
"""
fmap = []
# Compute the magnitude spectrogram of the input waveform
x = self.spectrogram(x)
# Add a channel dimension to the spectrogram tensor
x = x.unsqueeze(1)
# Apply the convolutional layers with leaky ReLU activation
for layer in self.convs:
x = layer(x.to(dtype=self.conv_post.weight.dtype))
x = F.leaky_relu(x, self.LRELU_SLOPE)
fmap.append(x)
# Apply the post-convolutional layer
x = self.conv_post(x)
fmap.append(x)
# Flatten the output tensor
x = torch.flatten(x, 1, -1)
return fmap, x
def spectrogram(self, x: torch.Tensor) -> torch.Tensor:
r"""Computes the magnitude spectrogram of the input waveform.
Args:
x (torch.Tensor): Input waveform tensor of shape [B, C, T].
Returns:
torch.Tensor: Magnitude spectrogram tensor of shape [B, F, TT], where F is the number of frequency bins and TT is the number of time frames.
"""
n_fft, hop_length, win_length = self.resolution
# Apply reflection padding to the input waveform
x = F.pad(
x,
(int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
mode="reflect",
)
# Squeeze the input waveform to remove the channel dimension
x = x.squeeze(1)
# Compute the short-time Fourier transform of the input waveform
x = torch.stft(
x,
n_fft=n_fft,
hop_length=hop_length,
win_length=win_length,
center=False,
return_complex=True,
window=torch.ones(win_length, device=x.device),
) # [B, F, TT, 2]
x = torch.view_as_real(x)
# Compute the magnitude spectrogram from the complex spectrogram
return torch.norm(x, p=2, dim=-1) # [B, F, TT]
|