Wrap inference in torch.no_grad() to avoid autograd-graph memory pressure

app.py

@@ -1,34 +1,53 @@
-"""Audio Interaction Model — Gradio demo.
+"""Audio Interaction Model — Gradio demo on ZeroGPU.
 
-Takes an audio clip (microphone or file) and an optional text instruction,
-streams a text response from the AudioInteraction model.
+Upload one or more audio clips (or record from microphone) and the model
+will listen, decide whether to respond, and generate a text answer.
+
+Based on the paper: "Audio Interaction Model" (arXiv:2606.05121)
+Model weights: https://huggingface.co/zhifeixie/AudioInteraction
 """
 
 import os
+
+# Avoid allocator fragmentation under memory spikes
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+# Prevent numba from initializing CUDA before spaces can patch torch
+os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
+
+# Install openai-whisper with --no-deps to avoid it downgrading torch.
+# Its real runtime deps (tiktoken, more-itertools, tqdm, scipy) are in requirements.txt.
+import subprocess, sys
+try:
+    import whisper
+except ImportError:
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install", "--no-deps", "openai-whisper"],
+        check=True,
+    )
 
-import spaces  # MUST come before torch / any CUDA-touching import
+import spaces  # MUST come before any CUDA-touching import
 import torch
 
-# Monkey-patch torch.load to allow loading the audio encoder checkpoint
-# (which pickles numpy/object globals and fails under torch 2.6+ weights_only=True default)
+# Patch torch.load to disable weights_only (the audio encoder checkpoint
+# was saved with older torch that pickles numpy/object globals).
 _orig_torch_load = torch.load
-def _patched_torch_load(*args, **kwargs):
-    return _orig_torch_load(*args, **{**kwargs, "weights_only": kwargs.get("weights_only", False)})
-torch.load = _patched_torch_load
+torch.load = lambda *a, **k: _orig_torch_load(*a, **{**k, "weights_only": k.get("weights_only", False)})
 
-import sys
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import json
+import sys as _sys
+import time
+from pathlib import Path
+from typing import List, Optional, Tuple
 
-import lightning as L
-import numpy as np
 import gradio as gr
-import whisper
-import tempfile
-import time
+import numpy as np
+from huggingface_hub import snapshot_download
 from transformers import AutoConfig
 from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
-from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+# Add the app root to PYTHONPATH so `src.*` and `utils` imports resolve
+_sys.path.insert(0, str(Path(__file__).resolve().parent))
 
 from src.audiointeraction.dataset.TOKENS import (
     ASSISTANT, AUDIO_BEGIN, ENGLISH, KEEP_SILENCE, ONLINE, PAD,
@@ -36,327 +55,335 @@ from src.audiointeraction.dataset.TOKENS import (
     HAPPY, SAD, ANGRY, SURPRISE, NORMAL, URGENT,
 )
 from src.audiointeraction.generate.base import (
-    AUDIO_TOKENS_PER_CHUNK, encode_audio_chunks, encode_silence_chunks, sample
+    AUDIO_TOKENS_PER_CHUNK,
+    sample,
+    encode_audio_chunks,
 )
-from src.audiointeraction.model import GPT, Config
+from src.audiointeraction.model import GPT
+from src.audiointeraction.config import Config
 from src.audiointeraction.tokenizer import Tokenizer
-from src.audiointeraction.utils import get_default_supported_precision
-
-from safetensors.torch import load_file
-from pathlib import Path
-import json
 
-# ── Constants ──────────────────────────────────────────────
+MODEL_REPO = "zhifeixie/AudioInteraction"
+EMOTION_EMOJI = {
+    HAPPY: "😊", SAD: "😢", ANGRY: "😠", SURPRISE: "😲",
+    NORMAL: "😐", URGENT: "⚠️",
+}
 
-MODEL_ID = "zhifeixie/AudioInteraction"
 SYSTEM_PROMPT = (
     "You are a helpful assistant. When there is no user text, if the audio contains a question, "
     "please answer it. If it is a sound effect, determine based on the sound whether help is needed."
 )
 
-EMOTION_EMOJI = {
-    HAPPY:    "😊",
-    SAD:      "😢",
-    ANGRY:    "😠",
-    SURPRISE: "😲",
-    NORMAL:   "😐",
-    URGENT:   "⚠️",
-}
-
 
-# ── Model loading (module scope, eager) ────────────────────
+# ---------------------------------------------------------------------------
+# Model loading (ZeroGPU-safe: no Lightning Fabric, no fabric.setup/init_module)
+# ---------------------------------------------------------------------------
 
-def _resolve_checkpoint_paths(checkpoint_dir: str):
-    ckpt = Path(checkpoint_dir)
-    return (
-        str(ckpt),
-        str(ckpt),
-        str(ckpt / "qwen25OmniConfig"),
-        str(ckpt / "audiointeraction_ChunkwisedEncoder.pth"),
+def _download_checkpoints() -> str:
+    """Download model weights from HuggingFace Hub and return local path."""
+    local_dir = "/home/user/checkpoints"
+    os.makedirs(local_dir, exist_ok=True)
+    snapshot_download(
+        repo_id=MODEL_REPO,
+        repo_type="model",
+        local_dir=local_dir,
     )
+    return local_dir
 
 
-def _load_model(model_config_dir, trained_checkpoint):
-    """Load GPT from sharded safetensors."""
-    config = Config.from_file(Path(model_config_dir) / "model_config.yaml")
-    model = GPT(config)
-    model.max_seq_length = config.block_size
+CHECKPOINT_DIR = _download_checkpoints()
+
+ckpt = Path(CHECKPOINT_DIR)
+MODEL_CONFIG_DIR = str(ckpt)
+TRAINED_CHECKPOINT = str(ckpt)
+QWEN_OMNI_CKPT = str(ckpt / "qwen25OmniConfig")
+AUDIO_TOWER_CKPT = str(ckpt / "audiointeraction_ChunkwisedEncoder.pth")
+
 
-    checkpoint_dir = Path(trained_checkpoint)
-    index_path = checkpoint_dir / "model.safetensors.index.json"
-    if not index_path.is_file():
-        raise FileNotFoundError(f"No model.safetensors.index.json under {checkpoint_dir}")
+def _load_model():
+    """Load the GPT model from sharded safetensors (ZeroGPU-safe, no Fabric)."""
+    config = Config.from_file(Path(MODEL_CONFIG_DIR) / "model_config.yaml")
+    # Force SDPA attention (flash_attn not available / sm_120 incompatible with FA3)
+    config.use_flash_attention = False
 
+    model = GPT(config)
+
+    index_path = Path(TRAINED_CHECKPOINT) / "model.safetensors.index.json"
     with open(index_path) as f:
         index = json.load(f)
     shard_files = sorted(set(index["weight_map"].values()))
-
     state_dict = {}
     for shard in shard_files:
-        state_dict.update(load_file(str(checkpoint_dir / shard), device="cpu"))
-
+        state_dict.update(load_file(str(Path(TRAINED_CHECKPOINT) / shard), device="cpu"))
     missing, unexpected = model.load_state_dict(state_dict, strict=True)
     if missing or unexpected:
-        print(f"[load_model] missing={missing[:3]}… unexpected={unexpected[:3]}…")
-    return model
+        print(f"[load_model] missing={missing[:3]}... unexpected={unexpected[:3]}...")
+
+    # Use bf16 precision for inference
+    model = model.to(torch.bfloat16)
+    # .to("cuda") is intercepted by ZeroGPU's spaces hijack — safe at module scope
+    model.to("cuda")
+    model.eval()
+    return model, config
 
 
-def _load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
-    cfg = AutoConfig.from_pretrained(qwen_omni_ckpt)
+def _load_audio_encoder():
+    """Load the Qwen2.5-Omni audio encoder."""
+    cfg = AutoConfig.from_pretrained(QWEN_OMNI_CKPT)
     audio_cfg = cfg.thinker_config.audio_config
     encoder = Qwen2_5OmniAudioEncoder._from_config(audio_cfg)
-    state_dict = torch.load(audio_tower_ckpt, map_location=device)
+    state_dict = torch.load(AUDIO_TOWER_CKPT, map_location="cpu")
     encoder.load_state_dict(state_dict)
-    encoder.to(device).requires_grad_(False).eval()
+    encoder = encoder.to(torch.bfloat16)
+    # .to("cuda") is intercepted by ZeroGPU's spaces hijack
+    encoder.to("cuda").requires_grad_(False).eval()
     return encoder
 
 
-# Download model weights at module scope
-print("Downloading model weights from HuggingFace...")
-ckpt_dir = snapshot_download(
-    repo_id=MODEL_ID,
-    repo_type="model",
-    local_dir="./checkpoints",
-    resume_download=True,
-)
-print(f"Model downloaded to {ckpt_dir}")
-
-model_config_dir, trained_checkpoint, qwen_omni_ckpt, audio_tower_ckpt = \
-    _resolve_checkpoint_paths(ckpt_dir)
-
-set_seed_called = False
-# We don't use Fabric on ZeroGPU — load directly and move to cuda
-device = "cuda"
-precision = "bf16-true"
-
-print("Loading language model...")
-model = _load_model(model_config_dir, trained_checkpoint)
-model = model.to(device).to(torch.bfloat16)
-model.eval()
-
-print("Loading audio encoder...")
-audio_encoder = _load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device)
-audio_encoder = audio_encoder.to(torch.bfloat16)
-
-tokenizer = Tokenizer(model_config_dir)
+print("Loading model and audio encoder...")
+model, model_config = _load_model()
+audio_encoder = _load_audio_encoder()
+tokenizer = Tokenizer(MODEL_CONFIG_DIR)
 
 system_ids = tokenizer.encode(SYSTEM_PROMPT).cpu().tolist()
-prefix_ids = torch.LongTensor(
-    [ONLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
-).to(device)
+_prefix_token_list = [ONLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
 
-print("Model loaded successfully!")
+# Set up KV cache (no fabric.init_tensor needed)
+model.set_kv_cache(batch_size=1)
+model.eval()
 
+print("Model and audio encoder loaded successfully!")
 
-# ── Inference ──────────────────────────────────────────────
 
-@spaces.GPU(duration=180)
-def interact(audio_file, text_instruction=""):
-    """Process an audio clip and optional text instruction, return the model's text response.
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
 
