Upload folder using huggingface_hub

.gitattributes

@@ -38,3 +38,4 @@ sample/02_translate/sample02_02.wav filter=lfs diff=lfs merge=lfs -text
 sample/02_translate/sample02_04.wav filter=lfs diff=lfs merge=lfs -text
 sample/03_cough_music/sample03_01.wav filter=lfs diff=lfs merge=lfs -text
 sample/03_cough_music/sample03_03.wav filter=lfs diff=lfs merge=lfs -text
+samples/sample03_01.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,28 +1,19 @@
 ---
 title: Audio Interaction Model
-emoji: 🎧
+emoji: 🎙️
 colorFrom: blue
 colorTo: red
 sdk: gradio
-sdk_version: 6.19.0
+sdk_version: 6.15.1
 app_file: app.py
 short_description: Streaming audio-language model that listens and responds
 python_version: "3.12"
-startup_duration_timeout: 30m
+startup_duration_timeout: 1h
 ---
 
-# Audio Interaction Model
+Audio Interaction Model demo — a streaming audio-language model that perceives audio
+and decides when to speak. Upload or record audio and get a text response.
 
-This Space demos 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 listen, understand, and respond with text.
-
-## Model
-
-- **Weights**: [zhifeixie/AudioInteraction](https://huggingface.co/zhifeixie/AudioInteraction)
-- **Architecture**: 3B parameter GPT based on Qwen2.5-Omni
-- **Paper**: [arXiv:2606.05121](https://arxiv.org/abs/2606.05121)
-- **Code**: [GitHub](https://github.com/xzf-thu/Audio-Interaction)
+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)

app.py

@@ -1,53 +1,34 @@
-"""Audio Interaction Model — Gradio demo on ZeroGPU.
+"""Audio Interaction Model — Gradio demo.
 
-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
+Takes an audio clip (microphone or file) and an optional text instruction,
+streams a text response from the AudioInteraction model.
 """
 
 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 any CUDA-touching import
+import spaces  # MUST come before torch / any CUDA-touching import
 import torch
 
-# Patch torch.load to disable weights_only (the audio encoder checkpoint
-# was saved with older torch that pickles numpy/object globals).
+# 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)
 _orig_torch_load = torch.load
-torch.load = lambda *a, **k: _orig_torch_load(*a, **{**k, "weights_only": k.get("weights_only", False)})
+def _patched_torch_load(*args, **kwargs):
+    return _orig_torch_load(*args, **{**kwargs, "weights_only": kwargs.get("weights_only", False)})
+torch.load = _patched_torch_load
 
-import json
-import sys as _sys
-import time
-from pathlib import Path
-from typing import List, Optional, Tuple
+import sys
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
-import gradio as gr
+import lightning as L
 import numpy as np
-from huggingface_hub import snapshot_download
+import gradio as gr
+import whisper
+import tempfile
+import time
 from transformers import AutoConfig
 from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
-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 huggingface_hub import snapshot_download
 
 from src.audiointeraction.dataset.TOKENS import (
     ASSISTANT, AUDIO_BEGIN, ENGLISH, KEEP_SILENCE, ONLINE, PAD,
@@ -55,334 +36,327 @@ from src.audiointeraction.dataset.TOKENS import (
     HAPPY, SAD, ANGRY, SURPRISE, NORMAL, URGENT,
 )
 from src.audiointeraction.generate.base import (
-    AUDIO_TOKENS_PER_CHUNK,
-    sample,
-    encode_audio_chunks,
+    AUDIO_TOKENS_PER_CHUNK, encode_audio_chunks, encode_silence_chunks, sample
 )
-from src.audiointeraction.model import GPT
-from src.audiointeraction.config import Config
+from src.audiointeraction.model import GPT, Config
 from src.audiointeraction.tokenizer import Tokenizer
+from src.audiointeraction.utils import get_default_supported_precision
 
-MODEL_REPO = "zhifeixie/AudioInteraction"
-EMOTION_EMOJI = {
-    HAPPY: "😊", SAD: "😢", ANGRY: "😠", SURPRISE: "😲",
-    NORMAL: "😐", URGENT: "⚠️",
-}
+from safetensors.torch import load_file
+from pathlib import Path
+import json
 
+# ── Constants ──────────────────────────────────────────────
+
+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 (ZeroGPU-safe: no Lightning Fabric, no fabric.setup/init_module)
-# ---------------------------------------------------------------------------
-
-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
-
-
-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")
+# ── Model loading (module scope, eager) ────────────────────
 
+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 _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
 
+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 = 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}")
 
-    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(Path(TRAINED_CHECKPOINT) / shard), device="cpu"))
+        state_dict.update(load_file(str(checkpoint_dir / 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]}...")
-
-    # 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
+        print(f"[load_model] missing={missing[:3]}… unexpected={unexpected[:3]}…")
+    return model
 
 
-def _load_audio_encoder():
-    """Load the Qwen2.5-Omni audio encoder."""
-    cfg = AutoConfig.from_pretrained(QWEN_OMNI_CKPT)
+def _load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
+    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="cpu")
+    state_dict = torch.load(audio_tower_ckpt, map_location=device)
     encoder.load_state_dict(state_dict)
-    encoder = encoder.to(torch.bfloat16)
-    # .to("cuda") is intercepted by ZeroGPU's spaces hijack
-    encoder.to("cuda").requires_grad_(False).eval()
+    encoder.to(device).requires_grad_(False).eval()
     return encoder
 
 
-print("Loading model and audio encoder...")
-model, model_config = _load_model()
-audio_encoder = _load_audio_encoder()
-tokenizer = Tokenizer(MODEL_CONFIG_DIR)
+# 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}")
 
-system_ids = tokenizer.encode(SYSTEM_PROMPT).cpu().tolist()
-_prefix_token_list = [ONLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
+model_config_dir, trained_checkpoint, qwen_omni_ckpt, audio_tower_ckpt = \
+    _resolve_checkpoint_paths(ckpt_dir)
 
-# Set up KV cache (no fabric.init_tensor needed)
-model.set_kv_cache(batch_size=1)
+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("Model and audio encoder loaded successfully!")
+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)
+
+system_ids = tokenizer.encode(SYSTEM_PROMPT).cpu().tolist()
+prefix_ids = torch.LongTensor(
+    [ONLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
+).to(device)
+
+print("Model loaded successfully!")
 
 
-# ---------------------------------------------------------------------------
-# Inference
-# ---------------------------------------------------------------------------
+# ── Inference ──────────────────────────────────────────────
 
-@spaces.GPU(duration=120)
-def interact_with_audio(
-    audio_file: str,
-    text_instruction: str = "",
-) -> str:
-    """Listen to an audio clip and generate a text response.
+@spaces.GPU(duration=180)
+def interact(audio_file, text_instruction=""):
+    """Process an audio clip and optional text instruction, return the model's text response.
 
-    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.
+    Args:
+        audio_file: Path to an audio file (wav, mp3, m4a, etc.).
+        text_instruction: Optional text instruction to guide the model's 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.
     """
     if audio_file is None:
-        return "Please provide an audio file."
-
-    device = next(model.parameters()).device
+        return "Please provide an audio file.", "No audio input"
 
-    # Build prefix tokens on the correct device
-    prefix_ids = torch.LongTensor(_prefix_token_list).to(device)
+    t0 = time.perf_counter()
 
-    # Reset KV cache for a fresh session
-    model.clear_kv_cache()
-    model.set_kv_cache(batch_size=1, device=device, dtype=torch.bfloat16)
+    # Set up KV cache
+    model.set_kv_cache(batch_size=1)
+    model.max_seq_length = model.config.block_size
 
-    token = prefix_ids
+    # Build prefix
+    token = prefix_ids.clone()
     input_pos = torch.arange(0, prefix_ids.size(0), device=device, dtype=torch.int64)
-    prompt_size = prefix_ids.size(0)
+    input_pos_maxp1 = torch.tensor(prefix_ids.size(0), device=device)
 
-    # Encode the audio file into chunks
-    audio_chunks = encode_audio_chunks(audio_file, audio_encoder, device)
-
-    # If there's a text instruction, prepend it as text tokens
+    # If text instruction provided, append it
     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])
