Spaces:
Running
Running
| from typing import List, Tuple, Union | |
| import torch | |
| import torch.nn as nn # pylint: disable=consider-using-from-import | |
| import torch.nn.functional as F | |
| from TTS.tts.layers.delightful_tts.conformer import ConformerMultiHeadedSelfAttention | |
| from TTS.tts.layers.delightful_tts.conv_layers import CoordConv1d | |
| from TTS.tts.layers.delightful_tts.networks import STL | |
| def get_mask_from_lengths(lengths: torch.Tensor) -> torch.Tensor: | |
| batch_size = lengths.shape[0] | |
| max_len = torch.max(lengths).item() | |
| ids = torch.arange(0, max_len, device=lengths.device).unsqueeze(0).expand(batch_size, -1) | |
| mask = ids >= lengths.unsqueeze(1).expand(-1, max_len) | |
| return mask | |
| def stride_lens(lens: torch.Tensor, stride: int = 2) -> torch.Tensor: | |
| return torch.ceil(lens / stride).int() | |
| class ReferenceEncoder(nn.Module): | |
| """ | |
| Referance encoder for utterance and phoneme prosody encoders. Reference encoder | |
| made up of convolution and RNN layers. | |
| Args: | |
| num_mels (int): Number of mel frames to produce. | |
| ref_enc_filters (list[int]): List of channel sizes for encoder layers. | |
| ref_enc_size (int): Size of the kernel for the conv layers. | |
| ref_enc_strides (List[int]): List of strides to use for conv layers. | |
| ref_enc_gru_size (int): Number of hidden features for the gated recurrent unit. | |
| Inputs: inputs, mask | |
| - **inputs** (batch, dim, time): Tensor containing mel vector | |
| - **lengths** (batch): Tensor containing the mel lengths. | |
| Returns: | |
| - **outputs** (batch, time, dim): Tensor produced by Reference Encoder. | |
| """ | |
| def __init__( | |
| self, | |
| num_mels: int, | |
| ref_enc_filters: List[Union[int, int, int, int, int, int]], | |
| ref_enc_size: int, | |
| ref_enc_strides: List[Union[int, int, int, int, int]], | |
| ref_enc_gru_size: int, | |
| ): | |
| super().__init__() | |
| n_mel_channels = num_mels | |
| self.n_mel_channels = n_mel_channels | |
| K = len(ref_enc_filters) | |
| filters = [self.n_mel_channels] + ref_enc_filters | |
| strides = [1] + ref_enc_strides | |
| # Use CoordConv at the first layer to better preserve positional information: https://arxiv.org/pdf/1811.02122.pdf | |
| convs = [ | |
| CoordConv1d( | |
| in_channels=filters[0], | |
| out_channels=filters[0 + 1], | |
| kernel_size=ref_enc_size, | |
| stride=strides[0], | |
| padding=ref_enc_size // 2, | |
| with_r=True, | |
| ) | |
| ] | |
| convs2 = [ | |
| nn.Conv1d( | |
| in_channels=filters[i], | |
| out_channels=filters[i + 1], | |
| kernel_size=ref_enc_size, | |
| stride=strides[i], | |
| padding=ref_enc_size // 2, | |
| ) | |
| for i in range(1, K) | |
| ] | |
| convs.extend(convs2) | |
| self.convs = nn.ModuleList(convs) | |
| self.norms = nn.ModuleList([nn.InstanceNorm1d(num_features=ref_enc_filters[i], affine=True) for i in range(K)]) | |
| self.gru = nn.GRU( | |
| input_size=ref_enc_filters[-1], | |
| hidden_size=ref_enc_gru_size, | |
| batch_first=True, | |
| ) | |
| def forward(self, x: torch.Tensor, mel_lens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """ | |
| inputs --- [N, n_mels, timesteps] | |
| outputs --- [N, E//2] | |
| """ | |
| mel_masks = get_mask_from_lengths(mel_lens).unsqueeze(1) | |
| x = x.masked_fill(mel_masks, 0) | |
| for conv, norm in zip(self.convs, self.norms): | |
| x = conv(x) | |
| x = F.leaky_relu(x, 0.3) # [N, 128, Ty//2^K, n_mels//2^K] | |
| x = norm(x) | |
| for _ in range(2): | |
| mel_lens = stride_lens(mel_lens) | |
| mel_masks = get_mask_from_lengths(mel_lens) | |
| x = x.masked_fill(mel_masks.unsqueeze(1), 0) | |
| x = x.permute((0, 2, 1)) | |
| x = torch.nn.utils.rnn.pack_padded_sequence(x, mel_lens.cpu().int(), batch_first=True, enforce_sorted=False) | |
| self.gru.flatten_parameters() | |
| x, memory = self.gru(x) # memory --- [N, Ty, E//2], out --- [1, N, E//2] | |
| x, _ = torch.nn.utils.rnn.pad_packed_sequence(x, batch_first=True) | |
| return x, memory, mel_masks | |
| def calculate_channels( # pylint: disable=no-self-use | |
| self, L: int, kernel_size: int, stride: int, pad: int, n_convs: int | |
| ) -> int: | |
| for _ in range(n_convs): | |
| L = (L - kernel_size + 2 * pad) // stride + 1 | |
| return L | |
| class UtteranceLevelProsodyEncoder(nn.Module): | |
| def __init__( | |
| self, | |
| num_mels: int, | |
| ref_enc_filters: List[Union[int, int, int, int, int, int]], | |
| ref_enc_size: int, | |
| ref_enc_strides: List[Union[int, int, int, int, int]], | |
| ref_enc_gru_size: int, | |
| dropout: float, | |
| n_hidden: int, | |
| bottleneck_size_u: int, | |
| token_num: int, | |
| ): | |
| """ | |
| Encoder to extract prosody from utterance. it is made up of a reference encoder | |
| with a couple of linear layers and style token layer with dropout. | |
| Args: | |
| num_mels (int): Number of mel frames to produce. | |
| ref_enc_filters (list[int]): List of channel sizes for ref encoder layers. | |
| ref_enc_size (int): Size of the kernel for the ref encoder conv layers. | |
| ref_enc_strides (List[int]): List of strides to use for teh ref encoder conv layers. | |
| ref_enc_gru_size (int): Number of hidden features for the gated recurrent unit. | |
| dropout (float): Probability of dropout. | |
| n_hidden (int): Size of hidden layers. | |
| bottleneck_size_u (int): Size of the bottle neck layer. | |
| Inputs: inputs, mask | |
| - **inputs** (batch, dim, time): Tensor containing mel vector | |
| - **lengths** (batch): Tensor containing the mel lengths. | |
| Returns: | |
| - **outputs** (batch, 1, dim): Tensor produced by Utterance Level Prosody Encoder. | |
| """ | |
| super().__init__() | |
| self.E = n_hidden | |
| self.d_q = self.d_k = n_hidden | |
| bottleneck_size = bottleneck_size_u | |
| self.encoder = ReferenceEncoder( | |
| ref_enc_filters=ref_enc_filters, | |
| ref_enc_gru_size=ref_enc_gru_size, | |
| ref_enc_size=ref_enc_size, | |
| ref_enc_strides=ref_enc_strides, | |
| num_mels=num_mels, | |
| ) | |
| self.encoder_prj = nn.Linear(ref_enc_gru_size, self.E // 2) | |
| self.stl = STL(n_hidden=n_hidden, token_num=token_num) | |
| self.encoder_bottleneck = nn.Linear(self.E, bottleneck_size) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, mels: torch.Tensor, mel_lens: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Shapes: | |
| mels: :math: `[B, C, T]` | |
| mel_lens: :math: `[B]` | |
| out --- [N, seq_len, E] | |
| """ | |
| _, embedded_prosody, _ = self.encoder(mels, mel_lens) | |
| # Bottleneck | |
| embedded_prosody = self.encoder_prj(embedded_prosody) | |
| # Style Token | |
| out = self.encoder_bottleneck(self.stl(embedded_prosody)) | |
| out = self.dropout(out) | |
| out = out.view((-1, 1, out.shape[3])) | |
| return out | |
| class PhonemeLevelProsodyEncoder(nn.Module): | |
| def __init__( | |
| self, | |
| num_mels: int, | |
| ref_enc_filters: List[Union[int, int, int, int, int, int]], | |
| ref_enc_size: int, | |
| ref_enc_strides: List[Union[int, int, int, int, int]], | |
| ref_enc_gru_size: int, | |
| dropout: float, | |
| n_hidden: int, | |
| n_heads: int, | |
| bottleneck_size_p: int, | |
| ): | |
| super().__init__() | |
| self.E = n_hidden | |
| self.d_q = self.d_k = n_hidden | |
| bottleneck_size = bottleneck_size_p | |
| self.encoder = ReferenceEncoder( | |
| ref_enc_filters=ref_enc_filters, | |
| ref_enc_gru_size=ref_enc_gru_size, | |
| ref_enc_size=ref_enc_size, | |
| ref_enc_strides=ref_enc_strides, | |
| num_mels=num_mels, | |
| ) | |
| self.encoder_prj = nn.Linear(ref_enc_gru_size, n_hidden) | |
| self.attention = ConformerMultiHeadedSelfAttention( | |
| d_model=n_hidden, | |
| num_heads=n_heads, | |
| dropout_p=dropout, | |
| ) | |
| self.encoder_bottleneck = nn.Linear(n_hidden, bottleneck_size) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| src_mask: torch.Tensor, | |
| mels: torch.Tensor, | |
| mel_lens: torch.Tensor, | |
| encoding: torch.Tensor, | |
| ) -> torch.Tensor: | |
| """ | |
| x --- [N, seq_len, encoder_embedding_dim] | |
| mels --- [N, Ty/r, n_mels*r], r=1 | |
| out --- [N, seq_len, bottleneck_size] | |
| attn --- [N, seq_len, ref_len], Ty/r = ref_len | |
| """ | |
| embedded_prosody, _, mel_masks = self.encoder(mels, mel_lens) | |
| # Bottleneck | |
| embedded_prosody = self.encoder_prj(embedded_prosody) | |
| attn_mask = mel_masks.view((mel_masks.shape[0], 1, 1, -1)) | |
| x, _ = self.attention( | |
| query=x, | |
| key=embedded_prosody, | |
| value=embedded_prosody, | |
| mask=attn_mask, | |
| encoding=encoding, | |
| ) | |
| x = self.encoder_bottleneck(x) | |
| x = x.masked_fill(src_mask.unsqueeze(-1), 0.0) | |
| return x | |