-    Args:
-        audio_file: Path to an audio file (wav, mp3, m4a, etc.).
-        text_instruction: Optional text instruction to guide the model's response.
+@spaces.GPU(duration=120)
+@torch.no_grad()
+def interact_with_audio(
+    audio_file: str,
+    text_instruction: str = "",
+) -> str:
+    """Listen to an audio clip and generate a text response.
 
-    Returns:
-        A tuple of (full_response_text, status_text) where full_response_text
-        is the model's text response and status_text provides processing info.
+    Upload an audio file (wav/mp3/m4a/flac/ogg) and optionally add a text
+    instruction. The model will listen to the audio and decide whether to
+    respond — it may answer a question, transcribe speech, translate,
+    describe a sound, or stay silent if the audio doesn't need a response.
     """
     if audio_file is None:
-        return "Please provide an audio file.", "No audio input"
+        return "Please provide an audio file."
+
+    device = next(model.parameters()).device
 
-    t0 = time.perf_counter()
+    # Build prefix tokens on the correct device
+    prefix_ids = torch.LongTensor(_prefix_token_list).to(device)
 
-    # Set up KV cache
-    model.set_kv_cache(batch_size=1)
-    model.max_seq_length = model.config.block_size
+    # Reset KV cache for a fresh session
+    model.clear_kv_cache()
+    model.set_kv_cache(batch_size=1, device=device, dtype=torch.bfloat16)
 
