vllm-omni/examples/online_serving/voxtral_tts/gradio_demo.py

"""
- Make sure to install the following for this example to function correctly:
- `pip install -e .`
- `pip install gradio==5.50 mistral_common=1.10.0`

Example use case:

python examples/online_serving/voxtral_tts/gradio_demo.py --host slurm-199-077 --port 8000

"""

import argparse
import io
import json
import logging
import socket
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any

import gradio as gr
import httpx
import numpy as np
import soundfile as sf
from text_preprocess import sanitize_tts_input_text_for_demo

logger = logging.getLogger()

LOGFORMAT = "%(asctime)s - %(levelname)s - %(message)s"
TIMEFORMAT = "%Y-%m-%d %H:%M:%S"

logging.basicConfig(level=logging.INFO, force=True, format=LOGFORMAT, datefmt=TIMEFORMAT)
logger.setLevel(logging.INFO)


_BASE_URL = f"{socket.gethostname()}:7860"

# Default fallback voices - comprehensive list
_DEFAULT_VOICES = [
    "casual_female",
    "casual_male",
    "cheerful_female",
    "neutral_female",
    "neutral_male",
    "ar_male",
    "de_female",
    "de_male",
    "es_female",
    "es_male",
    "fr_female",
    "fr_male",
    "hi_female",
    "hi_male",
    "it_female",
    "it_male",
    "nl_female",
    "nl_male",
    "pt_female",
    "pt_male",
]


def organize_voices_by_language(voices: list[str]) -> tuple[list[str], dict[str, list[str]]]:
    """Organize voices into language categories.

    Args:
        voices: List of voice names (e.g., ["neutral_male", "es_female", "fr_male"])

    Returns:
        Tuple of (sorted list of language categories, dictionary mapping language to voices)
    """
    # Define language prefixes and their display names
    LANGUAGE_PREFIXES = {
        "ar": "Arabic",
        "de": "German",
        "es": "Spanish",
        "fr": "French",
        "it": "Italian",
        "nl": "Dutch",
        "pt": "Portuguese",
        "hi": "Hindi",
    }

    # Initialize language to voices mapping
    language_voices: dict[str, list[str]] = {}

    for voice in voices:
        # Check for language prefix
        found_language = None
        for prefix, lang_name in LANGUAGE_PREFIXES.items():
            if voice.lower().startswith(f"{prefix}_"):
                found_language = lang_name
                break

        if found_language:
            # Add voice to its language category
            if found_language not in language_voices:
                language_voices[found_language] = []
            language_voices[found_language].append(voice)
        else:
            # Add to English (voices without language prefix)
            if "English" not in language_voices:
                language_voices["English"] = []
            language_voices["English"].append(voice)

    # Sort voices within each language category, with neutral_male first for English
    for lang in language_voices:
        if lang == "English":
            language_voices[lang].sort(key=lambda v: (0 if v == "neutral_male" else 1, v))
        else:
            language_voices[lang].sort()

    # Sort language categories (English first, then alphabetically)
    sorted_languages = sorted(language_voices.keys(), key=lambda x: (0 if x == "English" else 1, x.lower()))

    return sorted_languages, language_voices


# Configuration for server health check
_SERVER_CHECK_TIMEOUT = 300.0  # 5 minutes max wait
_SERVER_CHECK_INTERVAL = 5.0  # 5 seconds between retries


def wait_for_server(base_url: str, timeout: float = _SERVER_CHECK_TIMEOUT) -> bool:
    """Block until the server is available or timeout is reached.

    Args:
        base_url: Base URL of the server (e.g., "http://localhost:8091/v1")
        timeout: Maximum time to wait in seconds

    Returns:
        True if server became available, False if timeout reached
    """
    start_time = time.time()
    # Health endpoint is at /health (not /v1/health)
    health_url = base_url.replace("/v1", "") + "/health"

    logger.info(f"Waiting for server at {base_url} to become available...")

    with httpx.Client(timeout=5.0) as client:
        while time.time() - start_time < timeout:
            try:
                resp = client.get(health_url)
                if resp.status_code == 200:
                    logger.info("Server is now available!")
                    return True
            except Exception:
                pass
            elapsed = time.time() - start_time
            logger.info(f"Server not yet available ({elapsed:.1f}s elapsed), retrying in {_SERVER_CHECK_INTERVAL}s...")
            time.sleep(_SERVER_CHECK_INTERVAL)

    logger.warning(f"Server did not become available within {timeout}s timeout")
    return False


