mirror of https://github.com/coqui-ai/TTS.git
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
from torch import nn
|
|
from TTS.vocoder.layers.hifigan import MRF
|
|
|
|
|
|
class Generator(nn.Module):
|
|
|
|
def __init__(self, input_channel=80, hu=512, ku=[16, 16, 4, 4], kr=[3, 7, 11], Dr=[1, 3, 5]):
|
|
super(Generator, self).__init__()
|
|
self.input = nn.Sequential(
|
|
nn.ReflectionPad1d(3),
|
|
nn.utils.weight_norm(nn.Conv1d(input_channel, hu, kernel_size=7))
|
|
)
|
|
|
|
generator = []
|
|
|
|
for k in ku:
|
|
inp = hu
|
|
out = int(inp / 2)
|
|
generator += [
|
|
nn.LeakyReLU(0.2),
|
|
nn.ConvTranspose1d(inp, out, k, k // 2),
|
|
MRF(kr, out, Dr)
|
|
]
|
|
hu = out
|
|
self.generator = nn.Sequential(*generator)
|
|
|
|
self.output = nn.Sequential(
|
|
nn.LeakyReLU(0.2),
|
|
nn.ReflectionPad1d(3),
|
|
nn.utils.weight_norm(nn.Conv1d(hu, 1, kernel_size=7, stride=1)),
|
|
nn.Tanh()
|
|
|
|
)
|
|
|
|
def forward(self, x):
|
|
x1 = self.input(x)
|
|
x2 = self.generator(x1)
|
|
out = self.output(x2)
|
|
return out |