app.py

import spaces
import os
import sys
import tempfile
import subprocess
import numpy as np
import gradio as gr
import soundfile as sf
import random
import torch
from src.chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES

# Clone NeuTTS-Air repository if not present
NEUTTS_DIR = "neutts-air"
if not os.path.exists(NEUTTS_DIR):
    try:
        subprocess.run(["git", "clone", "https://github.com/neuphonic/neutts-air.git", NEUTTS_DIR], check=True)
        print(f"Successfully cloned NeuTTS-Air to {NEUTTS_DIR}")
    except Exception as e:
        print(f"Warning: Could not clone NeuTTS-Air: {e}")

# Add NeuTTS-Air to path - aligned with official implementation
sys.path.append(NEUTTS_DIR)

# Global variables for lazy loading
kokoro_pipe = None
neutts_model = None
chatterbox_model = None

# NeuTTS-Air configuration - aligned with official neutts-air/app.py
SAMPLES_PATH = os.path.join(os.getcwd(), NEUTTS_DIR, "samples")
DEFAULT_REF_TEXT = "So I'm live on radio. And I say, well, my dear friend James here clearly, and the whole room just froze. Turns out I'd completely misspoken and mentioned our other friend."
DEFAULT_REF_PATH = os.path.join(SAMPLES_PATH, "dave.wav")
DEFAULT_GEN_TEXT = "My name is Dave, and um, I'm from London."

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🚀 Running on device: {DEVICE}")

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频道观看次数达到了二十亿次，这绝对令人难以置信。"
    },
}

# ------------------------------------------------------------------
# 1. Lazy loaders
# ------------------------------------------------------------------
def load_kokoro():
    global kokoro_pipe
    if kokoro_pipe is None:
        from kokoro import KPipeline
        kokoro_pipe = KPipeline(lang_code='a')
    return kokoro_pipe

def load_neutts():
    """Initialize NeuTTS-Air model - aligned with official implementation"""
    global neutts_model
    if neutts_model is None:
        from neuttsair.neutts import NeuTTSAir
        # Configuration matches official neutts-air/app.py lines 14-19
        neutts_model = NeuTTSAir(
            backbone_repo="neuphonic/neutts-air",
            backbone_device="cuda",
            codec_repo="neuphonic/neucodec",
            codec_device="cuda"
        )
    return neutts_model

def get_or_load_chatterbox_model():
    """Loads the ChatterboxMultilingualTTS model if it hasn't been loaded already,
    and ensures it's on the correct device."""
    global chatterbox_model
    if chatterbox_model is None:
        print("Chatterbox model not loaded, initializing...")
        try:
            chatterbox_model = ChatterboxMultilingualTTS.from_pretrained(DEVICE)
            if hasattr(chatterbox_model, 'to') and str(chatterbox_model.device) != DEVICE:
                chatterbox_model.to(DEVICE)
            print(f"Chatterbox model loaded successfully. Internal device: {getattr(chatterbox_model, 'device', 'N/A')}")
        except Exception as e:
            print(f"Error loading Chatterbox model: {e}")
            raise
    return chatterbox_model

# Attempt to load the Chatterbox model at startup.
try:
    get_or_load_chatterbox_model()
except Exception as e:
    print(f"CRITICAL: Failed to load Chatterbox model on startup. Application may not function. Error: {e}")

# ------------------------------------------------------------------
# 2. Kokoro TTS inference
# ------------------------------------------------------------------
@spaces.GPU()
def kokoro_infer(text, voice, speed):
    if not text.strip():
        raise gr.Error("Please enter some text.")
    
    pipe = load_kokoro()
    generator = pipe(text, voice=voice, speed=speed)
    for gs, ps, audio in generator:
        # Save to temporary file
        fd, tmp = tempfile.mkstemp(suffix='.wav')
        os.close(fd)
        sf.write(tmp, audio, 24000)
        return tmp
    raise RuntimeError("Kokoro generation failed")

# ------------------------------------------------------------------
# 3. NeuTTS-Air inference - aligned with official implementation
# ------------------------------------------------------------------
@spaces.GPU()
def neutts_infer(ref_text: str, ref_audio_path: str, gen_text: str) -> tuple[int, np.ndarray]:
    """
    Generates speech using NeuTTS-Air given a reference audio and text, and new text to synthesize.
    
    Implementation aligned with official neutts-air/app.py lines 22-45.
    
    Args:
        ref_text (str): The text corresponding to the reference audio.
        ref_audio_path (str): The file path to the reference audio.
        gen_text (str): The new text to synthesize.
    Returns:
        tuple [int, np.ndarray]: A tuple containing the sample rate (24000) and the generated audio waveform as a numpy array.
    """
    if not gen_text.strip():
        raise gr.Error("Please enter text to generate.")
    if not ref_audio_path:
        raise gr.Error("Please provide reference audio.")
    if not ref_text.strip():
        raise gr.Error("Please provide reference text.")
    
    # Info messages aligned with official implementation
    gr.Info("Starting inference request!")
    gr.Info("Encoding reference...")
    
    tts = load_neutts()
    ref_codes = tts.encode_reference(ref_audio_path)
    
    gr.Info(f"Generating audio for input text: {gen_text}")
    wav = tts.infer(gen_text, ref_codes, ref_text)
    
    # Return format aligned with official implementation (line 45)
    return (24_000, wav)

