app.py

#!/usr/bin/env python3
"""
Chatterbox Indic LoRA — Gradio Web UI

Supports: Hindi, Telugu, Kannada, Bengali, Tamil, Malayalam, Marathi, Gujarati
Plus all 23 base Chatterbox-Multilingual languages.

Usage:
  python app.py                         # Auto-downloads everything
  python app.py --share                 # Public Gradio link
  python app.py --device cpu            # Force CPU
"""

import argparse
import random
from pathlib import Path

import numpy as np
import torch
import gradio as gr


# ── Config ──────────────────────────────────────────────────
INDIC_LANGUAGES = {
    "hi": "Hindi",
    "te": "Telugu",
    "kn": "Kannada",
    "bn": "Bengali",
    "ta": "Tamil",
    "ml": "Malayalam",
    "mr": "Marathi",
    "gu": "Gujarati",
}

EXAMPLE_SENTENCES = {
    "hi": "नमस्ते, आप कैसे हैं? आज मौसम बहुत अच्छा है।",
    "te": "నమస్కారం, మీరు ఎలా ఉన్నారు? ఈ రోజు వాతావరణం చాలా బాగుంది.",
    "kn": "ನಮಸ್ಕಾರ, ನೀವು ಹೇಗಿದ್ದೀರಿ? ಇಂದು ಹವಾಮಾನ ತುಂಬಾ ಚೆನ್ನಾಗಿದೆ.",
    "bn": "নমস্কার, আপনি কেমন আছেন? আজ আবহাওয়া খুব ভালো।",
    "ta": "வணக்கம், நீங்கள் எப்படி இருக்கிறீர்கள்? இன்று வானிலை மிகவும் நன்றாக உள்ளது.",
    "ml": "നമസ്കാരം, നിങ്ങൾ എങ്ങനെയുണ്ട്? ഇന്ന് കാലാവസ്ഥ വളരെ നല്ലതാണ്.",
    "mr": "नमस्कार, तुम्ही कसे आहात? आज हवामान खूप छान आहे.",
    "gu": "નમસ્તે, તમે કેમ છો? આજે હવામાન ખૂબ સારું છે.",
    "en": "Hello, how are you? The weather is very nice today.",
}

# ── Globals ─────────────────────────────────────────────────
MODEL = None
DEVICE = None
LORA_DIR = None


def detect_device(override=None):
    if override and override != "auto":
        return override
    if torch.cuda.is_available():
        return "cuda"
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def load_model():
    """Load base Chatterbox-Multilingual + apply Indic LoRA in one call."""
    global MODEL, LORA_DIR

    from chatterbox.mtl_tts import ChatterboxMultilingualTTS

    print(f"Loading Chatterbox + Indic LoRA on {DEVICE}...")
    MODEL = ChatterboxMultilingualTTS.from_indic_lora(device=DEVICE)
    LORA_DIR = Path(MODEL._lora_dir) if hasattr(MODEL, '_lora_dir') else None
    print("Model ready.")
    return MODEL


def set_seed(seed: int):
    torch.manual_seed(seed)
    if DEVICE == "cuda":
        torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)


def get_conds_choices():
    """List available speaker conditioning files."""
    if LORA_DIR is None:
        return []
    conds_dir = LORA_DIR / "conds"
    if not conds_dir.exists():
        return []
    choices = []
    for f in sorted(conds_dir.glob("*.pt")):
        name = f.stem  # e.g. "te_female"
        lang = name.split("_")[0]
        gender = name.split("_")[1] if "_" in name else "unknown"
        lang_name = INDIC_LANGUAGES.get(lang, lang.upper())
        label = f"{lang_name} {gender.title()} ({lang}_{gender})"
        choices.append((label, str(f)))
    return choices


