app.py

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

MODEL = None
# ZeroGPU supports CUDA placement at module load time via CUDA emulation.
TARGET_DEVICE = "cuda"

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')]
DEFAULT_LANGUAGE = 'en'
FIXED_LANGUAGE_ID = None
MAX_CHARS = 300

EXAMPLES = [
        ['Last month, we reached a new milestone with two billion views on our YouTube channel.', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/en_f1.flac', 'en', 0.5, 0.8, 0, 0.5],
        ["Bonjour, comment allez-vous? Aujourd'hui, il fait un temps magnifique pour se promener.", 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/fr_f1.flac', 'fr', 0.5, 0.8, 0, 0.5],
        ['Guten Tag! Heute ist ein wunderschöner Tag für einen Spaziergang im Park.', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/de_f1.flac', 'de', 0.5, 0.8, 0, 0.5],
        ['Hola, ¿cómo estás? Me encanta visitar nuevos lugares y conocer culturas diferentes.', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/es_f1.flac', 'es', 0.5, 0.8, 0, 0.5],
        ['Ciao! Oggi è una giornata perfetta per una passeggiata a Roma.', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/it_m1.flac', 'it', 0.5, 0.8, 0, 0.5],
        ['Olá, tudo bem? O Brasil é um país incrível com paisagens deslumbrantes.', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/pt_m1.flac', 'pt', 0.5, 0.8, 0, 0.5],
        ['नमस्ते, आज का दिन बहुत खूबसूरत है।', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/hi_f1.flac', 'hi', 0.5, 0.8, 0, 0.5],
        ['你好,今天天气真不错,我们一起去公园散步吧。', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/zh_f1.flac', 'zh', 0.5, 0.8, 0, 0.5],
        ['مرحبا، كيف حالك اليوم؟ الطقس جميل للنزهة في الحديقة.', 'https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ar_m1.flac', 'ar', 0.5, 0.8, 0, 0.5]
]


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(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 update_char_count(text: str):
    """Live character-count feedback for the synthesis box.

    Returns two values wired to outputs=[char_count, run_btn]:
      1. counter_md    -> "n / MAX_CHARS", turned red once over the limit
      2. button_update -> disables Generate while over the limit
    (The server still rejects over-limit text; disabling is the visible cue.)
    """
    n = len(text or "")
    over = n > MAX_CHARS
    counter_md = f"<span style='color:#e11'>{n} / {MAX_CHARS}</span>" if over else f"{n} / {MAX_CHARS}"
    return counter_md, gr.update(interactive=not over)


def get_or_load_model():
    global MODEL
    if MODEL is None:
        print(f"Model not loaded, initializing on {TARGET_DEVICE}...")
        try:
            MODEL = ChatterboxTTS.from_pretrained(TARGET_DEVICE)
        except Exception as exc:
            if TARGET_DEVICE != "cuda":
                raise
            print(f"CUDA model initialization failed, falling back to CPU: {exc}")
            MODEL = ChatterboxTTS.from_pretrained("cpu")
        print(f"Model loaded on {MODEL.device}.")
    return MODEL


def set_seed(seed: int, device: str):
    torch.manual_seed(seed)
    if device == "cuda" and torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)


@spaces.GPU
def generate_tts_audio(
    text_input: str,
    audio_prompt_path_input: str = None,
    language_id_input: str = None,
    exaggeration_input: float = 0.5,
    temperature_input: float = 0.8,
    seed_num_input: int = 0,
    cfgw_input: float = 0.5,
):
    """Generate speech from text with optional reference audio styling."""
    if len(text_input) > MAX_CHARS:
        raise gr.Error(f"Text exceeds {MAX_CHARS} characters ({len(text_input)}). Please shorten it.")
    current_model = get_or_load_model()
    device = current_model.device
    if seed_num_input != 0:
        set_seed(int(seed_num_input), device)
    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,
        "temperature": temperature_input,
        "cfg_weight": cfgw_input,
        "language_id": lang,
    }
    if chosen_prompt:
        generate_kwargs["audio_prompt_path"] = chosen_prompt
    wav = current_model.generate(text_input, **generate_kwargs)
    return (current_model.sr, wav.squeeze(0).cpu().numpy())


initialize_russian_stresser()
get_or_load_model()


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Chatterbox Multilingual TTS V3
        Chatterbox Multilingual TTS V3 (t3_mtl23ls_v3).

        Powered by model [`ResembleAI/chatterbox`](https://huggingface.co/ResembleAI/chatterbox).
        """
    )
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(
                value=default_text_for_ui(DEFAULT_LANGUAGE),
                label=f"Text to synthesize (max {MAX_CHARS} chars)",
                max_lines=5,
            )
            char_count = gr.Markdown(f"{len(default_text_for_ui(DEFAULT_LANGUAGE))} / {MAX_CHARS}")
            ref_wav = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="Reference Audio File (Optional)",
                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):
                seed_num = gr.Number(value=0, label="Random seed (0 for random)")
                temp = gr.Slider(0.05, 5, step=.05, label="Temperature", value=.8)
            run_btn = gr.Button("Generate", variant="primary")
        with gr.Column():
            audio_output = gr.Audio(label="Output Audio")

    inputs = [text, ref_wav, language, exaggeration, temp, seed_num, cfg_weight]
    run_btn.click(fn=generate_tts_audio, inputs=inputs, outputs=[audio_output])
    text.change(fn=update_char_count, inputs=[text], outputs=[char_count, run_btn])

    gr.Examples(
        examples=EXAMPLES,
        inputs=inputs,
        label="Examples",
    )

demo.launch(mcp_server=True)