Views
No views yet
exaggeration=0.5, cfg=0.5) work well for most prompts.exaggeration to around 0.7 or higher.exaggeration tends to speed up speech;1# !pip install --upgrade onnxruntime==1.22.1 huggingface_hub==0.34.4 transformers==4.46.3 numpy==2.2.6 tqdm==4.67.1 librosa==0.11.0 soundfile==0.13.1 resemble-perth==1.0.1
2# for Chinese, Japanese additionally pip install pkuseg==0.0.25 pykakasi==2.3.0
3
4import onnxruntime
5
6from huggingface_hub import hf_hub_download
7from transformers import AutoTokenizer
8
9import numpy as np
10from tqdm import tqdm
11import librosa
12import soundfile as sf
13from unicodedata import category
14import json
15
16S3GEN_SR = 24000
17START_SPEECH_TOKEN = 6561
18STOP_SPEECH_TOKEN = 6562
19SUPPORTED_LANGUAGES = {
20 "ar": "Arabic",
21 "da": "Danish",
22 "de": "German",
23 "el": "Greek",
24 "en": "English",
25 "es": "Spanish",
26 "fi": "Finnish",
27 "fr": "French",
28 "he": "Hebrew",
29 "hi": "Hindi",
30 "it": "Italian",
31 "ja": "Japanese",
32 "ko": "Korean",
33 "ms": "Malay",
34 "nl": "Dutch",
35 "no": "Norwegian",
36 "pl": "Polish",
37 "pt": "Portuguese",
38 "ru": "Russian",
39 "sv": "Swedish",
40 "sw": "Swahili",
41 "tr": "Turkish",
42 "zh": "Chinese",
43}
44
45
46class RepetitionPenaltyLogitsProcessor:
47 def __init__(self, penalty: float):
48 if not isinstance(penalty, float) or not (penalty > 0):
49 raise ValueError(f"`penalty` must be a strictly positive float, but is {penalty}")
50 self.penalty = penalty
51
52 def __call__(self, input_ids: np.ndarray, scores: np.ndarray) -> np.ndarray:
53 score = np.take_along_axis(scores, input_ids, axis=1)
54 score = np.where(score < 0, score * self.penalty, score / self.penalty)
55 scores_processed = scores.copy()
56 np.put_along_axis(scores_processed, input_ids, score, axis=1)
57 return scores_processed
58
59
60class ChineseCangjieConverter:
61 """Converts Chinese characters to Cangjie codes for tokenization."""
62
63 def __init__(self):
64 self.word2cj = {}
65 self.cj2word = {}
66 self.segmenter = None
67 self._load_cangjie_mapping()
68 self._init_segmenter()
69
70 def _load_cangjie_mapping(self):
71 """Load Cangjie mapping from HuggingFace model repository."""
72 try:
73 cangjie_file = hf_hub_download(
74 repo_id="onnx-community/chatterbox-multilingual-ONNX",
75 filename="Cangjie5_TC.json",
76 )
77
78 with open(cangjie_file, "r", encoding="utf-8") as fp:
79 data = json.load(fp)
80
81 for entry in data:
82 word, code = entry.split("\t")[:2]
83 self.word2cj[word] = code
84 if code not in self.cj2word:
85 self.cj2word[code] = [word]
86 else:
87 self.cj2word[code].append(word)
88
89 except Exception as e:
90 print(f"Could not load Cangjie mapping: {e}")
91
92 def _init_segmenter(self):
93 """Initialize pkuseg segmenter."""
94 try:
95 from pkuseg import pkuseg
96 self.segmenter = pkuseg()
97 except ImportError:
98 print("pkuseg not available - Chinese segmentation will be skipped")
99 self.segmenter = None
100
101 def _cangjie_encode(self, glyph: str):
102 """Encode a single Chinese glyph to Cangjie code."""
103 normed_glyph = glyph
104 code = self.word2cj.get(normed_glyph, None)
105 if code is None: # e.g. Japanese hiragana
106 return None
107 index = self.cj2word[code].index(normed_glyph)
108 index = str(index) if index > 0 else ""
109 return code + str(index)
110
111
112
113 def __call__(self, text):
114 """Convert Chinese characters in text to Cangjie tokens."""
115 output = []
116 if self.segmenter is not None:
117 segmented_words = self.segmenter.cut(text)
118 full_text = " ".join(segmented_words)
119 else:
120 full_text = text
121
122 for t in full_text:
123 if category(t) == "Lo":
124 cangjie = self._cangjie_encode(t)
125 if cangjie is None:
126 output.append(t)
127 continue
128 code = []
129 for c in cangjie:
130 code.append(f"[cj_{c}]")
131 code.append("[cj_.]")
132 code = "".join(code)
133 output.append(code)
134 else:
135 output.append(t)
136 return "".join(output)
137
138
139def is_kanji(c: str) -> bool:
140 """Check if character is kanji."""
141 return 19968 <= ord(c) <= 40959
142
143
144def is_katakana(c: str) -> bool:
145 """Check if character is katakana."""
146 return 12449 <= ord(c) <= 12538
147
148
149def hiragana_normalize(text: str) -> str:
150 """Japanese text normalization: converts kanji to hiragana; katakana remains the same."""
151 global _kakasi
152
153 try:
154 if _kakasi is None:
155 import pykakasi
156 _kakasi = pykakasi.kakasi()
157
158 result = _kakasi.convert(text)
159 out = []
160
161 for r in result:
162 inp = r['orig']
163 hira = r["hira"]
164
165 # Any kanji in the phrase
166 if any([is_kanji(c) for c in inp]):
167 if hira and hira[0] in ["は", "へ"]: # Safety check for empty hira
168 hira = " " + hira
169 out.append(hira)
170
171 # All katakana
172 elif all([is_katakana(c) for c in inp]) if inp else False: # Safety check for empty inp
173 out.append(r['orig'])
174
175 else:
176 out.append(inp)
177
178 normalized_text = "".join(out)
179
180 # Decompose Japanese characters for tokenizer compatibility
181 import unicodedata
182 normalized_text = unicodedata.normalize('NFKD', normalized_text)
183
184 return normalized_text
185
186 except ImportError:
187 print("pykakasi not available - Japanese text processing skipped")
188 return text
189
190
191def add_hebrew_diacritics(text: str) -> str:
192 """Hebrew text normalization: adds diacritics to Hebrew text."""
193 global _dicta
194
195 try:
196 if _dicta is None:
197 from dicta_onnx import Dicta
198 _dicta = Dicta()
199
200 return _dicta.add_diacritics(text)
201
202 except ImportError:
203 print("dicta_onnx not available - Hebrew text processing skipped")
204 return text
205 except Exception as e:
206 print(f"Hebrew diacritization failed: {e}")
207 return text
208
209
210def korean_normalize(text: str) -> str:
211 """Korean text normalization: decompose syllables into Jamo for tokenization."""
212
213 def decompose_hangul(char):
214 """Decompose Korean syllable into Jamo components."""
215 if not ('\uac00' <= char <= '\ud7af'):
216 return char
217
218 # Hangul decomposition formula
219 base = ord(char) - 0xAC00
220 initial = chr(0x1100 + base // (21 * 28))
221 medial = chr(0x1161 + (base % (21 * 28)) // 28)
222 final = chr(0x11A7 + base % 28) if base % 28 > 0 else ''
223
224 return initial + medial + final
225
226 # Decompose syllables and normalize punctuation
227 result = ''.join(decompose_hangul(char) for char in text)
228 return result.strip()
229
230
231def prepare_language(txt, language_id):
232 # Language-specific text processing
233 cangjie_converter = ChineseCangjieConverter()
234 if language_id == 'zh':
235 txt = cangjie_converter(txt)
236 elif language_id == 'ja':
237 txt = hiragana_normalize(txt)
238 elif language_id == 'he':
239 txt = add_hebrew_diacritics(txt)
240 elif language_id == 'ko':
241 txt = korean_normalize(txt)
242
243 # Prepend language token
244 if language_id:
245 txt = f"[{language_id.lower()}]{txt}"
246 return txt
247
248
249def run_inference(
250 text="The Lord of the Rings is the greatest work of literature.",
251 language_id="en",
252 target_voice_path=None,
253 max_new_tokens=256,
254 exaggeration=0.5,
255 output_dir="converted",
256 output_file_name="output.wav",
257 apply_watermark=True,
258):
259 # Validate language_id
260 if language_id and language_id.lower() not in SUPPORTED_LANGUAGES:
261 supported_langs = ", ".join(SUPPORTED_LANGUAGES.keys())
262 raise ValueError(
263 f"Unsupported language_id '{language_id}'. "
264 f"Supported languages: {supported_langs}"
265 )
266 model_id = "onnx-community/chatterbox-multilingual-ONNX"
267 if not target_voice_path:
268 target_voice_path = hf_hub_download(repo_id=model_id, filename="default_voice.wav", local_dir=output_dir)
269
270 ## Load model
271 speech_encoder_path = hf_hub_download(repo_id=model_id, filename="speech_encoder.onnx", local_dir=output_dir, subfolder='onnx')
272 hf_hub_download(repo_id=model_id, filename="speech_encoder.onnx_data", local_dir=output_dir, subfolder='onnx')
273 embed_tokens_path = hf_hub_download(repo_id=model_id, filename="embed_tokens.onnx", local_dir=output_dir, subfolder='onnx')
274 hf_hub_download(repo_id=model_id, filename="embed_tokens.onnx_data", local_dir=output_dir, subfolder='onnx')
275 conditional_decoder_path = hf_hub_download(repo_id=model_id, filename="conditional_decoder.onnx", local_dir=output_dir, subfolder='onnx')
276 hf_hub_download(repo_id=model_id, filename="conditional_decoder.onnx_data", local_dir=output_dir, subfolder='onnx')
277 language_model_path = hf_hub_download(repo_id=model_id, filename="language_model.onnx", local_dir=output_dir, subfolder='onnx')
278 hf_hub_download(repo_id=model_id, filename="language_model.onnx_data", local_dir=output_dir, subfolder='onnx')
279
280 # # Start inferense sessions
281 speech_encoder_session = onnxruntime.InferenceSession(speech_encoder_path)
282 embed_tokens_session = onnxruntime.InferenceSession(embed_tokens_path)
283 llama_with_past_session = onnxruntime.InferenceSession(language_model_path)
284 cond_decoder_session = onnxruntime.InferenceSession(conditional_decoder_path)
285
286 def execute_text_to_audio_inference(text):
287 print("Start inference script...")
288
289 audio_values, _ = librosa.load(target_voice_path, sr=S3GEN_SR)
290 audio_values = audio_values[np.newaxis, :].astype(np.float32)
291
292 ## Prepare input
293 tokenizer = AutoTokenizer.from_pretrained(model_id)
294 text = prepare_language(text, language_id)
295 input_ids = tokenizer(text, return_tensors="np")["input_ids"].astype(np.int64)
296
297 position_ids = np.where(
298 input_ids >= START_SPEECH_TOKEN,
299 0,
300 np.arange(input_ids.shape[1])[np.newaxis, :] - 1
301 )
302
303 ort_embed_tokens_inputs = {
304 "input_ids": input_ids,
305 "position_ids": position_ids.astype(np.int64),
306 "exaggeration": np.array([exaggeration], dtype=np.float32)
307 }
308
309 ## Instantiate the logits processors.
310 repetition_penalty = 1.2
311 repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)
312
313 num_hidden_layers = 30
314 num_key_value_heads = 16
315 head_dim = 64
316
317 generate_tokens = np.array([[START_SPEECH_TOKEN]])
318
319 # ---- Generation Loop using kv_cache ----
320 for i in tqdm(range(max_new_tokens), desc="Sampling", dynamic_ncols=True):
321
322 inputs_embeds = embed_tokens_session.run(None, ort_embed_tokens_inputs)[0]
323 if i == 0:
324 ort_speech_encoder_input = {
325 "audio_values": audio_values,
326 }
327 cond_emb, prompt_token, ref_x_vector, prompt_feat = speech_encoder_session.run(None, ort_speech_encoder_input)
328 inputs_embeds = np.concatenate((cond_emb, inputs_embeds), axis=1)
329
330 ## Prepare llm inputs
331 batch_size, seq_len, _ = inputs_embeds.shape
332 past_key_values = {
333 f"past_key_values.{layer}.{kv}": np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=np.float32)
334 for layer in range(num_hidden_layers)
335 for kv in ("key", "value")
336 }
337 attention_mask = np.ones((batch_size, seq_len), dtype=np.int64)
338 logits, *present_key_values = llama_with_past_session.run(None, dict(
339 inputs_embeds=inputs_embeds,
340 attention_mask=attention_mask,
341 **past_key_values,
342 ))
343
344 logits = logits[:, -1, :]
345 next_token_logits = repetition_penalty_processor(generate_tokens, logits)
346
347 next_token = np.argmax(next_token_logits, axis=-1, keepdims=True).astype(np.int64)
348 generate_tokens = np.concatenate((generate_tokens, next_token), axis=-1)
349 if (next_token.flatten() == STOP_SPEECH_TOKEN).all():
350 break
351
352 # Get embedding for the new token.
353 position_ids = np.full(
354 (input_ids.shape[0], 1),
355 i + 1,
356 dtype=np.int64,
357 )
358 ort_embed_tokens_inputs["input_ids"] = next_token
359 ort_embed_tokens_inputs["position_ids"] = position_ids
360
361 ## Update values for next generation loop
362 attention_mask = np.concatenate([attention_mask, np.ones((batch_size, 1), dtype=np.int64)], axis=1)
363 for j, key in enumerate(past_key_values):
364 past_key_values[key] = present_key_values[j]
365
366 speech_tokens = generate_tokens[:, 1:-1]
367 speech_tokens = np.concatenate([prompt_token, speech_tokens], axis=1)
368 return speech_tokens, ref_x_vector, prompt_feat
369
370 speech_tokens, speaker_embeddings, speaker_features = execute_text_to_audio_inference(text)
371 cond_incoder_input = {
372 "speech_tokens": speech_tokens,
373 "speaker_embeddings": speaker_embeddings,
374 "speaker_features": speaker_features,
375 }
376 wav = cond_decoder_session.run(None, cond_incoder_input)[0]
377 wav = np.squeeze(wav, axis=0)
378
379 # Optional: Apply watermark
380 if apply_watermark:
381 import perth
382 watermarker = perth.PerthImplicitWatermarker()
383 wav = watermarker.apply_watermark(wav, sample_rate=S3GEN_SR)
384
385 sf.write(output_file_name, wav, S3GEN_SR)
386 print(f"{output_file_name} was successfully saved")
387
388if __name__ == "__main__":
389 run_inference(
390 text="Bonjour, comment ça va? Ceci est le modèle de synthèse vocale multilingue Chatterbox, il prend en charge 23 langues.",
391 language_id="fr",
392 exaggeration=0.5,
393 output_file_name="output.wav",
394 apply_watermark=False,
395 )
396