initial commits

TTSInferencing.py

@@ -0,0 +1,267 @@
+
+import re
+import logging
+import torch
+import torchaudio
+import random
+import speechbrain
+from speechbrain.inference.interfaces import Pretrained
+from speechbrain.inference.text import GraphemeToPhoneme
+
+logger = logging.getLogger(__name__)
+
+class TTSInferencing(Pretrained):
+    """
+    A ready-to-use wrapper for TTS (text -> mel_spec).
+    Arguments
+    ---------
+    hparams
+        Hyperparameters (from HyperPyYAML)
+    """
+
+    HPARAMS_NEEDED = ["modules", "input_encoder"]
+
+    MODULES_NEEDED = ["encoder_prenet", "pos_emb_enc",
+                      "decoder_prenet", "pos_emb_dec",
+                      "Seq2SeqTransformer", "mel_lin",
+                      "stop_lin", "decoder_postnet"]
+
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        lexicon = self.hparams.lexicon
+        lexicon = ["@@"] + lexicon
+        self.input_encoder = self.hparams.input_encoder
+        self.input_encoder.update_from_iterable(lexicon, sequence_input=False)
+        self.input_encoder.add_unk()
+
+        self.modules = self.hparams.modules
+
+        self.g2p = GraphemeToPhoneme.from_hparams("speechbrain/soundchoice-g2p")
+
+
+
+
+    def generate_padded_phonemes(self, texts):
+        """Computes mel-spectrogram for a list of texts
+
+        Arguments
+        ---------
+        texts: List[str]
+            texts to be converted to spectrogram
+
+        Returns
+        -------
+        tensors of output spectrograms
+        """
+
+        # Preprocessing required at the inference time for the input text
+        # "label" below contains input text
+        # "phoneme_labels" contain the phoneme sequences corresponding to input text labels
+
+        phoneme_labels = list()
+
+        for label in texts:
+
+          phoneme_label = list()
+
+          label = self.custom_clean(label).upper()
+
+          words = label.split()
+          words = [word.strip() for word in words]
+          words_phonemes = self.g2p(words)
+
+          for i in range(len(words_phonemes)):
+              words_phonemes_seq = words_phonemes[i]
+              for phoneme in words_phonemes_seq:
+                  if not phoneme.isspace():
+                      phoneme_label.append(phoneme)
+          phoneme_labels.append(phoneme_label)
+
+
+        # encode the phonemes with input text encoder
+        encoded_phonemes = list()
+        for i in range(len(phoneme_labels)):
+            phoneme_label = phoneme_labels[i]
+            encoded_phoneme =  torch.LongTensor(self.input_encoder.encode_sequence(phoneme_label)).to(self.device)
+            encoded_phonemes.append(encoded_phoneme)
+
+
+        # Right zero-pad all one-hot text sequences to max input length
+        input_lengths, ids_sorted_decreasing = torch.sort(
+            torch.LongTensor([len(x) for x in encoded_phonemes]), dim=0, descending=True
+        )
+
+        max_input_len = input_lengths[0]
+
+        phoneme_padded = torch.LongTensor(len(encoded_phonemes), max_input_len).to(self.device)
+        phoneme_padded.zero_()
+
+        for seq_idx, seq in enumerate(encoded_phonemes):
+            phoneme_padded[seq_idx, : len(seq)] = seq
+
+
+        return phoneme_padded.to(self.device, non_blocking=True).float()
+
+
+    def encode_batch(self, texts):
+        """Computes mel-spectrogram for a list of texts
+
+        Texts must be sorted in decreasing order on their lengths
+
+        Arguments
+        ---------
+        texts: List[str]
+            texts to be encoded into spectrogram
+
+        Returns
+        -------
+        tensors of output spectrograms
+        """
+
+        # generate phonemes and padd the input texts
+        encoded_phoneme_padded = self.generate_padded_phonemes(texts)
+        phoneme_prenet_emb = self.modules['encoder_prenet'](encoded_phoneme_padded)
+        # Positional Embeddings
+        phoneme_pos_emb =  self.modules['pos_emb_enc'](encoded_phoneme_padded)
+        # Summing up embeddings
+        enc_phoneme_emb = phoneme_prenet_emb.permute(0,2,1)  + phoneme_pos_emb
+        enc_phoneme_emb = enc_phoneme_emb.to(self.device)
+
+
+        with torch.no_grad():
+
+          # generate sequential predictions via transformer decoder
+          start_token = torch.full((80, 1), fill_value= 0)
+          start_token[1] = 2
+          decoder_input = start_token.repeat(enc_phoneme_emb.size(0), 1, 1)
+          decoder_input = decoder_input.to(self.device, non_blocking=True).float()
+
+          num_itr = 0
+          stop_condition = [False] * decoder_input.size(0)
+          max_iter = 100
+
+          # while not all(stop_condition) and num_itr < max_iter:
+          while num_itr < max_iter:
+
+            # Decoder Prenet
+            mel_prenet_emb =  self.modules['decoder_prenet'](decoder_input).to(self.device).permute(0,2,1)
+
+            # Positional Embeddings
+            mel_pos_emb =  self.modules['pos_emb_dec'](mel_prenet_emb).to(self.device)
+            # Summing up Embeddings
+            dec_mel_spec = mel_prenet_emb + mel_pos_emb
+
+            # Getting the target mask to avoid looking ahead
+            tgt_mask = self.hparams.lookahead_mask(dec_mel_spec).to(self.device)
+
+            # Getting the source mask
+            src_mask = torch.zeros(enc_phoneme_emb.shape[1], enc_phoneme_emb.shape[1]).to(self.device)
+
+            # Padding masks for source and targets
+            src_key_padding_mask = self.hparams.padding_mask(enc_phoneme_emb, pad_idx = self.hparams.blank_index).to(self.device)
+            tgt_key_padding_mask = self.hparams.padding_mask(dec_mel_spec, pad_idx = self.hparams.blank_index).to(self.device)
+
+
+            # Running the Seq2Seq Transformer
+            decoder_outputs = self.modules['Seq2SeqTransformer'](src = enc_phoneme_emb, tgt = dec_mel_spec, src_mask = src_mask, tgt_mask = tgt_mask,
+                                                              src_key_padding_mask = src_key_padding_mask, tgt_key_padding_mask = tgt_key_padding_mask)
+
+            # Mel Linears
+            mel_linears =  self.modules['mel_lin'](decoder_outputs).permute(0,2,1)
+            mel_postnet = self.modules['decoder_postnet'](mel_linears) # mel tensor output
+            mel_pred = mel_linears + mel_postnet # mel tensor output
+
+            stop_token_pred =  self.modules['stop_lin'](decoder_outputs).squeeze(-1)
+
+            stop_condition_list = self.check_stop_condition(stop_token_pred)
+
+
+            # update the values of main stop conditions
+            stop_condition_update = [True if stop_condition_list[i] else stop_condition[i] for i in range(len(stop_condition))]
+            stop_condition = stop_condition_update
+
+
+            # Prepare input for the transformer input for next iteration
+            current_output = mel_pred[:, :, -1:]
+
+            decoder_input=torch.cat([decoder_input,current_output],dim=2)
+            num_itr = num_itr+1
+
+        mel_outputs =  decoder_input[:, :, 1:]
+
+        return mel_outputs
+
+
+
+    def encode_text(self, text):
+        """Runs inference for a single text str"""
+        return self.encode_batch([text])
+
+
+    def forward(self, text_list):
+        "Encodes the input texts."
+        return self.encode_batch(text_list)
+
+
+    def check_stop_condition(self, stop_token_pred):
+        """
+        check if stop token / EOS reached or not for mel_specs in the batch
+        """
+
+        # Applying sigmoid to perform binary classification
+        sigmoid_output = torch.sigmoid(stop_token_pred)
+        # Checking if the probability is greater than 0.5
+        stop_results = sigmoid_output > 0.8
+        stop_output = [all(result) for result in stop_results]
+
+        return stop_output
+
+
+
+    def custom_clean(self, text):
+        """
+        Uses custom criteria to clean text.
+
+        Arguments
+        ---------
+        text : str
+            Input text to be cleaned
+        model_name : str
+            whether to treat punctuations
+
+        Returns
+        -------
+        text : str
+            Cleaned text
+        """
+
+        _abbreviations = [
+            (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+            for x in [
+                ("mrs", "missus"),
+                ("mr", "mister"),
+                ("dr", "doctor"),
+                ("st", "saint"),
+                ("co", "company"),
+                ("jr", "junior"),
+                ("maj", "major"),
+                ("gen", "general"),
+                ("drs", "doctors"),
+                ("rev", "reverend"),
+                ("lt", "lieutenant"),
+                ("hon", "honorable"),
+                ("sgt", "sergeant"),
+                ("capt", "captain"),
+                ("esq", "esquire"),
+                ("ltd", "limited"),
+                ("col", "colonel"),
+                ("ft", "fort"),
+            ]
+        ]
+
+        text = re.sub(" +", " ", text)
+
+        for regex, replacement in _abbreviations:
+            text = re.sub(regex, replacement, text)
+        return text

hyperparams.yaml

@@ -0,0 +1,187 @@
+
+############################################################################
+# Model: TTS with attention-based mechanism
+# Tokens: g2p + possitional embeddings
+# losses: MSE & BCE
+# Training: LJSpeech
+# ############################################################################
+
+###################################
+# Experiment Parameters and setup #
+###################################
+seed: 1234
+__set_seed: !apply:torch.manual_seed [!ref <seed>]
+
+# Folder set up
+# output_folder: !ref .\\results\\tts\\<seed>
+# save_folder: !ref <output_folder>\\save
+
+output_folder: !ref ./results/<seed>
+save_folder: !ref <output_folder>/save
+
+
+################################
+# Model Parameters and model   #
+################################
+# Input parameters
+lexicon:
+    - AA
+    - AE
+    - AH
+    - AO
+    - AW
+    - AY
+    - B
+    - CH
+    - D
+    - DH
+    - EH
+    - ER
+    - EY
+    - F
+    - G
+    - HH
+    - IH
+    - IY
+    - JH
+    - K
+    - L
+    - M
+    - N
+    - NG
+    - OW
+    - OY
+    - P
+    - R
+    - S
+    - SH
+    - T
+    - TH
+    - UH
+    - UW
+    - V
+    - W
+    - Y
+    - Z
+    - ZH
+
+input_encoder: !new:speechbrain.dataio.encoder.TextEncoder
+
+
+
+################################
+# Model Parameters and model   #
+# Transformer Parameters
+################################
+d_model: 512
+nhead: 8
+num_encoder_layers: 3
+num_decoder_layers: 3
+dim_feedforward: 512
+dropout: 0.1
+
+
+# Decoder parameters
+# The number of frames in the target per encoder step
+n_frames_per_step: 1
+decoder_rnn_dim: 1024
+prenet_dim: 256
+max_decoder_steps: 1000
+gate_threshold: 0.5
+p_decoder_dropout: 0.1
+decoder_no_early_stopping: False
+
+blank_index: 0 # This special tokes is for padding
+
+
+# Masks
+lookahead_mask: !name:speechbrain.lobes.models.transformer.Transformer.get_lookahead_mask
+padding_mask: !name:speechbrain.lobes.models.transformer.Transformer.get_key_padding_mask
+
+
+################################
+# CNN 3-layers Prenet          #
+################################
+# Encoder Prenet
+encoder_prenet: !new:module_classes.CNNPrenet
+
+# Decoder Prenet
+decoder_prenet: !new:module_classes.CNNDecoderPrenet
+
+################################
+# Positional Encodings         #
+################################
+
+#encoder
+pos_emb_enc: !new:module_classes.ScaledPositionalEncoding
+    input_size: !ref <d_model>
+    max_len: 5000
+
+#decoder
+pos_emb_dec: !new:module_classes.ScaledPositionalEncoding
+    input_size: !ref <d_model>
+    max_len: 5000
+
+
+################################
+# S2S Transfomer               #
+################################
+
+Seq2SeqTransformer: !new:torch.nn.Transformer
+    d_model: !ref <d_model>
+    nhead: !ref <nhead>
+    num_encoder_layers: !ref <num_encoder_layers>
+    num_decoder_layers: !ref <num_decoder_layers>
+    dim_feedforward: !ref <dim_feedforward>
+    dropout: !ref <dropout>
+    batch_first: True
+
+
+################################
+# CNN 5-layers PostNet         #
+################################
+
+decoder_postnet: !new:speechbrain.lobes.models.Tacotron2.Postnet
+
+
+# Linear transformation on the top of the decoder.
+stop_lin: !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <d_model>
+    n_neurons: 1
+
+
+# Linear transformation on the top of the decoder.
+mel_lin: !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <d_model>
+    n_neurons: 80
+
+modules:
+    encoder_prenet: !ref <encoder_prenet>
+    pos_emb_enc: !ref <pos_emb_enc>
+    decoder_prenet: !ref <decoder_prenet>
+    pos_emb_dec: !ref <pos_emb_dec>
+    Seq2SeqTransformer: !ref <Seq2SeqTransformer>
+    mel_lin: !ref <mel_lin>
+    stop_lin: !ref <stop_lin>
+    decoder_postnet: !ref <decoder_postnet>
+
+
+model: !new:torch.nn.ModuleList
+    - [!ref <encoder_prenet>,!ref <pos_emb_enc>,
+       !ref <decoder_prenet>, !ref <pos_emb_dec>, !ref <Seq2SeqTransformer>,
+       !ref <mel_lin>, !ref <stop_lin>,  !ref <decoder_postnet>]
+
+
+pretrained_model_path: ./model.ckpt
+
+# The pretrainer allows a mapping between pretrained files and instances that
+# are declared in the yaml. E.g here, we will download the file model.ckpt
+# and it will be loaded into "model" which is pointing to the <model> defined
+# before.
+
+pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
+   collect_in: !ref <save_folder>
+   loadables:
+      model: !ref <model>
+   paths:
+      model: !ref <pretrained_model_path>

model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e5421fe987116817841652862ce070a421d7f5d7c8bbef68c83bec876b1eafb
+size 95804314

module_classes.py

@@ -0,0 +1,214 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+class CNNPrenet(torch.nn.Module):
+    def __init__(self):
+        super(CNNPrenet, self).__init__()
+
+        # Define the layers using Sequential container
+        self.conv_layers = nn.Sequential(
+            nn.Conv1d(in_channels=1, out_channels=512, kernel_size=3, padding=1),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+
+            nn.Conv1d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+
+            nn.Conv1d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.1)
+        )
+
+    def forward(self, x):
+
+        # Add a new dimension for the channel
+        x = x.unsqueeze(1)
+
+        # Pass input through convolutional layers
+        x = self.conv_layers(x)
+
+        # Remove the channel dimension
+        x = x.squeeze(1)
+
+        # Scale the output to the range [-1, 1]
+        x = torch.tanh(x)
+
+        return x
+
+
+
+class CNNDecoderPrenet(nn.Module):
+    def __init__(self, input_dim=80, hidden_dim=256, output_dim=256, final_dim=512, dropout_rate=0.5):
+        super(CNNDecoderPrenet, self).__init__()
+        self.layer1 = nn.Linear(input_dim, hidden_dim)
+        self.layer2 = nn.Linear(hidden_dim, output_dim)
+        self.linear_projection = nn.Linear(output_dim, final_dim) # Added linear projection
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+
+      # Transpose the input tensor to have the feature dimension as the last dimension
+      x = x.transpose(1, 2)
+      # Apply the linear layers
+      x = F.relu(self.layer1(x))
+      x = self.dropout(x)
+      x = F.relu(self.layer2(x))
+      x = self.dropout(x)
+      # Apply the linear projection
+      x = self.linear_projection(x)
+      x = x.transpose(1, 2)
+
+      return x
+
+
+
+
+class CNNPostNet(torch.nn.Module):
+    """
+    Conv Postnet
+    Arguments
+    ---------
+    n_mel_channels: int
+       input feature dimension for convolution layers
+    postnet_embedding_dim: int
+       output feature dimension for convolution layers
+    postnet_kernel_size: int
+       postnet convolution kernal size
+    postnet_n_convolutions: int
+       number of convolution layers
+    postnet_dropout: float
+        dropout probability fot postnet
+    """
+
+    def __init__(
+        self,
+        n_mel_channels=80,
+        postnet_embedding_dim=512,
+        postnet_kernel_size=5,
+        postnet_n_convolutions=5,
+        postnet_dropout=0.1,
+    ):
+        super(CNNPostNet, self).__init__()
+
+        self.conv_pre = nn.Conv1d(
+            in_channels=n_mel_channels,
+            out_channels=postnet_embedding_dim,
+            kernel_size=postnet_kernel_size,
+            padding="same",
+        )
+
+        self.convs_intermedite = nn.ModuleList()
+        for i in range(1, postnet_n_convolutions - 1):
+            self.convs_intermedite.append(
+                nn.Conv1d(
+                    in_channels=postnet_embedding_dim,
+                    out_channels=postnet_embedding_dim,
+                    kernel_size=postnet_kernel_size,
+                    padding="same",
+                ),
+            )
+
+        self.conv_post = nn.Conv1d(
+            in_channels=postnet_embedding_dim,
+            out_channels=n_mel_channels,
+            kernel_size=postnet_kernel_size,
+            padding="same",
+        )
+
+        self.tanh = nn.Tanh()
+        self.ln1 = nn.LayerNorm(postnet_embedding_dim)
+        self.ln2 = nn.LayerNorm(postnet_embedding_dim)
+        self.ln3 = nn.LayerNorm(n_mel_channels)
+        self.dropout1 = nn.Dropout(postnet_dropout)
+        self.dropout2 = nn.Dropout(postnet_dropout)
+        self.dropout3 = nn.Dropout(postnet_dropout)
+
+
+    def forward(self, x):
+        """Computes the forward pass
+        Arguments
+        ---------
+        x: torch.Tensor
+            a (batch, time_steps, features) input tensor
+        Returns
+        -------
+        output: torch.Tensor (the spectrogram predicted)
+        """
+        x = self.conv_pre(x)
+        x = self.ln1(x.permute(0, 2, 1)).permute(0, 2, 1)  # Transpose to [batch_size, feature_dim, sequence_length]
+        x = self.tanh(x)
+        x = self.dropout1(x)
+
+        for i in range(len(self.convs_intermedite)):
+            x = self.convs_intermedite[i](x)
+        x = self.ln2(x.permute(0, 2, 1)).permute(0, 2, 1)  # Transpose to [batch_size, feature_dim, sequence_length]
+        x = self.tanh(x)
+        x = self.dropout2(x)
+
+        x = self.conv_post(x)
+        x = self.ln3(x.permute(0, 2, 1)).permute(0, 2, 1)  # Transpose to [batch_size, feature_dim, sequence_length]
+        x = self.dropout3(x)
+
+        return x
+
+
+class ScaledPositionalEncoding(nn.Module):
+    """
+    This class implements the absolute sinusoidal positional encoding function
+    with an adaptive weight parameter alpha.
+
+    PE(pos, 2i)   = sin(pos/(10000^(2i/dmodel)))
+    PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel)))
+
+    Arguments
+    ---------
+    input_size: int
+        Embedding dimension.
+    max_len : int, optional
+        Max length of the input sequences (default 2500).
+    Example
+    -------
+    >>> a = torch.rand((8, 120, 512))
+    >>> enc = PositionalEncoding(input_size=a.shape[-1])
+    >>> b = enc(a)
+    >>> b.shape
+    torch.Size([1, 120, 512])
+    """
+
+    def __init__(self, input_size, max_len=2500):
+        super().__init__()
+        if input_size % 2 != 0:
+            raise ValueError(
+                f"Cannot use sin/cos positional encoding with odd channels (got channels={input_size})"
+            )
+        self.max_len = max_len
+        self.alpha = nn.Parameter(torch.ones(1))  # Define alpha as a trainable parameter
+        pe = torch.zeros(self.max_len, input_size, requires_grad=False)
+        positions = torch.arange(0, self.max_len).unsqueeze(1).float()
+        denominator = torch.exp(
+            torch.arange(0, input_size, 2).float()
+            * -(math.log(10000.0) / input_size)
+        )
+
+        pe[:, 0::2] = torch.sin(positions * denominator)
+        pe[:, 1::2] = torch.cos(positions * denominator)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        """
+        Arguments
+        ---------
+        x : tensor
+            Input feature shape (batch, time, fea)
+        """
+        pe_scaled = self.pe[:, :x.size(1)].clone().detach() * self.alpha  # Scale positional encoding by alpha
+        return pe_scaled
+