ru / infer_onnx.py
TeraSpace's picture
Update infer_onnx.py
8957f58
import scipy.io.wavfile
import os
import onnxruntime
import numpy as np
from huggingface_hub import snapshot_download
from num2words import num2words
import re
from transliterate import translit
import json
class TTS:
def __init__(self, model_name: str, save_path: str = "./model", add_time_to_end: float = 0.8) -> None:
if not os.path.exists(save_path):
os.mkdir(save_path)
model_dir = os.path.join(save_path, model_name)
if not os.path.exists(model_dir):
snapshot_download(repo_id=model_name,
allow_patterns=["*.txt", "*.onnx", "*.json"],
local_dir=model_dir,
local_dir_use_symlinks=False
)
self.model = onnxruntime.InferenceSession(os.path.join(model_dir, "exported/model.onnx"), providers=['CPUExecutionProvider'])
with open(os.path.join(model_dir, "exported/config.json")) as config_file:
self.config = json.load(config_file)["model_config"]
if os.path.exists(os.path.join(model_dir, "exported/dictionary.txt")):
from tokenizer import TokenizerG2P
print("Use g2p")
self.tokenizer = TokenizerG2P(os.path.join(model_dir, "exported"))
else:
from tokenizer import TokenizerGRUUT
print("Use gruut")
self.tokenizer = TokenizerGRUUT(os.path.join(model_dir, "exported"))
self.add_time_to_end = add_time_to_end
def _add_silent(self, audio, silence_duration: float = 1.0, sample_rate: int = 22050):
num_samples_silence = int(sample_rate * silence_duration)
silence_array = np.zeros(num_samples_silence, dtype=np.float32)
audio_with_silence = np.concatenate((audio, silence_array), axis=0)
return audio_with_silence
def save_wav(self, audio, path:str, sample_rate: int = 22050):
'''save audio to wav'''
scipy.io.wavfile.write(path, sample_rate, audio)
def _intersperse(self, lst, item):
result = [item] * (len(lst) * 2 + 1)
result[1::2] = lst
return result
def _get_seq(self, text):
phoneme_ids = self.tokenizer._get_seq(text)
phoneme_ids_inter = self._intersperse(phoneme_ids, 0)
return phoneme_ids_inter
def _num2wordsshor(self, match):
match = match.group()
ret = num2words(match, lang ='ru')
return ret
def __call__(self, text: str, length_scale=1.2):
text = translit(text, 'ru')
text = re.sub(r'\d+',self._num2wordsshor,text)
phoneme_ids = self._get_seq(text)
text = np.expand_dims(np.array(phoneme_ids, dtype=np.int64), 0)
text_lengths = np.array([text.shape[1]], dtype=np.int64)
scales = np.array(
[0.667, length_scale, 0.8],
dtype=np.float32,
)
audio = self.model.run(
None,
{
"input": text,
"input_lengths": text_lengths,
"scales": scales,
"sid": None,
},
)[0][0,0][0]
audio = self._add_silent(audio, silence_duration = self.add_time_to_end, sample_rate=self.config["samplerate"])
return audio