Spaces:
Build error
Build error
| import torch | |
| from torch import nn | |
| from .espnet_positional_embedding import RelPositionalEncoding | |
| from .espnet_transformer_attn import RelPositionMultiHeadedAttention, MultiHeadedAttention | |
| from .layers import Swish, ConvolutionModule, EncoderLayer, MultiLayeredConv1d | |
| from ..layers import Embedding | |
| def sequence_mask(length, max_length=None): | |
| if max_length is None: | |
| max_length = length.max() | |
| x = torch.arange(max_length, dtype=length.dtype, device=length.device) | |
| return x.unsqueeze(0) < length.unsqueeze(1) | |
| class ConformerLayers(nn.Module): | |
| def __init__(self, hidden_size, num_layers, kernel_size=9, dropout=0.0, num_heads=4, use_last_norm=True): | |
| super().__init__() | |
| self.use_last_norm = use_last_norm | |
| self.layers = nn.ModuleList() | |
| positionwise_layer = MultiLayeredConv1d | |
| positionwise_layer_args = (hidden_size, hidden_size * 4, 1, dropout) | |
| self.encoder_layers = nn.ModuleList([EncoderLayer( | |
| hidden_size, | |
| MultiHeadedAttention(num_heads, hidden_size, 0.0), | |
| positionwise_layer(*positionwise_layer_args), | |
| positionwise_layer(*positionwise_layer_args), | |
| ConvolutionModule(hidden_size, kernel_size, Swish()), | |
| dropout, | |
| ) for _ in range(num_layers)]) | |
| if self.use_last_norm: | |
| self.layer_norm = nn.LayerNorm(hidden_size) | |
| else: | |
| self.layer_norm = nn.Linear(hidden_size, hidden_size) | |
| def forward(self, x, x_mask): | |
| """ | |
| :param x: [B, T, H] | |
| :param padding_mask: [B, T] | |
| :return: [B, T, H] | |
| """ | |
| for l in self.encoder_layers: | |
| x, mask = l(x, x_mask) | |
| x = self.layer_norm(x) * x_mask | |
| return x | |
| class ConformerEncoder(ConformerLayers): | |
| def __init__(self, hidden_size, dict_size=0, in_size=0, strides=[2,2], num_layers=None): | |
| conformer_enc_kernel_size = 9 | |
| super().__init__(hidden_size, num_layers, conformer_enc_kernel_size) | |
| self.dict_size = dict_size | |
| if dict_size != 0: | |
| self.embed = Embedding(dict_size, hidden_size, padding_idx=0) | |
| else: | |
| self.seq_proj_in = torch.nn.Linear(in_size, hidden_size) | |
| self.seq_proj_out = torch.nn.Linear(hidden_size, in_size) | |
| self.mel_in = torch.nn.Linear(160, hidden_size) | |
| self.mel_pre_net = torch.nn.Sequential(*[ | |
| torch.nn.Conv1d(hidden_size, hidden_size, kernel_size=s * 2, stride=s, padding=s // 2) | |
| for i, s in enumerate(strides) | |
| ]) | |
| def forward(self, seq_out, mels_timbre, other_embeds=0): | |
| """ | |
| :param src_tokens: [B, T] | |
| :return: [B x T x C] | |
| """ | |
| x_lengths = (seq_out > 0).long().sum(-1) | |
| x = seq_out | |
| if self.dict_size != 0: | |
| x = self.embed(x) + other_embeds # [B, T, H] | |
| else: | |
| x = self.seq_proj_in(x) + other_embeds # [B, T, H] | |
| mels_timbre = self.mel_in(mels_timbre).transpose(1, 2) | |
| mels_timbre = self.mel_pre_net(mels_timbre).transpose(1, 2) | |
| T_out = x.size(1) | |
| if self.dict_size != 0: | |
| x_mask = torch.unsqueeze(sequence_mask(x_lengths + mels_timbre.size(1), x.size(1) + mels_timbre.size(1)), 2).to(x.dtype) | |
| else: | |
| x_mask = torch.cat((torch.ones(x.size(0), mels_timbre.size(1), 1).to(x.device), (x.abs().sum(2) > 0).float()[:, :, None]), dim=1) | |
| x = torch.cat((mels_timbre, x), 1) | |
| x = super(ConformerEncoder, self).forward(x, x_mask) | |
| if self.dict_size != 0: | |
| x = x[:, -T_out:, :] | |
| else: | |
| x = self.seq_proj_out(x[:, -T_out:, :]) | |
| return x | |
| class ConformerDecoder(ConformerLayers): | |
| def __init__(self, hidden_size, num_layers): | |
| conformer_dec_kernel_size = 9 | |
| super().__init__(hidden_size, num_layers, conformer_dec_kernel_size) | |