Create Audex 2B ZeroGPU demo

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.py[cod]
+.gradio/
+.env
+audex-cache/
+tmp/
+

README.md

@@ -1,13 +1,43 @@
 ---
-title: Nemotron Labs Audex
-emoji: 🐨
-colorFrom: red
-colorTo: purple
+title: Nemotron-Labs-Audex
 sdk: gradio
 sdk_version: 6.19.0
-python_version: '3.13'
 app_file: app.py
-pinned: false
+python_version: 3.12
+license: other
+short_description: Audio and text demo for NVIDIA Nemotron-Labs-Audex-2B.
+models:
+  - nvidia/Nemotron-Labs-Audex-2B
+tags:
+  - audio-understanding
+  - speech-recognition
+  - speech-translation
+  - text-generation
+  - nemotron-labs-audex
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Nemotron-Labs-Audex
+
+Public Gradio ZeroGPU demo for
+[nvidia/Nemotron-Labs-Audex-2B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B).
+
+Audex-2B is a unified audio-text LLM from NVIDIA. It supports audio
+understanding, speech recognition, speech translation, and text reasoning while
+using a single transformer decoder with an audio encoder and projected audio
+embeddings.
+
+Useful links:
+
+- Model card: https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B
+- Technical report: https://huggingface.co/papers/2607.05196
+
+This Space intentionally runs only the 2B model on ZeroGPU.
+
+The Space predownloads the model snapshot with the Space owner's `HF_TOKEN`
+secret when available, then keeps the loaded model resident across requests.
+Generation progress is split into input resolution, model snapshot, GPU model
+load, audio feature extraction, prompt preparation, token generation, and output
+decoding.
+
+Your use of this model is governed by the NVIDIA Oneway Noncommercial License
+linked from the model card.

app.py