def fetch_voices_and_languages(base_url: str, model: str) -> tuple[list[str], dict[str, list[str]]]:
    """Fetch available voices from the server API and organize by language.

    This function blocks until the server is available, then fetches the list
    of available voices from the /v1/audio/voices endpoint and organizes them.

    Args:
        base_url: Base URL of the server API (e.g., "http://localhost:8091/v1")
        model: Model name (used for logging)

    Returns:
        Tuple of (sorted language list, dictionary mapping language to voices)
    """
    # Always wait for server - if server is not up, demo cannot do anything
    if not wait_for_server(base_url):
        logger.warning("Server unavailable, using fallback voices")
        languages, language_voices = organize_voices_by_language(_DEFAULT_VOICES)
        return languages, language_voices

    try:
        with httpx.Client(timeout=10.0) as client:
            # The audio/voices endpoint is directly under /v1, not nested
            resp = client.get(
                f"{base_url}/audio/voices",
                headers={"Authorization": "Bearer EMPTY"},
            )
        if resp.status_code == 200:
            data = resp.json()
            voices = data.get("voices", [])
            if voices:
                logger.info(f"Fetched {len(voices)} voices from server: {voices}")
                languages, language_voices = organize_voices_by_language(sorted(voices))
                return languages, language_voices
            logger.warning("Server returned empty voices list")
    except Exception as e:
        logger.warning(f"Failed to fetch voices from server: {e}")

    # Fallback to default voices
    logger.info(f"Using fallback voices: {_DEFAULT_VOICES}")
    languages, language_voices = organize_voices_by_language(_DEFAULT_VOICES)
    return languages, language_voices


def make_update_voice_dropdown(language_voices: dict[str, list[str]]):
    """Return a callback that updates the voice dropdown when the user
    selects a different language."""

    def update_voice_dropdown(language: str) -> gr.Dropdown:
        voices = language_voices.get(language, [])
        return gr.Dropdown(choices=voices, value=voices[0] if voices else None, interactive=True)

    return update_voice_dropdown


def run_inference(
    voice_name: str,
    text_prompt: str,
    base_url: str,
    model: str,
) -> tuple[int, np.ndarray]:
    """Call /v1/audio/speech and return (sample_rate, audio_array)."""
    user_text_prompt = text_prompt.strip()
    if not user_text_prompt:
        raise gr.Error("Please enter a text prompt.")
    try:
        text_prompt = sanitize_tts_input_text_for_demo(user_text_prompt)
    except Exception as exc:
        raise gr.Error(f"Text preprocessing failed: {exc}") from exc

    payload: dict[str, Any] = {
        "input": text_prompt,
        "model": model,
        "response_format": "wav",
        "voice": voice_name,
    }

    response = httpx.post(
        f"{base_url}/audio/speech",
        json=payload,
        timeout=120.0,
    )
    response.raise_for_status()

    audio_array, sr = sf.read(io.BytesIO(response.content), dtype="float32")
    return sr, audio_array


def _save_example(
    outputs_dir: Path,
    voice_name: str,
    text_prompt: str,
    sr: int,
    audio_array: np.ndarray,
) -> tuple[str, str]:
    """
    Save inputs/outputs for sharing.
    Returns (share_id, saved_audio_path)
    """
    share_id = uuid.uuid4().hex

    # Save generated audio
    saved_audio_path = outputs_dir / f"{share_id}_gen.wav"
    sf.write(str(saved_audio_path), audio_array, sr)

    meta = {
        "id": share_id,
        "created_at": datetime.utcnow().isoformat(),
        "voice_name": voice_name,
        "text_prompt": text_prompt,
        "generated_audio_path": str(saved_audio_path),
    }

    with open(outputs_dir / f"{share_id}.json", "w") as f:
        json.dump(meta, f, ensure_ascii=False)

    return share_id, str(saved_audio_path)


def _load_from_share(
    outputs_dir: Path | None,
    request: gr.Request,
    languages: list[str] | None = None,
    language_voices: dict[str, list[str]] | None = None,
) -> tuple[str, str, str | None, str | None, dict[str, Any], str]:
    """
    Called on page load. If ?share_id=... is present, load stored example.
    Returns: (language, voice_name, text_prompt, output_audio, submit_btn_update, share_link_text)
    """
    fallback_language = (
        "English" if languages and "English" in languages else (languages[0] if languages else "English")
    )
    if language_voices:
        voices_list = language_voices.get(fallback_language, [])
        fallback_voice = voices_list[0] if voices_list else None
    else:
        fallback_voice = None

    if outputs_dir is None:
        return fallback_language, fallback_voice, "", None, gr.update(interactive=False), ""

    share_id = None
    if request and request.query_params:
        share_id = request.query_params.get("share_id")

    if not share_id:
        return fallback_language, fallback_voice, "", None, gr.update(interactive=False), ""

    meta_path = outputs_dir / f"{share_id}.json"
    if not meta_path.exists():
        logger.warning("No stored example for share_id=%s", share_id)
        return fallback_language, fallback_voice, "", None, gr.update(interactive=False), ""

    with open(meta_path) as f:
        meta = json.load(f)

    # Determine language and voice from the stored voice name
    voice = meta.get("voice_name", fallback_voice)
    text_prompt = meta.get("text_prompt", "")
    gen_path = meta.get("generated_audio_path")

    # Find which language this voice belongs to
    language = fallback_language
    if language_voices:
        for lang, voices in language_voices.items():
            if voice in voices:
                language = lang
                break

    if _BASE_URL:
        share_link = f"http://{_BASE_URL}?share_id={share_id}"
    else:
        share_link = f"Use this query on your current URL: ?share_id={share_id}"

    return language, voice, text_prompt, gen_path, gr.update(interactive=True), share_link


