v2 update: multilingual improvements

requirements.txt

@@ -13,6 +13,7 @@ safetensors
 
 # Optional language-specific dependencies
 # Uncomment the ones you need for specific languages:
- pkuseg                # For Chinese text segmentation (improves mixed text handling)
+ spacy_pkuseg          # For Chinese text segmentation
  pykakasi>=2.2.0       # For Japanese text processing (Kanji to Hiragana)
+ russian-text-stresser # For Russian stress labeling
 # dicta-onnx>=0.1.0     # For Hebrew diacritization

src/chatterbox/models/t3/modules/t3_config.py

@@ -28,7 +28,7 @@ class T3Config:
     
     @property
     def is_multilingual(self):
-        return self.text_tokens_dict_size == 2352
+        return self.text_tokens_dict_size == 2454
 
     @classmethod
     def english_only(cls):
@@ -38,4 +38,4 @@ class T3Config:
     @classmethod 
     def multilingual(cls):
         """Create configuration for multilingual TTS model."""
-        return cls(text_tokens_dict_size=2352)
+        return cls(text_tokens_dict_size=2454)

src/chatterbox/models/tokenizers/tokenizer.py

@@ -191,7 +191,7 @@ class ChineseCangjieConverter:
     def _init_segmenter(self):
         """Initialize pkuseg segmenter."""
         try:
-            from pkuseg import pkuseg
+            from spacy_pkuseg import pkuseg
             self.segmenter = pkuseg()
         except ImportError:
             logger.warning("pkuseg not available - Chinese segmentation will be skipped")
@@ -235,11 +235,53 @@ class ChineseCangjieConverter:
         return "".join(output)
 
 
+class RussianStressLabeler:
+    """Adds stress marks to Russian text when the optional dependency is available."""
+
+    def __init__(self):
+        self._stresser = None
+        self._available = False
+        self._error_logged = False
+        self._initialize()
+
+    def _initialize(self):
+        try:
+            from russian_text_stresser.text_stresser import RussianTextStresser
+        except ImportError:
+            logger.warning("russian_text_stresser not available - Russian stress labeling skipped")
+            self._error_logged = True
+            return
+        except Exception as exc:
+            logger.warning(f"Failed to import RussianTextStresser: {exc}")
+            self._error_logged = True
+            return
+
+        try:
+            self._stresser = RussianTextStresser()
+            self._available = True
+        except Exception as exc:
+            logger.warning(f"Failed to initialize RussianTextStresser: {exc}")
+            self._error_logged = True
+
+    def __call__(self, text: str) -> str:
+        if not text or not self._available:
+            return text
+
+        try:
+            return self._stresser.stress_text(text)
+        except Exception as exc:
+            if not self._error_logged:
+                logger.warning(f"Russian stress labeling failed: {exc}")
+                self._error_logged = True
+            return text
+        
+        
 class MTLTokenizer:
     def __init__(self, vocab_file_path):
         self.tokenizer: Tokenizer = Tokenizer.from_file(vocab_file_path)
         model_dir = Path(vocab_file_path).parent
         self.cangjie_converter = ChineseCangjieConverter(model_dir)
+        self.russian_stress_labeler = RussianStressLabeler()
         self.check_vocabset_sot_eot()
 
     def check_vocabset_sot_eot(self):
@@ -262,6 +304,8 @@ class MTLTokenizer:
             txt = add_hebrew_diacritics(txt)
         elif language_id == 'ko':
             txt = korean_normalize(txt)
+        elif language_id == 'ru':
+            txt = self.russian_stress_labeler(txt)
         
         # Prepend language token
         if language_id:

src/chatterbox/mtl_tts.py

@@ -168,7 +168,7 @@ class ChatterboxMultilingualTTS:
         ve.to(device).eval()
 
         t3 = T3(T3Config.multilingual())
-        t3_state = load_safetensors(ckpt_dir / "t3_23lang.safetensors")
+        t3_state = load_safetensors(ckpt_dir / "t3_mtl23ls_v2.safetensors")
         if "model" in t3_state.keys():
             t3_state = t3_state["model"][0]
         t3.load_state_dict(t3_state)
@@ -181,7 +181,7 @@ class ChatterboxMultilingualTTS:
         s3gen.to(device).eval()
 
         tokenizer = MTLTokenizer(
-            str(ckpt_dir / "mtl_tokenizer.json")
+            str(ckpt_dir / "grapheme_mtl_merged_expanded_v1.json")
         )
 
         conds = None
@@ -197,7 +197,7 @@ class ChatterboxMultilingualTTS:
                 repo_id=REPO_ID,
                 repo_type="model",
                 revision="main", 
-                allow_patterns=["ve.pt", "t3_23lang.safetensors", "s3gen.pt", "mtl_tokenizer.json", "conds.pt", "Cangjie5_TC.json"],
+                allow_patterns=["ve.pt", "t3_23ls_v2.safetensors", "s3gen.pt", "grapheme_mtl_merged_expanded_v1.json", "conds.pt", "Cangjie5_TC.json"],
                 token=os.getenv("HF_TOKEN"),
             )
         )