@@ -0,0 +1,860 @@
+from __future__ import annotations
+
+import gc
+import math
+import os
+import random
+import tempfile
+import threading
+import time
+import traceback
+import wave
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+from urllib.parse import urlparse
+
+import gradio as gr
+import numpy as np
+import requests
+import spaces
+import torch
+from huggingface_hub import snapshot_download
+from transformers import (
+    AutoConfig,
+    AutoFeatureExtractor,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+)
+
+
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+APP_TITLE = "Nemotron-Labs-Audex"
+MODEL_REPO_ID = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_CARD_URL = "https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B"
+PAPER_URL = "https://huggingface.co/papers/2607.05196"
+TEXT_VOCAB_SIZE = 131072
+MAX_AUDIO_BYTES = 80 * 1024 * 1024
+SAMPLE_RATE = 16000
+CACHE_DIR = os.getenv(
+    "AUDEX_CACHE_DIR",
+    "/data/audex-cache" if os.path.isdir("/data") else "/tmp/audex-cache",
+)
+APP_DIR = Path(__file__).resolve().parent
+ASSET_DIR = APP_DIR / "assets" / "examples"
+TMP_DIR = Path("/tmp/audex-inputs")
+
+ALLOW_PATTERNS = [
+    "checkpoint_folder_full/**",
+    "README.md",
+    "LICENSE",
+    "license/**",
+]
+
+SOUND_PLACEHOLDER = "<sound>"
+SOUND_TOKEN = "<so_embedding>"
+SOUND_START_TOKEN = "<so_start>"
+SOUND_END_TOKEN = "<so_end>"
+IM_END_TOKEN = "<|im_end|>"
+
+
+@dataclass
+class LoadedModel:
+    model_dir: Path
+    model: AutoModelForCausalLM
+    tokenizer: AutoTokenizer
+    feature_extractor: AutoFeatureExtractor
+    config: AutoConfig
+
+
+_MODEL_LOCK = threading.Lock()
+_LOADED: LoadedModel | None = None
+_SNAPSHOT_PATH: Path | None = None
+_PRELOAD_THREAD: threading.Thread | None = None
+_PRELOAD_ERROR: BaseException | None = None
+
+
+TASKS = {
+    "audio-understanding": {
+        "label": "Audio understanding",
+        "prompt": "Describe the audio in detail.",
+        "needs_audio": True,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "reasoning": False,
+    },
+    "speech-recognition": {
+        "label": "Speech recognition",
+        "prompt": "Transcribe the speech in the input audio.",
+        "needs_audio": True,
+        "temperature": 1.0,
+        "top_p": 1.0,
+        "reasoning": False,
+    },
+    "speech-translation": {
+        "label": "Speech translation",
+        "prompt": "Translate the speech in the input audio into English.",
+        "needs_audio": True,
+        "temperature": 1.0,
+        "top_p": 1.0,
+        "reasoning": False,
+    },
+    "text-reasoning": {
+        "label": "Text-only reasoning",
+        "prompt": "Explain why a unified audio-text model can preserve text reasoning while learning audio tasks.",
+        "needs_audio": False,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "reasoning": True,
+    },
+    "custom": {
+        "label": "Custom",
+        "prompt": "What is happening in this audio?",
+        "needs_audio": True,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "reasoning": False,
+    },
+}
+
+
+def _task_choices() -> list[tuple[str, str]]:
+    return [(cfg["label"], key) for key, cfg in TASKS.items()]
+
+
+def _hub_token() -> str | None:
+    return os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
+
+
+def _format_seconds(seconds: float) -> str:
+    seconds = max(0.0, float(seconds))
+    if seconds < 10:
+        return f"{seconds:.1f}s"
+    minutes, secs = divmod(int(round(seconds)), 60)
+    if minutes:
+        return f"{minutes}m {secs:02d}s"
+    return f"{secs}s"
+
+
+def _torch_cleanup() -> None:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
+def _friendly_error(exc: BaseException) -> str:
+    text = str(exc)
+    lowered = text.lower()
+    if "cuda out of memory" in lowered or "outofmemoryerror" in lowered:
+        return (
+            "CUDA ran out of memory. Try a shorter audio clip, fewer max tokens, "
+            "or wait for the Space to restart cleanly."
+        )
+    if "401" in lowered or "403" in lowered or "gated" in lowered:
+        return (
+            "The model download was rejected by Hugging Face. Configure the Space "
+            "owner `HF_TOKEN` secret with access to the model and restart the Space."
+        )
+    return f"{type(exc).__name__}: {text}"
+
+
+def _ensure_example_assets() -> None:
+    ASSET_DIR.mkdir(parents=True, exist_ok=True)
+    tone = ASSET_DIR / "tone_440hz.wav"
+    chirp = ASSET_DIR / "chirp_with_noise.wav"
+    if not tone.exists():
+        _write_wav(tone, _sine_wave(440.0, 2.0, 0.35))
+    if not chirp.exists():
+        _write_wav(chirp, _chirp_wave(220.0, 880.0, 3.0, 0.30))
+
+
+def _sine_wave(freq: float, seconds: float, amp: float) -> np.ndarray:
+    t = np.linspace(0, seconds, int(SAMPLE_RATE * seconds), endpoint=False)
+    return (amp * np.sin(2 * np.pi * freq * t)).astype(np.float32)
+
+
+def _chirp_wave(start_freq: float, end_freq: float, seconds: float, amp: float) -> np.ndarray:
+    t = np.linspace(0, seconds, int(SAMPLE_RATE * seconds), endpoint=False)
+    freqs = np.linspace(start_freq, end_freq, t.shape[0])
+    phase = 2 * np.pi * np.cumsum(freqs) / SAMPLE_RATE
+    noise = np.random.default_rng(7).normal(0, 0.02, t.shape[0])
+    return (amp * np.sin(phase) + noise).astype(np.float32)
+
+
+def _write_wav(path: Path, audio: np.ndarray) -> None:
+    pcm = np.clip(audio, -1, 1)
+    pcm = (pcm * 32767).astype("<i2")
+    with wave.open(str(path), "wb") as handle:
+        handle.setnchannels(1)
+        handle.setsampwidth(2)
+        handle.setframerate(SAMPLE_RATE)
+        handle.writeframes(pcm.tobytes())
+
+
+def _example_samples() -> list[list[object]]:
+    _ensure_example_assets()
+    return [
+        [
+            "audio-understanding",
+            str(ASSET_DIR / "tone_440hz.wav"),
+            "",
+            "Describe the audio in one sentence.",
+            False,
+            128,
+            0.7,
+            0.9,
+            0,
+        ],
+        [
+            "audio-understanding",
+            str(ASSET_DIR / "chirp_with_noise.wav"),
+            "",
+            "What kind of sound pattern do you hear?",
+            False,
+            160,
+            0.7,
+            0.9,
+            0,
+        ],
+        [
+            "text-reasoning",
+            None,
+            "",
+            "In three concise bullets, explain what makes Audex a unified audio-text model.",
+            True,
+            256,
+            0.7,
+            0.9,
+            0,
+        ],
+    ]
+
+
+def ensure_model_snapshot() -> Path:
+    global _SNAPSHOT_PATH
+    if _SNAPSHOT_PATH is not None:
+        return _SNAPSHOT_PATH
+
+    local_path = snapshot_download(
+        repo_id=MODEL_REPO_ID,
+        allow_patterns=ALLOW_PATTERNS,
+        cache_dir=CACHE_DIR,
+        token=_hub_token(),
+    )
+    _SNAPSHOT_PATH = Path(local_path)
+    return _SNAPSHOT_PATH
+
+
+def _preload_snapshot() -> None:
+    global _PRELOAD_ERROR
+    try:
+        ensure_model_snapshot()
+    except BaseException as exc:  # noqa: BLE001 - stored and surfaced to UI
+        _PRELOAD_ERROR = exc
+        traceback.print_exc()
+
+
+def start_preload() -> None:
+    global _PRELOAD_THREAD
+    if os.getenv("AUDEX_PRELOAD", "1") == "0":
+        return
+    if _PRELOAD_THREAD is not None:
+        return
+    _PRELOAD_THREAD = threading.Thread(target=_preload_snapshot, daemon=True)
+    _PRELOAD_THREAD.start()
+
+
+def wait_for_preload(progress: gr.Progress) -> float:
+    started = time.monotonic()
+    thread = _PRELOAD_THREAD
+    while thread is not None and thread.is_alive():
+        elapsed = time.monotonic() - started
+        progress(0.04, desc=f"Predownloading Audex-2B snapshot with Space HF_TOKEN | elapsed {_format_seconds(elapsed)}")
+        time.sleep(2)
+    if _PRELOAD_ERROR is not None:
+        raise _PRELOAD_ERROR
+    ensure_model_snapshot()
+    return time.monotonic() - started
+
+
+def refresh_remote_code_cache(model_dir: Path) -> None:
+    module_name = model_dir.resolve().name.replace("-", "_").replace(".", "_")
+    cache_root = Path(os.getenv("HF_MODULES_CACHE", Path.home() / ".cache/huggingface/modules"))
+    cache_path = cache_root / "transformers_modules" / module_name
+    if cache_path.exists():
+        import shutil
+
+        shutil.rmtree(cache_path)
+
+
+def resolve_audio_preprocessor_path(model_dir: Path, config) -> str:
+    path = getattr(config, "audio_preprocessor_path", None) or "audio_preprocessor"
+    candidate = Path(path)
+    if not candidate.is_absolute():
+        candidate = model_dir / candidate
+    return str(candidate)
+
+
+def load_model(progress: gr.Progress) -> tuple[LoadedModel, float]:
+    global _LOADED
+    with _MODEL_LOCK:
+        if _LOADED is not None:
+            progress(0.18, desc="Model already loaded on GPU")
+            return _LOADED, 0.0
+
+        started = time.monotonic()
+        snapshot_path = ensure_model_snapshot()
+        model_dir = snapshot_path / "checkpoint_folder_full"
+        refresh_remote_code_cache(model_dir)
+
+        progress(0.10, desc="Loading tokenizer and audio preprocessor")
+        tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+        config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+        feature_extractor = AutoFeatureExtractor.from_pretrained(
+            resolve_audio_preprocessor_path(model_dir, config)
+        )
+
+        if torch.cuda.is_available():
+            device = "cuda:0"
+            dtype = torch.bfloat16
+        else:
+            device = "cpu"
+            dtype = torch.float32
+
+        progress(0.14, desc=f"Loading Audex-2B weights to {device}")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_dir,
+            trust_remote_code=True,
+            torch_dtype=dtype,
+            device_map={"": device},
+        )
+        model.eval()
+        _LOADED = LoadedModel(
+            model_dir=model_dir,
+            model=model,
+            tokenizer=tokenizer,
+            feature_extractor=feature_extractor,
+            config=config,
+        )
+        return _LOADED, time.monotonic() - started
+
+
+def normalize_audio(audio: np.ndarray) -> np.ndarray:
+    audio = np.asarray(audio)
+    if audio.ndim == 2:
+        if audio.shape[1] <= 2:
+            audio = audio.mean(axis=1)
+        elif audio.shape[0] <= 2:
+            audio = audio.mean(axis=0)
+        else:
+            raise ValueError(f"Unsupported audio shape: {audio.shape}")
+
+    if audio.dtype == np.int16:
+        audio = audio.astype(np.float32) / 32768.0
+    elif audio.dtype != np.float32:
+        audio = audio.astype(np.float32)
+
+    max_abs = float(np.abs(audio).max()) if audio.size else 0.0
+    if max_abs > 1.0:
+        audio = audio / max_abs
+    return audio.astype(np.float32, copy=False)
+
+
+def load_audio(audio_path: str, target_sr: int = SAMPLE_RATE) -> tuple[np.ndarray, int]:
+    import librosa
+
+    audio_data, sr = librosa.load(audio_path, sr=target_sr, mono=True)
+    return normalize_audio(audio_data), sr
+
+
+def split_audio_into_clips(
+    audio: np.ndarray,
+    sample_rate: int = SAMPLE_RATE,
+    clip_duration: float = 30.0,
+) -> list[np.ndarray]:
+    audio = normalize_audio(audio)
+    clip_samples = int(round(sample_rate * clip_duration))
+    if clip_samples <= 0:
+        raise ValueError(f"Invalid clip duration: {clip_duration}")
+    if audio.size == 0:
+        audio = np.zeros(1, dtype=np.float32)
+
+    num_clips = max(1, math.ceil(audio.shape[0] / clip_samples))
+    clips: list[np.ndarray] = []
+    for idx in range(num_clips):
+        start = idx * clip_samples
+        clip = audio[start : start + clip_samples]
+        if clip.shape[0] < clip_samples:
+            clip = np.pad(clip, (0, clip_samples - clip.shape[0]))
+        clips.append(clip.astype(np.float32, copy=False))
+    return clips
+
+
+def extract_whisper_features(
+    feature_extractor,
+    audio: np.ndarray,
+    sample_rate: int,
+    clip_duration: float,
+) -> torch.Tensor:
+    clips = split_audio_into_clips(audio, sample_rate=sample_rate, clip_duration=clip_duration)
+    features = feature_extractor(
+        clips,
+        sampling_rate=sample_rate,
+        return_tensors="pt",
+        padding="max_length",
+        return_attention_mask=False,
+    )
+    input_features = features.input_features
+    if input_features.ndim != 3:
+        raise ValueError(f"Expected 3D Whisper features, got {tuple(input_features.shape)}")
+    return input_features
+
+
+def build_prompt_template(prompt: str, reasoning: bool, has_audio: bool) -> str:
+    audio_prefix = "<sound>\n" if has_audio else ""
+    if reasoning:
+        return f"<|im_start|>user\n{audio_prefix}{prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n"
+    return f"<|im_start|>user\n{audio_prefix}{prompt}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def expand_sound_placeholder(prompt: str, num_embeddings: int) -> str:
+    if prompt.count(SOUND_PLACEHOLDER) != 1:
+        raise ValueError(
+            f"Expected exactly one {SOUND_PLACEHOLDER}, found {prompt.count(SOUND_PLACEHOLDER)}"
+        )
+    replacement = SOUND_START_TOKEN + (SOUND_TOKEN * num_embeddings) + SOUND_END_TOKEN
+    return prompt.replace(SOUND_PLACEHOLDER, replacement)
+
+
+def split_thinking(response: str) -> tuple[str, str]:
+    if "</think>" not in response:
+        return "", response.strip()
+    thinking = response.rsplit("</think>", 1)[0].strip() + "</think>"
+    prediction = response.rsplit("</think>", 1)[1].strip()
+    return thinking, prediction
+
+
+def clean_response(tokenizer, output_ids, prompt_len: int) -> tuple[str, str, str]:
+    new_tokens = output_ids[0, prompt_len:]
+    response = tokenizer.decode(new_tokens, skip_special_tokens=False)
+    response = response.split(IM_END_TOKEN, 1)[0].strip()
+    thinking, answer = split_thinking(response)
+    return response, thinking, answer
+
+
+def validate_https_url(url: str) -> str:
+    parsed = urlparse(url)
+    if parsed.scheme != "https" or not parsed.netloc:
+        raise ValueError("Audio URL must be a valid HTTPS URL.")
+    return url
+
+
+def download_audio_url(url: str, progress: gr.Progress) -> tuple[str, float]:
+    started = time.monotonic()
+    TMP_DIR.mkdir(parents=True, exist_ok=True)
+    suffix = Path(urlparse(url).path).suffix[:10] or ".audio"
+    target = TMP_DIR / f"audex-url-{int(started * 1000)}-{random.randint(0, 9999)}{suffix}"
+
+    with requests.get(url, stream=True, timeout=(10, 120)) as resp:
+        resp.raise_for_status()
+        total = int(resp.headers.get("content-length") or 0)
+        downloaded = 0
+        with target.open("wb") as handle:
+            for chunk in resp.iter_content(chunk_size=1024 * 1024):
+                if not chunk:
+                    continue
+                downloaded += len(chunk)
+                if downloaded > MAX_AUDIO_BYTES:
+                    raise ValueError("Audio URL is too large. Limit is 80 MB.")
+                handle.write(chunk)
+                if total:
+                    progress(
+                        min(0.10, 0.02 + 0.08 * downloaded / total),
+                        desc=f"Downloading HTTPS audio {downloaded / 1_000_000:.1f}/{total / 1_000_000:.1f} MB",
+                    )
+                else:
+                    progress(
+                        0.04,
+                        desc=f"Downloading HTTPS audio {downloaded / 1_000_000:.1f} MB",
+                    )
+    return str(target), time.monotonic() - started
+
+
+def resolve_audio_input(
+    audio_file: str | None,
+    audio_url: str,
+    needs_audio: bool,
+    progress: gr.Progress,
+) -> tuple[str | None, float]:
+    if not needs_audio:
+        return None, 0.0
+    audio_url = audio_url.strip()
+    if audio_url:
+        progress(0.02, desc="Resolving HTTPS audio URL")
+        return download_audio_url(validate_https_url(audio_url), progress)
+    if audio_file:
+        return audio_file, 0.0
+    raise ValueError("Provide audio from the microphone, upload a file, or enter an HTTPS audio URL.")
+
+
+class GenerationProgress(StoppingCriteria):
+    def __init__(self, prompt_len: int, max_new_tokens: int, progress: gr.Progress) -> None:
+        self.prompt_len = int(prompt_len)
+        self.max_new_tokens = max(1, int(max_new_tokens))
+        self.progress = progress
+        self.started = time.monotonic()
+        self.last_reported = -1
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        new_tokens = max(0, int(input_ids.shape[-1]) - self.prompt_len)
+        if new_tokens != self.last_reported:
+            self.last_reported = new_tokens
+            elapsed = time.monotonic() - self.started
+            fraction = min(0.92, 0.62 + 0.30 * min(new_tokens, self.max_new_tokens) / self.max_new_tokens)
+            speed = new_tokens / elapsed if elapsed > 0 and new_tokens > 0 else 0.0
+            self.progress(
+                fraction,
+                desc=(
+                    f"Generating text tokens {new_tokens}/{self.max_new_tokens} | "
+                    f"elapsed {_format_seconds(elapsed)} | {speed:.1f} tok/s"
+                ),
+            )
+        return False
+
+
+def task_defaults(task: str):
+    cfg = TASKS[task]
+    return (
+        gr.update(value=cfg["prompt"]),
+        gr.update(value=cfg["reasoning"]),
+        gr.update(value=cfg["temperature"]),
+        gr.update(value=cfg["top_p"]),
+        gr.update(visible=bool(cfg["needs_audio"])),
+        gr.update(visible=bool(cfg["needs_audio"])),
+    )
+
+
+def estimate_duration(
+    task: str,
+    audio_file: str | None,
+    audio_url: str,
+    prompt: str,
+    reasoning: bool,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    *_args,
+    **_kwargs,
+) -> int:
+    del task, audio_file, audio_url, prompt, reasoning, temperature, top_p, top_k
+    loaded = _LOADED is not None
+    base = 40 if loaded else 90
+    return min(240, base + int(max_new_tokens) // 24)
+
+
+@spaces.GPU(duration=estimate_duration)
+def generate(
+    task: str,
+    audio_file: str | None,
+    audio_url: str,
+    prompt: str,
+    reasoning: bool,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    progress: gr.Progress = gr.Progress(track_tqdm=False),
+):
+    total_started = time.monotonic()
+    timings: dict[str, float] = {}
+    prompt = prompt.strip()
+    if not prompt:
+        return "", "", "Enter a prompt."
+
+    try:
+        task_cfg = TASKS[task]
+        needs_audio = bool(task_cfg["needs_audio"])
+        progress(0.01, desc="Resolving inputs")
+        resolved_audio, timings["audio_fetch"] = resolve_audio_input(
+            audio_file=audio_file,
+            audio_url=audio_url,
+            needs_audio=needs_audio,
+            progress=progress,
+        )
+
+        timings["snapshot_wait"] = wait_for_preload(progress)
+        loaded, timings["model_load"] = load_model(progress)
+        model = loaded.model
+        tokenizer = loaded.tokenizer
+        config = loaded.config
+
+        input_features = None
+        feature_started = time.monotonic()
+        if needs_audio:
+            progress(0.24, desc="Loading and resampling audio")
+            audio, sr = load_audio(str(resolved_audio), target_sr=SAMPLE_RATE)
+            duration = audio.shape[0] / float(sr)
+            progress(0.32, desc=f"Extracting Whisper audio features from {_format_seconds(duration)} audio")
+            input_features = extract_whisper_features(
+                loaded.feature_extractor,
+                audio,
+                sample_rate=sr,
+                clip_duration=float(getattr(config, "sound_clip_duration", 30.0)),
+            )
+            num_embeddings = input_features.shape[0] * int(getattr(config, "sound_embedding_size", 750))
+        else:
+            duration = 0.0
+            num_embeddings = 0
+        timings["audio_features"] = time.monotonic() - feature_started
+
+        prompt_started = time.monotonic()
+        progress(0.42, desc="Building ChatML prompt")
+        formatted_prompt = build_prompt_template(prompt, reasoning=reasoning, has_audio=needs_audio)
+        if needs_audio:
+            formatted_prompt = expand_sound_placeholder(formatted_prompt, num_embeddings)
+        tokenized = tokenizer(formatted_prompt, return_tensors="pt", add_special_tokens=False)
+        input_ids = tokenized.input_ids
+        attention_mask = tokenized.attention_mask if "attention_mask" in tokenized else torch.ones_like(input_ids)
+        prompt_len = int(input_ids.shape[-1])
+        device = model.device
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+        if input_features is not None:
+            input_features = input_features.to(device)
+        timings["prompt"] = time.monotonic() - prompt_started
+
+        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)
+
+        max_new_tokens = int(max_new_tokens)
+        temperature = float(temperature)
+        top_p = float(top_p)
+        top_k = int(top_k)
+        do_sample = temperature != 1.0 or 0.0 < top_p < 1.0 or top_k > 0
+        generation_kwargs = {
+            "do_sample": do_sample,
+            "eos_token_id": eos_token_id,
+            "pad_token_id": tokenizer.pad_token_id or getattr(config, "pad_token_id", 0),
+            "stopping_criteria": StoppingCriteriaList(
+                [GenerationProgress(prompt_len, max_new_tokens, progress)]
+            ),
+        }
+        if do_sample:
+            generation_kwargs["temperature"] = temperature
+            if top_p > 0.0:
+                generation_kwargs["top_p"] = top_p
+            if top_k > 0:
+                generation_kwargs["top_k"] = top_k
+
+        progress(0.60, desc="Starting model generation")
+        generation_started = time.monotonic()
+        with torch.inference_mode():
+            output_ids = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                input_features=input_features,
+                max_new_tokens=max_new_tokens,
+                **generation_kwargs,
+            )
+        timings["generation"] = time.monotonic() - generation_started
+
+        decode_started = time.monotonic()
+        progress(0.96, desc="Decoding output")
+        raw_response, thinking, answer = clean_response(tokenizer, output_ids, prompt_len)
+        timings["decode"] = time.monotonic() - decode_started
+        timings["total"] = time.monotonic() - total_started
+        _torch_cleanup()
+
+        if not answer and raw_response:
+            answer = raw_response
+
+        summary = build_summary(
+            task=TASKS[task]["label"],
+            prompt_tokens=prompt_len,
+            output_tokens=max(0, int(output_ids.shape[-1]) - prompt_len),
+            audio_duration=duration,
+            timings=timings,
+        )
+        return thinking, answer, summary
+    except Exception as exc:
+        traceback.print_exc()
+        return "", "", _friendly_error(exc)
+
+
+def build_summary(
+    *,
+    task: str,
+    prompt_tokens: int,
+    output_tokens: int,
+    audio_duration: float,
+    timings: dict[str, float],
+) -> str:
+    token_speed = output_tokens / timings["generation"] if timings.get("generation") else 0.0
+    audio_line = (
+        f"{_format_seconds(audio_duration)} audio · " if audio_duration > 0 else ""
+    )
+    return (
+        "### Run summary\n\n"
+        f"`{MODEL_REPO_ID}` · {task} · {audio_line}{prompt_tokens} prompt tokens · "
+        f"{output_tokens} output tokens\n\n"
+        "| Phase | Time |\n"
+        "| --- | ---: |\n"
+        f"| Total backend time | {_format_seconds(timings.get('total', 0))} |\n"
+        f"| HTTPS audio fetch | {_format_seconds(timings.get('audio_fetch', 0))} |\n"
+        f"| Snapshot wait/download | {_format_seconds(timings.get('snapshot_wait', 0))} |\n"
+        f"| Model load/reuse | {_format_seconds(timings.get('model_load', 0))} |\n"
+        f"| Audio feature extraction | {_format_seconds(timings.get('audio_features', 0))} |\n"
+        f"| Prompt preparation | {_format_seconds(timings.get('prompt', 0))} |\n"
+        f"| Token generation | {_format_seconds(timings.get('generation', 0))} |\n"
+        f"| Decode | {_format_seconds(timings.get('decode', 0))} |\n"
+        f"| Generation speed | {token_speed:.2f} tok/s |\n"
+    )
+
+
+APP_CSS = """
+.gradio-container {
+    max-width: 1280px !important;
+}
+
+#run_summary table {
+    width: 100%;
+}
+
+#run_summary th,
+#run_summary td {
+    padding: 6px 8px;
+}
+
+#run_summary th:last-child,
+#run_summary td:last-child {
+    text-align: right;
+    white-space: nowrap;
+}
+"""
+
+
+_ensure_example_assets()
+start_preload()
+
+
+with gr.Blocks(title=APP_TITLE, css=APP_CSS) as demo:
+    gr.Markdown(
+        f"""
+# {APP_TITLE}
+
+Audio and text demo for [`{MODEL_REPO_ID}`]({MODEL_CARD_URL}).
+Audex-2B supports audio understanding, speech recognition, speech translation,
+and text reasoning. See the [technical report]({PAPER_URL}).
+"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1, min_width=380):
+            task = gr.Radio(
+                choices=_task_choices(),
+                value="audio-understanding",
+                label="Task",
+            )
+            audio = gr.Audio(
+                label="Record or upload audio",
+                sources=["microphone", "upload"],
+                type="filepath",
+                visible=True,
+            )
+            audio_url = gr.Textbox(
+                label="HTTPS audio URL",
+                placeholder="https://example.com/audio.wav",
+                visible=True,
+            )
+            prompt = gr.Textbox(
+                label="Prompt",
+                value=TASKS["audio-understanding"]["prompt"],
+                lines=4,
+                max_lines=10,
+            )
+            with gr.Row():
+                reasoning = gr.Checkbox(label="Thinking mode", value=False)
+                top_k = gr.Number(label="Top-k", value=0, precision=0, minimum=0)
+            with gr.Row():
+                max_new_tokens = gr.Slider(
+                    minimum=32,
+                    maximum=2048,
+                    step=32,
+                    value=512,
+                    label="Max new tokens",
+                )
+            with gr.Row():
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.5,
+                    step=0.1,
+                    value=0.7,
+                    label="Temperature",
+                )
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.05,
+                    value=0.9,
+                    label="Top-p",
+                )
+            run = gr.Button("Generate", variant="primary")
+
+        with gr.Column(scale=1, min_width=420):
+            answer = gr.Textbox(label="Answer", lines=12, show_copy_button=True)
+            thinking = gr.Textbox(label="Thinking", lines=8, show_copy_button=True)
+            summary = gr.Markdown(
+                "The model snapshot starts predownloading with the Space owner's `HF_TOKEN` when the Space starts.",
+                elem_id="run_summary",
+            )
+
+    gr.Examples(
+        examples=_example_samples(),
+        inputs=[
+            task,
+            audio,
+            audio_url,
+            prompt,
+            reasoning,
+            max_new_tokens,
+            temperature,
+            top_p,
+            top_k,
+        ],
+        label="Examples",
+    )
+
+    task.change(
+        task_defaults,
+        inputs=task,
+        outputs=[prompt, reasoning, temperature, top_p, audio, audio_url],
+    )
+    run.click(
+        generate,
+        inputs=[
+            task,
+            audio,
+            audio_url,
+            prompt,
+            reasoning,
+            max_new_tokens,
+            temperature,
+            top_p,
+            top_k,
+        ],
+        outputs=[thinking, answer, summary],
+        api_name="generate",
+        concurrency_limit=1,
+    )
+
+demo.queue(default_concurrency_limit=1)
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=[str(ASSET_DIR)])

requirements.txt

@@ -0,0 +1,13 @@
+gradio==6.19.0
+spaces>=0.50.0
+torch==2.9.1
+transformers>=5.0.0
+accelerate>=1.12.0
+safetensors>=0.7.0
+huggingface_hub[hf_xet]>=1.22.0
+hf-transfer>=0.1.4
+librosa>=0.11.0
+soundfile>=0.13.0
+requests>=2.32.0
+numpy>=2.2.0
+