Polish unified Audex demo

Add unified 30B-A3B and 2B inference, streaming text and speech generation, reasoning controls, curated examples, local setup scripts, and ZeroGPU runtime support.

.gitattributes

@@ -35,5 +35,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 examples/mlk_speech.wav filter=lfs diff=lfs merge=lfs -text
 examples/sample_speech.wav filter=lfs diff=lfs merge=lfs -text
+wheels/*.whl filter=lfs diff=lfs merge=lfs -text
 wheels/causal_conv1d-1.6.2.post1-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text
 wheels/mamba_ssm-2.3.2.post1-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text
+examples/korean_speech.wav filter=lfs diff=lfs merge=lfs -text
+examples/question_1059.mp3 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+.cache/
+.local/
+.venv/
+__pycache__/
+_wheelout/
+*.pyc

README.md

@@ -1,32 +1,37 @@
 ---
-title: Nemotron Labs Audex 30B A3B
+title: Nemotron-Labs-Audex
 emoji: 🎧
 colorFrom: purple
 colorTo: indigo
 sdk: gradio
 sdk_version: 6.20.0
 app_file: app.py
-short_description: Audio understanding, speech recognition & translation
+short_description: Unified audio-text intelligence
 python_version: "3.10"
 startup_duration_timeout: 1h
 ---
 
-# Nemotron-Labs-Audex-30B-A3B
+# Nemotron-Labs-Audex
 
-A Gradio demo for [`nvidia/Nemotron-Labs-Audex-30B-A3B`](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B),
-a unified audio-text MoE (30B total / 3B active) built on Nemotron-Cascade-2.
-This demo exposes its **audio understanding, speech recognition, and speech
-translation** capabilities via the official Hugging Face inference path.
+A Gradio demo for
+[`nvidia/Nemotron-Labs-Audex-30B-A3B`](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B)
+and [`nvidia/Nemotron-Labs-Audex-2B`](https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B).
+The model selector defaults to the 30B-A3B and can switch to 2B model.
 
-## Correct numerics
+The demo includes **audio understanding, speech recognition, speech translation, text reasoning, text-to-speech, and speech-to-speech**.
 
-The model is a Nemotron-H hybrid (Mamba2 + attention + MLP). The Mamba layers
-require the compiled CUDA fast path from `mamba-ssm` (`selective_scan_cuda`) and
-`causal-conv1d`; the pure-torch fallback produces degenerate output. This Space
-builds both extensions from source against the runtime's exact toolchain
-(torch 2.11.0+cu130, CUDA 13.0, Python 3.10, sm_120) on first boot and caches
-the resulting wheels back into the repo (`wheels/`) for fast subsequent boots.
+Reasoning has no separate token cap by default for either model. `Max new tokens` sets a total cap shared by reasoning and the final answer.
 
-Runs on ZeroGPU `xlarge` (full 96 GB card) since the bf16 weights are ~65 GB.
+## Run locally
 
-*Licensed for non-commercial use only (NVIDIA One-Way Noncommercial License).*
+```bash
+git clone https://huggingface.co/spaces/L0SG/Nemotron-Labs-Audex
+cd Nemotron-Labs-Audex
+bash setup_local.sh
+bash run_local.sh
+```
+
+`setup_local.sh` creates an isolated `.venv` and installs PyTorch when needed.
+Use `AUDEX_TORCH_PACKAGE` or standard pip index environment variables to select a platform-specific PyTorch build.
+
+`run_local.sh` detects the selected GPU architecture and creates a matching local CUDA-extension cache.

app.py

@@ -1,41 +1,233 @@
 """
-Nemotron-Labs-Audex-30B-A3B — Audio Understanding / Speech Recognition & Translation demo.
-
-This unified audio-text MoE (30B total, 3B active) is a Nemotron-H hybrid
-(Mamba2 + attention + MLP). Correct numerics REQUIRE the compiled CUDA fast
-path from `mamba-ssm` (selective_scan_cuda) and `causal-conv1d`; the pure-torch
-fallback produces degenerate/repeated tokens. We therefore build both extensions
-from source against the Space's exact torch (2.11.0+cu130) / CUDA 13.0 / py3.10
-/ sm_120 toolchain, cache the resulting wheels back into this Space repo, and
-reuse them on subsequent boots.
+Nemotron-Labs-Audex — unified audio and text demo for the 30B-A3B and 2B models.
+
+The 30B-A3B model uses a Mamba2-Transformer Hybrid MoE backbone (30B total,
+3B active). Its correct numerics require the compiled CUDA fast path from
+`mamba-ssm` and `causal-conv1d`; the pure-PyTorch fallback produces degenerate
+or repeated tokens. The Space reuses its cached Blackwell wheels, while local
+runtimes can build and cache wheels for their CUDA architecture.
 """
 import os
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
 import glob
+import hashlib
+import importlib.util
+import re
 import subprocess
 import sys
 import time
+from collections.abc import Iterator
+from dataclasses import dataclass
+from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
+from threading import Event, Thread
+from types import ModuleType
 
 # Import spaces FIRST (before torch / any CUDA-touching import) so its
 # torch.cuda.* monkey-patch is installed. The kernel build below runs in
 # subprocesses, so it does not initialize CUDA in this process.
 import spaces  # noqa: E402
 
+
+@dataclass(frozen=True)
+class AudexRuntime:
+    model: object
+    tokenizer: object
+    config: object
+    feature_extractor: object
+    speech_decoder: object
+    cache_implementation: str | None = None
+
+
+def _max_reasoning_budget(max_new_tokens: int) -> int:
+    budget = max(16, (int(max_new_tokens) - 16) * 10 // 11)
+    return max(16, budget // 16 * 16)
+
+
+# Runtime and build configuration
 APP_DIR = Path(__file__).resolve().parent
-WHEELS_DIR = APP_DIR / "wheels"
-REPO_ID = os.environ.get("SPACE_ID", "L0SG/nemotron-labs-audex-30b-a3b")
+IS_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU") == "1"
+DEFAULT_WHEELS_DIR = APP_DIR / ("wheels" if IS_ZERO_GPU else ".local/wheels")
+DEFAULT_BUILD_DIR = APP_DIR / ("_wheelout" if IS_ZERO_GPU else ".local/build")
+WHEELS_DIR = Path(os.environ.get("AUDEX_WHEELS_DIR", str(DEFAULT_WHEELS_DIR)))
+BUILD_DIR = Path(os.environ.get("AUDEX_BUILD_DIR", str(DEFAULT_BUILD_DIR)))
+REPO_ID = os.environ.get("SPACE_ID", "L0SG/Nemotron-Labs-Audex")
+CAUSAL_CONV1D_VERSION = "1.6.2.post1"
+MAMBA_SSM_VERSION = "2.3.1"
+CAUSAL_CONV1D_SPEC = f"causal-conv1d=={CAUSAL_CONV1D_VERSION}"
+MAMBA_SSM_SPEC = f"mamba-ssm=={MAMBA_SSM_VERSION}"
+TORCH_ARCH = os.environ.get("AUDEX_TORCH_ARCH", "12.0" if IS_ZERO_GPU else "")
+
+# Model configuration
+MODEL_30B_ID = "nvidia/Nemotron-Labs-Audex-30B-A3B"
+MODEL_2B_ID = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_SUBFOLDER = "checkpoint_folder_full"
+DECODER_SUBFOLDER = "audex_causal_speech_decoder"
+DEFAULT_MODEL_NAME = "Nemotron-Labs-Audex-30B-A3B"
+MODEL_2B_NAME = "Nemotron-Labs-Audex-2B"
+
+# Shared generation configuration
+SAMPLE_RATE = 16000
+MAX_AUDIO_DURATION_SECONDS = float(os.environ.get("AUDEX_MAX_AUDIO_SECONDS", "120"))
+MAX_NEW_TOKENS = int(os.environ.get("AUDEX_MAX_NEW_TOKENS", "2048"))
+DEFAULT_MAX_NEW_TOKENS = min(
+    int(os.environ.get("AUDEX_DEFAULT_MAX_NEW_TOKENS", "256")),
+    MAX_NEW_TOKENS,
+)
+DEFAULT_REASONING_BUDGET = max(
+    0,
+    min(
+        int(os.environ.get("AUDEX_DEFAULT_REASONING_BUDGET", "0")),
+        _max_reasoning_budget(DEFAULT_MAX_NEW_TOKENS),
+    ),
+)
+STREAM_CHUNK_SIZE = max(1, int(os.environ.get("AUDEX_STREAM_CHUNK_SIZE", "8")))
+MAX_GPU_DURATION_SECONDS = 60
+MAX_TOKEN_WARNING = (
+    "⚠️ Maximum new-token limit reached; this output is incomplete. "
+    "Increase Max new tokens and run again."
+)
+TEXT_VOCAB_SIZE = 131072
+TEXT_SEED = 100
+TEXT_SYSTEM_PROMPT = (
+    "You are a helpful and harmless assistant.\n\n"
+    "You are not allowed to use any tools."
+)
+
+# Task names and prompts
+TEXT_TASK = "Text reasoning"
+TTS_TASK = "Text to speech (TTS)"
+S2S_TASK = "Speech to speech (S2S)"
+LEGACY_TASK_PROMPTS = {
+    "Describe the audio": "Describe the audio in detail.",
+    "Transcribe (ASR)": "Transcribe the speech in the input audio.",
+    "Translate speech to English": "Translate the spoken content in the audio to English.",
+    "Answer a question about the audio": "Where is the communication likely taking place?",
+}
+LEGACY_GREEDY_TASKS = {"Transcribe (ASR)", "Translate speech to English"}
+MMLU_PRO_EXAMPLE_PROMPT = (
+    "Question:\n"
+    "Which organelle is primarily responsible for ATP production in eukaryotic cells?\n\n"
+    "Answer Choices:\n"
+    "(A) Golgi apparatus\n"
+    "(B) Mitochondrion\n"
+    "(C) Lysosome\n"
+    "(D) Endoplasmic reticulum\n\n"
+    "Conclude your response with the sentence `The answer is \\boxed{{X}}.`, "
+    "in which X is the correct capital letter of your choice."
+)
+S2S_RESPONSE_PROMPT = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on "
+    "the Nemotron-Cascade-2 architecture. You may output your reasoning, followed "
+    "by your final response. You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each "
+    "sentence on its own separate line."
+)
+S2S_TASK_INSTRUCTION = (
+    "Use the response prompt above only to control the style of your final answer. "
+    "Do not quote, explain, or analyze the response prompt. "
+    "Answer the user request in the spoken transcript below. "
+    "The final answer must be one short, self-contained sentence of at most "
+    "20 words while preserving every requested name, number, and conclusion."
+)
+S2S_UI_PROMPT = f"{S2S_RESPONSE_PROMPT}\n\n{S2S_TASK_INSTRUCTION}"
+TTS_EXAMPLE_TEXT = (
+    "Artificial intelligence is helping people understand and create sound in new ways."
+)
 
-# Pinned versions recommended by the model card.
-CAUSAL_CONV1D_SPEC = "causal-conv1d==1.6.2.post1"
-MAMBA_SSM_SPEC = "mamba-ssm==2.3.2.post1"
-# sm_120 (RTX PRO 6000 Blackwell). Build env has no GPU, so arch must be explicit.
-TORCH_ARCH = "12.0"
+# Task-specific limits and feature switches
+TTS_MAX_NEW_TOKENS = min(
+    int(os.environ.get("AUDEX_TTS_MAX_NEW_TOKENS", "192")),
+    MAX_NEW_TOKENS,
+)
+S2S_TEXT_MAX_NEW_TOKENS = int(os.environ.get("AUDEX_S2S_TEXT_MAX_NEW_TOKENS", "2048"))
+S2S_REASONING_BUDGET = int(os.environ.get("AUDEX_S2S_REASONING_BUDGET", "0"))
+S2S_GPU_DURATION_SECONDS = int(os.environ.get("AUDEX_S2S_GPU_DURATION_SECONDS", "120"))
+S2S_SPOKEN_MAX_WORDS = 20
+S2S_TTS_MAX_NEW_TOKENS = int(os.environ.get("AUDEX_S2S_TTS_MAX_NEW_TOKENS", "2400"))
+S2S_TTS_SEGMENT_SILENCE_SECONDS = 0.2
+TTS_STREAMING_PLAYER_ENABLED = (
+    os.environ.get("AUDEX_TTS_STREAMING_PLAYER", "false").lower() == "true"
+)
+TASK_SPECS = {
+    "Speech recognition (ASR)": {
+        "modality": "audio",
+        "template": "<sound>\nTranscribe the speech in the input audio.",
+        "reasoning": False,
+        "temperature": 1.0,
+        "top_p": 1.0,
+        "greedy": True,
+    },
+    "Speech translation (AST)": {
+        "modality": "audio",
+        "template": "<sound>\nTranslate the spoken content in the audio to English.",
+        "reasoning": False,
+        "temperature": 1.0,
+        "top_p": 1.0,
+        "greedy": True,
+    },
+    "Audio description": {
+        "modality": "audio",
+        "template": "Describe the audio in detail.\n<sound>",
+        "reasoning": False,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "greedy": False,
+    },
+    "Audio question answering": {
+        "modality": "audio",
+        "template": "Where is the communication likely taking place?\n<sound>",
+        "reasoning": False,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "greedy": False,
+    },
+    TEXT_TASK: {
+        "modality": "text",
+        "template": MMLU_PRO_EXAMPLE_PROMPT,
+        "reasoning": True,
+        "temperature": 1.0,
+        "top_p": 0.95,
+        "greedy": False,
+    },
+    TTS_TASK: {
+        "modality": "tts",
+        "template": TTS_EXAMPLE_TEXT,
+        "reasoning": False,
+        "guidance_scale": 2.0,
+        "temperature": 0.8,
+        "top_p": 1.0,
+        "greedy": False,
+    },
+    S2S_TASK: {
+        "modality": "s2s",
+        "template": S2S_UI_PROMPT,
+        "reasoning": True,
+        "reasoning_budget": min(
+            S2S_REASONING_BUDGET,
+            _max_reasoning_budget(S2S_TEXT_MAX_NEW_TOKENS),
+        ),
+        "guidance_scale": 1.5,
+        "temperature": 1.0,
+        "top_p": 0.95,
+        "greedy": False,
+        "max_new_tokens": S2S_TEXT_MAX_NEW_TOKENS,
+        "max_new_tokens_limit": S2S_TEXT_MAX_NEW_TOKENS,
+    },
+}
 
 
-def _run(cmd, env=None, timeout=3000):
+def _run(
+    cmd: list[str],
+    env: dict[str, str] | None = None,
+    timeout: int = 3000,
+) -> subprocess.CompletedProcess[str]:
     print(f"[build] $ {' '.join(cmd)}", flush=True)
     p = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=timeout)
     if p.stdout:
@@ -50,45 +242,128 @@ def _run(cmd, env=None, timeout=3000):
     return p
 
 
-def _installed(mod):
+def _installed_version(package: str) -> str | None:
     try:
-        __import__(mod)
-        return True
-    except Exception:
+        return version(package)
+    except PackageNotFoundError:
+        return None
+
+
+def _kernels_ready() -> bool:
+    if (
+        _installed_version("causal-conv1d") != CAUSAL_CONV1D_VERSION
+        or _installed_version("mamba-ssm") != MAMBA_SSM_VERSION
+    ):
         return False
+    try:
+        from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+        from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+        from mamba_ssm.ops.triton.ssd_combined import (
+            mamba_chunk_scan_combined,
+            mamba_split_conv1d_scan_combined,
+        )
+    except ImportError:
+        return False
+    return all(
+        kernel is not None
+        for kernel in (
+            causal_conv1d_fn,
+            causal_conv1d_update,
+            selective_state_update,
+            mamba_chunk_scan_combined,
+            mamba_split_conv1d_scan_combined,
+        )
+    )
 
 
-def _pip(*args):
-    _run([sys.executable, "-m", "pip"] + list(args))
+def _pip(*args: str, env: dict[str, str] | None = None) -> None:
+    _run([sys.executable, "-m", "pip", *args], env=env)
 
 
-def _build_env():
+def _detect_torch_arch() -> str:
+    if TORCH_ARCH:
+        return TORCH_ARCH
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            (
+                "import torch; "
+                "assert torch.cuda.is_available(), 'CUDA is not available'; "
+                "major, minor = torch.cuda.get_device_capability(); "
+                "print(f'{major}.{minor}')"
+            ),
+        ],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(
+            "Unable to detect the GPU architecture. Set AUDEX_TORCH_ARCH "
+            f"explicitly. Details: {result.stderr.strip()}"
+        )
+    return result.stdout.strip()
+
+
+def _detect_cuda_home() -> str:
+    if cuda_home := os.environ.get("CUDA_HOME"):
+        if (Path(cuda_home) / "bin/nvcc").is_file():
+            return cuda_home
+        raise RuntimeError(f"CUDA_HOME does not contain bin/nvcc: {cuda_home}")
+    if IS_ZERO_GPU:
+        return "/cuda-image/usr/local/cuda-13.0"
+
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "from torch.utils.cpp_extension import CUDA_HOME; print(CUDA_HOME or '')",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    candidates = [result.stdout.strip(), "/usr/local/cuda"]
+    for candidate in candidates:
+        if candidate and (Path(candidate) / "bin/nvcc").is_file():
+            return candidate
+    raise RuntimeError(
+        "A CUDA toolkit with nvcc is required to build the 30B Mamba kernels. "
+        "Install the toolkit or set CUDA_HOME."
+    )
+
+
+def _build_env() -> dict[str, str]:
     env = dict(os.environ)
     env["MAMBA_FORCE_BUILD"] = "TRUE"
     env["CAUSAL_CONV1D_FORCE_BUILD"] = "TRUE"
-    env["TORCH_CUDA_ARCH_LIST"] = TORCH_ARCH
+    env["TORCH_CUDA_ARCH_LIST"] = _detect_torch_arch()
     env["MAX_JOBS"] = env.get("MAX_JOBS", "4")
-    # Ensure nvcc is discoverable.
-    cuda_home = env.get("CUDA_HOME") or "/cuda-image/usr/local/cuda-13.0"
+    cuda_home = _detect_cuda_home()
     env["CUDA_HOME"] = cuda_home
     env["PATH"] = f"{cuda_home}/bin:" + env.get("PATH", "")
     return env
 
 
-def ensure_kernels():
+def ensure_kernels() -> None:
     """Install compiled causal-conv1d + mamba-ssm. Use cached wheels if present,
     otherwise build from source and cache the wheels back into the repo."""
-    if _installed("causal_conv1d") and _installed("selective_scan_cuda"):
+    if _kernels_ready():
         print("[build] kernels already importable", flush=True)
         return
 
-    WHEELS_DIR.mkdir(exist_ok=True)
-    cached = sorted(glob.glob(str(WHEELS_DIR / "*.whl")))
-    if cached:
+    WHEELS_DIR.mkdir(parents=True, exist_ok=True)
+    cached_causal = sorted(
+        glob.glob(str(WHEELS_DIR / f"causal_conv1d-{CAUSAL_CONV1D_VERSION}-*.whl"))
+    )
+    cached_mamba = sorted(
+        glob.glob(str(WHEELS_DIR / f"mamba_ssm-{MAMBA_SSM_VERSION}-*.whl"))
+    )
+    if cached_causal and cached_mamba:
+        cached = [cached_causal[-1], cached_mamba[-1]]
         print(f"[build] installing cached wheels: {cached}", flush=True)
         try:
             _pip("install", "--no-deps", "--no-build-isolation", *cached)
-            if _installed("causal_conv1d") and _installed("selective_scan_cuda"):
+            if _kernels_ready():
                 print("[build] cached wheels installed OK", flush=True)
                 return
             print("[build] cached wheels imported incompletely; rebuilding", flush=True)
@@ -96,25 +371,47 @@ def ensure_kernels():
             print(f"[build] cached wheel install failed ({e!r}); rebuilding", flush=True)
 
     env = _build_env()
-    out_dir = APP_DIR / "_wheelout"
-    out_dir.mkdir(exist_ok=True)
+    BUILD_DIR.mkdir(parents=True, exist_ok=True)
 
     # Build wheels (no-build-isolation => uses the preinstalled torch).
     t0 = time.time()
     print("[build] building causal-conv1d + mamba-ssm from source (this can take ~20 min)", flush=True)
-    _pip("wheel", "--no-build-isolation", "--no-deps", "-w", str(out_dir),
-         CAUSAL_CONV1D_SPEC)
+    _pip(
+        "wheel",
+        "--no-build-isolation",
+        "--no-deps",
+        "-w",
+        str(BUILD_DIR),
+        CAUSAL_CONV1D_SPEC,
+        env=env,
+    )
     # mamba-ssm needs causal-conv1d importable during its own build; install it first.
-    built = sorted(glob.glob(str(out_dir / "causal_conv1d*.whl")))
-    _pip("install", "--no-deps", *built)
-    _pip("wheel", "--no-build-isolation", "--no-deps", "-w", str(out_dir),
-         MAMBA_SSM_SPEC)
+    built_causal = sorted(
+        glob.glob(str(BUILD_DIR / f"causal_conv1d-{CAUSAL_CONV1D_VERSION}-*.whl"))
+    )
+    if not built_causal:
+        raise RuntimeError(f"build produced no wheel for {CAUSAL_CONV1D_SPEC}")
+    _pip("install", "--no-deps", built_causal[-1])
+    _pip(
+        "wheel",
+        "--no-build-isolation",
+        "--no-deps",
+        "-w",
+        str(BUILD_DIR),
+        MAMBA_SSM_SPEC,
+        env=env,
+    )
     print(f"[build] source build finished in {time.time()-t0:.0f}s", flush=True)
 
-    all_wheels = sorted(glob.glob(str(out_dir / "*.whl")))
+    built_mamba = sorted(
+        glob.glob(str(BUILD_DIR / f"mamba_ssm-{MAMBA_SSM_VERSION}-*.whl"))
+    )
+    if not built_mamba:
+        raise RuntimeError(f"build produced no wheel for {MAMBA_SSM_SPEC}")
+    all_wheels = [built_causal[-1], built_mamba[-1]]
     _pip("install", "--no-deps", *all_wheels)
 
-    if not (_installed("causal_conv1d") and _installed("selective_scan_cuda")):
+    if not _kernels_ready():
         raise RuntimeError("kernel build completed but imports still fail")
     print("[build] kernels built + installed OK", flush=True)
 
@@ -125,6 +422,10 @@ def ensure_kernels():
             if not dst.exists():
                 import shutil
                 shutil.copy(w, dst)
+        if not IS_ZERO_GPU:
+            print(f"[build] kernels cached locally in {WHEELS_DIR}", flush=True)
+            return
+
         from huggingface_hub import HfApi
         tok = os.environ.get("HF_TOKEN")
         if tok:
@@ -147,8 +448,14 @@ ensure_kernels()
 # ---- Now safe to bring in torch / model ----
 import torch
 import gradio as gr
+import numpy as np
 from huggingface_hub import snapshot_download
-from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForCausalLM, AutoTokenizer
+from transformers import (
+    AutoConfig,
+    AutoFeatureExtractor,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+)
 
 from audio_utils import (
     IM_END_TOKEN,
@@ -160,41 +467,37 @@ from audio_utils import (
     resolve_audio_preprocessor_path,
     split_thinking,
 )
-
-MODEL_ID = "nvidia/Nemotron-Labs-Audex-30B-A3B"
-SUBFOLDER = "checkpoint_folder_full"
-SAMPLE_RATE = 16000
-
-print("[load] downloading checkpoint…", flush=True)
-LOCAL_DIR = snapshot_download(
-    MODEL_ID,
-    allow_patterns=[f"{SUBFOLDER}/*"],
-    token=os.environ.get("HF_TOKEN"),
+from reasoning_utils import ReasoningBudgetLogitsProcessor
+from tts_player import TTS_PLAYER_CSS, TTS_PLAYER_JS, TTS_PLAYER_TEMPLATE, player_value
+from tts_utils import (
+    EventStoppingCriteria,
+    TokenIdStreamer,
+    encode_pcm_chunk,
+    load_speech_decoder,
+    stream_tts,
+    write_wav,
 )
-MODEL_PATH = os.path.join(LOCAL_DIR, SUBFOLDER)
-print(f"[load] checkpoint at {MODEL_PATH}", flush=True)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
-config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True)
-feature_extractor = AutoFeatureExtractor.from_pretrained(
-    resolve_audio_preprocessor_path(MODEL_PATH, config)
-)
-print("[load] tokenizer/config/feature_extractor loaded", flush=True)
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_PATH,
-    trust_remote_code=True,
-    dtype=torch.bfloat16,
-    low_cpu_mem_usage=True,
-).eval()
-model = model.to("cuda")
-print("[load] model loaded and moved to cuda (packed by ZeroGPU)", flush=True)
+def _load_native_module(model_path: str) -> ModuleType:
+    module_path = Path(model_path).resolve() / "modeling_nemotron_h_audio_native.py"
+    if not module_path.is_file():
+        raise FileNotFoundError(
+            f"Native Audex model adapter not found in checkpoint: {module_path}"
+        )
+
+    path_hash = hashlib.sha256(str(module_path).encode()).hexdigest()[:12]
+    module_name = f"audex_native_{path_hash}"
+    if module_name in sys.modules:
+        return sys.modules[module_name]
 
-# Verify fast path is active (import-time flag inside the remote modeling module).
-try:
-    import transformers_modules  # noqa
-except Exception:
-    pass
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Unable to load native Audex model adapter: {module_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
 
 
 def _find_fast_path_flag() -> bool | None:
@@ -204,16 +507,168 @@ def _find_fast_path_flag() -> bool | None:
     return None
 
 
-TASK_PROMPTS = {
-    "Describe the audio": "Describe the audio in detail.",
-    "Transcribe (ASR)": "Transcribe the speech in the input audio.",
-    "Translate speech to English": "Translate the speech in the input audio into English.",
-    "Answer a question about the audio": "What is being said in this audio, and what is the tone?",
+def _resolve_model_paths(
+    model_id: str,
+    model_name: str,
+    model_path_env: str,
+    decoder_path_env: str,
+) -> tuple[str, str]:
+    model_path_override = os.environ.get(model_path_env)
+    if model_path_override:
+        model_path = Path(model_path_override).expanduser().resolve()
+        if not model_path.is_dir():
+            raise FileNotFoundError(f"{model_path_env} does not exist: {model_path}")
+        model_root = model_path.parent
+        print(f"[load] using local {model_name} checkpoint at {model_path}", flush=True)
+    else:
+        print(f"[load] downloading {model_name} checkpoint…", flush=True)
+        model_root = Path(
+            snapshot_download(
+                model_id,
+                allow_patterns=[
+                    f"{MODEL_SUBFOLDER}/*",
+                    f"{DECODER_SUBFOLDER}/*",
+                ],
+                token=os.environ.get("HF_TOKEN"),
+            )
+        )
+        model_path = model_root / MODEL_SUBFOLDER
+
+    decoder_path = (
+        Path(
+            os.environ.get(
+                decoder_path_env,
+                str(model_root / DECODER_SUBFOLDER),
+            )
+        )
+        .expanduser()
+        .resolve()
+    )
+    print(f"[load] {model_name} checkpoint at {model_path}", flush=True)
+    print(f"[load] {model_name} speech decoder at {decoder_path}", flush=True)
+    return str(model_path), str(decoder_path)
+
+
+def _load_30b_runtime() -> AudexRuntime:
+    model_path, decoder_path = _resolve_model_paths(
+        MODEL_30B_ID,
+        DEFAULT_MODEL_NAME,
+        "AUDEX_MODEL_PATH",
+        "AUDEX_DECODER_PATH",
+    )
+    native_module = _load_native_module(model_path)
+    config_class = native_module.NemotronHAudexConfig
+    model_class = native_module.NemotronHAudexForConditionalGeneration
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    chat_template_path = Path(model_path) / "chat_template.jinja"
+    if chat_template_path.is_file():
+        tokenizer.chat_template = chat_template_path.read_text()
+    config = config_class.from_pretrained(model_path)
+    feature_extractor = AutoFeatureExtractor.from_pretrained(
+        resolve_audio_preprocessor_path(model_path, config)
+    )
+    print(f"[load] {DEFAULT_MODEL_NAME} tokenizer/config/feature extractor loaded", flush=True)
+
+    model_load_kwargs = {
+        "pretrained_model_name_or_path": model_path,
+        "dtype": torch.bfloat16,
+        "low_cpu_mem_usage": True,
+        "output_loading_info": True,
+    }
+    if not IS_ZERO_GPU:
+        model_load_kwargs["device_map"] = {"": "cuda:0"}
+    model, loading_info = model_class.from_pretrained(**model_load_kwargs)
+    invalid_keys = {
+        key: loading_info[key]
+        for key in ("missing_keys", "unexpected_keys", "mismatched_keys")
+        if loading_info[key]
+    }
+    if invalid_keys:
+        raise RuntimeError(f"Native checkpoint loading was incomplete: {invalid_keys}")
+    model = model.eval()
+    if IS_ZERO_GPU:
+        model = model.to("cuda")
+    runtime_name = "packed by ZeroGPU" if IS_ZERO_GPU else "local CUDA"
+    print(f"[load] {DEFAULT_MODEL_NAME} loaded on {runtime_name}", flush=True)
+
+    if not _find_fast_path_flag():
+        raise RuntimeError(
+            "Transformers loaded without the native Mamba2-Transformer Hybrid fast path."
+        )
+    print("[load] native Mamba2-Transformer Hybrid fast path active", flush=True)
+    speech_decoder = load_speech_decoder(decoder_path)
+    print(f"[load] {DEFAULT_MODEL_NAME} speech decoder loaded", flush=True)
+    return AudexRuntime(
+        model=model,
+        tokenizer=tokenizer,
+        config=config,
+        feature_extractor=feature_extractor,
+        speech_decoder=speech_decoder,
+    )
+
+
+def _load_2b_runtime() -> AudexRuntime:
+    model_path, decoder_path = _resolve_model_paths(
+        MODEL_2B_ID,
+        MODEL_2B_NAME,
+        "AUDEX_2B_MODEL_PATH",
+        "AUDEX_2B_DECODER_PATH",
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    chat_template_path = Path(model_path) / "chat_template.jinja"
+    if chat_template_path.is_file():
+        tokenizer.chat_template = chat_template_path.read_text()
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    feature_extractor = AutoFeatureExtractor.from_pretrained(
+        resolve_audio_preprocessor_path(model_path, config)
+    )
+    model_load_kwargs = {
+        "pretrained_model_name_or_path": model_path,
+        "trust_remote_code": True,
+        "dtype": torch.bfloat16,
+        "low_cpu_mem_usage": True,
+    }
+    if not IS_ZERO_GPU:
+        model_load_kwargs["device_map"] = {"": "cuda:0"}
+    model = AutoModelForCausalLM.from_pretrained(**model_load_kwargs).eval()
+    if IS_ZERO_GPU:
+        model = model.to("cuda")
+    runtime_name = "packed by ZeroGPU" if IS_ZERO_GPU else "local CUDA"
+    print(f"[load] {MODEL_2B_NAME} loaded on {runtime_name}", flush=True)
+    speech_decoder = load_speech_decoder(decoder_path)
+    print(f"[load] {MODEL_2B_NAME} speech decoder loaded", flush=True)
+    return AudexRuntime(
+        model=model,
+        tokenizer=tokenizer,
+        config=config,
+        feature_extractor=feature_extractor,
+        speech_decoder=speech_decoder,
+        cache_implementation="static",
+    )
+
+
+model_runtimes = {
+    DEFAULT_MODEL_NAME: _load_30b_runtime(),
+    MODEL_2B_NAME: _load_2b_runtime(),
 }
-GREEDY_TASKS = {"Transcribe (ASR)", "Translate speech to English"}
-MAX_AUDIO_DURATION_SECONDS = 120.0
-MAX_NEW_TOKENS = 512
-MAX_GPU_DURATION_SECONDS = 60
+
+
+def _blocked_output_token_ids(runtime: AudexRuntime) -> list[int]:
+    sound_token_ids = {
+        int(getattr(runtime.config, name))
+        for name in ("sound_token_id", "sound_start_token_id", "sound_end_token_id")
+        if getattr(runtime.config, name, None) is not None
+    }
+    return sorted(
+        set(range(TEXT_VOCAB_SIZE, int(runtime.config.vocab_size))) | sound_token_ids
+    )
+
+
+def _get_runtime(model_name: str) -> AudexRuntime:
+    try:
+        return model_runtimes[model_name]
+    except KeyError as error:
+        raise gr.Error(f"Unknown model: {model_name}") from error
 
 
 def _probe_audio_duration(audio: str | None) -> float:
@@ -248,7 +703,138 @@ def _estimate(
     return MAX_GPU_DURATION_SECONDS
 
 
+def _stream_fields(response: str, reasoning: bool) -> tuple[str, str]:
+    visible = response
+    for marker in (IM_END_TOKEN, "<|end_of_text|>", "</s>"):
+        visible = visible.split(marker, 1)[0]
+
+    if not reasoning:
+        _, answer = split_thinking(visible)
+        return answer, ""
+    if "</think>" not in visible:
+        return "", visible.removeprefix("<think>").strip()
+
+    thinking, answer = visible.rsplit("</think>", 1)
+    return answer.strip(), thinking.removeprefix("<think>").strip()
+
+
+def _stream_model_generate(
+    runtime: AudexRuntime,
+    generation_kwargs: dict[str, object],
+    reasoning: bool,
+    reasoning_budget: int | None = None,
+) -> Iterator[tuple[str, str]]:
+    model_kwargs = dict(generation_kwargs)
+    eos_token_id = model_kwargs.get("eos_token_id")
+    eos_token_ids = {eos_token_id} if isinstance(eos_token_id, int) else set(eos_token_id or [])
+    logits_processors = list(model_kwargs.pop("logits_processor", []))
+    if reasoning and reasoning_budget is not None and reasoning_budget > 0:
+        prompt_length = int(model_kwargs["input_ids"].shape[-1])
+        logits_processors.append(
+            ReasoningBudgetLogitsProcessor(
+                runtime.tokenizer,
+                prompt_length=prompt_length,
+                reasoning_budget=int(reasoning_budget),
+            )
+        )
+    streamer = TokenIdStreamer()
+    cancel_event = Event()
+    generation_error: list[BaseException] = []
+    generated_tokens = 0
+    buffered_token_ids: list[int] = []
+    response = ""
+    last_fields: tuple[str, str] | None = None
+    generation_finished = False
+
+    def generate() -> None:
+        try:
+            stopping_criteria = list(model_kwargs.pop("stopping_criteria", []))
+            cache_kwargs = (
+                {
+                    "cache_implementation": runtime.cache_implementation,
+                    "disable_compile": True,
+                }
+                if runtime.cache_implementation
+                else {}
+            )
+            with torch.inference_mode():
+                runtime.model.generate(
+                    **model_kwargs,
+                    **cache_kwargs,
+                    logits_processor=logits_processors,
+                    stopping_criteria=[
+                        *stopping_criteria,
+                        EventStoppingCriteria(cancel_event),
+                    ],
+                    streamer=streamer,
+                    use_cache=True,
+                )
+        except BaseException as error:
+            generation_error.append(error)
+            streamer.fail(error)
+
+    def flush() -> tuple[str, str]:
+        nonlocal response
+        response += runtime.tokenizer.decode(
+            buffered_token_ids,
+            skip_special_tokens=False,
+            clean_up_tokenization_spaces=False,
+        )
+        buffered_token_ids.clear()
+        return _stream_fields(response, reasoning)
+
+    thread = Thread(target=generate, daemon=True)
+    thread.start()
+    try:
+        for token_id in streamer:
+            generated_tokens += 1
+            buffered_token_ids.append(token_id)
+            if token_id in eos_token_ids:
+                generation_finished = True
+            if len(buffered_token_ids) < STREAM_CHUNK_SIZE and not generation_finished:
+                continue
+
+            fields = flush()
+            if fields != last_fields:
+                yield fields
+                last_fields = fields
+            if generation_finished:
+                break
+
+        if generation_error:
+            raise generation_error[0]
+        if buffered_token_ids:
+            fields = flush()
+            if fields != last_fields:
+                yield fields
+                last_fields = fields
+
+        hit_max_tokens = (
+            not generation_finished
+            and generated_tokens >= int(model_kwargs["max_new_tokens"])
+        )
+        final_fields = _stream_fields(response, reasoning)
+        if reasoning and generation_finished and "</think>" not in response:
+            final_fields = _stream_fields(response, reasoning=False)
+        if hit_max_tokens:
+            answer, thinking = final_fields
+            answer = f"{answer}\n\n{MAX_TOKEN_WARNING}".strip()
+            final_fields = answer, thinking
+        if final_fields != last_fields:
+            yield final_fields
+        print(
+            f"[gpu] generated_tokens={generated_tokens} "
+            f"reasoning_closed={'</think>' in response} "
+            f"hit_max_tokens={hit_max_tokens}",
+            flush=True,
+        )
+    finally:
+        cancel_event.set()
+        thread.join(timeout=10)
+
+
 def _generate(
+    runtime: AudexRuntime,
     audio: str,
     prompt: str,
     reasoning: bool,
@@ -256,10 +842,11 @@ def _generate(
     temperature: float,
     top_p: float,
     greedy: bool,
-) -> tuple[str, str]:
+    reasoning_budget: int | None = None,
+) -> Iterator[tuple[str, str]]:
     started_at = time.perf_counter()
-    ff = _find_fast_path_flag()
-    print(f"[gpu] is_fast_path_available={ff}", flush=True)
+    if runtime is model_runtimes[DEFAULT_MODEL_NAME]:
+        print(f"[gpu] is_fast_path_available={_find_fast_path_flag()}", flush=True)
 
     wav, sr = load_audio(audio, target_sr=SAMPLE_RATE)
     audio_duration = wav.shape[-1] / sr
@@ -270,24 +857,41 @@ def _generate(
         )
     if max_new_tokens > MAX_NEW_TOKENS:
         raise gr.Error(f"This demo supports up to {MAX_NEW_TOKENS} output tokens.")
+    if (
+        reasoning
+        and reasoning_budget is not None
+        and reasoning_budget > 0
+        and (
+            reasoning_budget >= max_new_tokens
+            or reasoning_budget > _max_reasoning_budget(max_new_tokens)
+        )
+    ):
+        raise gr.Error(
+            "Reasoning budget is too high for max new tokens after reserving "
+            "the newline grace window and final answer."
+        )
 
     input_features = extract_whisper_features(
-        feature_extractor, wav, sample_rate=sr,
-        clip_duration=float(getattr(config, "sound_clip_duration", 30.0)),
+        runtime.feature_extractor,
+        wav,
+        sample_rate=sr,
+        clip_duration=float(getattr(runtime.config, "sound_clip_duration", 30.0)),
+    )
+    num_embeddings = input_features.shape[0] * int(
+        getattr(runtime.config, "sound_embedding_size", 750)
     )
-    num_embeddings = input_features.shape[0] * int(getattr(config, "sound_embedding_size", 750))
     formatted = build_prompt_template(prompt.strip(), reasoning=bool(reasoning),
                                       prompt_repitition="none")
     expanded = expand_sound_placeholder(formatted, num_embeddings)
-    tok = tokenizer(expanded, return_tensors="pt", add_special_tokens=False)
+    tok = runtime.tokenizer(expanded, return_tensors="pt", add_special_tokens=False)
     input_ids = tok.input_ids.to("cuda")
     attention_mask = (tok.attention_mask if "attention_mask" in tok
                       else build_attention_mask(input_ids)).to("cuda")
     input_features = input_features.to("cuda")
 
-    eos_token_id = tokenizer.convert_tokens_to_ids(IM_END_TOKEN)
-    if eos_token_id is None or eos_token_id == tokenizer.unk_token_id:
-        eos_token_id = getattr(config, "eos_token_id", None)
+    eos_token_id = runtime.tokenizer.convert_tokens_to_ids(IM_END_TOKEN)
+    if eos_token_id is None or eos_token_id == runtime.tokenizer.unk_token_id:
+        eos_token_id = getattr(runtime.config, "eos_token_id", None)
 
     temperature = max(float(temperature), 1e-4)
     top_p = float(top_p)
@@ -295,31 +899,36 @@ def _generate(
     gen_kwargs = dict(
         do_sample=do_sample,
         eos_token_id=eos_token_id,
-        pad_token_id=tokenizer.pad_token_id or getattr(config, "pad_token_id", 0),
+        pad_token_id=runtime.tokenizer.pad_token_id
+        or getattr(runtime.config, "pad_token_id", 0),
         max_new_tokens=int(max_new_tokens),
+        suppress_tokens=_blocked_output_token_ids(runtime),
     )
+    # The official top-k value is 0, which means no top-k filter in Transformers.
     if do_sample:
         gen_kwargs["temperature"] = temperature
         if top_p > 0.0:
             gen_kwargs["top_p"] = top_p
 
-    with torch.inference_mode():
-        out = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            input_features=input_features,
-            **gen_kwargs,
-        )
-    new_tokens = out[0, input_ids.shape[-1]:]
-    response = tokenizer.decode(new_tokens, skip_special_tokens=False)
-    response = response.split(IM_END_TOKEN, 1)[0].strip()
-    thinking, prediction = split_thinking(response)
+    torch.manual_seed(0)
+    torch.cuda.manual_seed_all(0)
+    generation_kwargs = {
+        **gen_kwargs,
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "input_features": input_features,
+    }
+    yield from _stream_model_generate(
+        runtime,
+        generation_kwargs,
+        reasoning,
+        reasoning_budget,
+    )
     print(
         f"[gpu] audio_seconds={audio_duration:.1f} greedy={greedy} "
         f"elapsed_seconds={time.perf_counter() - started_at:.1f}",
         flush=True,
     )
-    return thinking, prediction
 
 
 @spaces.GPU(duration=_estimate, size="xlarge")
@@ -336,7 +945,7 @@ def run(
 
     Predefined ASR and translation tasks use the model's recommended greedy
     decoding. Other tasks and custom instructions use the sampling controls.
-    Audio may be up to two minutes long, with up to 512 output tokens.
+    Input duration and output length are limited by the current runtime.
 
     Returns:
         The final answer and, when enabled, the reasoning trace.
@@ -345,9 +954,14 @@ def run(
         return "Please provide an audio input.", ""
 
     custom_instruction = custom_prompt.strip() if custom_prompt else ""
-    prompt = custom_instruction or TASK_PROMPTS.get(task, TASK_PROMPTS["Describe the audio"])
-    greedy = not custom_instruction and task in GREEDY_TASKS
-    thinking, answer = _generate(
+    prompt = custom_instruction or LEGACY_TASK_PROMPTS.get(
+        task,
+        LEGACY_TASK_PROMPTS["Describe the audio"],
+    )
+    greedy = not custom_instruction and task in LEGACY_GREEDY_TASKS
+    answer, thinking = "", ""
+    for answer, thinking in _generate(
+        model_runtimes[DEFAULT_MODEL_NAME],
         audio,
         prompt,
         reasoning,
@@ -355,77 +969,762 @@ def run(
         temperature=float(temperature),
         top_p=float(top_p),
         greedy=greedy,
+    ):
+        pass
+    return answer, thinking
+
+
+def _estimate_text(
+    prompt: str,
+    reasoning: bool,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    *args: object,
+    **kwargs: object,
+) -> int:
+    if not prompt.strip() or int(max_new_tokens) > MAX_NEW_TOKENS:
+        return 10
+    return MAX_GPU_DURATION_SECONDS
+
+
+def _generate_text(
+    runtime: AudexRuntime,
+    prompt: str,
+    reasoning: bool,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    reasoning_budget: int | None = None,
+    seed: int = TEXT_SEED,
+) -> Iterator[tuple[str, str]]:
+    """Run text-only inference with the official Audex chat template."""
+    if not prompt.strip():
+        yield "Please provide a text prompt.", ""
+        return
+    if max_new_tokens > MAX_NEW_TOKENS:
+        raise gr.Error(f"This demo supports up to {MAX_NEW_TOKENS} output tokens.")
+    if (
+        reasoning
+        and reasoning_budget is not None
+        and reasoning_budget > 0
+        and (
+            reasoning_budget >= max_new_tokens
+            or reasoning_budget > _max_reasoning_budget(max_new_tokens)
+        )
+    ):
+        raise gr.Error(
+            "Reasoning budget is too high for max new tokens after reserving "
+            "the newline grace window and final answer."
+        )
+
+    formatted = runtime.tokenizer.apply_chat_template(
+        [
+            {"role": "system", "content": TEXT_SYSTEM_PROMPT},
+            {"role": "user", "content": prompt.strip()},
+        ],
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=reasoning,
+    )
+    tok = runtime.tokenizer(formatted, return_tensors="pt", add_special_tokens=False)
+    input_ids = tok.input_ids.to("cuda")
+    attention_mask = (
+        tok.attention_mask if "attention_mask" in tok else build_attention_mask(input_ids)
+    ).to("cuda")
+
+    temperature = max(float(temperature), 1e-4)
+    top_p = float(top_p)
+    do_sample = (temperature != 1.0) or (0.0 < top_p < 1.0)
+    gen_kwargs = {
+        "do_sample": do_sample,
+        "eos_token_id": runtime.tokenizer.convert_tokens_to_ids(IM_END_TOKEN),
+        "pad_token_id": runtime.tokenizer.pad_token_id
+        or getattr(runtime.config, "pad_token_id", 0),
+        "max_new_tokens": int(max_new_tokens),
+        "suppress_tokens": _blocked_output_token_ids(runtime),
+    }
+    if do_sample:
+        gen_kwargs["temperature"] = temperature
+        if top_p > 0.0:
+            gen_kwargs["top_p"] = top_p
+
+    started_at = time.perf_counter()
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    generation_kwargs = {
+        **gen_kwargs,
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+    }
+    yield from _stream_model_generate(
+        runtime,
+        generation_kwargs,
+        reasoning,
+        reasoning_budget,
     )
+    print(
+        f"[gpu] text_input_tokens={input_ids.shape[-1]} "
+        f"elapsed_seconds={time.perf_counter() - started_at:.1f}",
+        flush=True,
+    )
+
+
+def _generate_tts(
+    runtime: AudexRuntime,
+    text: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    *,
+    top_k: int = 0,
+    guidance_scale: float = 2.0,
+    token_limit: int | None = None,
+    segment_sentences: bool = False,
+) -> Iterator[tuple[str, str | None, dict[str, object]]]:
+    text = text.strip()
+    if not text:
+        yield "Please provide text to synthesize.", None, player_value(0, reset=True)
+        return
+    max_tokens = TTS_MAX_NEW_TOKENS if token_limit is None else token_limit
+    if max_new_tokens > max_tokens:
+        raise gr.Error(f"TTS supports up to {max_tokens} speech tokens.")
+
+    segments = _split_sentence_segments(text) if segment_sentences else [text]
+    if not segments:
+        raise gr.Error("Text response had no final answer for TTS.")
+
+    sequence = 0
+    chunks: list[np.ndarray] = []
+    first_token_seconds: float | None = None
+    total_tokens = 0
+    request_tag = hashlib.sha256(text.encode()).hexdigest()[:8]
+    started_at = time.perf_counter()
+    torch.cuda.reset_peak_memory_stats()
+    print(f"[gpu] tts_request={request_tag} started", flush=True)
+    initial_player = (
+        player_value(sequence, reset=True) if TTS_STREAMING_PLAYER_ENABLED else gr.skip()
+    )
+    yield "Generating speech tokens…", None, initial_player
+
+    try:
+        for segment_index, segment in enumerate(segments, start=1):
+            segment_tokens = 0
+            for event in stream_tts(
+                model=runtime.model,
+                tokenizer=runtime.tokenizer,
+                decoder=runtime.speech_decoder,
+                text=segment,
+                max_new_tokens=int(max_new_tokens),
+                temperature=float(temperature),
+                top_p=float(top_p),
+                top_k=top_k,
+                guidance_scale=guidance_scale,
+            ):
+                segment_tokens = event.token_count
+                if event.token_count and first_token_seconds is None:
+                    first_token_seconds = time.perf_counter() - started_at
+                if event.pcm is not None:
+                    chunks.append(event.pcm)
+                    sequence += 1
+                    player_update = (
+                        player_value(
+                            sequence,
+                            pcm=encode_pcm_chunk(event.pcm),
+                            token_count=total_tokens + event.token_count,
+                        )
+                        if TTS_STREAMING_PLAYER_ENABLED
+                        else gr.skip()
+                    )
+                    yield (
+                        f"Generating speech segment {segment_index}/{len(segments)} "
+                        f"· {total_tokens + event.token_count} tokens",
+                        None,
+                        player_update,
+                    )
+
+            total_tokens += segment_tokens
+            if segment_index < len(segments):
+                chunks.append(
+                    np.zeros(
+                        round(SAMPLE_RATE * S2S_TTS_SEGMENT_SILENCE_SECONDS),
+                        dtype=np.float32,
+                    )
+                )
+
+        if not chunks:
+            raise RuntimeError("TTS completed without producing audio")
+        waveform = np.concatenate(chunks)
+        wav_path = write_wav(waveform)
+        sequence += 1
+        elapsed_seconds = time.perf_counter() - started_at
+        peak_gib = torch.cuda.max_memory_allocated() / (1024**3)
+        token_rate = total_tokens / max(elapsed_seconds, 1e-6)
+        print(
+            f"[gpu] tts_request={request_tag} tts_tokens={total_tokens} "
+            f"ttfc_seconds={first_token_seconds:.2f} "
+            f"tokens_per_second={token_rate:.2f} audio_seconds={waveform.size / SAMPLE_RATE:.2f} "
+            f"elapsed_seconds={elapsed_seconds:.2f} peak_memory_gib={peak_gib:.2f}",
+            flush=True,
+        )
+        yield (
+            f"Complete · {waveform.size / SAMPLE_RATE:.1f}s audio",
+            wav_path,
+            (
+                player_value(sequence, token_count=total_tokens, done=True)
+                if TTS_STREAMING_PLAYER_ENABLED
+                else gr.skip()
+            ),
+        )
+    except GeneratorExit:
+        print(f"[gpu] tts_request={request_tag} cancelled", flush=True)
+        raise
+
+
+def _split_sentence_segments(text: str) -> list[str]:
+    return [
+        part.strip()
+        for line in text.splitlines()
+        for part in re.findall(r"[^.!?]+[.!?]+[\"')\]]*|[^.!?]+$", line.strip())
+        if part.strip()
+    ]
+
+
+def _clean_transcription(text: str) -> str:
+    first_quote = text.find("'")
+    last_quote = text.rfind("'")
+    if first_quote != -1 and last_quote > first_quote:
+        return text[first_quote + 1 : last_quote].strip()
+    return text.strip()
+
+
+def _compose_s2s_input(response_prompt: str, transcript: str) -> str:
+    return (
+        f"{response_prompt.strip()}\n\n"
+        f"<spoken_transcript>\n{transcript.strip()}\n</spoken_transcript>"
+    )
+
+
+def _limit_spoken_answer(text: str) -> str:
+    segments = _split_sentence_segments(text)
+    sentence = segments[-1] if segments else text.strip()
+    words = sentence.split()
+    if len(words) <= S2S_SPOKEN_MAX_WORDS:
+        return sentence
+    return " ".join(words[:S2S_SPOKEN_MAX_WORDS]).rstrip(",;:") + "."
+
+
+def _generate_s2s(
+    runtime: AudexRuntime,
+    audio: str,
+    response_prompt: str,
+    reasoning: bool,
+    reasoning_budget: int,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    guidance_scale: float,
+) -> Iterator[tuple[object, object, object, object]]:
+    yield "Transcribing input speech…", "", gr.skip(), gr.skip()
+    transcript = ""
+    for transcription, _ in _generate(
+        runtime,
+        audio,
+        "Transcribe the input speech.",
+        False,
+        min(256, MAX_NEW_TOKENS),
+        temperature=1.0,
+        top_p=1.0,
+        greedy=True,
+    ):
+        transcript = _clean_transcription(transcription)
+        yield f"Transcript: {transcript}", "", gr.skip(), gr.skip()
+    if not transcript:
+        raise gr.Error("Speech-to-speech transcription produced no text.")
+
+    text_input = _compose_s2s_input(response_prompt, transcript)
+    candidate_answer = ""
+    thinking = ""
+    for current_answer, thinking in _generate_text(
+        runtime,
+        text_input,
+        reasoning,
+        int(max_new_tokens),
+        temperature,
+        top_p,
+        reasoning_budget=reasoning_budget,
+    ):
+        candidate_answer = current_answer or candidate_answer
+        yield "", thinking, gr.skip(), gr.skip()
+    if MAX_TOKEN_WARNING in candidate_answer:
+        yield candidate_answer, thinking, gr.skip(), gr.skip()
+        return
+    if not candidate_answer and not thinking:
+        raise gr.Error("Speech-to-speech response generation produced no text.")
+
+    answer = _limit_spoken_answer(candidate_answer)
+    if not answer:
+        raise gr.Error("Speech-to-speech response generation produced no final answer.")
+    yield answer, thinking, gr.skip(), gr.skip()
+
+    for _, wav_path, player in _generate_tts(
+        runtime,
+        answer,
+        S2S_TTS_MAX_NEW_TOKENS,
+        0.1,
+        1.0,
+        top_k=80,
+        guidance_scale=guidance_scale,
+        token_limit=S2S_TTS_MAX_NEW_TOKENS,
+        segment_sentences=True,
+    ):
+        yield answer, thinking, wav_path or gr.skip(), player
+
+
+@spaces.GPU(duration=_estimate_text, size="xlarge")
+def run_text(
+    prompt: str,
+    reasoning: bool,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+) -> tuple[str, str]:
+    """Run the legacy text-only API endpoint."""
+    answer, thinking = "", ""
+    for answer, thinking in _generate_text(
+        model_runtimes[DEFAULT_MODEL_NAME],
+        prompt,
+        reasoning,
+        max_new_tokens,
+        temperature,
+        top_p,
+    ):
+        pass
     return answer, thinking
 
 
-THEME = gr.themes.Citrus()
+def _estimate_unified(
+    model_name: str,
+    task: str,
+    audio: str | None,
+    prompt: str,
+    reasoning: bool,
+    reasoning_budget: int,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    guidance_scale: float,
+    *args: object,
+    **kwargs: object,
+) -> int:
+    if task == S2S_TASK:
+        return S2S_GPU_DURATION_SECONDS
+    if task in {TEXT_TASK, TTS_TASK}:
+        if task == TTS_TASK and int(max_new_tokens) > TTS_MAX_NEW_TOKENS:
+            return 10
+        return _estimate_text(
+            prompt,
+            reasoning,
+            max_new_tokens,
+            temperature,
+            top_p,
+        )
+    return _estimate(
+        audio,
+        task,
+        prompt,
+        reasoning,
+        max_new_tokens,
+        temperature,
+        top_p,
+    )
+
+
+@spaces.GPU(duration=_estimate_unified, size="xlarge")
+def run_unified(
+    model_name: str,
+    task: str,
+    audio: str | None,
+    prompt: str,
+    reasoning: bool,
+    reasoning_budget: int,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    guidance_scale: float,
+) -> Iterator[tuple[object, object, object, object]]:
+    """Run a selected Audex model for audio, text, or speech generation."""
+    yield "", "", gr.skip(), gr.skip()
+    runtime = _get_runtime(model_name)
+    settings = TASK_SPECS[task]
+    if settings["modality"] == "s2s":
+        if audio is None:
+            yield "Please provide an audio input.", "", gr.skip(), gr.skip()
+            return
+        yield from _generate_s2s(
+            runtime,
+            audio,
+            prompt,
+            reasoning,
+            int(reasoning_budget),
+            int(max_new_tokens),
+            temperature,
+            top_p,
+            float(guidance_scale),
+        )
+        return
+    if settings["modality"] == "tts":
+        for status, wav_path, player in _generate_tts(
+            runtime,
+            prompt,
+            int(max_new_tokens),
+            temperature,
+            top_p,
+            guidance_scale=float(guidance_scale),
+        ):
+            yield status, "", wav_path or gr.skip(), player
+        return
+    if settings["modality"] == "text":
+        for answer, thinking in _generate_text(
+            runtime,
+            prompt,
+            reasoning,
+            max_new_tokens,
+            temperature,
+            top_p,
+            reasoning_budget=int(reasoning_budget),
+        ):
+            yield answer, thinking, gr.skip(), gr.skip()
+        return
+    if audio is None:
+        yield "Please provide an audio input.", "", gr.skip(), gr.skip()
+        return
+
+    instruction = (prompt.strip() or str(settings["template"])).replace("<sound>", "").strip()
+    for answer, thinking in _generate(
+        runtime,
+        audio,
+        instruction,
+        reasoning,
+        int(max_new_tokens),
+        temperature=float(temperature),
+        top_p=float(top_p),
+        greedy=bool(settings["greedy"]),
+        reasoning_budget=int(reasoning_budget),
+    ):
+        yield answer, thinking, gr.skip(), gr.skip()
+
+
+def unified_task_defaults(
+    task: str,
+) -> tuple[dict[str, object], ...]:
+    settings = TASK_SPECS[task]
+    is_audio = settings["modality"] in {"audio", "s2s"}
+    is_tts = settings["modality"] == "tts"
+    has_speech_output = settings["modality"] in {"tts", "s2s"}
+    reasoning = bool(settings["reasoning"])
+    default_max_new_tokens = int(settings.get("max_new_tokens", DEFAULT_MAX_NEW_TOKENS))
+    max_new_tokens_limit = int(settings.get("max_new_tokens_limit", MAX_NEW_TOKENS))
+    reasoning_budget_limit = _max_reasoning_budget(default_max_new_tokens)
+    return (
+        gr.update(visible=is_audio),
+        gr.update(
+            value=str(settings["template"]),
+            lines=8 if settings["modality"] == "s2s" else (3 if is_audio else 8),
+            label=(
+                "Text to synthesize"
+                if is_tts
+                else (
+                    "Response instruction / prompt"
+                    if settings["modality"] == "s2s"
+                    else "Task template / prompt"
+                )
+            ),
+        ),
+        gr.update(value=reasoning, visible=not is_tts),
+        gr.update(
+            value=min(
+                int(settings.get("reasoning_budget", DEFAULT_REASONING_BUDGET)),
+                reasoning_budget_limit,
+            ),
+            maximum=reasoning_budget_limit,
+            visible=not is_tts,
+            interactive=reasoning,
+        ),
+        float(settings["temperature"]),
+        float(settings["top_p"]),
+        gr.update(
+            value=TTS_MAX_NEW_TOKENS if is_tts else default_max_new_tokens,
+            maximum=TTS_MAX_NEW_TOKENS if is_tts else max_new_tokens_limit,
+        ),
+        gr.update(
+            value=float(settings.get("guidance_scale", 2.0)),
+            visible=has_speech_output,
+        ),
+        gr.update(label="Status" if is_tts else "Answer"),
+        gr.update(visible=not is_tts),
+        gr.update(value=None, visible=has_speech_output),
+        gr.update(
+            value=player_value(0, reset=True),
+            visible=is_tts and TTS_STREAMING_PLAYER_ENABLED,
+        ),
+    )
+
+
+def update_reasoning_budget_limit(
+    max_new_tokens: int,
+    reasoning_budget: int,
+) -> dict[str, object]:
+    maximum = _max_reasoning_budget(max_new_tokens)
+    return gr.update(maximum=maximum, value=min(int(reasoning_budget), maximum))
 
-with gr.Blocks(title="Nemotron-Labs-Audex-30B-A3B") as demo:
+
+theme = gr.themes.Citrus()
+
+with gr.Blocks(title="Nemotron-Labs-Audex") as demo:
     gr.Markdown(
-        "# 🎧 Nemotron-Labs-Audex-30B-A3B\n"
-        "Unified audio-text MoE (30B total / 3B active) for **audio understanding, "
-        "speech recognition, and speech translation**. Runs the official Hugging Face "
-        "inference path with compiled `mamba-ssm` + `causal-conv1d` CUDA kernels for "
-        "correct numerics.\n\n"
-        "*Non-commercial use only (NVIDIA One-Way Noncommercial License).*"
+        "# 🎧 Nemotron-Labs-Audex\n"
+        '<div style="display: flex; align-items: center; gap: 0.4rem; flex-wrap: wrap;">'
+        '<a href="https://arxiv.org/abs/2607.05196">'
+        '<img alt="Technical Report" src="https://img.shields.io/badge/2607.05196-Technical_Report-blue"></a>'
+        '<a href="https://huggingface.co/collections/nvidia/nemotron-labs-audex">'
+        '<img alt="Models" src="https://img.shields.io/badge/🤗-Models-blue"></a>'
+        "</div>\n\n"
+        "This is an interactive demo of [Nemotron-Labs-Audex-30B-A3B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B) and [Nemotron-Labs-Audex-2B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B). \n\n"
+        "Audex extends the vocabulary for discrete audio tokens, as well as an audio encoder for audio inputs. Audex delivers strong abilities on audio tasks (audio understanding, speech recognition and translation, text-to-speech, audio generation, and speech-to-speech generation) while preserving very compelling reasoning, alignment, knowledge, long-context, and agentic capabilities of its text-only LLM backbone with marginal or no regression. Audex operates in both **thinking** and **instruct** (non-thinking) modes. \n\n"
+        "Task selection loads a default prompt template and inference recipe.\n\n"
+        "Your use of this model is governed by the [NVIDIA OneWay Noncommercial License](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B/blob/main/license/NVIDIA-OneWay-Noncommercial-License.docx/)."
     )
     with gr.Row():
         with gr.Column():
-            audio_in = gr.Audio(
+            unified_model = gr.Radio(
+                choices=list(model_runtimes),
+                value=DEFAULT_MODEL_NAME,
+                label="Model",
+            )
+            unified_task = gr.Dropdown(
+                choices=list(TASK_SPECS),
+                value="Speech recognition (ASR)",
+                label="Task",
+            )
+            unified_audio = gr.Audio(
                 type="filepath",
-                label="Input audio (up to 2 minutes)",
+                label=f"Input audio (up to {MAX_AUDIO_DURATION_SECONDS / 60:g} minutes)",
                 sources=["upload", "microphone"],
             )
-            task = gr.Radio(
-                choices=list(TASK_PROMPTS.keys()),
-                value="Transcribe (ASR)",
-                label="Task",
+            unified_prompt = gr.Textbox(
+                value=str(TASK_SPECS["Speech recognition (ASR)"]["template"]),
+                label="Task template / prompt",
+                lines=3,
             )
-            custom_prompt = gr.Textbox(
-                label="Custom instruction (optional — overrides Task)",
-                placeholder="e.g. Summarize what the speaker is saying.",
-                lines=2,
+            gr.Markdown(
+                "`<sound>` marks where audio embeddings enter the prompt; the app "
+                "replaces it automatically. Text-only tasks have no input-modality token."
             )
-            run_btn = gr.Button("Run", variant="primary")
+            with gr.Row():
+                unified_run_btn = gr.Button("Run", variant="primary", elem_id="audex-run")
+                unified_stop_btn = gr.Button("Stop", variant="stop", elem_id="audex-stop")
             with gr.Accordion("Advanced options", open=False):
-                reasoning = gr.Checkbox(value=False, label="Enable reasoning (<think>) mode")
-                max_new_tokens = gr.Slider(
+                unified_reasoning = gr.Checkbox(
+                    value=False,
+                    label="Enable reasoning (<think>) mode",
+                )
+                unified_reasoning_budget = gr.Slider(
+                    0,
+                    _max_reasoning_budget(MAX_NEW_TOKENS),
+                    value=DEFAULT_REASONING_BUDGET,
+                    step=16,
+                    label="Reasoning budget (0 = unlimited)",
+                    info=(
+                        "No reasoning cap by default. Set a positive threshold to enable "
+                        "a decoding-time cap with a 10% newline grace."
+                    ),
+                    interactive=False,
+                )
+                unified_max_new_tokens = gr.Slider(
                     16,
                     MAX_NEW_TOKENS,
-                    value=256,
+                    value=DEFAULT_MAX_NEW_TOKENS,
                     step=16,
                     label="Max new tokens",
+                    info=(
+                        "Total cap shared by reasoning and the final answer. "
+                        "If reached, the displayed output is incomplete."
+                    ),
+                )
+                unified_temperature = gr.Slider(
+                    0.1,
+                    1.5,
+                    value=1.0,
+                    step=0.05,
+                    label="Temperature",
+                )
+                unified_top_p = gr.Slider(
+                    0.1,
+                    1.0,
+                    value=1.0,
+                    step=0.05,
+                    label="Top-p",
+                )
+                unified_guidance_scale = gr.Slider(
+                    1.0,
+                    3.0,
+                    value=2.0,
+                    step=0.1,
+                    label="TTS CFG scale",
+                    visible=False,
                 )
-                temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature")
-                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
                 gr.Markdown(
-                    "ASR and speech translation use greedy decoding. "
-                    "Sampling controls apply to other tasks and custom instructions."
+                    "Task selection restores the official defaults. ASR and AST use "
+                    "greedy decoding; audio understanding and text reasoning use sampling. "
+                    "Reasoning has no separate cap unless you set one. TTS uses CFG 2.0; "
+                    "speech-to-speech uses CFG 1.5."
                 )
         with gr.Column():
-            answer_out = gr.Textbox(label="Answer", lines=10)
-            thinking_out = gr.Textbox(label="Reasoning trace (if enabled)", lines=6)
+            unified_answer_out = gr.Textbox(label="Answer", lines=10)
+            unified_thinking_out = gr.Textbox(
+                label="Reasoning trace (if enabled)",
+                lines=10,
+            )
+            unified_tts_audio_out = gr.Audio(
+                label="Generated speech",
+                type="filepath",
+                autoplay=False,
+                visible=False,
+            )
+            unified_tts_player = gr.HTML(
+                value=player_value(0, reset=True),
+                html_template=TTS_PLAYER_TEMPLATE,
+                css_template=TTS_PLAYER_CSS,
+                js_on_load=TTS_PLAYER_JS,
+                apply_default_css=False,
+                visible=False,
+            )
+
+    unified_task.change(
+        unified_task_defaults,
+        inputs=unified_task,
+        outputs=[
+            unified_audio,
+            unified_prompt,
+            unified_reasoning,
+            unified_reasoning_budget,
+            unified_temperature,
+            unified_top_p,
+            unified_max_new_tokens,
+            unified_guidance_scale,
+            unified_answer_out,
+            unified_thinking_out,
+            unified_tts_audio_out,
+            unified_tts_player,
+        ],
+        api_name=False,
+    )
+    unified_reasoning.change(
+        lambda enabled: gr.update(interactive=enabled),
+        inputs=unified_reasoning,
+        outputs=unified_reasoning_budget,
+        api_name=False,
+    )
+    unified_max_new_tokens.change(
+        update_reasoning_budget_limit,
+        inputs=[unified_max_new_tokens, unified_reasoning_budget],
+        outputs=unified_reasoning_budget,
+        api_name=False,
+    )
 
-    run_btn.click(
-        run,
-        inputs=[audio_in, task, custom_prompt, reasoning, max_new_tokens, temperature, top_p],
-        outputs=[answer_out, thinking_out],
-        api_name="run",
+    unified_run_event = unified_run_btn.click(
+        run_unified,
+        inputs=[
+            unified_model,
+            unified_task,
+            unified_audio,
+            unified_prompt,
+            unified_reasoning,
+            unified_reasoning_budget,
+            unified_max_new_tokens,
+            unified_temperature,
+            unified_top_p,
+            unified_guidance_scale,
+        ],
+        outputs=[
+            unified_answer_out,
+            unified_thinking_out,
+            unified_tts_audio_out,
+            unified_tts_player,
+        ],
+        api_name="run_unified",
+        concurrency_id="audex-gpu",
+        concurrency_limit=1,
     )
+    unified_stop_btn.click(fn=None, cancels=[unified_run_event], api_name=False)
 
+    gr.Markdown(
+        "### Curated examples\n"
+        "Audio and text examples share the same task-driven interface."
+    )
     gr.Examples(
         examples=[
-            ["examples/mlk_speech.wav", "Transcribe (ASR)", "", False, 256, 0.7, 0.9],
-            ["examples/sample_speech.wav", "Transcribe (ASR)", "", False, 256, 0.7, 0.9],
-            ["examples/mlk_speech.wav", "Describe the audio", "", False, 256, 0.7, 0.9],
+            [
+                "Speech recognition (ASR)",
+                "examples/mlk_speech.wav",
+                TASK_SPECS["Speech recognition (ASR)"]["template"],
+            ],
+            [
+                "Speech recognition (ASR)",
+                "examples/sample_speech.wav",
+                TASK_SPECS["Speech recognition (ASR)"]["template"],
+            ],
+            [
+                "Speech translation (AST)",
+                "examples/korean_speech.wav",
+                TASK_SPECS["Speech translation (AST)"]["template"],
+            ],
+            [
+                "Audio description",
+                "examples/mlk_speech.wav",
+                TASK_SPECS["Audio description"]["template"],
+            ],
+            [
+                "Audio description",
+                "examples/sample_speech.wav",
+                TASK_SPECS["Audio description"]["template"],
+            ],
+            [
+                "Audio question answering",
+                "examples/mlk_speech.wav",
+                TASK_SPECS["Audio question answering"]["template"],
+            ],
+            [
+                TEXT_TASK,
+                None,
+                MMLU_PRO_EXAMPLE_PROMPT,
+            ],
+            [
+                S2S_TASK,
+                "examples/question_1059.mp3",
+                S2S_UI_PROMPT,
+            ],
+            [
+                TTS_TASK,
+                None,
+                TTS_EXAMPLE_TEXT,
+            ],
+        ],
+        inputs=[
+            unified_task,
+            unified_audio,
+            unified_prompt,
         ],
-        inputs=[audio_in, task, custom_prompt, reasoning, max_new_tokens, temperature, top_p],
-        outputs=[answer_out, thinking_out],
-        fn=run,
         cache_examples=False,
-        run_on_click=True,
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=8).launch(theme=THEME, ssr_mode=False, mcp_server=True)
+    demo.queue(max_size=8).launch(
+        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
+        server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
+        share=os.environ.get("GRADIO_SHARE", "false").lower() == "true",
+        theme=theme,
+        ssr_mode=False,
+        mcp_server=True,
+    )

examples/korean_speech.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65857d13573b3f2f793092865cc4da3ddd46d921c6f9ea6eb77e263aba00779c
+size 255532

examples/question_1059.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4615906395ee1dc6fdcb33dac1be954b211d32b7096132c8d110530d0d385d5
+size 228000

reasoning_utils.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from transformers.generation.logits_process import LogitsProcessor
+
+
+class ReasoningBudgetLogitsProcessor(LogitsProcessor):
+    """Apply an optional reasoning-token budget for one sequence."""
+
+    def __init__(
+        self,
+        tokenizer: Any,
+        prompt_length: int,
+        reasoning_budget: int,
+        grace_fraction: float = 0.1,
+        early_exit: str = ".\n</think>\n\n",
+    ) -> None:
+        if reasoning_budget < 1:
+            raise ValueError("reasoning_budget must be positive")
+
+        self.tokenizer = tokenizer
+        self.prompt_length = prompt_length
+        self.reasoning_budget = reasoning_budget
+        self.hard_limit = reasoning_budget + max(
+            1, math.ceil(reasoning_budget * grace_fraction)
+        )
+        self.reasoning_end_ids = tokenizer.encode(
+            "</think>", add_special_tokens=False
+        )
+        self.early_exit_ids = tokenizer.encode(
+            early_exit, add_special_tokens=False
+        )
+        self.forced_index: int | None = None
+        self.done = False
+
+    def _force(self, scores: torch.FloatTensor, token_id: int) -> torch.FloatTensor:
+        scores.fill_(-float("inf"))
+        scores[:, token_id] = 0
+        return scores
+
+    def __call__(
+        self,
+        input_ids: torch.LongTensor,
+        scores: torch.FloatTensor,
+    ) -> torch.FloatTensor:
+        if self.done:
+            return scores
+        if input_ids.shape[0] != 1:
+            raise ValueError("Reasoning budget control requires batch size 1")
+
+        generated_ids = input_ids[0, self.prompt_length :].tolist()
+        if self.forced_index is not None:
+            self.forced_index += 1
+            if self.forced_index >= len(self.early_exit_ids):
+                self.done = True
+                return scores
+            return self._force(scores, self.early_exit_ids[self.forced_index])
+
+        if generated_ids[-len(self.reasoning_end_ids) :] == self.reasoning_end_ids:
+            self.done = True
+            return scores
+
+        generated_tokens = len(generated_ids)
+        if generated_tokens < self.reasoning_budget:
+            return scores
+
+        ended_line = bool(
+            generated_ids
+            and "\n"
+            in self.tokenizer.decode(
+                [generated_ids[-1]],
+                skip_special_tokens=False,
+                clean_up_tokenization_spaces=False,
+            )
+        )
+        if not ended_line and generated_tokens < self.hard_limit:
+            return scores
+
+        self.forced_index = 0
+        return self._force(scores, self.early_exit_ids[0])

requirements.txt

@@ -1,4 +1,4 @@
-transformers>=4.53.3
+transformers==5.14.0
 accelerate
 sentencepiece
 librosa

run_local.sh

@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+APP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+VENV_DIR="${AUDEX_VENV_DIR:-$APP_DIR/.venv}"
+PYTHON="$VENV_DIR/bin/python"
+
+if [[ ! -x "$PYTHON" ]]; then
+    echo "Missing local environment. Run: bash $APP_DIR/setup_local.sh" >&2
+    exit 1
+fi
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
+
+gpu_info="$("$PYTHON" - <<'PY'
+import re
+import sys
+
+import torch
+
+if not torch.cuda.is_available():
+    raise SystemExit("CUDA is unavailable in this PyTorch environment.")
+
+major, minor = torch.cuda.get_device_capability()
+cuda = (torch.version.cuda or "unknown").replace(".", "")
+torch_version = re.sub(r"[^0-9A-Za-z]+", "", torch.__version__.split("+", 1)[0])
+python_version = f"{sys.version_info.major}{sys.version_info.minor}"
+memory_gib = torch.cuda.get_device_properties(0).total_memory / 1024**3
+print(f"{major}.{minor}\tsm{major}{minor}-cu{cuda}-torch{torch_version}-py{python_version}\t{memory_gib:.0f}")
+PY
+)"
+IFS=$'\t' read -r detected_arch runtime_tag memory_gib <<< "$gpu_info"
+
+export AUDEX_TORCH_ARCH="${AUDEX_TORCH_ARCH:-$detected_arch}"
+export AUDEX_WHEELS_DIR="${AUDEX_WHEELS_DIR:-$APP_DIR/.local/wheels/$runtime_tag}"
+export AUDEX_BUILD_DIR="${AUDEX_BUILD_DIR:-$APP_DIR/.local/build/$runtime_tag}"
+export AUDEX_MAX_AUDIO_SECONDS="${AUDEX_MAX_AUDIO_SECONDS:-900}"
+export AUDEX_MAX_NEW_TOKENS="${AUDEX_MAX_NEW_TOKENS:-2048}"
+export AUDEX_DEFAULT_MAX_NEW_TOKENS="${AUDEX_DEFAULT_MAX_NEW_TOKENS:-1024}"
+export AUDEX_TTS_MAX_NEW_TOKENS="${AUDEX_TTS_MAX_NEW_TOKENS:-1024}"
+export AUDEX_S2S_TEXT_MAX_NEW_TOKENS="${AUDEX_S2S_TEXT_MAX_NEW_TOKENS:-2048}"
+export AUDEX_S2S_TTS_MAX_NEW_TOKENS="${AUDEX_S2S_TTS_MAX_NEW_TOKENS:-2400}"
+export GRADIO_SERVER_NAME="${GRADIO_SERVER_NAME:-0.0.0.0}"
+export GRADIO_SERVER_PORT="${GRADIO_SERVER_PORT:-7860}"
+export GRADIO_SHARE="${GRADIO_SHARE:-false}"
+
+echo "GPU architecture: $AUDEX_TORCH_ARCH (${memory_gib} GiB)"
+echo "Models: Audex-30B-A3B and Audex-2B (downloaded from the Hub if uncached)"
+
+exec "$PYTHON" "$APP_DIR/app.py"

setup_local.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+APP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+VENV_DIR="${AUDEX_VENV_DIR:-$APP_DIR/.venv}"
+PYTHON="$VENV_DIR/bin/python"
+
+if [[ ! -x "$PYTHON" ]]; then
+    python3 -m venv "$VENV_DIR"
+fi
+
+"$PYTHON" -m pip install --upgrade pip
+if ! "$PYTHON" -c "import torch" >/dev/null 2>&1; then
+    "$PYTHON" -m pip install "${AUDEX_TORCH_PACKAGE:-torch}"
+fi
+"$PYTHON" -m pip install \
+    --requirement "$APP_DIR/requirements.txt" \
+    "gradio==6.20.0" \
+    "huggingface_hub" \
+    "spaces>=0.51.0"
+
+echo "Local environment ready: $VENV_DIR"
+echo "Launch with: bash $APP_DIR/run_local.sh"

tts_player.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+
+_WORKLET_SOURCE = (Path(__file__).parent / "web" / "PCMPlayerWorklet.js").read_text()
+
+TTS_PLAYER_TEMPLATE = """
+<div class="audex-player">
+  <div class="audex-player-status">Ready</div>
+  <button class="audex-player-enable" type="button">Enable playback</button>
+  <audio class="audex-player-audio" controls></audio>
+  <a class="audex-player-download" hidden download="audex-tts.wav">Download WAV</a>
+</div>
+"""
+
+TTS_PLAYER_CSS = """
+.audex-player {
+  display: grid;
+  gap: 0.75rem;
+  padding: 1rem;
+  border: 1px solid var(--border-color-primary);
+  border-radius: var(--radius-lg);
+}
+.audex-player-status {
+  color: var(--body-text-color-subdued);
+}
+.audex-player-audio {
+  width: 100%;
+}
+.audex-player-download {
+  color: var(--link-text-color);
+  font-weight: 600;
+}
+"""
+
+TTS_PLAYER_JS = """
+const WORKLET_SOURCE = __WORKLET_SOURCE__;
+const status = element.querySelector(".audex-player-status");
+const enable = element.querySelector(".audex-player-enable");
+const audio = element.querySelector(".audex-player-audio");
+const download = element.querySelector(".audex-player-download");
+let context = null;
+let player = null;
+let chunks = [];
+let wavUrl = null;
+let lastSequence = -1;
+
+const ensurePlayer = async () => {
+  if (context) {
+    await context.resume();
+    return;
+  }
+  context = new AudioContext();
+  const blob = new Blob([WORKLET_SOURCE], { type: "text/javascript" });
+  const moduleUrl = URL.createObjectURL(blob);
+  await context.audioWorklet.addModule(moduleUrl);
+  URL.revokeObjectURL(moduleUrl);
+  player = new AudioWorkletNode(
+    context,
+    "audex-pcm-player",
+    { outputChannelCount: [1] },
+  );
+  player.connect(context.destination);
+  player.port.onmessage = (event) => {
+    if (event.data.type === "underrun") {
+      status.textContent = "Buffering generated speech…";
+    }
+  };
+  await context.resume();
+};
+
+const decodePcm = (encoded) => {
+  const binary = atob(encoded);
+  const bytes = new Uint8Array(binary.length);
+  for (let index = 0; index < binary.length; index += 1) {
+    bytes[index] = binary.charCodeAt(index);
+  }
+  return new Float32Array(bytes.buffer);
+};
+
+const resample = (input, sourceRate, targetRate) => {
+  if (sourceRate === targetRate) return input;
+  const output = new Float32Array(Math.max(1, Math.round(input.length * targetRate / sourceRate)));
+  const ratio = sourceRate / targetRate;
+  for (let index = 0; index < output.length; index += 1) {
+    const position = index * ratio;
+    const left = Math.floor(position);
+    const right = Math.min(left + 1, input.length - 1);
+    const fraction = position - left;
+    output[index] = input[left] * (1 - fraction) + input[right] * fraction;
+  }
+  return output;
+};
+
+const buildWav = (parts, sampleRate) => {
+  const length = parts.reduce((total, part) => total + part.length, 0);
+  const buffer = new ArrayBuffer(44 + length * 2);
+  const view = new DataView(buffer);
+  const write = (offset, text) => {
+    for (let index = 0; index < text.length; index += 1) {
+      view.setUint8(offset + index, text.charCodeAt(index));
+    }
+  };
+  write(0, "RIFF");
+  view.setUint32(4, 36 + length * 2, true);
+  write(8, "WAVE");
+  write(12, "fmt ");
+  view.setUint32(16, 16, true);
+  view.setUint16(20, 1, true);
+  view.setUint16(22, 1, true);
+  view.setUint32(24, sampleRate, true);
+  view.setUint32(28, sampleRate * 2, true);
+  view.setUint16(32, 2, true);
+  view.setUint16(34, 16, true);
+  write(36, "data");
+  view.setUint32(40, length * 2, true);
+  let offset = 44;
+  for (const part of parts) {
+    for (const sample of part) {
+      const clipped = Math.max(-1, Math.min(1, sample));
+      view.setInt16(offset, clipped < 0 ? clipped * 32768 : clipped * 32767, true);
+      offset += 2;
+    }
+  }
+  return new Blob([buffer], { type: "audio/wav" });
+};
+
+const reset = async () => {
+  chunks = [];
+  lastSequence = -1;
+  audio.removeAttribute("src");
+  audio.load();
+  download.hidden = true;
+  if (wavUrl) URL.revokeObjectURL(wavUrl);
+  wavUrl = null;
+  await ensurePlayer();
+  player.port.postMessage({ type: "reset" });
+  status.textContent = "Waiting for speech tokens…";
+};
+
+const consume = async () => {
+  const value = props.value || {};
+  if (value.sequence === lastSequence) return;
+  lastSequence = value.sequence;
+  if (value.reset) await reset();
+  if (value.pcm) {
+    await ensurePlayer();
+    const pcm = decodePcm(value.pcm);
+    chunks.push(pcm.slice());
+    const playback = resample(pcm, value.sample_rate || 16000, context.sampleRate);
+    player.port.postMessage({ type: "audio", samples: playback.buffer }, [playback.buffer]);
+    status.textContent = `Streaming ${value.token_count || 0} speech tokens…`;
+  }
+  if (value.done && chunks.length) {
+    const wav = buildWav(chunks, value.sample_rate || 16000);
+    wavUrl = URL.createObjectURL(wav);
+    audio.src = wavUrl;
+    download.href = wavUrl;
+    download.hidden = false;
+    status.textContent = `Complete · ${value.token_count || 0} speech tokens`;
+  }
+};
+
+const runButton = document.getElementById("audex-run");
+if (runButton) runButton.addEventListener("click", () => ensurePlayer());
+enable.addEventListener("click", async () => {
+  await ensurePlayer();
+  status.textContent = "Playback enabled";
+});
+const stopButton = document.getElementById("audex-stop");
+if (stopButton) {
+  stopButton.addEventListener("click", () => {
+    if (player) player.port.postMessage({ type: "reset" });
+    status.textContent = "Stopped";
+  });
+}
+watch("value", consume);
+consume();
+""".replace("__WORKLET_SOURCE__", json.dumps(_WORKLET_SOURCE))
+
+
+def player_value(
+    sequence: int,
+    *,
+    pcm: str | None = None,
+    token_count: int = 0,
+    done: bool = False,
+    reset: bool = False,
+) -> dict[str, object]:
+    return {
+        "sequence": sequence,
+        "sample_rate": 16000,
+        "pcm": pcm,
+        "token_count": token_count,
+        "done": done,
+        "reset": reset,
+    }

tts_utils.py

@@ -0,0 +1,314 @@
+from __future__ import annotations
+
+import base64
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from queue import Queue
+from threading import Event, Thread
+from typing import Any, Iterator
+
+import numpy as np
+import soundfile as sf
+import torch
+from transformers import AutoModel
+from transformers.generation.logits_process import LogitsProcessor
+from transformers.generation.stopping_criteria import StoppingCriteria
+from transformers.generation.streamers import BaseStreamer
+
+
+SAMPLE_RATE = 16000
+TTS_PREFIX = "<|text to speech|> Generate speech for this transcription. "
+DEFAULT_SYSTEM_PROMPT = (
+    "You are a helpful and harmless assistant.\n\n"
+    "You are not allowed to use any tools."
+)
+
+
+@dataclass(frozen=True)
+class SpeechTokenMap:
+    start: int
+    end: int
+    codec_start: int
+    codec_end: int
+    eos: int
+
+
+@dataclass(frozen=True)
+class TTSStreamEvent:
+    pcm: np.ndarray | None
+    token_count: int
+    elapsed_seconds: float
+    done: bool = False
+
+
+class SpeechTokenLogitsProcessor(LogitsProcessor):
+    def __init__(self, token_map: SpeechTokenMap):
+        self.token_map = token_map
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        eos_scores = scores[:, self.token_map.eos].clone()
+        scores[:, : self.token_map.end] = -float("inf")
+        scores[:, self.token_map.codec_end + 1 :] = -float("inf")
+        scores[:, self.token_map.eos] = eos_scores
+        return scores
+
+
+class EventStoppingCriteria(StoppingCriteria):
+    def __init__(self, event: Event):
+        self.event = event
+
+    def __call__(
+        self,
+        input_ids: torch.LongTensor,
+        scores: torch.FloatTensor,
+        **kwargs: Any,
+    ) -> torch.BoolTensor:
+        return torch.full(
+            (input_ids.shape[0],),
+            self.event.is_set(),
+            dtype=torch.bool,
+            device=input_ids.device,
+        )
+
+
+class TokenIdStreamer(BaseStreamer):
+    _END = object()
+
+    def __init__(self) -> None:
+        self.queue: Queue[int | BaseException | object] = Queue()
+        self.is_prompt = True
+
+    def put(self, value: torch.Tensor) -> None:
+        if self.is_prompt:
+            self.is_prompt = False
+            return
+        token_ids = value.reshape(-1).tolist()
+        if len(token_ids) != 1:
+            raise ValueError(f"TTS streaming requires batch size 1, got {len(token_ids)} tokens")
+        self.queue.put(int(token_ids[0]))
+
+    def end(self) -> None:
+        self.queue.put(self._END)
+
+    def fail(self, error: BaseException) -> None:
+        self.queue.put(error)
+        self.end()
+
+    def __iter__(self) -> Iterator[int]:
+        while True:
+            item = self.queue.get()
+            if item is self._END:
+                return
+            if isinstance(item, BaseException):
+                raise item
+            yield int(item)
+
+
+def build_tts_prompt(text: str, tokenizer: Any) -> str:
+    messages = [
+        {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
+        {"role": "user", "content": f"{TTS_PREFIX}{text}"},
+    ]
+    return (
+        tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+        + "<speechgen_start>"
+    )
+
+
+def build_tts_null_prompt(cond_prompt: str, tokenizer: Any, max_iters: int = 64) -> str:
+    target_len = len(tokenizer.encode(cond_prompt))
+
+    def template(null_text: str) -> str:
+        return build_tts_prompt(null_text, tokenizer)
+
+    base_len = len(tokenizer.encode(template("")))
+    n_unk = max(1, target_len - base_len)
+    for _ in range(max_iters):
+        prompt = template("<unk>" * n_unk)
+        current_len = len(tokenizer.encode(prompt))
+        if current_len == target_len:
+            return prompt
+        n_unk += 1 if current_len < target_len else -1
+        n_unk = max(1, n_unk)
+    raise ValueError("Unable to construct an equal-length TTS CFG null prompt")
+
+
+def build_speech_token_map(tokenizer: Any) -> SpeechTokenMap:
+    token_map = SpeechTokenMap(
+        start=tokenizer.convert_tokens_to_ids("<speechgen_start>"),
+        end=tokenizer.convert_tokens_to_ids("<speechgen_end>"),
+        codec_start=tokenizer.convert_tokens_to_ids("<speechcodec_0>"),
+        codec_end=tokenizer.convert_tokens_to_ids("<speechcodec_65535>"),
+        eos=int(tokenizer.eos_token_id),
+    )
+    expected = (131075, 131076, 131077, 196612)
+    actual = (token_map.start, token_map.end, token_map.codec_start, token_map.codec_end)
+    if actual != expected:
+        raise ValueError(f"Unexpected Audex speech token layout: expected={expected}, actual={actual}")
+    return token_map
+
+
+def load_speech_decoder(model_path: str, device: str = "cuda") -> Any:
+    path = Path(model_path)
+    for filename in ("config.json", "model.safetensors"):
+        if not (path / filename).is_file():
+            raise FileNotFoundError(f"Speech decoder file not found: {path / filename}")
+
+    decoder, loading_info = AutoModel.from_pretrained(
+        str(path),
+        trust_remote_code=True,
+        output_loading_info=True,
+    )
+    invalid_keys = {
+        key: loading_info[key]
+        for key in ("missing_keys", "unexpected_keys", "mismatched_keys")
+        if loading_info[key]
+    }
+    if invalid_keys:
+        raise RuntimeError(f"Speech decoder checkpoint loading was incomplete: {invalid_keys}")
+
+    for module in decoder.modules():
+        if {"theta", "cache"} <= getattr(module, "_non_persistent_buffers_set", set()):
+            module.rope_init(device=module.theta.device)
+            module._rope_ready = False
+
+    sample_rate = int(getattr(decoder.config, "sample_rate", SAMPLE_RATE))
+    if sample_rate != SAMPLE_RATE:
+        raise ValueError(f"Expected a {SAMPLE_RATE} Hz speech decoder, got {sample_rate} Hz")
+    decoder = decoder.to(device).eval()
+    for parameter in decoder.parameters():
+        parameter.requires_grad = False
+    return decoder
+
+
+def stream_tts(
+    model: Any,
+    tokenizer: Any,
+    decoder: Any,
+    text: str,
+    max_new_tokens: int,
+    *,
+    temperature: float = 0.8,
+    top_p: float = 1.0,
+    top_k: int = 0,
+    guidance_scale: float = 2.0,
+    seed: int = 0,
+    chunk_frames: int = 5,
+) -> Iterator[TTSStreamEvent]:
+    text = text.strip()
+    if not text:
+        raise ValueError("Text to synthesize must not be empty")
+
+    token_map = build_speech_token_map(tokenizer)
+    cond_prompt = build_tts_prompt(text, tokenizer)
+    uncond_prompt = build_tts_null_prompt(cond_prompt, tokenizer)
+    cond_ids = tokenizer.encode(cond_prompt, return_tensors="pt").to(model.device)
+    uncond_ids = tokenizer.encode(uncond_prompt, return_tensors="pt").to(model.device)
+    if cond_ids.shape != uncond_ids.shape:
+        raise ValueError(
+            f"TTS CFG prompt lengths differ: conditional={cond_ids.shape[-1]}, "
+            f"unconditional={uncond_ids.shape[-1]}"
+        )
+
+    streamer = TokenIdStreamer()
+    cancel_event = Event()
+    session = decoder.create_session(
+        chunk_frames=chunk_frames,
+        sample_rate=SAMPLE_RATE,
+        return_numpy=True,
+    )
+    generation_error: list[BaseException] = []
+    started_at = time.perf_counter()
+
+    def generate() -> None:
+        try:
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            model.generate(
+                input_ids=cond_ids,
+                attention_mask=torch.ones_like(cond_ids),
+                negative_prompt_ids=uncond_ids,
+                negative_prompt_attention_mask=torch.ones_like(uncond_ids),
+                guidance_scale=guidance_scale,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+                **({"top_k": top_k} if top_k > 0 else {}),
+                max_new_tokens=max_new_tokens,
+                eos_token_id=[token_map.end, token_map.eos],
+                pad_token_id=tokenizer.pad_token_id or token_map.eos,
+                logits_processor=[SpeechTokenLogitsProcessor(token_map)],
+                stopping_criteria=[EventStoppingCriteria(cancel_event)],
+                streamer=streamer,
+                use_cache=True,
+            )
+        except BaseException as error:
+            generation_error.append(error)
+            streamer.fail(error)
+
+    thread = Thread(target=generate, daemon=True)
+    thread.start()
+    token_count = 0
+    completed = False
+    try:
+        for token_id in streamer:
+            if token_id == token_map.end or token_id == token_map.eos:
+                completed = True
+                break
+            if not token_map.codec_start <= token_id <= token_map.codec_end:
+                raise RuntimeError(f"Unexpected token in TTS output: {token_id}")
+
+            token_count += 1
+            codec_index = token_id - token_map.codec_start
+            if token_count == 1:
+                yield TTSStreamEvent(
+                    pcm=None,
+                    token_count=token_count,
+                    elapsed_seconds=time.perf_counter() - started_at,
+                )
+            for _, pcm in session.push([[codec_index]]):
+                yield TTSStreamEvent(
+                    pcm=np.asarray(pcm, dtype=np.float32),
+                    token_count=token_count,
+                    elapsed_seconds=time.perf_counter() - started_at,
+                )
+        else:
+            completed = True
+
+        if generation_error:
+            raise generation_error[0]
+        if completed:
+            for _, pcm in session.flush():
+                yield TTSStreamEvent(
+                    pcm=np.asarray(pcm, dtype=np.float32),
+                    token_count=token_count,
+                    elapsed_seconds=time.perf_counter() - started_at,
+                )
+            yield TTSStreamEvent(
+                pcm=None,
+                token_count=token_count,
+                elapsed_seconds=time.perf_counter() - started_at,
+                done=True,
+            )
+    finally:
+        cancel_event.set()
+        thread.join(timeout=10)
+
+
+def encode_pcm_chunk(pcm: np.ndarray) -> str:
+    samples = np.asarray(pcm, dtype="<f4")
+    return base64.b64encode(samples.tobytes()).decode("ascii")
+
+
+def write_wav(pcm: np.ndarray, sample_rate: int = SAMPLE_RATE) -> str:
+    output = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    output.close()
+    sf.write(output.name, np.asarray(pcm, dtype=np.float32), sample_rate, subtype="PCM_16")
+    return output.name

web/PCMPlayerWorklet.js

@@ -0,0 +1,55 @@
+class PCMPlayerProcessor extends AudioWorkletProcessor {
+  constructor() {
+    super();
+    this.queue = [];
+    this.offset = 0;
+    this.bufferedSamples = 0;
+    this.started = false;
+    this.startThreshold = Math.round(sampleRate * 0.25);
+    this.port.onmessage = (event) => {
+      if (event.data.type === "reset") {
+        this.queue = [];
+        this.offset = 0;
+        this.bufferedSamples = 0;
+        this.started = false;
+        return;
+      }
+      if (event.data.type === "audio") {
+        const samples = new Float32Array(event.data.samples);
+        this.queue.push(samples);
+        this.bufferedSamples += samples.length;
+      }
+    };
+  }
+
+  process(_inputs, outputs) {
+    const output = outputs[0][0];
+    output.fill(0);
+    if (!this.started) {
+      this.started = this.bufferedSamples >= this.startThreshold;
+      if (!this.started) {
+        return true;
+      }
+    }
+
+    let outputOffset = 0;
+    while (outputOffset < output.length && this.queue.length) {
+      const current = this.queue[0];
+      const count = Math.min(output.length - outputOffset, current.length - this.offset);
+      output.set(current.subarray(this.offset, this.offset + count), outputOffset);
+      outputOffset += count;
+      this.offset += count;
+      this.bufferedSamples -= count;
+      if (this.offset === current.length) {
+        this.queue.shift();
+        this.offset = 0;
+      }
+    }
+    if (this.started && outputOffset < output.length) {
+      this.port.postMessage({ type: "underrun" });
+    }
+    return true;
+  }
+}
+
+registerProcessor("audex-pcm-player", PCMPlayerProcessor);