Spaces:
Paused
Paused
| 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 | |