Add opt-in v3 multilingual checkpoint, skip analyzer for v3 (#516)

* Add opt-in v3 multilingual checkpoint, skip analyzer for v3

* Remove alignment analyzer; lower rep_penalty default to 1.2; trim final speech token artifact

Author: ZihanJin

README.md

@@ -84,6 +84,8 @@ ta.save("test-english.wav", wav, model.sr)
 
 # Multilingual examples
 multilingual_model = ChatterboxMultilingualTTS.from_pretrained(device=device)
+# v2 remains the default. To use the v3 multilingual checkpoint:
+# multilingual_model = ChatterboxMultilingualTTS.from_pretrained(device=device, t3_model="v3")
 
 french_text = "Bonjour, comment ça va? Ceci est le modèle de synthèse vocale multilingue Chatterbox, il prend en charge 23 langues."
 wav_french = multilingual_model.generate(french_text, language_id="fr")

example_tts.py

@@ -21,6 +21,7 @@
 ta.save("test-1.wav", wav, model.sr)
 
 multilingual_model = ChatterboxMultilingualTTS.from_pretrained(device=device)
+# v2 is the default. Pass t3_model="v3" to use the v3 multilingual checkpoint.
 text = "Bonjour, comment ça va? Ceci est le modèle de synthèse vocale multilingue Chatterbox, il prend en charge 23 langues."
 wav = multilingual_model.generate(text, language_id="fr")
 ta.save("test-2.wav", wav, multilingual_model.sr)

multilingual_app.py

@@ -1,11 +1,14 @@
 import random
+import os
 import numpy as np
 import torch
 from chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
 import gradio as gr
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+T3_MODEL = os.getenv("CHATTERBOX_MULTILINGUAL_T3_MODEL", "v2")
 print(f"🚀 Running on device: {DEVICE}")
+print(f"Using multilingual T3 model: {T3_MODEL}")
 
 # --- Global Model Initialization ---
 MODEL = None
@@ -140,7 +143,7 @@ def get_or_load_model():
     if MODEL is None:
         print("Model not loaded, initializing...")
         try:
-            MODEL = ChatterboxMultilingualTTS.from_pretrained(DEVICE)
+            MODEL = ChatterboxMultilingualTTS.from_pretrained(DEVICE, t3_model=T3_MODEL)
             if hasattr(MODEL, 'to') and str(MODEL.device) != DEVICE:
                 MODEL.to(DEVICE)
             print(f"Model loaded successfully. Internal device: {getattr(MODEL, 'device', 'N/A')}")

src/chatterbox/models/t3/inference/alignment_stream_analyzer.py

@@ -23,14 +23,12 @@ def __init__(
         speech_head,
         latents_queue=None,
         logits_queue=None,
-        alignment_stream_analyzer: 'AlignmentStreamAnalyzer'=None,
     ):
         super().__init__(config)
         self.model = llama
         self.speech_enc = speech_enc
         self.speech_head = speech_head
         self._added_cond = False
-        self.alignment_stream_analyzer = alignment_stream_analyzer
 
     @torch.inference_mode()
     def prepare_inputs_for_generation(
@@ -105,9 +103,6 @@ def forward(
         logits = self.speech_head(hidden_states)
         # assert inputs_embeds.size(0) == 1 # (disabled for CFG)
 
-        # NOTE: hallucination handler may modify logits to force emit an EOS token
-        # logits = self.alignment_stream_analyzer.step(logits)
-
         return CausalLMOutputWithCrossAttentions(
             logits=logits,
             past_key_values=tfmr_out.past_key_values,

src/chatterbox/models/t3/inference/t3_hf_backend.py

@@ -24,7 +24,6 @@
 from .modules.t3_config import T3Config
 from .llama_configs import LLAMA_CONFIGS
 from .inference.t3_hf_backend import T3HuggingfaceBackend
-from .inference.alignment_stream_analyzer import AlignmentStreamAnalyzer
 from ..utils import AttrDict
 
 
@@ -275,24 +274,11 @@ def inference(
         # TODO? synchronize the expensive compile function
         # with self.compile_lock:
         if not self.compiled:
-            # Default to None for English models, only create for multilingual
-            alignment_stream_analyzer = None
-            if self.hp.is_multilingual:
-                alignment_stream_analyzer = AlignmentStreamAnalyzer(
-                    self.tfmr,
-                    None,
-                    text_tokens_slice=(len_cond, len_cond + text_tokens.size(-1)),
-                    alignment_layer_idx=9, # TODO: hparam or something?
-                    eos_idx=self.hp.stop_speech_token,
-                )
-                assert alignment_stream_analyzer.eos_idx == self.hp.stop_speech_token
-
             patched_model = T3HuggingfaceBackend(
                 config=self.cfg,
                 llama=self.tfmr,
                 speech_enc=self.speech_emb,
                 speech_head=self.speech_head,
-                alignment_stream_analyzer=alignment_stream_analyzer,
             )
             self.patched_model = patched_model
             self.compiled = True
@@ -341,7 +327,7 @@ def inference(
             inputs_embeds=inputs_embeds,
             past_key_values=None,
             use_cache=True,
-            output_attentions=True,
+            output_attentions=False,
             output_hidden_states=True,
             return_dict=True,
         )
@@ -357,14 +343,6 @@ def inference(
             cfg = torch.as_tensor(cfg_weight, device=cond.device, dtype=cond.dtype)
             logits = cond + cfg * (cond - uncond)
 
-            # Apply alignment stream analyzer integrity checks
-            if self.patched_model.alignment_stream_analyzer is not None:
-                if logits.dim() == 1:            # guard in case something upstream squeezed
-                    logits = logits.unsqueeze(0) # (1, V)
-                # Pass the last generated token for repetition tracking
-                last_token = generated_ids[0, -1].item() if len(generated_ids[0]) > 0 else None
-                logits = self.patched_model.alignment_stream_analyzer.step(logits, next_token=last_token)  # (1, V)
-
             # Apply repetition penalty
             ids_for_proc = generated_ids[:1, ...]   # batch = 1
             logits = repetition_penalty_processor(ids_for_proc, logits)  # expects (B,V)
@@ -400,7 +378,7 @@ def inference(
             output = self.patched_model(
                 inputs_embeds=next_token_embed,
                 past_key_values=past,
-                output_attentions=True,
+                output_attentions=False,
                 output_hidden_states=True,
                 return_dict=True,
             )