# ------------------------------------------------------------------
# 4. Chatterbox Multilingual inference
# ------------------------------------------------------------------
def default_audio_for_ui(lang: str) -> str | None:
    return LANGUAGE_CONFIG.get(lang, {}).get("audio")

def default_text_for_ui(lang: str) -> str:
    return LANGUAGE_CONFIG.get(lang, {}).get("text", "")

def get_supported_languages_display() -> str:
    """Generate a formatted display of all supported languages."""
    language_items = []
    for code, name in sorted(SUPPORTED_LANGUAGES.items()):
        language_items.append(f"**{name}** (`{code}`)")
    
    # Split into 2 lines
    mid = len(language_items) // 2
    line1 = " • ".join(language_items[:mid])
    line2 = " • ".join(language_items[mid:])
    
    return f"""
### 🌍 Supported Languages ({len(SUPPORTED_LANGUAGES)} total)
{line1}
{line2}
"""

def set_seed(seed: int):
    """Sets the random seed for reproducibility across torch, numpy, and random."""
    torch.manual_seed(seed)
    if DEVICE == "cuda":
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)

@spaces.GPU
def chatterbox_infer(
    text_input: str,
    language_id: str,
    audio_prompt_path_input: str = None,
    exaggeration_input: float = 0.5,
    temperature_input: float = 0.8,
    seed_num_input: int = 0,
    cfgw_input: float = 0.5
) -> tuple[int, np.ndarray]:
    """
    Generate high-quality speech audio from text using Chatterbox Multilingual model with optional reference audio styling.
    """
    current_model = get_or_load_chatterbox_model()

    if current_model is None:
        raise RuntimeError("TTS model is not loaded.")

    if seed_num_input != 0:
        set_seed(int(seed_num_input))

    print(f"Generating audio for text: '{text_input[:50]}...'")
    
    # Handle optional audio prompt
    chosen_prompt = audio_prompt_path_input or default_audio_for_ui(language_id)

    generate_kwargs = {
        "exaggeration": exaggeration_input,
        "temperature": temperature_input,
        "cfg_weight": cfgw_input,
    }
    if chosen_prompt:
        generate_kwargs["audio_prompt_path"] = chosen_prompt
        print(f"Using audio prompt: {chosen_prompt}")
    else:
        print("No audio prompt provided; using default voice.")
        
    wav = current_model.generate(
        text_input[:300],  # Truncate text to max chars
        language_id=language_id,
        **generate_kwargs
    )
    print("Audio generation complete.")
    return (current_model.sr, wav.squeeze(0).numpy())

def on_language_change(lang, current_ref, current_text):
    return default_audio_for_ui(lang), default_text_for_ui(lang)

# ------------------------------------------------------------------
# 5. Gradio UI with model selection
# ------------------------------------------------------------------
css = """footer {visibility: hidden}"""