-    # Build prefix
-    token = prefix_ids.clone()
+    token = prefix_ids
     input_pos = torch.arange(0, prefix_ids.size(0), device=device, dtype=torch.int64)
-    input_pos_maxp1 = torch.tensor(prefix_ids.size(0), device=device)
+    prompt_size = prefix_ids.size(0)
 
-    # If text instruction provided, append it
-    if text_instruction.strip():
-        text_ids = tokenizer.encode(text_instruction).cpu().tolist()
-        text_tokens = torch.LongTensor(
-            [TEXT_BEGIN] + text_ids + [TEXT_END]
-        ).to(device)
-        token = torch.cat([token, text_tokens])
-        input_pos = torch.cat([input_pos, torch.arange(
-            prefix_ids.size(0), prefix_ids.size(0) + len(text_tokens),
-            device=device, dtype=torch.int64
-        )])
-
-    # Encode audio
+    # Encode the audio file into chunks
     audio_chunks = encode_audio_chunks(audio_file, audio_encoder, device)
-    n_chunks = len(audio_chunks)
 
-    # Run streaming inference
-    response_parts = []
-    current_text = []
+    # If there's a text instruction, prepend it as text tokens
+    if text_instruction.strip():
+        # Insert text instruction after the system prompt
+        instruction_ids = tokenizer.encode(text_instruction.strip()).cpu().tolist()
+        text_tokens = [TEXT_BEGIN] + instruction_ids + [TEXT_END]
+        text_tensor = torch.LongTensor(text_tokens).to(device)
+        token = torch.cat([token, text_tensor])
+
+    turns: List[List[int]] = []
     listening = True
     audio_idx = -1
