Spaces:
Running
Running
| import torch | |
| import torch.nn as nn # pylint: disable=consider-using-from-import | |
| from TTS.tts.layers.delightful_tts.conv_layers import ConvTransposed | |
| class PhonemeProsodyPredictor(nn.Module): | |
| """Non-parallel Prosody Predictor inspired by: https://arxiv.org/pdf/2102.00851.pdf | |
| It consists of 2 layers of 1D convolutions each followed by a relu activation, layer norm | |
| and dropout, then finally a linear layer. | |
| Args: | |
| hidden_size (int): Size of hidden channels. | |
| kernel_size (int): Kernel size for the conv layers. | |
| dropout: (float): Probability of dropout. | |
| bottleneck_size (int): bottleneck size for last linear layer. | |
| lrelu_slope (float): Slope of the leaky relu. | |
| """ | |
| def __init__( | |
| self, | |
| hidden_size: int, | |
| kernel_size: int, | |
| dropout: float, | |
| bottleneck_size: int, | |
| lrelu_slope: float, | |
| ): | |
| super().__init__() | |
| self.d_model = hidden_size | |
| self.layers = nn.ModuleList( | |
| [ | |
| ConvTransposed( | |
| self.d_model, | |
| self.d_model, | |
| kernel_size=kernel_size, | |
| padding=(kernel_size - 1) // 2, | |
| ), | |
| nn.LeakyReLU(lrelu_slope), | |
| nn.LayerNorm(self.d_model), | |
| nn.Dropout(dropout), | |
| ConvTransposed( | |
| self.d_model, | |
| self.d_model, | |
| kernel_size=kernel_size, | |
| padding=(kernel_size - 1) // 2, | |
| ), | |
| nn.LeakyReLU(lrelu_slope), | |
| nn.LayerNorm(self.d_model), | |
| nn.Dropout(dropout), | |
| ] | |
| ) | |
| self.predictor_bottleneck = nn.Linear(self.d_model, bottleneck_size) | |
| def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Shapes: | |
| x: :math: `[B, T, D]` | |
| mask: :math: `[B, T]` | |
| """ | |
| mask = mask.unsqueeze(2) | |
| for layer in self.layers: | |
| x = layer(x) | |
| x = x.masked_fill(mask, 0.0) | |
| x = self.predictor_bottleneck(x) | |
| return x | |