import os
os.environ.setdefault("CHATTERBOX_FLASH_ENGINE", "torch")
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")

import spaces
import torch
import numpy as np
import gradio as gr

import subprocess
import sys

subprocess.run(
    [sys.executable, "-m", "pip", "install", "--no-deps",
     "chatterbox-tts==0.1.7", "chatterbox-flash==0.1.0"],
    check=True,
)

from chatterbox_flash import ChatterboxFlashTTS

MODEL_ID = "ResembleAI/chatterbox-flash"
# Chatterbox-Flash has no built-in voice: `generate()` raises
# "Conditioning not prepared" unless conditionals exist. Ship a default
# reference clip so zero-shot generation works without a user upload.
DEFAULT_AUDIO_PROMPT = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "female_shadowheart4.flac"
)
_tts = None


def get_tts():
    global _tts
    if _tts is None:
        print(f"Loading Chatterbox-Flash from {MODEL_ID}...")
        _tts = ChatterboxFlashTTS.from_pretrained(
            MODEL_ID, device="cuda", dtype=torch.bfloat16,
        )
        print("Model loaded successfully.")
    return _tts


@spaces.GPU(duration=60)
def generate_tts(
    text_input: str,
    audio_prompt_path: str | None = None,
    exaggeration: float = 0.5,
    temperature: float = 0.6,
    cfg_scale: float = 1.0,
    num_steps: int = 10,
    seed_num: int = 0,
):
    """Generate speech from text using Chatterbox-Flash block-diffusion TTS."""
    tts = get_tts()

    if seed_num != 0:
        torch.manual_seed(int(seed_num))
        torch.cuda.manual_seed(int(seed_num))
        np.random.seed(int(seed_num))

    generate_kwargs = {
        "exaggeration": exaggeration,
        "temperature": temperature,
        "cfg_scale": cfg_scale,
        "num_steps": num_steps,
        "backend": "torch",
    }

    # Fall back to the bundled reference clip when the user supplies none, so
    # conditionals are always prepared before generation.
    prompt_path = audio_prompt_path or DEFAULT_AUDIO_PROMPT
    generate_kwargs["audio_prompt_path"] = prompt_path

    if tts.conds is None:
        tts.prepare_conditionals(prompt_path, exaggeration=exaggeration)

    wav = tts.generate(text_input[:300], **generate_kwargs)
    return (tts.sr, wav.squeeze(0).cpu().numpy())


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Chatterbox-Flash TTS
        Prior-calibrated block-diffusion zero-shot TTS by Resemble AI.
        Provide a reference audio clip to clone a voice, or generate with the default voice.

        [Paper](https://huggingface.co/papers/2605.30748) · [Model](https://huggingface.co/ResembleAI/chatterbox-flash) · [GitHub](https://github.com/resemble-ai/chatterbox-flash)
        """
    )

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=3):
            text = gr.Textbox(
                value="Sometimes it's better to just let things slide, you know?",
                label="Text to synthesize (max 300 chars)",
                max_lines=5,
            )
            ref_wav = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                value=DEFAULT_AUDIO_PROMPT,
                label="Reference Audio (for voice cloning) — defaults to the bundled voice",
            )

            with gr.Accordion("Advanced settings", open=False):
                exaggeration = gr.Slider(
                    0.25, 2.0, step=0.05,
                    label="Exaggeration (0.5=neutral, higher=more expressive)",
                    value=0.5,
                )
                temperature = gr.Slider(
                    0.05, 2.0, step=0.05,
                    label="Temperature",
                    value=0.6,
                )
                cfg_scale = gr.Slider(
                    0.2, 1.0, step=0.05,
                    label="CFG Scale",
                    value=1.0,
                )
                num_steps = gr.Slider(
                    1, 30, step=1,
                    label="Denoising Steps",
                    value=10,
                )
                seed_num = gr.Number(
                    value=0, label="Seed (0=random)", precision=0,
                )

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

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

    gr.Examples(
        examples=[
            ["Sometimes it's better to just let things slide, you know?"],
            ["The quick brown fox jumps over the lazy dog. Pack my box with five dozen liquor jugs."],
            ["In the depths of winter, I finally learned that within me lay an invincible summer."],
        ],
        inputs=[text],
        outputs=[audio_output],
        fn=generate_tts,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=generate_tts,
        inputs=[
            text, ref_wav, exaggeration,
            temperature, cfg_scale, num_steps, seed_num,
        ],
        outputs=[audio_output],
        api_name="generate",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)