[Voxtral TTS] Add multilingual support in gradio demo (#2151)

Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

examples/offline_inference/voxtral_tts/README.md

@@ -58,7 +58,7 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
 | `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
 | `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
 | `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
-| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) |
+| `--voice NAME` | Voice preset to use instead of reference audio file. Check Huggingface `mistralai/Voxtral-4B-TTS-2603` to get the list of available voices |
 | `--write-audio` | Write generated audio to WAV files |
 | `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
 | `--log-stats` | Enable detailed statistics logging |

examples/online_serving/voxtral_tts/gradio_demo.py

@@ -14,6 +14,7 @@ import io
 import json
 import logging
 import socket
+import time
 import uuid
 from datetime import datetime
 from pathlib import Path
@@ -35,6 +36,178 @@ 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,
@@ -99,39 +272,59 @@ def _save_example(
 def _load_from_share(
     outputs_dir: Path | None,
     request: gr.Request,
-) -> tuple[str | None, str, str | None, dict[str, Any], str]:
+    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: (voice_name, text_prompt, output_audio, submit_btn_update, share_link_text)
+    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 "Neutral_Male", "", None, gr.update(interactive=False), ""
+        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 "Neutral_Male", "", None, gr.update(interactive=False), ""
+        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 "Neutral_Male", "", None, gr.update(interactive=False), ""
+        return fallback_language, fallback_voice, "", None, gr.update(interactive=False), ""
 
     with open(meta_path) as f:
         meta = json.load(f)
 
-    voice = meta.get("voice_name", "Neutral_Male")
+    # 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 voice, text_prompt, gen_path, gr.update(interactive=True), share_link
+    return language, voice, text_prompt, gen_path, gr.update(interactive=True), share_link
 
 
 def main(
@@ -148,15 +341,28 @@ def main(
         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(
-                    ["Neutral_Male", "Neutral_Female", "Cheerful_Female", "Casual_Male", "Casual_Female"],
+                    choices=voices_for_default_lang,
                     label="Voice",
-                    value="Neutral_Male",
+                    value=voices_for_default_lang[0] if voices_for_default_lang else None,
                 )
                 text_prompt = gr.Textbox(
                     label="Text prompt",
@@ -182,6 +388,13 @@ def main(
                     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())
@@ -216,25 +429,38 @@ def main(
         )
 
         # --- Clear everything and disable submit again ---
-        def _on_reset():
-            return (
-                "Neutral_Male",  # voice_name
-                "",  # text_prompt
-                None,  # output_audio
-                gr.update(interactive=False),  # submit_btn
-                "",  # share_link_box
-            )
+        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=_on_reset,
+            fn=make_on_reset(languages, language_voices),
             inputs=[],
-            outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
+            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=lambda request: _load_from_share(outputs_dir, request),
+            fn=make_load_from_share(outputs_dir, languages, language_voices),
             inputs=[],
-            outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
+            outputs=[language_dropdown, voice_name, text_prompt, output_audio, submit_btn, share_link_box],
         )
 
     launch_kwargs: dict[str, Any] = {