Add multilingual text processors

app.py

@@ -1,6 +1,7 @@
 import random
 import numpy as np
 import torch
+from chatterbox.src.chatterbox.models.tokenizers.tokenizer import initialize_russian_stresser
 from chatterbox.src.chatterbox.tts import ChatterboxTTS
 import gradio as gr
 import spaces
@@ -9,9 +10,99 @@ MODEL = None
 # ZeroGPU supports CUDA placement at module load time via CUDA emulation.
 TARGET_DEVICE = "cuda"
 
-DEFAULT_CONFIG = {
-    "audio": 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/en_f1.flac',
-    "text": 'Last month, we reached a new milestone with two billion views on our YouTube channel.',
+LANGUAGE_CONFIG = {
+    "ar": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ar_f/ar_prompts2.flac",
+        "text": "في الشهر الماضي، وصلنا إلى معلم جديد بمليارين من المشاهدات على قناتنا على يوتيوب.",
+    },
+    "da": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/da_m1.flac",
+        "text": "Sidste måned nåede vi en ny milepæl med to milliarder visninger på vores YouTube-kanal.",
+    },
+    "de": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/de_f1.flac",
+        "text": "Letzten Monat haben wir einen neuen Meilenstein erreicht: zwei Milliarden Aufrufe auf unserem YouTube-Kanal.",
+    },
+    "el": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/el_m.flac",
+        "text": "Τον περασμένο μήνα, φτάσαμε σε ένα νέο ορόσημο με δύο δισεκατομμύρια προβολές στο κανάλι μας στο YouTube.",
+    },
+    "en": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/en_f1.flac",
+        "text": "Last month, we reached a new milestone with two billion views on our YouTube channel.",
+    },
+    "es": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/es_f1.flac",
+        "text": "El mes pasado alcanzamos un nuevo hito: dos mil millones de visualizaciones en nuestro canal de YouTube.",
+    },
+    "fi": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/fi_m.flac",
+        "text": "Viime kuussa saavutimme uuden virstanpylvään kahden miljardin katselukerran kanssa YouTube-kanavallamme.",
+    },
+    "fr": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/fr_f1.flac",
+        "text": "Le mois dernier, nous avons atteint un nouveau jalon avec deux milliards de vues sur notre chaîne YouTube.",
+    },
+    "he": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/he_m1.flac",
+        "text": "בחודש שעבר הגענו לאבן דרך חדשה עם שני מיליארד צפיות בערוץ היוטיוב שלנו.",
+    },
+    "hi": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/hi_f1.flac",
+        "text": "पिछले महीने हमने एक नया मील का पत्थर छुआ: हमारे YouTube चैनल पर दो अरब व्यूज़।",
+    },
+    "it": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/it_m1.flac",
+        "text": "Il mese scorso abbiamo raggiunto un nuovo traguardo: due miliardi di visualizzazioni sul nostro canale YouTube.",
+    },
+    "ja": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ja/ja_prompts1.flac",
+        "text": "先月、私たちのYouTubeチャンネルで二十億回の再生回数という新たなマイルストーンに到達しました。",
+    },
+    "ko": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ko_f.flac",
+        "text": "지난달 우리는 유튜브 채널에서 이십억 조회수라는 새로운 이정표에 도달했습니다.",
+    },
+    "ms": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ms_f.flac",
+        "text": "Bulan lepas, kami mencapai pencapaian baru dengan dua bilion tontonan di saluran YouTube kami.",
+    },
+    "nl": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/nl_m.flac",
+        "text": "Vorige maand bereikten we een nieuwe mijlpaal met twee miljard weergaven op ons YouTube-kanaal.",
+    },
+    "no": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/no_f1.flac",
+        "text": "Forrige måned nådde vi en ny milepæl med to milliarder visninger på YouTube-kanalen vår.",
+    },
+    "pl": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/pl_m.flac",
+        "text": "W zeszłym miesiącu osiągnęliśmy nowy kamień milowy z dwoma miliardami wyświetleń na naszym kanale YouTube.",
+    },
+    "pt": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/pt_m1.flac",
+        "text": "No mês passado, alcançámos um novo marco: dois mil milhões de visualizações no nosso canal do YouTube.",
+    },
+    "ru": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ru_m.flac",
+        "text": "В прошлом месяце мы достигли нового рубежа: два миллиарда просмотров на нашем YouTube-канале.",
+    },
+    "sv": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/sv_f.flac",
+        "text": "Förra månaden nådde vi en ny milstolpe med två miljarder visningar på vår YouTube-kanal.",
+    },
+    "sw": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/sw_m.flac",
+        "text": "Mwezi uliopita, tulifika hatua mpya ya maoni ya bilioni mbili kweny kituo chetu cha YouTube.",
+    },
+    "tr": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/tr_m.flac",
+        "text": "Geçen ay YouTube kanalımızda iki milyar görüntüleme ile yeni bir dönüm noktasına ulaştık.",
+    },
+    "zh": {
+        "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/zh_f2.flac",
+        "text": "上个月，我们达到了一个新的里程碑。 我们的YouTube频道观看次数达到了二十亿次，这绝对令人难以置信。",
+    },
 }
 
 LANGUAGE_CHOICES = [('en', 'English'), ('ar', 'Arabic'), ('da', 'Danish'), ('de', 'German'), ('el', 'Greek'), ('es', 'Spanish'), ('fi', 'Finnish'), ('fr', 'French'), ('he', 'Hebrew'), ('hi', 'Hindi'), ('it', 'Italian'), ('ja', 'Japanese'), ('ko', 'Korean'), ('ms', 'Malay'), ('nl', 'Dutch'), ('no', 'Norwegian'), ('pl', 'Polish'), ('pt', 'Portuguese'), ('ru', 'Russian'), ('sv', 'Swedish'), ('sw', 'Swahili'), ('tr', 'Turkish'), ('zh', 'Chinese')]
