Spaces:
Running
Running
File size: 2,189 Bytes
10f957b |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm, spectral_norm
class DiscriminatorR(torch.nn.Module):
def __init__(self, hp, resolution):
super(DiscriminatorR, self).__init__()
self.resolution = resolution
self.LRELU_SLOPE = hp.mpd.lReLU_slope
norm_f = weight_norm if hp.mrd.use_spectral_norm == False else spectral_norm
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):
fmap = []
x = self.spectrogram(x)
x = x.unsqueeze(1)
for l in self.convs:
x = l(x)
x = F.leaky_relu(x, self.LRELU_SLOPE)
fmap.append(x)
x = self.conv_post(x)
fmap.append(x)
x = torch.flatten(x, 1, -1)
return fmap, x
def spectrogram(self, x):
n_fft, hop_length, win_length = self.resolution
x = F.pad(x, (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)), mode='reflect')
x = x.squeeze(1)
x = torch.stft(x, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False, return_complex=False) #[B, F, TT, 2]
mag = torch.norm(x, p=2, dim =-1) #[B, F, TT]
return mag
class MultiResolutionDiscriminator(torch.nn.Module):
def __init__(self, hp):
super(MultiResolutionDiscriminator, self).__init__()
self.resolutions = eval(hp.mrd.resolutions)
self.discriminators = nn.ModuleList(
[DiscriminatorR(hp, resolution) for resolution in self.resolutions]
)
def forward(self, x):
ret = list()
for disc in self.discriminators:
ret.append(disc(x))
return ret # [(feat, score), (feat, score), (feat, score)]
|