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)