-    emotion_tag = ""
-    max_tokens = 4096
-
-    try:
-        with torch.inference_mode():
-            for _ in range(max_tokens - input_pos.numel()):
-                if listening:
-                    audio_idx += 1
-                    if audio_idx >= n_chunks:
-                        break
-                    # Append [AUDIO_BEGIN, PAD*N, ASSISTANT]
-                    new_tokens = torch.LongTensor(
-                        [AUDIO_BEGIN] + [PAD] * AUDIO_TOKENS_PER_CHUNK + [ASSISTANT]
-                    ).to(device)
-                    new_positions = input_pos[-1] + torch.arange(1, len(new_tokens) + 1, device=device)
-                    token = torch.cat([token, new_tokens])
-                    input_pos = torch.cat([input_pos, new_positions])
-                    input_pos_maxp1 = input_pos_maxp1 + len(new_tokens)
-
-                    logits = model(
-                        token.view(1, -1), None, 1,
-                        audio_chunks[audio_idx].to(device).to(torch.bfloat16),
-                        input_pos,
-                        input_pos_maxp1=input_pos_maxp1,
-                        audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
-                    )
-                else:
-                    logits = model(
-                        token.view(1, -1), None, 1,
-                        None,
-                        input_pos,
-                        input_pos_maxp1=input_pos_maxp1,
-                        audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
-                    )
-
-                token = sample(logits, temperature=0.0, top_p=0.0).to(torch.int64)
-                int_token = token.item()
-                input_pos = input_pos[-1].unsqueeze(0).add_(1)
-                input_pos_maxp1 = input_pos_maxp1 + 1
-
-                if listening:
-                    if int_token == TEXT_BEGIN:
-                        listening = False
-                        current_text = [int_token]
-                    elif int_token == KEEP_SILENCE:
-                        pass  # model stays silent
+    current_turn: List[int] = []
+    text_started = False
+    emotion_prefix = ""
+
+    max_steps = 4096
+
+    for _ in range(max_steps - input_pos.numel()):
+        if listening:
+            audio_idx += 1
+            if audio_idx >= len(audio_chunks):
+                break
+            # Append listening block: [AUDIO_BEGIN, PAD*N, ASSISTANT]
+            new_tokens = torch.LongTensor(
+                [AUDIO_BEGIN] + [PAD] * AUDIO_TOKENS_PER_CHUNK + [ASSISTANT]
+            ).to(device)
+            new_positions = input_pos[-1] + torch.arange(1, len(new_tokens) + 1, device=device)
+            token = torch.cat([token, new_tokens])
+            input_pos = torch.cat([input_pos, new_positions])
+
+            logits = model(
+                token.view(1, -1), None, 1,
+                audio_chunks[audio_idx].to(device),
+                input_pos,
+                input_pos_maxp1=None,
+                audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
+            )
+        else:
+            logits = model(
+                token.view(1, -1), None, 1,
+                None,
+                input_pos,
+                input_pos_maxp1=None,
+                audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
+            )
+
+        token = sample(logits, temperature=0.0, top_p=0.0).to(torch.int64)
+        int_token = token.item()
+        # Advance input_pos by one -- REPLACE, not append: once the KV cache
+        # holds everything before it, the next forward call only needs the
+        # position of the single new token (see the reference `_advance_one`
+        # in src/audiointeraction/generate/base.py, which returns just
+        # `new_pos`). Concatenating here made input_pos grow every step and
+        # blow past the cache length.
+        new_pos = input_pos[-1].unsqueeze(0).add_(1)
+        input_pos = new_pos
+
+        if listening:
+            if int_token == TEXT_BEGIN:
+                listening = False
+                current_turn = [int_token]
+                text_started = False
+                emotion_prefix = ""
+            elif int_token == KEEP_SILENCE:
+                turns.append([int_token])
+            else:
+                # Unexpected token
+                break
+        else:
+            current_turn.append(int_token)
+            if int_token == TEXT_END:
+                turns.append(current_turn)
+                current_turn = []
+                text_started = False
+                listening = True
+            else:
+                n = len(current_turn)
+                if n == 2:
+                    text_started = True
+                    if int_token in EMOTION_EMOJI:
+                        emotion_prefix = EMOTION_EMOJI[int_token] + " "
                     else:
