"""
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
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", "nvidia/Nemotron-Labs-Audex")
CAUSAL_CONV1D_VERSION = "1.6.2.post1"
MAMBA_SSM_VERSION = "2.3.2.post1"
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", "900"))
MAX_NEW_TOKENS = int(os.environ.get("AUDEX_MAX_NEW_TOKENS", "4096"))
DEFAULT_MAX_NEW_TOKENS = min(
    int(os.environ.get("AUDEX_DEFAULT_MAX_NEW_TOKENS", "1024")),
    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."
)
TTS_TOKEN_WARNING = (
    "⚠️ Maximum speech-token limit reached; generated audio may be incomplete."
)
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."
)

# Task-specific limits and feature switches
TTS_MAX_NEW_TOKENS = min(
    int(os.environ.get("AUDEX_TTS_MAX_NEW_TOKENS", "512")),
    MAX_NEW_TOKENS,
)
TTS_DEFAULT_NEW_TOKENS = min(
    int(os.environ.get("AUDEX_TTS_DEFAULT_NEW_TOKENS", "256")),
    TTS_MAX_NEW_TOKENS,
)
TTS_30B_GPU_DURATION_SECONDS = int(
    os.environ.get("AUDEX_TTS_30B_GPU_DURATION_SECONDS", "120")
)
S2S_TEXT_MAX_NEW_TOKENS = int(os.environ.get("AUDEX_S2S_TEXT_MAX_NEW_TOKENS", "4096"))
S2S_REASONING_BUDGET = int(os.environ.get("AUDEX_S2S_REASONING_BUDGET", "3584"))
S2S_30B_GPU_DURATION_SECONDS = int(
    os.environ.get("AUDEX_S2S_30B_GPU_DURATION_SECONDS", "320")
)
S2S_2B_GPU_DURATION_SECONDS = int(
    os.environ.get("AUDEX_S2S_2B_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: 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:
        print(p.stdout[-4000:], flush=True)
    if p.returncode != 0:
        print(p.stderr[-8000:], flush=True)
        raise RuntimeError(f"command failed ({p.returncode}): {' '.join(cmd)}")
    else:
        # surface tail of stderr (nvcc warnings etc.) but not as failure
        if p.stderr:
            print(p.stderr[-2000:], flush=True)
    return p


def _installed_version(package: str) -> str | None:
    try:
        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: str, env: dict[str, str] | None = None) -> None:
    _run([sys.executable, "-m", "pip", *args], env=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"] = _detect_torch_arch()
    env["MAX_JOBS"] = env.get("MAX_JOBS", "4")
    cuda_home = _detect_cuda_home()
    env["CUDA_HOME"] = cuda_home
    env["PATH"] = f"{cuda_home}/bin:" + env.get("PATH", "")
    return env


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 _kernels_ready():
        print("[build] kernels already importable", flush=True)
        return

    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 _kernels_ready():
                print("[build] cached wheels installed OK", flush=True)
                return
            print("[build] cached wheels imported incompletely; rebuilding", flush=True)
        except Exception as e:
            print(f"[build] cached wheel install failed ({e!r}); rebuilding", flush=True)

    env = _build_env()
    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(BUILD_DIR),
        CAUSAL_CONV1D_SPEC,
        env=env,
    )
    # mamba-ssm needs causal-conv1d importable during its own build; install it first.
    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)

    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 _kernels_ready():
        raise RuntimeError("kernel build completed but imports still fail")
    print("[build] kernels built + installed OK", flush=True)

    # Cache wheels back into the repo for fast subsequent boots.
    try:
        for w in all_wheels:
            dst = WHEELS_DIR / Path(w).name
            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:
            HfApi(token=tok).upload_folder(
                folder_path=str(WHEELS_DIR),
                path_in_repo="wheels",
                repo_id=REPO_ID,
                repo_type="space",
                commit_message="cache compiled mamba-ssm + causal-conv1d wheels",
            )
            print("[build] cached wheels uploaded to repo", flush=True)
        else:
            print("[build] no HF_TOKEN; skipping wheel cache upload", flush=True)
    except Exception as e:
        print(f"[build] wheel cache upload skipped ({e!r})", flush=True)


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 audio_utils import (
    IM_END_TOKEN,
    build_attention_mask,
    build_prompt_template,
    expand_sound_placeholder,
    extract_whisper_features,
    load_audio,
    resolve_audio_preprocessor_path,
    split_thinking,
)
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,
)


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]

    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:
    for name, mod in list(sys.modules.items()):
        if name.endswith("modeling_nemotron_h") and hasattr(mod, "is_fast_path_available"):
            return bool(mod.is_fast_path_available)
    return None


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(),
}


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:
    if not audio:
        return 0.0
    try:
        import soundfile as sf
        return float(sf.info(audio).duration)
    except Exception:
        return 0.0


def _estimate(
    audio: str | None,
    task: str,
    custom_prompt: str | None,
    reasoning: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    *args: object,
    **kwargs: object,
) -> int:
    audio_duration = _probe_audio_duration(audio)
    if (
        audio is None
        or audio_duration > MAX_AUDIO_DURATION_SECONDS
        or int(max_new_tokens) > MAX_NEW_TOKENS
    ):
        return 10

    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()


def _generate(
    runtime: AudexRuntime,
    audio: str,
    prompt: str,
    reasoning: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    greedy: bool,
    reasoning_budget: int | None = None,
) -> Iterator[tuple[str, str]]:
    started_at = time.perf_counter()
    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
    if audio_duration > MAX_AUDIO_DURATION_SECONDS:
        raise gr.Error(
            f"Audio is {audio_duration / 60:.1f} minutes long; "
            f"this demo supports up to {MAX_AUDIO_DURATION_SECONDS / 60:.0f} minutes."
        )
    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(
        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)
    )
    formatted = build_prompt_template(prompt.strip(), reasoning=bool(reasoning),
                                      prompt_repitition="none")
    expanded = expand_sound_placeholder(formatted, num_embeddings)
    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 = 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)
    do_sample = not greedy and ((temperature != 1.0) or (0.0 < top_p < 1.0))
    gen_kwargs = dict(
        do_sample=do_sample,
        eos_token_id=eos_token_id,
        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

    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,
    )


@spaces.GPU(duration=_estimate, size="xlarge")
def run(
    audio: str | None,
    task: str,
    custom_prompt: str | None,
    reasoning: bool,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
) -> tuple[str, str]:
    """Run audio understanding, speech recognition, or speech translation.

    Predefined ASR and translation tasks use the model's recommended greedy
    decoding. Other tasks and custom instructions use the sampling controls.
    Input duration and output length are limited by the current runtime.

    Returns:
        The final answer and, when enabled, the reasoning trace.
    """
    if audio is None:
        return "Please provide an audio input.", ""

    custom_instruction = custom_prompt.strip() if custom_prompt else ""
    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,
        int(max_new_tokens),
        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
    truncated = False
    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
            segment_truncated = False
            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.done:
                    segment_truncated = event.truncated
                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_truncated:
                truncated = True
                break
            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} "
            f"truncated={truncated}",
            flush=True,
        )
        status = (
            f"{TTS_TOKEN_WARNING} · {waveform.size / SAMPLE_RATE:.1f}s audio"
            if truncated
            else f"Complete · {waveform.size / SAMPLE_RATE:.1f}s audio"
        )
        yield (
            status,
            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:
    text = text.strip()
    quoted = re.fullmatch(
        r"(?is)(?:(?:(?:(?:source\s+)?language)\s*:[^.\n]+\.\s*)?"
        r"(?:(?:the\s+)?(?:transcription|transcript)(?:\s+(?:is|reads))?"
        r"|(?:the\s+)?(?:spoken\s+)?content\s+of\s+the\s+(?:input\s+)?audio\s+is)"
        r"\s*:?\s*)?(['\"])(.*)\1[.!]?",
        text,
    )
    return quoted.group(2).strip() if quoted else text


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 status, 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,
    ):
        visible_answer = (
            f"{answer}\n\n{TTS_TOKEN_WARNING}"
            if TTS_TOKEN_WARNING in status
            else answer
        )
        yield visible_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


def _estimate_unified(
    model_name: str,
    task: str,
    audio: str | None,
    prompt: str,
    reasoning: bool,
    reasoning_budget: int,
    max_new_tokens: int,
    tts_max_new_tokens: int,
    temperature: float,
    top_p: float,
    guidance_scale: float,
    *args: object,
    **kwargs: object,
) -> int:
    if task == S2S_TASK:
        if model_name == DEFAULT_MODEL_NAME:
            return S2S_30B_GPU_DURATION_SECONDS
        return S2S_2B_GPU_DURATION_SECONDS
    if task == TTS_TASK:
        if model_name == DEFAULT_MODEL_NAME:
            return TTS_30B_GPU_DURATION_SECONDS
        return MAX_GPU_DURATION_SECONDS
    if task == TEXT_TASK:
        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,
    tts_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 "", "", None, 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(tts_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"])
    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))
    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=default_max_new_tokens,
            visible=not is_tts,
        ),
        gr.update(value=TTS_DEFAULT_NEW_TOKENS, visible=is_tts),
        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))


theme = gr.themes.Citrus()

with gr.Blocks(title="Nemotron-Labs-Audex") as demo:
    gr.Markdown(
        "# 🎧 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():
            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=f"Input audio (up to {MAX_AUDIO_DURATION_SECONDS / 60:g} minutes)",
                sources=["upload", "microphone"],
            )
            unified_prompt = gr.Textbox(
                value=str(TASK_SPECS["Speech recognition (ASR)"]["template"]),
                label="Task template / prompt",
                lines=3,
            )
            gr.Markdown(
                "`<sound>` marks where audio embeddings enter the prompt; the app "
                "replaces it automatically. Text-only tasks have no input-modality token."
            )
            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):
                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=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_tts_max_new_tokens = gr.Slider(
                    16,
                    TTS_MAX_NEW_TOKENS,
                    value=TTS_DEFAULT_NEW_TOKENS,
                    step=16,
                    label="Max speech tokens",
                    info=(
                        "Speech-token cap. If reached, generated audio may be incomplete."
                    ),
                    visible=False,
                )
                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,
                )
                gr.Markdown(
                    "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; speech-to-speech "
                    "defaults to a 3,584-token reasoning budget. TTS uses CFG 2.0; "
                    "speech-to-speech uses CFG 1.5. TTS defaults to 256 speech tokens "
                    "and supports up to 512."
                )
        with gr.Column():
            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_tts_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,
    )

    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_tts_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=[
            [
                "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,
        ],
        cache_examples=False,
    )

if __name__ == "__main__":
    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,
    )