app_base.py

import random
import os
import numpy as np
import torch
from src.chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
import gradio as gr
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import torch

# 🔥 GLOBAL PATCH: force all torch.load to CPU
_original_torch_load = torch.load

def _cpu_only_torch_load(*args, **kwargs):
    kwargs.setdefault("map_location", torch.device("cpu"))
    return _original_torch_load(*args, **kwargs)

torch.load = _cpu_only_torch_load




LANGUAGE_CONFIG = {
   
    "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."
    },
   
    "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 चैनल पर दो अरब व्यूज़।"
    },
   
}

# --- UI Helpers ---
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}
"""



DEVICE = "cpu"
MODEL = None

def get_or_load_model():
    """Loads ChatterboxMultilingualTTS strictly on CPU (library-safe)."""
    global MODEL

    if MODEL is None:
        print("Model not loaded, initializing on CPU only...")

        try:
            # ❌ Do NOT pass map_location (not supported)
            MODEL = ChatterboxMultilingualTTS.from_pretrained("cpu")

            # ✅ Force CPU after load
            if hasattr(MODEL, "to"):
                MODEL = MODEL.to("cpu")

            MODEL.eval()

            # Disable grads for CPU inference
            for p in MODEL.parameters():
                p.requires_grad = False

            print("✅ Model loaded successfully on CPU")

        except Exception as e:
            print(f"❌ Error loading model on CPU: {e}")
            raise

    return MODEL

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

def set_seed(seed: int):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)


    
def resolve_audio_prompt(language_id: str, provided_path: str | None) -> str | None:
    """
    Decide which audio prompt to use:
    - If user provided a path (upload/mic/url), use it.
    - Else, fall back to language-specific default (if any).
    """
    if provided_path and str(provided_path).strip():
        return provided_path
    return LANGUAGE_CONFIG.get(language_id, {}).get("audio")



# ===============================
# Singing formatter (TEXT ONLY)
# ===============================
def format_for_singing(lyrics: str) -> str:
    """
    Encode melody directly into text for Chatterbox.
    NO instructions. ONLY singable text.
    """
    lines = []

    for line in lyrics.splitlines():
        line = line.strip()
        if not line:
            continue

        # Light vowel stretching (safe, readable)
        line = (
            line.replace("a", "aa")
                .replace("e", "ee")
                .replace("i", "ii")
                .replace("o", "oo")
                .replace("u", "uu")
        )

        # Add rhythm + pause
        lines.append(f"{line} ♪ ...")

    return "\n".join(lines)


# ===============================
# TTS generator (FIXED)
# ===============================
def generate_tts_audio(
    text_input: str,
    lyrics_input: str,
    mode: 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]:

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

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

    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

    # ===============================
    # STRICT MODE TOGGLE (IMPORTANT)
    # ===============================
    if mode == "Sing 🎵":
        if not lyrics_input.strip():
            raise gr.Error("Please enter lyrics for Sing mode.")
        final_text = format_for_singing(lyrics_input)
    else:
        if not text_input.strip():
            raise gr.Error("Please enter text for Speak mode.")
        final_text = text_input

    # ===============================
    # CPU-safe inference
    # ===============================
    with torch.no_grad():
        wav = current_model.generate(
            final_text[:300],
            language_id=language_id,
            **generate_kwargs
        )

    wav = wav.squeeze(0).detach().cpu().numpy()
    return current_model.sr, wav


# ===============================
# GRADIO UI
# ===============================
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Chatterbox Multilingual Demo
        Speak or sing text using Chatterbox (CPU-only).
        """
    )

    gr.Markdown(get_supported_languages_display())

    with gr.Row():
        with gr.Column():
            initial_lang = "hi"

            mode = gr.Radio(
                choices=["Speak 🗣️", "Sing 🎵"],
                value="Speak 🗣️",
                label="Output Mode"
            )

            text = gr.Textbox(
                value="धृतराष्ट्र ने कहा—हे संजय! धर्मभूमि कुरुक्षेत्र में एकत्रित, युद्ध की इच्छा वाले मेरे और पाण्डु के पुत्रों ने क्या किया?संजय ने कहा—उस समय राजा दुर्योधन ने व्यूह रचनायुक्त पाण्डवों की सेना को देखकर और द्रोणाचार्य के पास जाकर यह वचन कहा।",
                label="Text (Speak mode only)",
                placeholder="Add text here",
                lines=6,
                max_lines=10
            )

            lyrics = gr.Textbox(
                label="Lyrics (Sing mode only)",
                placeholder="Paste singable lyrics (one line per phrase)",
                max_lines=10
            )

            language_id = gr.Dropdown(
                choices=list(ChatterboxMultilingualTTS.get_supported_languages().keys()),
                value=initial_lang,
                label="Language"
            )

            ref_wav = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="Reference Audio (Optional)",
                value=default_audio_for_ui(initial_lang)
            )

            exaggeration = gr.Slider(
                0.25, 2.0, step=0.05,
                label="Exaggeration",
                value=0.5
            )

            cfg_weight = gr.Slider(
                0.2, 1.0, step=0.05,
                label="CFG / Pace",
                value=0.5
            )

            with gr.Accordion("More options", open=False):
                seed_num = gr.Number(value=0, label="Random seed (0 = random)")
                temp = gr.Slider(0.05, 5.0, step=0.05, label="Temperature", value=0.8)

            run_btn = gr.Button("Generate", variant="primary")

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


    # ===============================
    # AUTO-TUNE FOR SING MODE
    # ===============================
    def on_mode_change(mode):
        if mode == "Sing 🎵":
            return (
                gr.update(visible=False),  # hide text
                gr.update(visible=True),   # show lyrics
                1.3,                       # exaggeration
                1.0,                       # temperature
                0.45                       # cfg
            )
        else:
            return (
                gr.update(visible=True),
                gr.update(visible=False),
                0.5,
                0.8,
                0.5
            )

    mode.change(
        fn=on_mode_change,
        inputs=mode,
        outputs=[text, lyrics, exaggeration, temp, cfg_weight],
        show_progress=False
    )

    run_btn.click(
        fn=generate_tts_audio,
        inputs=[
            text,
            lyrics,
            mode,
            language_id,
            ref_wav,
            exaggeration,
            temp,
            seed_num,
            cfg_weight,
        ],
        outputs=[audio_output],
    )

demo.launch(mcp_server=True)