-                        break
-                else:
-                    current_text.append(int_token)
-                    if int_token == TEXT_END:
-                        # Decode the turn (skip TEXT_BEGIN and optional emotion)
-                        text_tokens = current_text[1:-1]  # strip TEXT_BEGIN and TEXT_END
-                        if text_tokens:
-                            # Check if first token is an emotion tag
-                            if text_tokens[0] in EMOTION_EMOJI:
-                                emotion_tag = EMOTION_EMOJI[text_tokens[0]]
-                                text_tokens = text_tokens[1:]
-                            decoded = tokenizer.decode(torch.tensor(text_tokens))
-                            if decoded:
-                                response_parts.append((emotion_tag, decoded))
-                        current_text = []
-                        listening = True
-                    elif len(current_text) >= 2:
-                        # Could be streaming text, but we collect at TEXT_END
-                        pass
-
-        # Handle any incomplete turn
-        if not listening and len(current_text) > 1:
-            text_tokens = current_text[1:]  # skip TEXT_BEGIN
-            if text_tokens[0] in EMOTION_EMOJI if text_tokens else False:
-                emotion_tag = EMOTION_EMOJI[text_tokens[0]]
-                text_tokens = text_tokens[1:]
-            decoded = tokenizer.decode(torch.tensor(text_tokens))
+                        piece = tokenizer.decode(torch.tensor([int_token]))
+                        if piece:
+                            # Start output
+                            pass
+                elif n >= 3:
+                    pass
+
+    # Decode all turns into text
+    responses = []
+    for turn in turns:
+        if turn[0] == KEEP_SILENCE:
+            continue
+        # Extract text tokens: skip TEXT_BEGIN and optional emotion tag
+        text_tokens_list = turn[1:]  # skip TEXT_BEGIN
+        if text_tokens_list and text_tokens_list[-1] == TEXT_END:
+            text_tokens_list = text_tokens_list[:-1]
+        # Check if first token is an emotion tag
+        emotion = ""
+        if text_tokens_list and text_tokens_list[0] in EMOTION_EMOJI:
+            emotion = EMOTION_EMOJI[text_tokens_list[0]] + " "
+            text_tokens_list = text_tokens_list[1:]
+
+        if text_tokens_list:
+            decoded = tokenizer.decode(torch.tensor(text_tokens_list))
             if decoded:
-                response_parts.append((emotion_tag, decoded))
+                responses.append(emotion + decoded)
 
-    finally:
-        model.clear_kv_cache()
+    model.clear_kv_cache()
 
-    elapsed = time.perf_counter() - t0
+    if not responses:
+        return "🔇 The model listened to the audio and chose to stay silent. (This is expected for non-question sounds — the model decides for itself when to speak.)"
 
-    if not response_parts:
-        return "The model listened to the audio but chose to stay silent. (This is expected for non-speech sounds when no response is warranted.)", \
-               f"Processed {n_chunks} audio chunks in {elapsed:.1f}s — model stayed silent"
+    return "\n\n".join(responses)
 
-    # Format response
-    formatted = []
-    for emoji, text in response_parts:
-        if emoji:
-            formatted.append(f"{emoji} {text}")
-        else:
-            formatted.append(text)
-    full_response = "\n".join(formatted)
 