@@ -31,12 +122,17 @@ EXAMPLES = [
 ]
 
 
-def default_audio_for_ui():
-    return DEFAULT_CONFIG.get("audio")
+def default_audio_for_ui(language_id: str):
+    return LANGUAGE_CONFIG.get(language_id, LANGUAGE_CONFIG[DEFAULT_LANGUAGE]).get("audio")
 
 
-def default_text_for_ui():
-    return DEFAULT_CONFIG.get("text", "")
+def default_text_for_ui(language_id: str):
+    return LANGUAGE_CONFIG.get(language_id, LANGUAGE_CONFIG[DEFAULT_LANGUAGE]).get("text", "")
+
+
+def on_language_change(language_id: str, current_ref_wav: str | None, current_text: str):
+    print(f"Switching UI defaults to language: {language_id}")
+    return default_audio_for_ui(language_id), default_text_for_ui(language_id)
 
 
 def get_or_load_model():
@@ -78,8 +174,8 @@ def generate_tts_audio(
     device = current_model.device
     if seed_num_input != 0:
         set_seed(int(seed_num_input), device)
-    chosen_prompt = audio_prompt_path_input or default_audio_for_ui()
     lang = language_id_input or DEFAULT_LANGUAGE
+    chosen_prompt = audio_prompt_path_input or default_audio_for_ui(lang)
     print(f"Generating on {device} (lang={lang}) for text: '{text_input[:50]}...'")
     generate_kwargs = {
         "exaggeration": exaggeration_input,
@@ -93,6 +189,7 @@ def generate_tts_audio(
     return (current_model.sr, wav.squeeze(0).cpu().numpy())
 
 
+initialize_russian_stresser()
 get_or_load_model()
 
 
@@ -101,13 +198,15 @@ with gr.Blocks() as demo:
         """
         # Chatterbox Multilingual TTS V3
         Chatterbox Multilingual TTS V3 (t3_mtl23ls_v3).
+        UI revision: `multilingual-defaults-v1`
+
         Powered by model [`ResembleAI/chatterbox`](https://huggingface.co/ResembleAI/chatterbox).
         """
     )
     with gr.Row():
         with gr.Column():
             text = gr.Textbox(
-                value=default_text_for_ui(),
+                value=default_text_for_ui(DEFAULT_LANGUAGE),
                 label="Text to synthesize (max chars 300)",
                 max_lines=5,
             )
@@ -115,13 +214,19 @@ with gr.Blocks() as demo:
                 sources=["upload", "microphone"],
                 type="filepath",
                 label="Reference Audio File (Optional)",
-                value=default_audio_for_ui(),
+                value=default_audio_for_ui(DEFAULT_LANGUAGE),
             )
             language = gr.Dropdown(
                 choices=[(name, code) for code, name in LANGUAGE_CHOICES],
                 value=DEFAULT_LANGUAGE,
                 label="Language",
             )
+            language.change(
+                fn=on_language_change,
+                inputs=[language, ref_wav, text],
+                outputs=[ref_wav, text],
+                show_progress=False,
+            )
             exaggeration = gr.Slider(0.25, 2, step=.05, label="Exaggeration (Neutral = 0.5)", value=.5)
             cfg_weight = gr.Slider(0.2, 1, step=.05, label="CFG/Pace", value=0.5)
             with gr.Accordion("More options", open=False):

chatterbox/src/chatterbox/models/tokenizers/tokenizer.py

@@ -1,6 +1,14 @@
+import importlib
+import json
 import logging
+from pathlib import Path
+import subprocess
+import sys
+import threading
+from unicodedata import category, normalize
 
 import torch
+from huggingface_hub import hf_hub_download
 from tokenizers import Tokenizer
 
 
@@ -13,9 +21,132 @@ SPECIAL_TOKENS = [SOT, EOT, UNK, SPACE, "[PAD]", "[SEP]", "[CLS]", "[MASK]"]
 
 logger = logging.getLogger(__name__)
 
+CANGJIE_FILENAME = "Cangjie5_TC.json"
+CANGJIE_REPO_ID = "ResembleAI/chatterbox"
+RUSSIAN_STRESSER_PACKAGE = (
+    "git+https://github.com/Vuizur/add-stress-to-epub.git"
+    "@9e064c815373ef6f0360dd97b0d0fc37b52d256f"
+)
+RUSSIAN_STRESSER_INSTALL_DIR = Path.home() / ".cache" / "chatterbox" / "russian-text-stresser"
+_russian_stresser_lock = threading.Lock()
+_russian_stresser_local = threading.local()
+
+
+class ChineseCangjieConverter:
+    """Convert Chinese characters to Cangjie tokens used by the model."""
+
+    def __init__(self, model_dir: Path):
+        self.word2cj = {}
+        self.cj2word = {}
+        self.segmenter = None
+        self._load_cangjie_mapping(model_dir)
+        self._init_segmenter()
+
+    def _load_cangjie_mapping(self, model_dir: Path) -> None:
+        cangjie_path = model_dir / CANGJIE_FILENAME
+        if not cangjie_path.exists():
+            cangjie_path = Path(
+                hf_hub_download(
+                    repo_id=CANGJIE_REPO_ID,
+                    filename=CANGJIE_FILENAME,
+                )
+            )
+
+        with cangjie_path.open("r", encoding="utf-8") as fp:
+            for entry in json.load(fp):
+                word, code = entry.split("\t")[:2]
+                self.word2cj[word] = code
+                self.cj2word.setdefault(code, []).append(word)
+
+    def _init_segmenter(self) -> None:
+        from spacy_pkuseg import pkuseg
+
+        self.segmenter = pkuseg()
+
+    def _cangjie_encode(self, glyph: str) -> str | None:
+        code = self.word2cj.get(glyph)
+        if code is None:
+            return None
+        index = self.cj2word[code].index(glyph)
+        return code + (str(index) if index > 0 else "")
+
+    def __call__(self, text: str) -> str:
+        full_text = " ".join(self.segmenter.cut(text))
+        output = []
+        for glyph in full_text:
+            if category(glyph) != "Lo":
+                output.append(glyph)
+                continue
+
+            cangjie = self._cangjie_encode(glyph)
+            if cangjie is None:
+                output.append(glyph)
+                continue
+
+            output.extend(f"[cj_{code}]" for code in cangjie)
+            output.append("[cj_.]")
+        return "".join(output)
+
+
+def _load_russian_stresser_class():
+    install_dir = str(RUSSIAN_STRESSER_INSTALL_DIR)
+    if RUSSIAN_STRESSER_INSTALL_DIR.exists() and install_dir not in sys.path:
+        sys.path.insert(0, install_dir)
+
+    try:
+        from russian_text_stresser.text_stresser import RussianTextStresser
+        return RussianTextStresser
+    except ImportError:
+        logger.info("Installing russian_text_stresser for the first Russian request")
+        RUSSIAN_STRESSER_INSTALL_DIR.mkdir(parents=True, exist_ok=True)
+        subprocess.check_call([
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "--no-cache-dir",
+            "--no-deps",
+            "--target",
+            install_dir,
+            RUSSIAN_STRESSER_PACKAGE,
+        ])
+        if install_dir not in sys.path:
+            sys.path.insert(0, install_dir)
+        importlib.invalidate_caches()
+        from russian_text_stresser.text_stresser import RussianTextStresser
+        return RussianTextStresser
+
+
+def add_russian_stress(text: str) -> str:
+    """Add stress marks to Russian text."""
+    try:
+        stresser = getattr(_russian_stresser_local, "instance", None)
+        if stresser is None:
+            with _russian_stresser_lock:
+                RussianTextStresser = _load_russian_stresser_class()
+                stresser = RussianTextStresser()
+                _russian_stresser_local.instance = stresser
+
+        return stresser.stress_text(text)
+    except Exception as exc:
+        raise RuntimeError("Russian stress labeling failed") from exc
+
+
+def initialize_russian_stresser() -> None:
+    """Install, initialize, and verify Russian stress labeling."""
+    health_check_text = normalize("NFKD", "твои слова ничего не значат.")
+    stressed_text = add_russian_stress(health_check_text)
+    if "\u0301" not in stressed_text:
+        raise RuntimeError(
+            "Russian stresser health check failed: no stress marks were produced"
+        )
+    logger.info("Russian stresser health check passed")
+
+
 class MTLTokenizer:
     def __init__(self, vocab_file_path):
         self.tokenizer: Tokenizer = Tokenizer.from_file(vocab_file_path)
+        self.cangjie_converter = ChineseCangjieConverter(Path(vocab_file_path).parent)
         self.check_vocabset_sot_eot()
 
     def check_vocabset_sot_eot(self):
@@ -23,24 +154,70 @@ class MTLTokenizer:
         assert SOT in voc
         assert EOT in voc
 
-    def text_to_tokens(self, text: str, language_id: str = None):
-        text_tokens = self.encode(text, language_id=language_id)
+    def preprocess_text(
+        self,
+        raw_text: str,
+        language_id: str = None,
+        lowercase: bool = True,
+        nfkd_normalize: bool = True,
+    ):
+        """Apply the multilingual model's shared text preprocessing."""
+        preprocessed_text = raw_text
+        if language_id == "zh":
+            return preprocessed_text
+        if lowercase:
+            preprocessed_text = preprocessed_text.lower()
+        if nfkd_normalize:
+            preprocessed_text = normalize("NFKD", preprocessed_text)
+        return preprocessed_text
+
+    def text_to_tokens(
+        self,
+        text: str,
+        language_id: str = None,
+        lowercase: bool = True,
+        nfkd_normalize: bool = True,
+    ):
+        text_tokens = self.encode(
+            text,
+            language_id=language_id,
+            lowercase=lowercase,
+            nfkd_normalize=nfkd_normalize,
+        )
         text_tokens = torch.IntTensor(text_tokens).unsqueeze(0)
         return text_tokens
 
-    def encode(self, txt: str, language_id: str = None, verbose=False):
+    def encode(
+        self,
+        txt: str,
+        language_id: str = None,
+        lowercase: bool = True,
+        nfkd_normalize: bool = True,
+    ):
         """
-        clean_text > (prepend lang_id) > replace SPACE > encode text using Tokenizer
-        
+        preprocess text > language-specific processing > prepend lang_id > encode
+
         The multilingual model expects language tokens like [en], [fr] to be prepended
         to condition the synthesis for the appropriate language.
         """
+        txt = self.preprocess_text(
+            txt,
+            language_id=language_id,
+            lowercase=lowercase,
+            nfkd_normalize=nfkd_normalize,
+        )
+
+        if language_id == "zh":
+            txt = self.cangjie_converter(txt)
+        elif language_id == "ru":
+            txt = add_russian_stress(txt)
+
         # Prepend language token if provided
         if language_id:
             lang_token = f"[{language_id.lower()}]"
             txt = lang_token + txt
-        
-        txt = txt.replace(' ', SPACE)
+
+        txt = txt.replace(" ", SPACE)
         code = self.tokenizer.encode(txt)
         ids = code.ids
         return ids
@@ -49,10 +226,12 @@ class MTLTokenizer:
         if isinstance(seq, torch.Tensor):
             seq = seq.cpu().numpy()
 
-        txt: str = self.tokenizer.decode(seq,
-        skip_special_tokens=False)
-        txt = txt.replace(' ', '')
-        txt = txt.replace(SPACE, ' ')
-        txt = txt.replace(EOT, '')
-        txt = txt.replace(UNK, '')
+        txt: str = self.tokenizer.decode(
+            seq,
+            skip_special_tokens=False
+        )
+        txt = txt.replace(" ", "")
+        txt = txt.replace(SPACE, " ")
+        txt = txt.replace(EOT, "")
+        txt = txt.replace(UNK, "")
         return txt

chatterbox/src/chatterbox/tts.py

@@ -22,6 +22,7 @@ REPO_ID = "ResembleAI/chatterbox"
 BASE_REPO_ID = "ResembleAI/chatterbox"
 T3_FILENAME = "t3_mtl23ls_v3.safetensors"
 TOKENIZER_FILENAME = "grapheme_mtl_merged_expanded_v1.json"
+CANGJIE_FILENAME = "Cangjie5_TC.json"
 T3_TEXT_VOCAB_SIZE = 2454
 
 
@@ -60,7 +61,7 @@ def punc_norm(text: str) -> str:
 
     # Add full stop if no ending punc
     text = text.rstrip(" ")
-    sentence_enders = {".", "!", "?", "-", ","}
+    sentence_enders = {".", "!", "?", "-", ",", "、", "，", "。", "？", "！"}
     if not any(text.endswith(p) for p in sentence_enders):
         text += "."
 
@@ -183,7 +184,7 @@ class ChatterboxTTS:
     @classmethod
     def from_pretrained(cls, device: torch.device) -> 'ChatterboxTTS':
         token = os.getenv("HF_TOKEN")
-        base_files = ["ve.pt", "s3gen.pt", TOKENIZER_FILENAME, "conds.pt"]
+        base_files = ["ve.pt", "s3gen.pt", TOKENIZER_FILENAME, CANGJIE_FILENAME, "conds.pt"]
         base_dir = Path(
             snapshot_download(
                 repo_id=BASE_REPO_ID,

requirements.txt

@@ -15,3 +15,14 @@ resemble-perth==1.0.1
 silero-vad==5.1.2
 conformer==0.3.2
 safetensors
+spacy_pkuseg==1.0.1
+
+# Runtime dependencies for russian-text-stresser. The package itself is
+# installed without its incompatible spaCy 3.6 pin on the first Russian request.
+beautifulsoup4>=4.11.1
+lxml>=4.9.1
+pymorphy2>=0.9.1
+pymorphy2-dicts-ru>=2.4.417127.4579844
+spacy==3.8.11
+stressed-cyrillic-tools>=0.1.10
+transliterate>=1.10.2