Views
No views yet
1# minimal onnx inference extracted from coqui-tts
2import json
3import re
4from typing import Callable, List
5
6import numpy as np
7import onnxruntime as ort
8import scipy
9
10# Regular expression matching whitespace:
11_whitespace_re = re.compile(r"\s+")
12
13
14class Graphemes:
15 def __init__(
16 self,
17 characters: str = None,
18 punctuations: str = None,
19 pad: str = None,
20 eos: str = None,
21 bos: str = None,
22 blank: str = "<BLNK>",
23 is_unique: bool = False,
24 is_sorted: bool = True,
25 ) -> None:
26 self._characters = characters
27 self._punctuations = punctuations
28 self._pad = pad
29 self._eos = eos
30 self._bos = bos
31 self._blank = blank
32 self.is_unique = is_unique
33 self.is_sorted = is_sorted
34 self._create_vocab()
35
36 @property
37 def pad_id(self) -> int:
38 return self.char_to_id(self.pad) if self.pad else len(self.vocab)
39
40 @property
41 def blank_id(self) -> int:
42 return self.char_to_id(self.blank) if self.blank else len(self.vocab)
43
44 @property
45 def eos_id(self) -> int:
46 return self.char_to_id(self.eos) if self.eos else len(self.vocab)
47
48 @property
49 def bos_id(self) -> int:
50 return self.char_to_id(self.bos) if self.bos else len(self.vocab)
51
52 @property
53 def characters(self):
54 return self._characters
55
56 @characters.setter
57 def characters(self, characters):
58 self._characters = characters
59 self._create_vocab()
60
61 @property
62 def punctuations(self):
63 return self._punctuations
64
65 @punctuations.setter
66 def punctuations(self, punctuations):
67 self._punctuations = punctuations
68 self._create_vocab()
69
70 @property
71 def pad(self):
72 return self._pad
73
74 @pad.setter
75 def pad(self, pad):
76 self._pad = pad
77 self._create_vocab()
78
79 @property
80 def eos(self):
81 return self._eos
82
83 @eos.setter
84 def eos(self, eos):
85 self._eos = eos
86 self._create_vocab()
87
88 @property
89 def bos(self):
90 return self._bos
91
92 @bos.setter
93 def bos(self, bos):
94 self._bos = bos
95 self._create_vocab()
96
97 @property
98 def blank(self):
99 return self._blank
100
101 @blank.setter
102 def blank(self, blank):
103 self._blank = blank
104 self._create_vocab()
105
106 @property
107 def vocab(self):
108 return self._vocab
109
110 @vocab.setter
111 def vocab(self, vocab):
112 self._vocab = vocab
113 self._char_to_id = {char: idx for idx, char in enumerate(self.vocab)}
114 self._id_to_char = {
115 idx: char for idx, char in enumerate(self.vocab) # pylint: disable=unnecessary-comprehension
116 }
117
118 @property
119 def num_chars(self):
120 return len(self._vocab)
121
122 def _create_vocab(self):
123 self._vocab = [self._pad] + list(self._punctuations) + list(self._characters) + [self._blank]
124 self._char_to_id = {char: idx for idx, char in enumerate(self.vocab)}
125 # pylint: disable=unnecessary-comprehension
126 self._id_to_char = {idx: char for idx, char in enumerate(self.vocab)}
127
128 def char_to_id(self, char: str) -> int:
129 try:
130 return self._char_to_id[char]
131 except KeyError as e:
132 raise KeyError(f" [!] {repr(char)} is not in the vocabulary.") from e
133
134 def id_to_char(self, idx: int) -> str:
135 return self._id_to_char[idx]
136
137
138class TTSTokenizer:
139 """🐸TTS tokenizer to convert input characters to token IDs and back.
140
141 Token IDs for OOV chars are discarded but those are stored in `self.not_found_characters` for later.
142
143 Args:
144 characters (Characters):
145 A Characters object to use for character-to-ID and ID-to-character mappings.
146
147 text_cleaner (callable):
148 A function to pre-process the text before tokenization and phonemization. Defaults to None.
149 """
150
151 def __init__(
152 self,
153 text_cleaner: Callable = None,
154 characters: Graphemes = None,
155 add_blank: bool = False,
156 use_eos_bos=False,
157 ):
158 self.text_cleaner = text_cleaner
159 self.add_blank = add_blank
160 self.use_eos_bos = use_eos_bos
161 self.characters = characters
162 self.not_found_characters = []
163
164 @property
165 def characters(self):
166 return self._characters
167
168 @characters.setter
169 def characters(self, new_characters):
170 self._characters = new_characters
171 self.pad_id = self.characters.char_to_id(self.characters.pad) if self.characters.pad else None
172 self.blank_id = self.characters.char_to_id(self.characters.blank) if self.characters.blank else None
173
174 def encode(self, text: str) -> List[int]:
175 """Encodes a string of text as a sequence of IDs."""
176 token_ids = []
177 for char in text:
178 try:
179 idx = self.characters.char_to_id(char)
180 token_ids.append(idx)
181 except KeyError:
182 # discard but store not found characters
183 if char not in self.not_found_characters:
184 self.not_found_characters.append(char)
185 print(text)
186 print(f" [!] Character {repr(char)} not found in the vocabulary. Discarding it.")
187 return token_ids
188
189 def text_to_ids(self, text: str) -> List[int]: # pylint: disable=unused-argument
190 """Converts a string of text to a sequence of token IDs.
191
192 Args:
193 text(str):
194 The text to convert to token IDs.
195
196 1. Text normalization
197 3. Add blank char between characters
198 4. Add BOS and EOS characters
199 5. Text to token IDs
200 """
201 if self.text_cleaner is not None:
202 text = self.text_cleaner(text)
203 text = self.encode(text)
204 if self.add_blank:
205 text = self.intersperse_blank_char(text, True)
206 if self.use_eos_bos:
207 text = self.pad_with_bos_eos(text)
208 return text
209
210 def pad_with_bos_eos(self, char_sequence: List[str]):
211 """Pads a sequence with the special BOS and EOS characters."""
212 return [self.characters.bos_id] + list(char_sequence) + [self.characters.eos_id]
213
214 def intersperse_blank_char(self, char_sequence: List[str], use_blank_char: bool = False):
215 """Intersperses the blank character between characters in a sequence.
216
217 Use the ```blank``` character if defined else use the ```pad``` character.
218 """
219 char_to_use = self.characters.blank_id if use_blank_char else self.characters.pad
220 result = [char_to_use] * (len(char_sequence) * 2 + 1)
221 result[1::2] = char_sequence
222 return result
223
224
225class VitsOnnxInference:
226 def __init__(self, onnx_model_path: str, config_path: str, cuda=False):
227 self.config = {}
228 if config_path:
229 with open(config_path) as f:
230 self.config = json.load(f)
231 providers = [
232 "CPUExecutionProvider"
233 if cuda is False
234 else ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"})
235 ]
236 sess_options = ort.SessionOptions()
237 self.onnx_sess = ort.InferenceSession(
238 onnx_model_path,
239 sess_options=sess_options,
240 providers=providers,
241 )
242
243 _pad = self.config.get("characters", {}).get("pad", "_")
244 _punctuations = self.config.get("characters", {}).get("punctuations", "!\"(),-.:;?\u00a1\u00bf ")
245 _letters = self.config.get("characters", {}).get("characters",
246 "ABCDEFGHIJKLMNOPQRSTUVXYZabcdefghijklmnopqrstuvwxyz\u00c1\u00c9\u00cd\u00d3\u00da\u00e1\u00e9\u00ed\u00f1\u00f3\u00fa\u00fc")
247
248 vocab = Graphemes(characters=_letters,
249 punctuations=_punctuations,
250 pad=_pad)
251
252 self.tokenizer = TTSTokenizer(
253 text_cleaner=self.normalize_text,
254 characters=vocab,
255 add_blank=self.config.get("add_blank", True),
256 use_eos_bos=False,
257 )
258
259 @staticmethod
260 def normalize_text(text: str) -> str:
261 """Basic pipeline that lowercases and collapses whitespace without transliteration."""
262 text = text.lower()
263 text = text.replace(";", ",")
264 text = text.replace("-", " ")
265 text = text.replace(":", ",")
266 text = re.sub(r"[\<\>\(\)\[\]\"]+", "", text)
267 text = re.sub(_whitespace_re, " ", text).strip()
268 return text
269
270 def inference_onnx(self, text: str):
271 """ONNX inference"""
272 x = np.asarray(
273 self.tokenizer.text_to_ids(text),
274 dtype=np.int64,
275 )[None, :]
276
277 x_lengths = np.array([x.shape[1]], dtype=np.int64)
278
279 scales = np.array(
280 [self.config.get("inference_noise_scale", 0.667),
281 self.config.get("length_scale", 1.0),
282 self.config.get("inference_noise_scale_dp", 1.0), ],
283 dtype=np.float32,
284 )
285 input_params = {"input": x, "input_lengths": x_lengths, "scales": scales}
286
287 audio = self.onnx_sess.run(
288 ["output"],
289 input_params,
290 )
291 return audio[0][0]
292
293 @staticmethod
294 def save_wav(wav: np.ndarray, path: str, sample_rate: int = 16000) -> None:
295 """Save float waveform to a file using Scipy.
296
297 Args:
298 wav (np.ndarray): Waveform with float values in range [-1, 1] to save.
299 path (str): Path to a output file.
300 """
301 wav_norm = wav * (32767 / max(0.01, np.max(np.abs(wav))))
302 wav_norm = wav_norm.astype(np.int16)
303 scipy.io.wavfile.write(path, sample_rate, wav_norm)
304
305 def synth(self, text: str, path: str):
306 wavs = self.inference_onnx(text)
307 self.save_wav(wavs[0], path, self.config.get("sample_rate", 16000))
308```