-    status = f"Processed {n_chunks} audio chunks in {elapsed:.1f}s — {len(response_parts)} response(s)"
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+DESCRIPTION = """
+# Audio Interaction Model
 
-    return full_response, status
+This is a demo of the **Audio Interaction Model** from the paper
+[Audio Interaction Model](https://arxiv.org/abs/2606.05121).
 
+The model is an always-on streaming audio language model that listens to
+audio and **decides for itself when to speak**. Upload an audio clip
+(or record from microphone) and the model will:
 
-# ── Gradio UI ──────────────────────────────────────────────
+- 🎧 **Listen** to the audio and understand its content
+- 🧠 **Decide** whether a response is needed
+- 💬 **Respond** with text (transcription, answer, translation, description, etc.)
 
-CSS = """
-#col-container { max-width: 900px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
+The model may also choose to stay silent (⟨Silent⟩) if the audio doesn't
+warrant a response — this is a core feature of the Audio Interaction Model,
+not a bug!
 """
 
-with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
-    gr.Markdown("# 🎙️ Audio Interaction Model")
-    gr.Markdown(
-        "Upload or record an audio clip — the model listens and responds with text. "
-        "This is a demo of [AudioInteraction](https://huggingface.co/zhifeixie/AudioInteraction), "
-        "a streaming audio-language model that perceives audio and decides when to speak."
-    )
+with gr.Blocks(title="Audio Interaction Model") as demo:
+    gr.Markdown(DESCRIPTION)
 
-    with gr.Column(elem_id="col-container"):
-        with gr.Row():
+    with gr.Row():
+        with gr.Column():
             audio_input = gr.Audio(
                 label="Audio Input",
                 type="filepath",
-                sources=["microphone", "upload"],
+                sources=["upload", "microphone"],
             )
+            text_instruction = gr.Textbox(
+                label="Text Instruction (optional)",
+                placeholder="e.g. 'Translate this to English' or leave empty",
+                lines=2,
+            )
+            submit_btn = gr.Button("Interact", variant="primary")
 
-        text_input = gr.Textbox(
-            label="Text Instruction (optional)",
-            placeholder="e.g., 'Translate this to English' or leave blank for auto-response",
-            lines=2,
-        )
-
-        run_btn = gr.Button("Run", variant="primary")
-
-        with gr.Row():
-            response_output = gr.Textbox(
+        with gr.Column():
+            output_text = gr.Textbox(
                 label="Model Response",
-                lines=8,
+                lines=10,
                 interactive=False,
             )
 
-        status_output = gr.Textbox(
-            label="Status",
-            interactive=False,
-        )
-
-    run_btn.click(
-        fn=interact,
-        inputs=[audio_input, text_input],
-        outputs=[response_output, status_output],
+    submit_btn.click(
+        fn=interact_with_audio,
+        inputs=[audio_input, text_instruction],
+        outputs=output_text,
     )
 
     gr.Examples(
         examples=[
-            ["samples/sample01_01.mp3", ""],
-            ["samples/sample02_01.mp3", "Translate this to English"],
-            ["samples/sample03_01.wav", ""],
+            ["sample/01_count_bark/sample01_01.mp3", ""],
+            ["sample/01_count_bark/sample01_02.wav", ""],
+            ["sample/02_translate/sample02_01.mp3", "Translate to English"],
+            ["sample/03_cough_music/sample03_01.wav", ""],
         ],
-        inputs=[audio_input, text_input],
-        outputs=[response_output, status_output],
-        fn=interact,
+        inputs=[audio_input, text_instruction],
+        outputs=output_text,
+        fn=interact_with_audio,
         cache_examples=True,
         cache_mode="lazy",
     )
 
+    gr.Markdown("""
+    ---
+    **Model**: [zhifeixie/AudioInteraction](https://huggingface.co/zhifeixie/AudioInteraction)
+    | **Paper**: [arXiv:2606.05121](https://arxiv.org/abs/2606.05121)
+    | **Code**: [GitHub](https://github.com/xzf-thu/Audio-Interaction)
+
+    The model is a 3B parameter streaming audio language model based on
+    Qwen2.5-Omni. It processes audio in 400ms chunks and generates text
+    responses when it decides intervention is needed.
+    """)
+
 demo.launch(mcp_server=True)