def generate_speech(text, language, speaker, audio_ref, exaggeration, temperature, cfg_weight, seed_num):
    """Generate speech with the Indic LoRA model."""
    if MODEL is None:
        raise gr.Error("Model not loaded yet. Please wait...")

    if not text or not text.strip():
        raise gr.Error("Please enter some text.")

    if seed_num and int(seed_num) != 0:
        set_seed(int(seed_num))

    # Determine language code
    lang_code = language.split("(")[-1].strip(")") if "(" in language else language

    generate_kwargs = {
        "language_id": lang_code,
        "exaggeration": exaggeration,
        "temperature": temperature,
        "cfg_weight": cfg_weight,
    }

    # Use uploaded audio ref or pre-extracted conds
    if audio_ref:
        generate_kwargs["audio_prompt_path"] = audio_ref
    elif speaker:
        # Load pre-extracted speaker conditioning
        from chatterbox.mtl_tts import Conditionals
        conds_path = speaker
        MODEL.conds = Conditionals.load(conds_path).to(DEVICE)

    wav = MODEL.generate(text[:500], **generate_kwargs)
    return (MODEL.sr, wav.squeeze(0).cpu().numpy())


def on_language_change(lang_choice):
    """Update example text when language changes."""
    lang_code = lang_choice.split("(")[-1].strip(")") if "(" in lang_choice else lang_choice
    return EXAMPLE_SENTENCES.get(lang_code, "")


# ── Build Gradio UI ────────────────────────────────────────
def create_ui():
    lang_choices = [f"{name} ({code})" for code, name in INDIC_LANGUAGES.items()]
    lang_choices.append("English (en)")
    speaker_choices = get_conds_choices()

    with gr.Blocks(title="Chatterbox Indic LoRA", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
# Chatterbox Indic LoRA
**8 Indian languages via LoRA fine-tuning on [Chatterbox-Multilingual](https://github.com/resemble-ai/chatterbox)**

Hindi | Telugu | Kannada | Bengali | Tamil | Malayalam | Marathi | Gujarati
        """)

        with gr.Row():
            with gr.Column(scale=3):
                language = gr.Dropdown(
                    choices=lang_choices,
                    value=lang_choices[0],
                    label="Language",
                )
                text = gr.Textbox(
                    value=EXAMPLE_SENTENCES["hi"],
                    label="Text (max 500 chars)",
                    lines=3,
                    max_lines=5,
                )
                speaker = gr.Dropdown(
                    choices=speaker_choices,
                    value=speaker_choices[0][1] if speaker_choices else None,
                    label="Speaker Voice",
                    info="Pre-extracted speaker conditionings (male/female per language)",
                )
                audio_ref = gr.Audio(
                    sources=["upload", "microphone"],
                    type="filepath",
                    label="Or upload your own reference audio (overrides speaker)",
                )

                with gr.Accordion("Advanced", open=False):
                    exaggeration = gr.Slider(0.25, 2.0, value=0.5, step=0.05, label="Exaggeration")
                    temperature = gr.Slider(0.05, 2.0, value=0.8, step=0.05, label="Temperature")
                    cfg_weight = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="CFG Weight")
                    seed_num = gr.Number(value=0, label="Seed (0 = random)")

                generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")

            with gr.Column(scale=2):
                audio_output = gr.Audio(label="Generated Speech", type="numpy")

        # Wire events
        language.change(fn=on_language_change, inputs=[language], outputs=[text])
        generate_btn.click(
            fn=generate_speech,
            inputs=[text, language, speaker, audio_ref, exaggeration, temperature, cfg_weight, seed_num],
            outputs=[audio_output],
        )

        gr.Markdown("""
---
**Model:** [reenigne314/chatterbox-indic-lora](https://huggingface.co/reenigne314/chatterbox-indic-lora)
| **Base:** [ResembleAI/chatterbox](https://github.com/resemble-ai/chatterbox)
        """)

    return demo


# ── Main ────────────────────────────────────────────────────
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chatterbox Indic LoRA Web UI")
    parser.add_argument("--device", default="auto", help="Device: auto, cuda, mps, cpu")
    parser.add_argument("--share", action="store_true", help="Create public Gradio link")
    parser.add_argument("--port", type=int, default=7860, help="Port number")
    args = parser.parse_args()

    DEVICE = detect_device(args.device)
    print(f"Using device: {DEVICE}")

    load_model()

    demo = create_ui()
    demo.launch(server_name="0.0.0.0", server_port=args.port, share=args.share)