Upload folder using huggingface_hub

README.md

@@ -72,23 +72,77 @@ Chatterbox is provided in an exported ONNX format, enabling fast and portable in
 - Easy voice conversion script using onnxruntime
 - [Outperforms ElevenLabs](https://podonos.com/resembleai/chatterbox)
 
-# Changelog
+# What's fixed in this repo
+
+This is a corrected ONNX export of [ResembleAI/chatterbox](https://huggingface.co/ResembleAI/chatterbox) multilingual TTS. The original community ONNX exports ([onnx-community/chatterbox-multilingual-ONNX](https://huggingface.co/onnx-community/chatterbox-multilingual-ONNX)) had multiple issues that caused **incorrect pronunciation for Polish and other diacritics-heavy languages** (sounding like Russian instead of the target language). This repo fixes all of them. Polish speech generation is now on par with the PyTorch reference model.
+
+### Issues found and fixed
+
+**1. Wrong model weights (root cause of "Russian accent")**
+
+The original ONNX export used weights from an older checkpoint (`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.
+
+Additionally, the `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.
+
+**2. Truncated embedding vocabulary**
+
+`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.
+
+**3. Missing text preprocessing**
+
+The ONNX inference code was missing two critical preprocessing steps that the PyTorch `MTLTokenizer.preprocess_text()` always applies:
+- `text.lower()` — the model was trained on lowercase text
+- `unicodedata.normalize("NFKD", text)` — decomposes diacritics into base character + combining mark (e.g., `ą` → `a` + `̨`), which is how the tokenizer was trained
+
+Without NFKD normalization, composed Polish characters hit the wrong (or missing) token embeddings.
+
+### Recommended settings
 
-- **2025-03-30**: Fixed Polish and other diacritics-heavy languages (German, French, etc.):
-  - **embed_tokens.onnx**: Expanded `text_emb.weight` from `(2352, 1024)` to `(2454, 1024)` to match the full tokenizer vocabulary. The previous export truncated 102 embedding rows, causing out-of-bounds errors for Polish characters (ą, ę, ć, ś, ź, ż, ń, ó).
-  - **Preprocessing fix**: Added mandatory `text.lower()` + NFKD Unicode normalization before tokenization, matching the original PyTorch `MTLTokenizer.preprocess_text()`. Without NFKD, composed diacritics (e.g. `ą` U+0105) are tokenized as single unknown tokens instead of being decomposed into base character + combining mark, which is how the model 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) with `cfg_weight=0.5` and `temperature=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.
 
 # Tips
 - **General Use (TTS and Voice Agents):**
-  - The default settings (`exaggeration=0.5`, `cfg=0.5`) work well for most prompts.
+  - The default settings (`exaggeration=0.5`, `temperature=0.3`) work well for most prompts.
+  - Use `temperature=0` for fully deterministic (greedy) output.
 
 - **Expressive or Dramatic Speech:**
-  - Try increase `exaggeration` to around `0.7` or higher.
-  - Higher `exaggeration` tends to speed up speech;
+  - Try increasing `exaggeration` to around `0.7` or higher.
+  - Higher `exaggeration` tends to speed up speech.
 
+- **If you hear wrong accent/language:**
+  - Make sure text preprocessing includes `text.lower()` + NFKD normalization.
+  - Keep temperature low (0.0–0.3) until CFG is implemented.
+  - Verify you are using the ONNX models from this repo, not the original community exports.
 
 # Usage
-[Link to GitHub ONNX Export and Inference script](https://github.com/VladOS95-cyber/onnx_conversion_scripts/tree/main/chatterbox)
+
+Included `test_polish.py` can be used as a standalone inference script:
+
+```bash
+pip install onnxruntime transformers numpy tqdm librosa soundfile
+
+# Default settings (Polish, temp=0.3)
+python test_polish.py
+
+# Custom text and settings
+python test_polish.py --text "Cześć, jak się masz?" --temperature 0.5 --exaggeration 0.7 -o output.wav
+
+# Greedy (deterministic)
+python test_polish.py --temperature 0 -o output_greedy.wav
+```
+
+### Full inference code
+
+[Link to original GitHub ONNX conversion scripts](https://github.com/VladOS95-cyber/onnx_conversion_scripts/tree/main/chatterbox)
 
 ```python
 # !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
@@ -483,8 +537,9 @@ def run_inference(
 
 if __name__ == "__main__":
     run_inference(
-        text="Bonjour, comment ça va? Ceci est le modèle de synthèse vocale multilingue Chatterbox, il prend en charge 23 langues.",
-        language_id="fr",
+        text="Dzień dobry, nazywam się Weronika. Dziś jest piękna pogoda i chciałabym zaprosić wszystkich na spacer po parku.",
+        language_id="pl",
+        target_voice_path="polish-female-common-voice.wav",
         exaggeration=0.5,
         output_file_name="output.wav",
         apply_watermark=False,
@@ -493,10 +548,21 @@ if __name__ == "__main__":
 ```
 
 
+# Changelog
+
+- **2026-03-31**: Full model re-export and Polish language fix
+  - Re-exported `language_model.onnx` from the correct checkpoint (`t3_mtl23ls_v2.safetensors`) using `transformers==4.46.3` for correct RoPE computation
+  - Patched all 4 embedding weights in `embed_tokens.onnx` from the correct checkpoint
+  - Added NFKD normalization + lowercase preprocessing to match PyTorch training pipeline
+  - Expanded `text_emb.weight` from `(2352, 1024)` to `(2454, 1024)` to cover full tokenizer vocabulary
+  - Added `test_polish.py` with tuned defaults and CLI interface
+  - Included Polish reference voice and example output
+
 # Acknowledgements
 - [Xenova](https://huggingface.co/Xenova)
 - [Vladislav Bronzov](https://github.com/VladOS95-cyber)
 - [Resemble AI](https://github.com/resemble-ai/chatterbox)
+- [Folx](https://huggingface.co/Folx) — Polish language fixes and correct model re-export
 
 # Built-in PerTh Watermarking for Responsible AI
 

example_output_polish.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:812239fdc8d163d71adf4101029c0294dbc12d458a78c732dddf2732f40935db
+size 358124

onnx/embed_tokens.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:aff4b0863dec868c9ccc589d7d7617ea54c6b2377822e138bb023b83acf9f794
-size 13288
+oid sha256:cd3449f120c1738c34c1908de07ee4fccbc12f2723e88b89dba5f037b254a6f3
+size 13200

onnx/embed_tokens.onnx_data

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3222d642b7f3ece8d1b12072b20e7f8e6fda53d70641cbe0383ad6f963cc9dfc
+oid sha256:d7c50b7240475f19e94675869460a1204f6d8a72348bf72c1814e5a004b491c9
 size 68808704

onnx/language_model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:861a34585605e8ad671051788afc495dcbeaee833a41523a1b33aded9c3babc7
-size 171387
+oid sha256:3a38c047d205122aa36e43c8a6e5ac027b391ce90f522cfa206f41ab39a4bdea
+size 657688

onnx/language_model.onnx_data

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b3556d41085196c122b7197e4d44ec4475b6d7cfe0971a70faa95caa38ad787a
-size 2080632832
+oid sha256:87a4e23357d396d42ef959165eaecfb95108963d2c2147b46b01cca130af734d
+size 2047078400

onnx/speech_encoder.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8f1c8a0f89b77bf9cd5dd8f2e034eb2c79dc00fe70d41196b28c257643b00ccb
-size 1184608
+oid sha256:36f42eb070f7f90423b51bea3475c784b4cc88756514f527bcfb26067b59ee01
+size 592665681

onnx_fixed/embed_tokens.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aff4b0863dec868c9ccc589d7d7617ea54c6b2377822e138bb023b83acf9f794
+size 13288

onnx_fixed/embed_tokens.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae9479b9f126b2ca973ab4057425a3820d36571b2133d453527bab05d6c5e359
+size 68808704

polish-female-common-voice.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49acc23c4513cb957f9591a38370db089d2121fa77c01d740ffd66ce3ae3bd94
+size 491322

test_polish.py

@@ -0,0 +1,288 @@
+"""Test Polish language TTS with chatterbox-ONNX-polish.
+Verifies that Polish diacritics (ą, ę, ć, ś, ź, ż, ł, ń, ó) are handled correctly.
+
+All generation parameters match the original PyTorch ChatterboxMultilingualTTS.generate() defaults.
+"""
+
+import argparse
+import os
+import onnxruntime
+from transformers import AutoTokenizer
+import numpy as np
+from tqdm import tqdm
+import librosa
+import soundfile as sf
+from unicodedata import normalize
+
+S3GEN_SR = 24000
+START_SPEECH_TOKEN = 6561
+STOP_SPEECH_TOKEN = 6562
+
+# Path to this script's directory (the model repo)
+MODEL_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+class RepetitionPenaltyLogitsProcessor:
+    def __init__(self, penalty: float):
+        self.penalty = penalty
+
+    def __call__(self, input_ids, scores):
+        score = np.take_along_axis(scores, input_ids, axis=1)
+        score = np.where(score < 0, score * self.penalty, score / self.penalty)
+        scores_processed = scores.copy()
+        np.put_along_axis(scores_processed, input_ids, score, axis=1)
+        return scores_processed
+
+
+def top_p_filter(logits, top_p):
+    """Nucleus (top-p) filtering: zero out tokens outside the top-p cumulative probability."""
+    if top_p >= 1.0:
+        return logits
+    sorted_indices = np.argsort(-logits, axis=-1)
+    sorted_logits = np.take_along_axis(logits, sorted_indices, axis=-1)
+    probs = np.exp(sorted_logits - np.max(sorted_logits, axis=-1, keepdims=True))
+    probs = probs / probs.sum(axis=-1, keepdims=True)
+    cumulative_probs = np.cumsum(probs, axis=-1)
+    # Remove tokens with cumulative probability above top_p (keep at least 1 token)
+    sorted_mask = cumulative_probs - probs > top_p
+    sorted_logits[sorted_mask] = -np.inf
+    # Unsort
+    result = np.empty_like(logits)
+    np.put_along_axis(result, sorted_indices, sorted_logits, axis=-1)
+    return result
+
+
+def min_p_filter(logits, min_p):
+    """Min-p filtering: remove tokens with probability below min_p * max_prob."""
+    if min_p <= 0.0:
+        return logits
+    probs = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
+    probs = probs / probs.sum(axis=-1, keepdims=True)
+    max_prob = probs.max(axis=-1, keepdims=True)
+    mask = probs < (min_p * max_prob)
+    logits = logits.copy()
+    logits[mask] = -np.inf
+    return logits
+
+
+def prepare_language(txt, language_id):
+    """Match PyTorch MTLTokenizer.preprocess_text(): lowercase + NFKD normalize."""
+    txt = txt.lower()
+    txt = normalize("NFKD", txt)
+    if language_id:
+        txt = f"[{language_id.lower()}]{txt}"
+    return txt
+
+
+def run_inference(
+    text,
+    language_id="pl",
+    target_voice_path=None,
+    # Generation parameters — tuned for ONNX multilingual quality
+    max_new_tokens=500,
+    exaggeration=0.5,        # 0.0-1.0: emotion/intensity control (unique to Chatterbox)
+    cfg_weight=0.0,          # CFG not yet supported in ONNX, keep at 0
+    temperature=0.3,         # low temp avoids language drift without CFG
+    top_p=1.0,               # nucleus sampling threshold (1.0 = disabled)
+    min_p=0.05,              # min probability filter (relative to top token)
+    repetition_penalty=2.0,  # higher penalty prevents repetition loops
+    seed=None,               # random seed for reproducibility
+    model_dir=MODEL_DIR,
+    output_file_name="example_output_polish.wav",
+):
+    if seed is not None:
+        np.random.seed(seed)
+
+    onnx_dir = os.path.join(model_dir, "onnx")
+
+    if not target_voice_path:
+        target_voice_path = os.path.join(model_dir, "default_voice.wav")
+
+    print(f"Text: {text}")
+    print(f"Settings: exaggeration={exaggeration}, cfg_weight={cfg_weight}, "
+          f"temperature={temperature}, top_p={top_p}, min_p={min_p}, "
+          f"repetition_penalty={repetition_penalty}, seed={seed}")
+    print(f"Voice: {os.path.basename(target_voice_path)}")
+
+    speech_encoder_session = onnxruntime.InferenceSession(os.path.join(onnx_dir, "speech_encoder.onnx"))
+    embed_tokens_session = onnxruntime.InferenceSession(os.path.join(onnx_dir, "embed_tokens.onnx"))
+    llama_with_past_session = onnxruntime.InferenceSession(os.path.join(onnx_dir, "language_model.onnx"))
+    cond_decoder_session = onnxruntime.InferenceSession(os.path.join(onnx_dir, "conditional_decoder.onnx"))
+
+    audio_values, _ = librosa.load(target_voice_path, sr=S3GEN_SR)
+    audio_values = audio_values[np.newaxis, :].astype(np.float32)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    text_prepared = prepare_language(text, language_id)
+    input_ids = tokenizer(text_prepared, return_tensors="np")["input_ids"].astype(np.int64)
+
+    position_ids = np.where(
+        input_ids >= START_SPEECH_TOKEN, 0,
+        np.arange(input_ids.shape[1])[np.newaxis, :] - 1
+    )
+
+    ort_embed_tokens_inputs = {
+        "input_ids": input_ids,
+        "position_ids": position_ids.astype(np.int64),
+        "exaggeration": np.array([exaggeration], dtype=np.float32)
+    }
+
+    repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)
+    num_hidden_layers = 30
+    num_key_value_heads = 16
+    head_dim = 64
+
+    use_cfg = cfg_weight > 0.0
+    # Batch size: 2 for CFG (conditional + unconditional), 1 without
+    B = 2 if use_cfg else 1
+
+    # Load text position embeddings for CFG uncond branch
+    # PyTorch zeros text content but KEEPS position embeddings in the uncond branch
+    text_pos_emb_weight = None
+    if use_cfg:
+        import onnx
+        from onnx import numpy_helper
+        onnx_model = onnx.load(os.path.join(onnx_dir, "embed_tokens.onnx"), load_external_data=True)
+        for init in onnx_model.graph.initializer:
+            if init.name == "text_pos_emb.weight":
+                text_pos_emb_weight = numpy_helper.to_array(init)  # (2050, 1024)
+                break
+        del onnx_model
+
+    generate_tokens = np.array([[START_SPEECH_TOKEN]])
+
+    # --- Initial step: build full context embeddings ---
+    inputs_embeds = embed_tokens_session.run(None, ort_embed_tokens_inputs)[0]
+    cond_emb, prompt_token, ref_x_vector, prompt_feat = speech_encoder_session.run(
+        None, {"audio_values": audio_values}
+    )
+    cond_inputs = np.concatenate((cond_emb, inputs_embeds), axis=1)
+
+    if use_cfg:
+        # CFG unconditional: replace text embeddings with ONLY position embeddings
+        # (matching PyTorch: text_emb[1].zero_() then += text_pos_emb)
+        uncond_embeds = inputs_embeds.copy()
+        text_mask = input_ids[0] < START_SPEECH_TOKEN  # which tokens are text
+        for j in range(len(text_mask)):
+            if text_mask[j]:
+                pos = int(position_ids[0, j])
+                uncond_embeds[0, j, :] = text_pos_emb_weight[pos]  # position only, no content
+        uncond_inputs = np.concatenate((cond_emb, uncond_embeds), axis=1)
+        current_embeds = np.concatenate([cond_inputs, uncond_inputs], axis=0)  # (2, seq, dim)
+    else:
+        current_embeds = cond_inputs  # (1, seq, dim)
+
+    _, seq_len, _ = current_embeds.shape
+    past_key_values = {
+        f"past_key_values.{layer}.{kv}": np.zeros([B, num_key_value_heads, 0, head_dim], dtype=np.float32)
+        for layer in range(num_hidden_layers)
+        for kv in ("key", "value")
+    }
+    attention_mask = np.ones((B, seq_len), dtype=np.int64)
+
+    # --- Generation loop ---
+    for i in tqdm(range(max_new_tokens), desc="Generating", dynamic_ncols=True):
+        logits, *present_key_values = llama_with_past_session.run(None, dict(
+            inputs_embeds=current_embeds,
+            attention_mask=attention_mask,
+            **past_key_values,
+        ))
+
+        logits = logits[:, -1, :]  # (B, vocab)
+
+        if use_cfg:
+            # CFG: logits = cond + cfg_weight * (cond - uncond)
+            cond_logits = logits[0:1, :]
+            uncond_logits = logits[1:2, :]
+            logits = cond_logits + cfg_weight * (cond_logits - uncond_logits)
+
+        # Apply repetition penalty
+        logits = repetition_penalty_processor(generate_tokens, logits)
+
+        # Sampling or greedy
+        if temperature > 0:
+            logits = logits / temperature
+            logits = min_p_filter(logits, min_p)
+            logits = top_p_filter(logits, top_p)
+            probs = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
+            probs = probs / probs.sum(axis=-1, keepdims=True)
+            probs = np.nan_to_num(probs, nan=0.0)
+            if probs.sum() == 0:
+                probs[0, np.argmax(logits)] = 1.0
+            next_token = np.array([[np.random.choice(probs.shape[-1], p=probs[0])]], dtype=np.int64)
+        else:
+            next_token = np.argmax(logits, axis=-1, keepdims=True).astype(np.int64)
+
+        generate_tokens = np.concatenate((generate_tokens, next_token), axis=-1)
+        if (next_token.flatten() == STOP_SPEECH_TOKEN).all():
+            break
+
+        # Embed next speech token for next step
+        next_pos = np.full((1, 1), i + 1, dtype=np.int64)
+        next_embeds = embed_tokens_session.run(None, {
+            "input_ids": next_token,
+            "position_ids": next_pos,
+            "exaggeration": np.array([exaggeration], dtype=np.float32),
+        })[0]  # (1, 1, dim)
+
+        # For CFG: duplicate token embedding for both branches
+        if use_cfg:
+            current_embeds = np.concatenate([next_embeds, next_embeds], axis=0)
+        else:
+            current_embeds = next_embeds
+
+        attention_mask = np.concatenate([attention_mask, np.ones((B, 1), dtype=np.int64)], axis=1)
+        # Map present.X.key/value outputs -> past_key_values.X.key/value inputs
+        for layer in range(num_hidden_layers):
+            for ki, kv in enumerate(("key", "value")):
+                out_idx = layer * 2 + ki
+                past_key_values[f"past_key_values.{layer}.{kv}"] = present_key_values[out_idx]
+
+    speech_tokens = generate_tokens[:, 1:-1]
+    speech_tokens = np.concatenate([prompt_token, speech_tokens], axis=1)
+
+    wav = cond_decoder_session.run(None, {
+        "speech_tokens": speech_tokens,
+        "speaker_embeddings": ref_x_vector,
+        "speaker_features": prompt_feat,
+    })[0]
+    wav = np.squeeze(wav, axis=0)
+    sf.write(output_file_name, wav, S3GEN_SR)
+    print(f"\n{output_file_name} saved!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Chatterbox Polish TTS test")
+    parser.add_argument("--text", type=str,
+        default="Dzień dobry, nazywam się Weronika. Dziś jest piękna pogoda i chciałabym zaprosić wszystkich na spacer po parku.")
+    parser.add_argument("--lang", type=str, default="pl")
+    parser.add_argument("--voice", type=str, default=os.path.join(MODEL_DIR, "polish-female-common-voice.wav"))
+    parser.add_argument("--output", "-o", type=str, default=os.path.join(MODEL_DIR, "example_output_polish.wav"))
+    parser.add_argument("--exaggeration", type=float, default=0.5, help="Emotion intensity 0.0-1.0 (default: 0.5)")
+    parser.add_argument("--cfg-weight", type=float, default=0.0, help="Classifier-free guidance weight (default: 0.0, not yet supported)")
+    parser.add_argument("--temperature", type=float, default=0.3, help="Sampling temperature, 0=greedy (default: 0.3)")
+    parser.add_argument("--top-p", type=float, default=1.0, help="Nucleus sampling threshold (default: 1.0)")
+    parser.add_argument("--min-p", type=float, default=0.05, help="Min probability filter (default: 0.05)")
+    parser.add_argument("--repetition-penalty", type=float, default=2.0, help="Repetition penalty (default: 2.0)")
+    parser.add_argument("--max-tokens", type=int, default=500, help="Max speech tokens to generate (default: 500)")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("POLISH TTS TEST")
+    print("=" * 60)
+
+    run_inference(
+        text=args.text,
+        language_id=args.lang,
+        target_voice_path=args.voice,
+        output_file_name=args.output,
+        exaggeration=args.exaggeration,
+        cfg_weight=args.cfg_weight,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        min_p=args.min_p,
+        repetition_penalty=args.repetition_penalty,
+        max_new_tokens=args.max_tokens,
+        seed=args.seed,
+    )