Upload 3 files

TTSModel.py

@@ -0,0 +1,173 @@
+
+import re
+import logging
+import torch
+import torchaudio
+import random
+import speechbrain as sb
+import torch as nn
+from speechbrain.utils.fetching import fetch
+from speechbrain.inference.interfaces import Pretrained
+from speechbrain.inference.text import GraphemeToPhoneme
+
+
+logger = logging.getLogger(__name__)
+
+
+class TTSModel(Pretrained):
+    """
+    A ready-to-use wrapper for Transformer TTS (text -> mel_spec).
+    Arguments
+    ---------
+    hparams
+        Hyperparameters (from HyperPyYAML)"""
+
+    HPARAMS_NEEDED = ["model", "blank_index", "padding_mask", "lookahead_mask", "mel_spec_feats", "label_encoder"]
+    MODULES_NEEDED = ["modules"]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.label_encoder = self.hparams.label_encoder
+        #self.label_encoder.update_from_iterable(self.hparams["lexicon"], sequence_input=False)
+        self.g2p = GraphemeToPhoneme.from_hparams("speechbrain/soundchoice-g2p")
+
+
+    def text_to_phoneme(self, text):
+        """
+        Generates phoneme sequences for the given text using a Grapheme-to-Phoneme (G2P) model.
+
+        Args:
+            text (str): The input text.
+
+        Returns:
+            list: List of phoneme sequences for the words in the text.
+        """
+        abbreviation_expansions = {
+            "Mr.": "Mister",
+            "Mrs.": "Misess",
+            "Dr.": "Doctor",
+            "No.": "Number",
+            "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"
+        }
+
+        # Expand abbreviations
+        for abbreviation, expansion in abbreviation_expansions.items():
+            text = text.replace(abbreviation, expansion)
+
+        phonemes = self.g2p(text)
+        phonemes = self.label_encoder.encode_sequence(phonemes)
+        phoneme_seq = torch.LongTensor(phonemes)
+
+        return phoneme_seq, len(phoneme_seq)
+
+    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, output lengths and alignments
+        """
+        with torch.no_grad():
+            phoneme_seqs = [self.text_to_phoneme(text)[0] for text in texts]
+            phoneme_seqs_padded, input_lengths = self.pad_sequences(phoneme_seqs)
+
+            encoded_phoneme = self.mods.encoder_emb(phoneme_seqs_padded)
+            encoder_emb = self.mods.enc_pre_net(encoded_phoneme)
+            pos_emb_enc = self.mods.pos_emb_enc(encoder_emb)
+            encoder_emb = encoder_emb + pos_emb_enc
+
+
+            stop_generated = False
+            decoder_input = torch.zeros(1, 80, 1, device=self.device)
+            stop_tokens_logits = []
+            max_generation_length = 1000
+            sequence_length = 0
+
+            result = []
+            result.append(decoder_input)
+
+            src_mask = torch.zeros(encoder_emb.size(1), encoder_emb.size(1), device=self.device)
+            src_key_padding_mask = self.hparams.padding_mask(encoder_emb, self.hparams.blank_index)
+
+
+            while not stop_generated and sequence_length < max_generation_length:
+                encoded_mel = self.mods.dec_pre_net(decoder_input)
+                pos_emb_dec = self.mods.pos_emb_dec(encoded_mel)
+                decoder_emb = encoded_mel + pos_emb_dec
+
+                decoder_output = self.mods.Seq2SeqTransformer(
+                    encoder_emb, decoder_emb, src_mask=src_mask,
+                    src_key_padding_mask=src_key_padding_mask)
+
+                mel_output = self.mods.mel_lin(decoder_output)
+
+                stop_token_logit = self.mods.stop_lin(decoder_output).squeeze(-1)
+
+                post_mel_outputs = self.mods.postnet(mel_output.to(self.device))
+                refined_mel_output = mel_output + post_mel_outputs.to(self.device)
+                refined_mel_output = refined_mel_output.transpose(1, 2)
+
+                stop_tokens_logits.append(stop_token_logit)
+                stop_token_probs = torch.sigmoid(stop_token_logit)
+
+                if torch.any(stop_token_probs[:, -1] >= self.hparams.stop_threshold):
+                    stop_generated = True
+
+                decoder_input = refined_mel_output
+                result.append(decoder_input)
+                sequence_length += 1
+
+            results = torch.cat(result, dim=2)
+            stop_tokens_logits = torch.cat(stop_tokens_logits, dim=1)
+
+        return results
+
+    def pad_sequences(self, sequences):
+      """Pad sequences to the maximum length sequence in the batch.
+
+      Arguments
+      ---------
+      sequences: List[torch.Tensor]
+          The sequences to pad
+
+      Returns
+      -------
+      Padded sequences and original lengths
+      """
+      max_length = max([len(seq) for seq in sequences])
+      padded_seqs = torch.zeros(len(sequences), max_length, dtype=torch.long)
+      lengths = []
+      for i, seq in enumerate(sequences):
+          length = len(seq)
+          padded_seqs[i, :length] = seq
+          lengths.append(length)
+      return padded_seqs, torch.tensor(lengths)
+
+    def encode_text(self, text):
+        """Runs inference for a single text str"""
+        return self.encode_batch(text)
+
+    def forward(self, texts):
+        "Encodes the input texts."
+        return self.encode_batch(texts)

hyperparams.yaml

@@ -0,0 +1,173 @@
+
+
+################################
+# Audio Parameters             #
+################################
+sample_rate: 22050
+hop_length: 256
+win_length: 1024
+n_mel_channels: 80
+n_fft: 1024
+mel_fmin: 0.0
+mel_fmax: 8000.0
+power: 1
+normalized: False
+min_max_energy_norm: True
+norm: "slaney"
+mel_scale: "slaney"
+dynamic_range_compression: True
+mel_normalized: False
+min_f0: 65  #(torchaudio pyin values)
+max_f0: 2093 #(torchaudio pyin values)
+
+positive_weight: 5.0
+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
+    - ' '
+n_symbols: 42 #fixed depending on symbols in the lexicon +1 for a dummy symbol used for padding
+padding_idx: 0
+
+# Define model architecture
+d_model: 512
+nhead: 8
+num_encoder_layers: 6
+num_decoder_layers: 6
+dim_feedforward: 2048
+dropout: 0.2
+blank_index: 0 # This special token is for padding
+bos_index: 1
+eos_index: 2
+stop_weight: 0.45
+stop_threshold: 0.5
+
+
+###################PRENET#######################
+enc_pre_net: !new:models.EncoderPrenet
+dec_pre_net: !new:models.DecoderPrenet
+
+
+encoder_emb: !new:torch.nn.Embedding
+    num_embeddings: 128
+    embedding_dim: !ref <d_model>
+    padding_idx: !ref <blank_index>
+
+pos_emb_enc: !new:models.ScaledPositionalEncoding
+    d_model: !ref <d_model>
+
+decoder_emb: !new:torch.nn.Embedding
+    num_embeddings: 128
+    embedding_dim: !ref <d_model>
+    padding_idx: !ref <blank_index>
+
+pos_emb_dec: !new:models.ScaledPositionalEncoding
+    d_model: !ref <d_model>
+
+
+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
+
+postnet: !new:models.PostNet
+    mel_channels: !ref <n_mel_channels>
+    postnet_channels: 512
+    kernel_size: 5
+    postnet_layers: 5
+
+mel_lin: !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <d_model>
+    n_neurons: !ref <n_mel_channels>
+
+stop_lin: !new:speechbrain.nnet.linear.Linear
+    input_size: !ref <d_model>
+    n_neurons: 1
+
+mel_spec_feats: !name:speechbrain.lobes.models.FastSpeech2.mel_spectogram
+    sample_rate: !ref <sample_rate>
+    hop_length: !ref <hop_length>
+    win_length: !ref <win_length>
+    n_fft: !ref <n_fft>
+    n_mels: !ref <n_mel_channels>
+    f_min: !ref <mel_fmin>
+    f_max: !ref <mel_fmax>
+    power: !ref <power>
+    normalized: !ref <normalized>
+    min_max_energy_norm: !ref <min_max_energy_norm>
+    norm: !ref <norm>
+    mel_scale: !ref <mel_scale>
+    compression: !ref <dynamic_range_compression>
+
+modules:
+    enc_pre_net: !ref <enc_pre_net>
+    encoder_emb: !ref <encoder_emb>
+    pos_emb_enc: !ref <pos_emb_enc>
+
+    dec_pre_net: !ref <dec_pre_net>
+    #decoder_emb: !ref <decoder_emb>
+    pos_emb_dec: !ref <pos_emb_dec>
+
+    Seq2SeqTransformer: !ref <Seq2SeqTransformer>
+    postnet: !ref <postnet>
+    mel_lin: !ref <mel_lin>
+    stop_lin: !ref <stop_lin>
+    model: !ref <model>
+
+lookahead_mask: !name:speechbrain.lobes.models.transformer.Transformer.get_lookahead_mask
+padding_mask: !name:speechbrain.lobes.models.transformer.Transformer.get_key_padding_mask
+
+model: !new:torch.nn.ModuleList
+    - [!ref <enc_pre_net>, !ref <encoder_emb>, !ref <pos_emb_enc>, !ref <dec_pre_net>, !ref <pos_emb_dec>, !ref <Seq2SeqTransformer>, !ref <postnet>, !ref <mel_lin>, !ref <stop_lin>]
+
+label_encoder: !new:speechbrain.dataio.encoder.TextEncoder
+
+pretrained_path: /content/
+
+pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
+    loadables:
+        model: !ref <model>
+        label_encoder: !ref <label_encoder>
+    paths:
+        model: !ref <pretrained_path>/model.ckpt
+        label_encoder: !ref <pretrained_path>/label_encoder.txt

label_encoder.txt

@@ -0,0 +1,46 @@
+'AA' => 0
+'AE' => 40
+'AH' => 41
+'AO' => 3
+'AW' => 4
+'AY' => 5
+'B' => 6
+'CH' => 7
+'D' => 8
+'DH' => 9
+'EH' => 10
+'ER' => 11
+'EY' => 12
+'F' => 13
+'G' => 14
+'HH' => 15
+'IH' => 16
+'IY' => 17
+'JH' => 18
+'K' => 19
+'L' => 20
+'M' => 21
+'N' => 22
+'NG' => 23
+'OW' => 24
+'OY' => 25
+'P' => 26
+'R' => 27
+'S' => 28
+'SH' => 29
+'T' => 30
+'TH' => 31
+'UH' => 32
+'UW' => 33
+'V' => 34
+'W' => 35
+'Y' => 36
+'Z' => 37
+'ZH' => 38
+' ' => 39
+'<bos>' => 1
+'<eos>' => 2
+================
+'starting_index' => 0
+'bos_label' => '<bos>'
+'eos_label' => '<eos>'