with gr.Blocks(css=css, title="Text2Audio - Kokoro, NeuTTS-Air & Chatterbox") as demo:
    gr.Markdown("# 🎙️ Text-to-Audio Generation")
    gr.Markdown("Choose between **Kokoro TTS** (fast English TTS), **NeuTTS-Air** (voice cloning with reference audio), or **Chatterbox Multilingual** (23 languages with styling)")
    
    # Model selection
    model_choice = gr.Radio(
        choices=["Kokoro TTS", "NeuTTS-Air", "Chatterbox Multilingual"],
        value="Kokoro TTS",
        label="Select TTS Engine",
        interactive=True
    )
    
    # Kokoro TTS Interface
    with gr.Group(visible=True) as kokoro_group:
        gr.Markdown("### 🎧 Kokoro TTS Settings")
        with gr.Row():
            with gr.Column():
                kokoro_voice = gr.Dropdown(
                    label="Voice",
                    choices=['af_heart', 'af_sky', 'af_mist', 'af_dusk'],
                    value='af_heart'
                )
                kokoro_speed = gr.Slider(0.5, 2.0, 1.0, step=0.1, label="Speed")
            
            with gr.Column(scale=3):
                kokoro_text = gr.Textbox(
                    label="Text to speak",
                    placeholder="Type or paste text here…",
                    lines=6,
                    max_lines=12
                )
                kokoro_btn = gr.Button("🎧 Synthesise with Kokoro", variant="primary")
                kokoro_audio_out = gr.Audio(label="Generated speech", type="filepath")
        
        gr.Markdown("**Kokoro** – fast, high-quality English TTS. Audio is returned as 24 kHz WAV.")
    
    # NeuTTS-Air Interface - aligned with official implementation
    with gr.Group(visible=False) as neutts_group:
        gr.Markdown("### ☁️ NeuTTS-Air Settings")
        # Interface structure aligned with official neutts-air/app.py lines 47-57
        neutts_ref_text = gr.Textbox(
            label="Reference Text",
            value=DEFAULT_REF_TEXT,
            lines=3
        )
        neutts_ref_audio = gr.Audio(
            type="filepath",
            label="Reference Audio",
            value=DEFAULT_REF_PATH if os.path.exists(DEFAULT_REF_PATH) else None
        )
        neutts_gen_text = gr.Textbox(
            label="Text to Generate",
            value=DEFAULT_GEN_TEXT,
            lines=3
        )
        neutts_btn = gr.Button("☁️ Generate with NeuTTS-Air", variant="primary")
        neutts_audio_out = gr.Audio(label="Generated Speech", type="numpy")
        
        gr.Markdown("**NeuTTS-Air** – Upload a reference audio sample, provide the reference text, and enter new text to synthesize.")
    
    # Chatterbox Multilingual Interface
    with gr.Group(visible=False) as chatterbox_group:
        gr.Markdown("### 🗣️ Chatterbox Multilingual Settings")
        gr.Markdown(
            """
            Generate high-quality multilingual speech from text with reference audio styling, supporting 23 languages.
            
            For a hosted version of Chatterbox Multilingual and for finetuning, please visit [resemble.ai](https://app.resemble.ai)
            """
        )
        # Display supported languages
        gr.Markdown(get_supported_languages_display())
        with gr.Row():
            with gr.Column():
                initial_lang = "en"
                chatterbox_text = gr.Textbox(
                    value=default_text_for_ui(initial_lang),
                    label="Text to synthesize (max chars 300)",
                    max_lines=5
                )
                
                chatterbox_language_id = gr.Dropdown(
                    choices=list(SUPPORTED_LANGUAGES.keys()),
                    value=initial_lang,
                    label="Language",
                    info="Select the language for text-to-speech synthesis"
                )
                
                chatterbox_ref_wav = gr.Audio(
                    sources=["upload", "microphone"],
                    type="filepath",
                    label="Reference Audio File (Optional)",
                    value=default_audio_for_ui(initial_lang)
                )
                
                gr.Markdown(
                    "💡 **Note**: Ensure that the reference clip matches the specified language tag. Otherwise, language transfer outputs may inherit the accent of the reference clip's language. To mitigate this, set the CFG weight to 0.",
                    elem_classes=["audio-note"]
                )
                
                chatterbox_exaggeration = gr.Slider(
                    0.25, 2, step=.05, label="Exaggeration (Neutral = 0.5, extreme values can be unstable)", value=.5
                )
                chatterbox_cfg_weight = gr.Slider(
                    0.2, 1, step=.05, label="CFG/Pace", value=0.5
                )

                with gr.Accordion("More options", open=False):
                    chatterbox_seed_num = gr.Number(value=0, label="Random seed (0 for random)")
                    chatterbox_temp = gr.Slider(0.05, 5, step=.05, label="Temperature", value=.8)

                chatterbox_btn = gr.Button("🗣️ Generate with Chatterbox", variant="primary")

            with gr.Column():
                chatterbox_audio_output = gr.Audio(label="Output Audio")

        chatterbox_language_id.change(
            fn=on_language_change,
            inputs=[chatterbox_language_id, chatterbox_ref_wav, chatterbox_text],
            outputs=[chatterbox_ref_wav, chatterbox_text],
            show_progress=False
        )

        chatterbox_btn.click(
            fn=chatterbox_infer,
            inputs=[
                chatterbox_text,
                chatterbox_language_id,
                chatterbox_ref_wav,
                chatterbox_exaggeration,
                chatterbox_temp,
                chatterbox_seed_num,
                chatterbox_cfg_weight,
            ],
            outputs=[chatterbox_audio_output],
        )
    
    # Event handlers
    def toggle_interface(choice):
        if choice == "Kokoro TTS":
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
        elif choice == "NeuTTS-Air":
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
    
    model_choice.change(
        fn=toggle_interface,
        inputs=[model_choice],
        outputs=[kokoro_group, neutts_group, chatterbox_group]
    )
    
    kokoro_btn.click(
        kokoro_infer,
        inputs=[kokoro_text, kokoro_voice, kokoro_speed],
        outputs=kokoro_audio_out
    )
    
    neutts_btn.click(
        neutts_infer,
        inputs=[neutts_ref_text, neutts_ref_audio, neutts_gen_text],
        outputs=neutts_audio_out
    )

if __name__ == "__main__":
    # Launch configuration aligned with official implementation (line 60)
    demo.launch(allowed_paths=[SAMPLES_PATH] if os.path.exists(SAMPLES_PATH) else None, mcp_server=True, inbrowser=True)