Views
No views yet
t3_23lang.safetensors), not the latest t3_mtl23ls_v2.safetensors that the PyTorch ChatterboxMultilingualTTS uses. This caused 61 weight mismatches across the language model (including a 0.42 max diff in the final layer norm) and 3 mismatched embedding weights. We verified this numerically — every weight now exactly matches the PyTorch model.language_model.onnx was re-exported from the correct Llama backbone using transformers==4.46.3 to ensure the RoPE (rotary position embeddings) cos/sin caches are computed correctly. Different transformers versions produce different RoPE implementations (32-dim vs 64-dim), causing completely divergent attention outputs even with identical weights.embed_tokens.onnx had text_emb.weight with shape (2352, 1024) but the tokenizer produces IDs up to 2454. The 102 missing rows caused out-of-bounds crashes for Polish diacritics (ą, ę, ć, ś, ź, ż, ń, ó) and other non-ASCII characters. Fixed by extracting the full (2454, 1024) embeddings from the correct PyTorch checkpoint.MTLTokenizer.preprocess_text() always applies:text.lower() — the model was trained on lowercase textunicodedata.normalize("NFKD", text) — decomposes diacritics into base character + combining mark (e.g., ą → a + ̨), which is how the tokenizer was trained| Parameter | Value | Notes |
|---|---|---|
temperature | 0.3 | Low temp prevents language drift without CFG |
repetition_penalty | 2.0 | Matches PyTorch multilingual default |
exaggeration | 0.5 | Default emotion intensity |
cfg_weight | 0.0 | CFG not yet implemented for ONNX inference |
min_p | 0.05 | Filters low-probability tokens |
max_tokens | 500 | Enough for most sentences |
Note: The PyTorch model uses Classifier-Free Guidance (CFG) withcfg_weight=0.5andtemperature=0.8, which allows higher temperature without language drift. CFG requires running the language model with batch=2 (conditional + unconditional branches). This is not yet implemented in the ONNX inference script, so we compensate with lower temperature.
exaggeration=0.5, temperature=0.3) work well for most prompts.temperature=0 for fully deterministic (greedy) output.exaggeration to around 0.7 or higher.exaggeration tends to speed up speech.text.lower() + NFKD normalization.test_polish.py can be used as a standalone inference script:1pip install onnxruntime transformers numpy tqdm librosa soundfile
2
3# Default settings (Polish, temp=0.3)
4python test_polish.py
5
6# Custom text and settings
7python test_polish.py --text "Cześć, jak się masz?" --temperature 0.5 --exaggeration 0.7 -o output.wav
8
9# Greedy (deterministic)
10python test_polish.py --temperature 0 -o output_greedy.wav1# !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, normalize
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="Folx/chatterbox-ONNX-polish",
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 def __call__(self, text):
112 """Convert Chinese characters in text to Cangjie tokens."""
113 output = []
114 if self.segmenter is not None:
115 segmented_words = self.segmenter.cut(text)
116 full_text = " ".join(segmented_words)
117 else:
118 full_text = text
119
120 for t in full_text:
121 if category(t) == "Lo":
122 cangjie = self._cangjie_encode(t)
123 if cangjie is None:
124 output.append(t)
125 continue
126 code = []
127 for c in cangjie:
128 code.append(f"[cj_{c}]")
129 code.append("[cj_.]")
130 code = "".join(code)
131 output.append(code)
132 else:
133 output.append(t)
134 return "".join(output)
135
136
137def is_kanji(c: str) -> bool:
138 """Check if character is kanji."""
139 return 19968 <= ord(c) <= 40959
140
141
142def is_katakana(c: str) -> bool:
143 """Check if character is katakana."""
144 return 12449 <= ord(c) <= 12538
145
146
147def hiragana_normalize(text: str) -> str:
148 """Japanese text normalization: converts kanji to hiragana; katakana remains the same."""
149 global _kakasi
150
151 try:
152 if _kakasi is None:
153 import pykakasi
154 _kakasi = pykakasi.kakasi()
155
156 result = _kakasi.convert(text)
157 out = []
158
159 for r in result:
160 inp = r['orig']
161 hira = r["hira"]
162
163 # Any kanji in the phrase
164 if any([is_kanji(c) for c in inp]):
165 if hira and hira[0] in ["は", "へ"]: # Safety check for empty hira
166 hira = " " + hira
167 out.append(hira)
168
169 # All katakana
170 elif all([is_katakana(c) for c in inp]) if inp else False: # Safety check for empty inp
171 out.append(r['orig'])
172
173 else:
174 out.append(inp)
175
176 normalized_text = "".join(out)
177
178 # Decompose Japanese characters for tokenizer compatibility
179 import unicodedata
180 normalized_text = unicodedata.normalize('NFKD', normalized_text)
181
182 return normalized_text
183
184 except ImportError:
185 print("pykakasi not available - Japanese text processing skipped")
186 return text
187
188
189def add_hebrew_diacritics(text: str) -> str:
190 """Hebrew text normalization: adds diacritics to Hebrew text."""
191 global _dicta
192
193 try:
194 if _dicta is None:
195 from dicta_onnx import Dicta
196 _dicta = Dicta()
197
198 return _dicta.add_diacritics(text)
199
200 except ImportError:
201 print("dicta_onnx not available - Hebrew text processing skipped")
202 return text
203 except Exception as e:
204 print(f"Hebrew diacritization failed: {e}")
205 return text
206
207
208def korean_normalize(text: str) -> str:
209 """Korean text normalization: decompose syllables into Jamo for tokenization."""
210
211 def decompose_hangul(char):
212 """Decompose Korean syllable into Jamo components."""
213 if not ('\uac00' <= char <= '\ud7af'):
214 return char
215
216 # Hangul decomposition formula
217 base = ord(char) - 0xAC00
218 initial = chr(0x1100 + base // (21 * 28))
219 medial = chr(0x1161 + (base % (21 * 28)) // 28)
220 final = chr(0x11A7 + base % 28) if base % 28 > 0 else ''
221
222 return initial + medial + final
223
224 # Decompose syllables and normalize punctuation
225 result = ''.join(decompose_hangul(char) for char in text)
226 return result.strip()
227
228
229def prepare_language(txt, language_id):
230 # IMPORTANT: lowercase and NFKD-normalize text to match PyTorch training preprocessing.
231 # Without this, diacritics in Polish, German, French, etc. produce wrong token IDs.
232 txt = txt.lower()
233 txt = normalize("NFKD", txt)
234
235 # Language-specific text processing
236 cangjie_converter = ChineseCangjieConverter()
237 if language_id == 'zh':
238 txt = cangjie_converter(txt)
239 elif language_id == 'ja':
240 txt = hiragana_normalize(txt)
241 elif language_id == 'he':
242 txt = add_hebrew_diacritics(txt)
243 elif language_id == 'ko':
244 txt = korean_normalize(txt)
245
246 # Prepend language token
247 if language_id:
248 txt = f"[{language_id.lower()}]{txt}"
249 return txt
250
251
252def run_inference(
253 text="The Lord of the Rings is the greatest work of literature.",
254 language_id="en",
255 target_voice_path=None,
256 max_new_tokens=256,
257 exaggeration=0.5,
258 output_dir="converted",
259 output_file_name="output.wav",
260 apply_watermark=True,
261):
262 # Validate language_id
263 if language_id and language_id.lower() not in SUPPORTED_LANGUAGES:
264 supported_langs = ", ".join(SUPPORTED_LANGUAGES.keys())
265 raise ValueError(
266 f"Unsupported language_id '{language_id}'. "
267 f"Supported languages: {supported_langs}"
268 )
269 model_id = "Folx/chatterbox-ONNX-polish"
270 if not target_voice_path:
271 target_voice_path = hf_hub_download(repo_id=model_id, filename="default_voice.wav", local_dir=output_dir)
272
273 ## Load model
274 speech_encoder_path = hf_hub_download(repo_id=model_id, filename="speech_encoder.onnx", local_dir=output_dir, subfolder='onnx')
275 hf_hub_download(repo_id=model_id, filename="speech_encoder.onnx_data", local_dir=output_dir, subfolder='onnx')
276 embed_tokens_path = hf_hub_download(repo_id=model_id, filename="embed_tokens.onnx", local_dir=output_dir, subfolder='onnx')
277 hf_hub_download(repo_id=model_id, filename="embed_tokens.onnx_data", local_dir=output_dir, subfolder='onnx')
278 conditional_decoder_path = hf_hub_download(repo_id=model_id, filename="conditional_decoder.onnx", local_dir=output_dir, subfolder='onnx')
279 hf_hub_download(repo_id=model_id, filename="conditional_decoder.onnx_data", local_dir=output_dir, subfolder='onnx')
280 language_model_path = hf_hub_download(repo_id=model_id, filename="language_model.onnx", local_dir=output_dir, subfolder='onnx')
281 hf_hub_download(repo_id=model_id, filename="language_model.onnx_data", local_dir=output_dir, subfolder='onnx')
282
283 # # Start inferense sessions
284 speech_encoder_session = onnxruntime.InferenceSession(speech_encoder_path)
285 embed_tokens_session = onnxruntime.InferenceSession(embed_tokens_path)
286 llama_with_past_session = onnxruntime.InferenceSession(language_model_path)
287 cond_decoder_session = onnxruntime.InferenceSession(conditional_decoder_path)
288
289 def execute_text_to_audio_inference(text):
290 print("Start inference script...")
291
292 audio_values, _ = librosa.load(target_voice_path, sr=S3GEN_SR)
293 audio_values = audio_values[np.newaxis, :].astype(np.float32)
294
295 ## Prepare input
296 tokenizer = AutoTokenizer.from_pretrained(model_id)
297 text = prepare_language(text, language_id)
298 input_ids = tokenizer(text, return_tensors="np")["input_ids"].astype(np.int64)
299
300 position_ids = np.where(
301 input_ids >= START_SPEECH_TOKEN,
302 0,
303 np.arange(input_ids.shape[1])[np.newaxis, :] - 1
304 )
305
306 ort_embed_tokens_inputs = {
307 "input_ids": input_ids,
308 "position_ids": position_ids.astype(np.int64),
309 "exaggeration": np.array([exaggeration], dtype=np.float32)
310 }
311
312 ## Instantiate the logits processors.
313 repetition_penalty = 1.2
314 repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)
315
316 num_hidden_layers = 30
317 num_key_value_heads = 16
318 head_dim = 64
319
320 generate_tokens = np.array([[START_SPEECH_TOKEN]])
321
322 # ---- Generation Loop using kv_cache ----
323 for i in tqdm(range(max_new_tokens), desc="Sampling", dynamic_ncols=True):
324
325 inputs_embeds = embed_tokens_session.run(None, ort_embed_tokens_inputs)[0]
326 if i == 0:
327 ort_speech_encoder_input = {
328 "audio_values": audio_values,
329 }
330 cond_emb, prompt_token, ref_x_vector, prompt_feat = speech_encoder_session.run(None, ort_speech_encoder_input)
331 inputs_embeds = np.concatenate((cond_emb, inputs_embeds), axis=1)
332
333 ## Prepare llm inputs
334 batch_size, seq_len, _ = inputs_embeds.shape
335 past_key_values = {
336 f"past_key_values.{layer}.{kv}": np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=np.float32)
337 for layer in range(num_hidden_layers)
338 for kv in ("key", "value")
339 }
340 attention_mask = np.ones((batch_size, seq_len), dtype=np.int64)
341 logits, *present_key_values = llama_with_past_session.run(None, dict(
342 inputs_embeds=inputs_embeds,
343 attention_mask=attention_mask,
344 **past_key_values,
345 ))
346
347 logits = logits[:, -1, :]
348 next_token_logits = repetition_penalty_processor(generate_tokens, logits)
349
350 next_token = np.argmax(next_token_logits, axis=-1, keepdims=True).astype(np.int64)
351 generate_tokens = np.concatenate((generate_tokens, next_token), axis=-1)
352 if (next_token.flatten() == STOP_SPEECH_TOKEN).all():
353 break
354
355 # Get embedding for the new token.
356 position_ids = np.full(
357 (input_ids.shape[0], 1),
358 i + 1,
359 dtype=np.int64,
360 )
361 ort_embed_tokens_inputs["input_ids"] = next_token
362 ort_embed_tokens_inputs["position_ids"] = position_ids
363
364 ## Update values for next generation loop
365 attention_mask = np.concatenate([attention_mask, np.ones((batch_size, 1), dtype=np.int64)], axis=1)
366 for j, key in enumerate(past_key_values):
367 past_key_values[key] = present_key_values[j]
368
369 speech_tokens = generate_tokens[:, 1:-1]
370 speech_tokens = np.concatenate([prompt_token, speech_tokens], axis=1)
371 return speech_tokens, ref_x_vector, prompt_feat
372
373 speech_tokens, speaker_embeddings, speaker_features = execute_text_to_audio_inference(text)
374 cond_incoder_input = {
375 "speech_tokens": speech_tokens,
376 "speaker_embeddings": speaker_embeddings,
377 "speaker_features": speaker_features,
378 }
379 wav = cond_decoder_session.run(None, cond_incoder_input)[0]
380 wav = np.squeeze(wav, axis=0)
381
382 # Optional: Apply watermark
383 if apply_watermark:
384 import perth
385 watermarker = perth.PerthImplicitWatermarker()
386 wav = watermarker.apply_watermark(wav, sample_rate=S3GEN_SR)
387
388 sf.write(output_file_name, wav, S3GEN_SR)
389 print(f"{output_file_name} was successfully saved")
390
391if __name__ == "__main__":
392 run_inference(
393 text="Dzień dobry, nazywam się Weronika. Dziś jest piękna pogoda i chciałabym zaprosić wszystkich na spacer po parku.",
394 language_id="pl",
395 target_voice_path="polish-female-common-voice.wav",
396 exaggeration=0.5,
397 output_file_name="output.wav",
398 apply_watermark=False,
399 )
400language_model.onnx from the correct checkpoint (t3_mtl23ls_v2.safetensors) using transformers==4.46.3 for correct RoPE computationembed_tokens.onnx from the correct checkpointtext_emb.weight from (2352, 1024) to (2454, 1024) to cover full tokenizer vocabularytest_polish.py with tuned defaults and CLI interface