def main(
    model: str,
    host: str,
    port: str,
    output_dir: str | None = None,
) -> None:
    base_url = f"http://{host}:{port}/v1"
    logger.info(f"Using speech API at: {base_url}/audio/speech")

    outputs_dir: Path | None = None
    if output_dir is not None:
        outputs_dir = Path(output_dir)
        outputs_dir.mkdir(parents=True, exist_ok=True)

    # Load available voices and organize by language
    # This will block until the server is available
    languages, language_voices = fetch_voices_and_languages(base_url, model)

    with gr.Blocks(title="Voxtral TTS", fill_height=True) as demo:
        gr.Markdown("## Voxtral TTS")

        with gr.Row():
            with gr.Column():
                # Language dropdown (first level)
                default_language = "English" if "English" in languages else languages[0]
                language_dropdown = gr.Dropdown(
                    choices=languages,
                    label="Language",
                    value=default_language,
                )
                # Voice dropdown (second level, updates based on language)
                voices_for_default_lang = language_voices.get(default_language, [])
                voice_name = gr.Dropdown(
                    choices=voices_for_default_lang,
                    label="Voice",
                    value=voices_for_default_lang[0] if voices_for_default_lang else None,
                )
                text_prompt = gr.Textbox(
                    label="Text prompt",
                    placeholder="Enter the text you want to synthesize...",
                    lines=4,
                )
                with gr.Row():
                    reset_btn = gr.Button("Clear")
                    submit_btn = gr.Button("Generate audio", interactive=False)

            with gr.Column():
                output_audio = gr.Audio(
                    label="Generated audio",
                    show_download_button=True,
                    interactive=False,
                    autoplay=True,
                    type="filepath",
                )
                share_link_box = gr.Textbox(
                    label="Shareable link",
                    interactive=False,
                    show_copy_button=True,
                    visible=outputs_dir is not None,
                )

        # --- Cascading dropdown: update voice list when language changes ---
        language_dropdown.change(
            fn=make_update_voice_dropdown(language_voices),
            inputs=[language_dropdown],
            outputs=voice_name,
        )

        # --- UI logic: disable submit until text is non-empty ---
        def _toggle_submit(text: str):
            enabled = bool(text.strip())
            return gr.update(interactive=enabled)

        text_prompt.change(
            fn=_toggle_submit,
            inputs=[text_prompt],
            outputs=submit_btn,
        )

        # --- Wiring inference + persistence to the button ---
        def _on_submit(voice: str, text: str):
            assert text.strip() != ""
            sr, audio_array = run_inference(voice, text, base_url, model)
            if outputs_dir is not None:
                share_id, saved_audio_path = _save_example(
                    outputs_dir,
                    voice_name=voice,
                    text_prompt=text,
                    sr=sr,
                    audio_array=audio_array,
                )
                share_link = f"{_BASE_URL}?share_id={share_id}"
                return saved_audio_path, share_link
            return (sr, audio_array), ""

        submit_btn.click(
            fn=_on_submit,
            inputs=[voice_name, text_prompt],
            outputs=[output_audio, share_link_box],
        )

        # --- Clear everything and disable submit again ---
        def make_on_reset(languages: list[str], language_voices: dict[str, list[str]]):
            def _on_reset():
                language = "English" if "English" in languages else languages[0]
                voices_list = language_voices.get(language, [])
                voice = voices_list[0] if voices_list else None
                return (
                    language,  # language_dropdown
                    voice,  # voice_name
                    "",  # text_prompt
                    None,  # output_audio
                    gr.update(interactive=False),  # submit_btn
                    "",  # share_link_box
                )

            return _on_reset

        reset_btn.click(
            fn=make_on_reset(languages, language_voices),
            inputs=[],
            outputs=[language_dropdown, voice_name, text_prompt, output_audio, submit_btn, share_link_box],
        )

        def make_load_from_share(outputs_dir: Path | None, languages: list[str], language_voices: dict[str, list[str]]):
            def _load(request: gr.Request):
                return _load_from_share(outputs_dir, request, languages, language_voices)

            return _load

        demo.load(
            fn=make_load_from_share(outputs_dir, languages, language_voices),
            inputs=[],
            outputs=[language_dropdown, voice_name, text_prompt, output_audio, submit_btn, share_link_box],
        )

    launch_kwargs: dict[str, Any] = {
        "server_name": "0.0.0.0",
        "share": True,
    }
    if outputs_dir is not None:
        launch_kwargs["allowed_paths"] = [str(outputs_dir)]
    demo.launch(**launch_kwargs)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Voxtral TTS Gradio Demo")
    parser.add_argument("--model", type=str, default="mistralai/Voxtral-4B-TTS-2603", help="Name of model repo on HF")
    parser.add_argument("--host", type=str, default="localhost", help="Name of host")
    parser.add_argument("--port", type=str, default="8091", help="port number")
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Directory to save generated audio and share links. "
        "If not provided, save/share functionality is disabled.",
    )

    args = parser.parse_args()

    main(
        model=args.model,
        host=args.host,
        port=args.port,
        output_dir=args.output_dir,
    )