+        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
+    audio_chunks = encode_audio_chunks(audio_file, audio_encoder, device)
+    n_chunks = len(audio_chunks)
 
-    turns: List[List[int]] = []
+    # Run streaming inference
+    response_parts = []
+    current_text = []
     listening = True
     audio_idx = -1
-    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] + " "
+    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
                     else:
-                        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))
+                        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))
             if decoded:
-                responses.append(emotion + decoded)
-
-    model.clear_kv_cache()
+                response_parts.append((emotion_tag, decoded))
 
-    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.)"
+    finally:
+        model.clear_kv_cache()
 
-    return "\n\n".join(responses)
+    elapsed = time.perf_counter() - t0
 
+    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"
 
-# ---------------------------------------------------------------------------
-# Gradio UI
-# ---------------------------------------------------------------------------
+    # 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)
 
-DESCRIPTION = """
-# Audio Interaction Model
+    status = f"Processed {n_chunks} audio chunks in {elapsed:.1f}s — {len(response_parts)} response(s)"
 
-This is a demo of the **Audio Interaction Model** from the paper
-[Audio Interaction Model](https://arxiv.org/abs/2606.05121).
+    return full_response, status
 
-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:
 
-- 🎧 **Listen** to the audio and understand its content
-- 🧠 **Decide** whether a response is needed
-- 💬 **Respond** with text (transcription, answer, translation, description, etc.)
+# ── Gradio UI ──────────────────────────────────────────────
 
-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!
+CSS = """
+#col-container { max-width: 900px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
 """
 
-with gr.Blocks(title="Audio Interaction Model") as demo:
-    gr.Markdown(DESCRIPTION)
+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.Row():
-        with gr.Column():
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
             audio_input = gr.Audio(
                 label="Audio Input",
                 type="filepath",
-                sources=["upload", "microphone"],
-            )
-            text_instruction = gr.Textbox(
-                label="Text Instruction (optional)",
-                placeholder="e.g. 'Translate this to English' or leave empty",
-                lines=2,
+                sources=["microphone", "upload"],
             )
-            submit_btn = gr.Button("Interact", variant="primary")
 
-        with gr.Column():
-            output_text = gr.Textbox(
+        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(
                 label="Model Response",
-                lines=10,
+                lines=8,
                 interactive=False,
             )
 
-    submit_btn.click(
-        fn=interact_with_audio,
-        inputs=[audio_input, text_instruction],
-        outputs=output_text,
+        status_output = gr.Textbox(
+            label="Status",
+            interactive=False,
+        )
+
+    run_btn.click(
+        fn=interact,
+        inputs=[audio_input, text_input],
+        outputs=[response_output, status_output],
     )
 
     gr.Examples(
         examples=[
-            ["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", ""],
+            ["samples/sample01_01.mp3", ""],
+            ["samples/sample02_01.mp3", "Translate this to English"],
+            ["samples/sample03_01.wav", ""],
         ],
-        inputs=[audio_input, text_instruction],
-        outputs=output_text,
-        fn=interact_with_audio,
+        inputs=[audio_input, text_input],
+        outputs=[response_output, status_output],
+        fn=interact,
         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)

requirements.txt

@@ -1,15 +1,10 @@
+transformers==4.57.6
 lightning
-transformers>=5.0.0
+openai-whisper
 safetensors
-tokenizers
 sentencepiece
-torchaudio
-librosa
-numpy
-pyyaml
+tokenizers
 einops
-tiktoken
-more-itertools
-tqdm
-scipy
-numba
+numpy
+torchaudio
+pyyaml

samples/sample03_01.wav

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

src/audiointeraction/dataset/utils/extract_online_feature.py

@@ -0,0 +1,51 @@
+"""Encode raw audio samples with Qwen2.5-Omni's audio tower; save to AudioFeat.pt.
+
+Pure library — paths come in as arguments. The audio encoder is loaded once
+per (qwen_omni_ckpt, audio_tower_ckpt, device) tuple and cached for reuse.
+"""
+
+import os
+
+import numpy as np
+import torch
+import whisper
+from transformers import AutoConfig, Qwen2_5OmniForConditionalGeneration
+
+
+_encoder_cache = {}  # keyed by (qwen_omni_ckpt, audio_tower_ckpt, str(device))
+
+
+def _load_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
+    cfg = AutoConfig.from_pretrained(qwen_omni_ckpt)
+    enc = Qwen2_5OmniForConditionalGeneration._from_config(cfg).thinker.audio_tower
+    enc.load_state_dict(torch.load(audio_tower_ckpt, map_location=device))
+    return enc.to(device).requires_grad_(False).eval()
+
+
+def _split_into_chunks(n, chunk_size):
+    chunks = [chunk_size] * (n // chunk_size)
+    if n % chunk_size:
+        chunks.append(n % chunk_size)
+    return chunks
+
+
+def extract_audio_features(audio_samples, save_dir, *,
+                           qwen_omni_ckpt, audio_tower_ckpt, device="cuda"):
+    """Encode raw audio samples; save the feature tensor to `<save_dir>/AudioFeat.pt`."""
+    key = (qwen_omni_ckpt, audio_tower_ckpt, str(device))
+    if key not in _encoder_cache:
+        _encoder_cache[key] = _load_encoder(qwen_omni_ckpt, audio_tower_ckpt, device)
+    encoder = _encoder_cache[key]
+
+    mel = whisper.log_mel_spectrogram(np.array(audio_samples, dtype=np.float32), n_mels=128)
+    len_feature = mel.shape[1]
+
+    with torch.no_grad():
+        feat = encoder(
+            torch.tensor(mel).to(device),
+            torch.tensor(_split_into_chunks(len_feature, 40)).to(device),
+            torch.tensor((len_feature - 1) // 2 + 1).to(device),
+        ).last_hidden_state
+
+    os.makedirs(save_dir, exist_ok=True)
+    torch.save(feat.detach().cpu(), os.path.join(save_dir, "AudioFeat.pt"))

src/audiointeraction/dataset/utils/load_audio.py

@@ -0,0 +1,56 @@
+"""
+Shared audio loading helpers for the dataset preprocessing scripts.
+
+The audio_tower in Qwen2.5-Omni does two stride-2 conv downsamples, so a mel
+spectrogram of length L produces ((L-1)//2 + 1 - 2)//2 + 1 encoder-output
+frames. We refer to these encoder-output frames just as "frames" throughout.
+"""
+
+import librosa
+import numpy as np
+import whisper
+
+
+# Number of raw audio samples (@ 16 kHz) per encoder-output frame.
+# 16000 Hz * 40 ms/frame = 640 samples / frame.
+SAMPLES_PER_FRAME = 640
+
+
+def _count_output_lengths(input_lengths: int):
+    """Two-step downsample matching the audio_tower's conv stack."""
+    input_lengths = (input_lengths - 1) // 2 + 1
+    output_lengths = (input_lengths - 2) // 2 + 1
+    return input_lengths, output_lengths
+
+
+def _load_mel(audio_path: str, max_seconds: float = None):
+    """Load audio @ 16 kHz, compute log-mel + downsampled lengths.
+
+    Args:
+        audio_path: path to audio file readable by librosa.
+        max_seconds: optional truncation of the raw audio (offline samples
+            cap at 20 s; online uses the full clip).
+
+    Returns:
+        audio        : list[float], raw audio samples (post-truncation, pre-mel-padding).
+        mel          : torch.Tensor of shape (128, len_feature), log-mel spectrogram.
+        len_feature  : int, mel.shape[1] (the mel time axis length).
+        input_len    : int, length after the first conv downsample.
+        output_len   : int, length after both conv downsamples (encoder frames).
+    """
+    audio_np, _ = librosa.load(audio_path, sr=16000)
+    audio = audio_np.tolist()
+
+    if max_seconds is not None:
+        audio = audio[: int(max_seconds * 16000)]
+
+    # Pad to a 160-sample multiple for whisper's mel-hop alignment.
+    audio_for_mel = (
+        audio if len(audio) % 160 == 0
+        else audio + [0] * (160 - len(audio) % 160)
+    )
+    mel = whisper.log_mel_spectrogram(np.array(audio_for_mel, dtype=np.float32), n_mels=128)
+    len_feature = mel.shape[1]
+    input_len, output_len = _count_output_lengths(len_feature)
+
+    return audio, mel, len_feature, input_len, output_len