Spaces:
Running
Running
| import torch | |
| from torch import nn | |
| from torch.nn.utils.parametrizations import weight_norm | |
| from TTS.utils.io import load_fsspec | |
| from TTS.vocoder.layers.melgan import ResidualStack | |
| class MelganGenerator(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels=80, | |
| out_channels=1, | |
| proj_kernel=7, | |
| base_channels=512, | |
| upsample_factors=(8, 8, 2, 2), | |
| res_kernel=3, | |
| num_res_blocks=3, | |
| ): | |
| super().__init__() | |
| # assert model parameters | |
| assert (proj_kernel - 1) % 2 == 0, " [!] proj_kernel should be an odd number." | |
| # setup additional model parameters | |
| base_padding = (proj_kernel - 1) // 2 | |
| act_slope = 0.2 | |
| self.inference_padding = 2 | |
| # initial layer | |
| layers = [] | |
| layers += [ | |
| nn.ReflectionPad1d(base_padding), | |
| weight_norm(nn.Conv1d(in_channels, base_channels, kernel_size=proj_kernel, stride=1, bias=True)), | |
| ] | |
| # upsampling layers and residual stacks | |
| for idx, upsample_factor in enumerate(upsample_factors): | |
| layer_in_channels = base_channels // (2**idx) | |
| layer_out_channels = base_channels // (2 ** (idx + 1)) | |
| layer_filter_size = upsample_factor * 2 | |
| layer_stride = upsample_factor | |
| layer_output_padding = upsample_factor % 2 | |
| layer_padding = upsample_factor // 2 + layer_output_padding | |
| layers += [ | |
| nn.LeakyReLU(act_slope), | |
| weight_norm( | |
| nn.ConvTranspose1d( | |
| layer_in_channels, | |
| layer_out_channels, | |
| layer_filter_size, | |
| stride=layer_stride, | |
| padding=layer_padding, | |
| output_padding=layer_output_padding, | |
| bias=True, | |
| ) | |
| ), | |
| ResidualStack(channels=layer_out_channels, num_res_blocks=num_res_blocks, kernel_size=res_kernel), | |
| ] | |
| layers += [nn.LeakyReLU(act_slope)] | |
| # final layer | |
| layers += [ | |
| nn.ReflectionPad1d(base_padding), | |
| weight_norm(nn.Conv1d(layer_out_channels, out_channels, proj_kernel, stride=1, bias=True)), | |
| nn.Tanh(), | |
| ] | |
| self.layers = nn.Sequential(*layers) | |
| def forward(self, c): | |
| return self.layers(c) | |
| def inference(self, c): | |
| c = c.to(self.layers[1].weight.device) | |
| c = torch.nn.functional.pad(c, (self.inference_padding, self.inference_padding), "replicate") | |
| return self.layers(c) | |
| def remove_weight_norm(self): | |
| for _, layer in enumerate(self.layers): | |
| if len(layer.state_dict()) != 0: | |
| try: | |
| nn.utils.parametrize.remove_parametrizations(layer, "weight") | |
| except ValueError: | |
| layer.remove_weight_norm() | |
| def load_checkpoint( | |
| self, config, checkpoint_path, eval=False, cache=False | |
| ): # pylint: disable=unused-argument, redefined-builtin | |
| state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache) | |
| self.load_state_dict(state["model"]) | |
| if eval: | |
| self.eval() | |
| assert not self.training | |
| self.remove_weight_norm() | |