|
import torch |
|
|
|
from Layers.ResidualBlock import HiFiGANResidualBlock as ResidualBlock |
|
|
|
|
|
class HiFiGANGenerator(torch.nn.Module): |
|
|
|
def __init__(self, |
|
path_to_weights, |
|
in_channels=80, |
|
out_channels=1, |
|
channels=512, |
|
kernel_size=7, |
|
upsample_scales=(8, 6, 4, 4), |
|
upsample_kernel_sizes=(16, 12, 8, 8), |
|
resblock_kernel_sizes=(3, 7, 11), |
|
resblock_dilations=[(1, 3, 5), (1, 3, 5), (1, 3, 5)], |
|
use_additional_convs=True, |
|
bias=True, |
|
nonlinear_activation="LeakyReLU", |
|
nonlinear_activation_params={"negative_slope": 0.1}, |
|
use_weight_norm=True, ): |
|
super().__init__() |
|
assert kernel_size % 2 == 1, "Kernal size must be odd number." |
|
assert len(upsample_scales) == len(upsample_kernel_sizes) |
|
assert len(resblock_dilations) == len(resblock_kernel_sizes) |
|
self.num_upsamples = len(upsample_kernel_sizes) |
|
self.num_blocks = len(resblock_kernel_sizes) |
|
self.input_conv = torch.nn.Conv1d(in_channels, |
|
channels, |
|
kernel_size, |
|
1, |
|
padding=(kernel_size - 1) // 2, ) |
|
self.upsamples = torch.nn.ModuleList() |
|
self.blocks = torch.nn.ModuleList() |
|
for i in range(len(upsample_kernel_sizes)): |
|
self.upsamples += [ |
|
torch.nn.Sequential(getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params), |
|
torch.nn.ConvTranspose1d(channels // (2 ** i), |
|
channels // (2 ** (i + 1)), |
|
upsample_kernel_sizes[i], |
|
upsample_scales[i], |
|
padding=(upsample_kernel_sizes[i] - upsample_scales[i]) // 2, ), )] |
|
for j in range(len(resblock_kernel_sizes)): |
|
self.blocks += [ResidualBlock(kernel_size=resblock_kernel_sizes[j], |
|
channels=channels // (2 ** (i + 1)), |
|
dilations=resblock_dilations[j], |
|
bias=bias, |
|
use_additional_convs=use_additional_convs, |
|
nonlinear_activation=nonlinear_activation, |
|
nonlinear_activation_params=nonlinear_activation_params, )] |
|
self.output_conv = torch.nn.Sequential( |
|
torch.nn.LeakyReLU(), |
|
torch.nn.Conv1d(channels // (2 ** (i + 1)), |
|
out_channels, |
|
kernel_size, |
|
1, |
|
padding=(kernel_size - 1) // 2, ), |
|
torch.nn.Tanh(), ) |
|
if use_weight_norm: |
|
self.apply_weight_norm() |
|
self.load_state_dict(torch.load(path_to_weights, map_location='cpu')["generator"]) |
|
|
|
def forward(self, c, normalize_before=False): |
|
if normalize_before: |
|
c = (c - self.mean) / self.scale |
|
c = self.input_conv(c.unsqueeze(0)) |
|
for i in range(self.num_upsamples): |
|
c = self.upsamples[i](c) |
|
cs = 0.0 |
|
for j in range(self.num_blocks): |
|
cs = cs + self.blocks[i * self.num_blocks + j](c) |
|
c = cs / self.num_blocks |
|
c = self.output_conv(c) |
|
return c.squeeze(0).squeeze(0) |
|
|
|
def remove_weight_norm(self): |
|
def _remove_weight_norm(m): |
|
try: |
|
torch.nn.utils.remove_weight_norm(m) |
|
except ValueError: |
|
return |
|
|
|
self.apply(_remove_weight_norm) |
|
|
|
def apply_weight_norm(self): |
|
def _apply_weight_norm(m): |
|
if isinstance(m, torch.nn.Conv1d) or isinstance(m, torch.nn.ConvTranspose1d): |
|
torch.nn.utils.weight_norm(m) |
|
|
|
self.apply(_apply_weight_norm) |
|
|
