| import torch | |
| import os | |
| class TTSInference: | |
| def __init__(self, device=None): | |
| print("Initializing TTS model to %s" % device) | |
| from .tasks.tts.tts_utils import load_data_preprocessor | |
| from .utils.commons.hparams import set_hparams | |
| if device is None: | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| self.hparams = set_hparams("text_to_speech/checkpoints/ljspeech/ps_adv_baseline/config.yaml") | |
| self.device = device | |
| self.data_dir = 'text_to_speech/checkpoints/ljspeech/data_info' | |
| self.preprocessor, self.preprocess_args = load_data_preprocessor() | |
| self.ph_encoder, self.word_encoder = self.preprocessor.load_dict(self.data_dir) | |
| self.spk_map = self.preprocessor.load_spk_map(self.data_dir) | |
| self.model = self.build_model() | |
| self.model.eval() | |
| self.model.to(self.device) | |
| self.vocoder = self.build_vocoder() | |
| self.vocoder.eval() | |
| self.vocoder.to(self.device) | |
| print("TTS loaded!") | |
| def build_model(self): | |
| from .utils.commons.ckpt_utils import load_ckpt | |
| from .modules.tts.portaspeech.portaspeech import PortaSpeech | |
| ph_dict_size = len(self.ph_encoder) | |
| word_dict_size = len(self.word_encoder) | |
| model = PortaSpeech(ph_dict_size, word_dict_size, self.hparams) | |
| load_ckpt(model, 'text_to_speech/checkpoints/ljspeech/ps_adv_baseline', 'model') | |
| model.to(self.device) | |
| with torch.no_grad(): | |
| model.store_inverse_all() | |
| model.eval() | |
| return model | |
| def forward_model(self, inp): | |
| sample = self.input_to_batch(inp) | |
| with torch.no_grad(): | |
| output = self.model( | |
| sample['txt_tokens'], | |
| sample['word_tokens'], | |
| ph2word=sample['ph2word'], | |
| word_len=sample['word_lengths'].max(), | |
| infer=True, | |
| forward_post_glow=True, | |
| spk_id=sample.get('spk_ids') | |
| ) | |
| mel_out = output['mel_out'] | |
| wav_out = self.run_vocoder(mel_out) | |
| wav_out = wav_out.cpu().numpy() | |
| return wav_out[0] | |
| def build_vocoder(self): | |
| from .utils.commons.hparams import set_hparams | |
| from .modules.vocoder.hifigan.hifigan import HifiGanGenerator | |
| from .utils.commons.ckpt_utils import load_ckpt | |
| base_dir = 'text_to_speech/checkpoints/hifi_lj' | |
| config_path = f'{base_dir}/config.yaml' | |
| config = set_hparams(config_path, global_hparams=False) | |
| vocoder = HifiGanGenerator(config) | |
| load_ckpt(vocoder, base_dir, 'model_gen') | |
| return vocoder | |
| def run_vocoder(self, c): | |
| c = c.transpose(2, 1) | |
| y = self.vocoder(c)[:, 0] | |
| return y | |
| def preprocess_input(self, inp): | |
| """ | |
| :param inp: {'text': str, 'item_name': (str, optional), 'spk_name': (str, optional)} | |
| :return: | |
| """ | |
| preprocessor, preprocess_args = self.preprocessor, self.preprocess_args | |
| text_raw = inp['text'] | |
| item_name = inp.get('item_name', '<ITEM_NAME>') | |
| spk_name = inp.get('spk_name', '<SINGLE_SPK>') | |
| ph, txt, word, ph2word, ph_gb_word = preprocessor.txt_to_ph( | |
| preprocessor.txt_processor, text_raw, preprocess_args) | |
| word_token = self.word_encoder.encode(word) | |
| ph_token = self.ph_encoder.encode(ph) | |
| spk_id = self.spk_map[spk_name] | |
| item = {'item_name': item_name, 'text': txt, 'ph': ph, 'spk_id': spk_id, | |
| 'ph_token': ph_token, 'word_token': word_token, 'ph2word': ph2word, | |
| 'ph_words':ph_gb_word, 'words': word} | |
| item['ph_len'] = len(item['ph_token']) | |
| return item | |
| def input_to_batch(self, item): | |
| item_names = [item['item_name']] | |
| text = [item['text']] | |
| ph = [item['ph']] | |
| txt_tokens = torch.LongTensor(item['ph_token'])[None, :].to(self.device) | |
| txt_lengths = torch.LongTensor([txt_tokens.shape[1]]).to(self.device) | |
| word_tokens = torch.LongTensor(item['word_token'])[None, :].to(self.device) | |
| word_lengths = torch.LongTensor([txt_tokens.shape[1]]).to(self.device) | |
| ph2word = torch.LongTensor(item['ph2word'])[None, :].to(self.device) | |
| spk_ids = torch.LongTensor(item['spk_id'])[None, :].to(self.device) | |
| batch = { | |
| 'item_name': item_names, | |
| 'text': text, | |
| 'ph': ph, | |
| 'txt_tokens': txt_tokens, | |
| 'txt_lengths': txt_lengths, | |
| 'word_tokens': word_tokens, | |
| 'word_lengths': word_lengths, | |
| 'ph2word': ph2word, | |
| 'spk_ids': spk_ids, | |
| } | |
| return batch | |
| def postprocess_output(self, output): | |
| return output | |
| def infer_once(self, inp): | |
| inp = self.preprocess_input(inp) | |
| output = self.forward_model(inp) | |
| output = self.postprocess_output(output) | |
| return output | |