[OmniVoice] Add two-stage TTS serving support (#2463)

Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-04-03 20:41:20 +00:00
1 parent 10db95f9a9
commit 0e83ebe1d4
@@ -320,6 +320,29 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

+  - label: "OmniVoice E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_omnivoice.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
  - label: "Voxtral-TTS E2E Test"
    timeout_in_minutes: 20
    depends_on: upload-ready-pipeline
@@ -0,0 +1,73 @@
+# OmniVoice
+
+This directory contains an offline demo for running OmniVoice TTS models with vLLM Omni. It generates speech from text and saves WAV files locally.
+
+## Model Overview
+
+[OmniVoice](https://huggingface.co/k2-fsa/OmniVoice) is a zero-shot multilingual TTS model supporting 600+ languages. It uses a diffusion language model (Qwen3-0.6B backbone) with iterative masked unmasking to generate speech.
+
+Three inference modes are supported:
+
+- **Auto Voice**: Generate speech without any reference — the model picks a voice automatically.
+- **Voice Clone**: Clone a voice from a reference audio + transcription.
+- **Voice Design**: Control voice style via natural language instruction (e.g., "female, low pitch, british accent").
+
+## Setup
+
+Ensure the model is downloaded:
+
+```bash
+huggingface-cli download k2-fsa/OmniVoice
+```
+
+> **Note:** Voice cloning requires `transformers>=5.3.0` for `HiggsAudioV2TokenizerModel`. Auto voice and voice design modes work with `transformers>=4.57.0`.
+
+## Quick Start
+
+Auto voice (text only):
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice --text "Hello, this is a test."
+```
+
+Voice design (with style instruction):
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice \
+    --text "Hello, this is a test." \
+    --instruct "female, low pitch, british accent"
+```
+
+Voice clone (with reference audio):
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice \
+    --text "Hello, this is a test." \
+    --ref-audio ref.wav \
+    --ref-text "This is the reference transcription."
+```
+
+## Language Support
+
+Specify a language for improved quality:
+
+```bash
+python end2end.py --model k2-fsa/OmniVoice \
+    --text "你好，这是一个测试。" \
+    --lang zh
+```
+
+## Architecture
+
+OmniVoice uses a two-stage pipeline:
+
+- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text.
+- **Stage 1 (Decoder)**: HiggsAudioV2 RVQ quantizer + DAC acoustic decoder. Converts tokens to 24kHz waveform.
+
+Both stages use `GPUGenerationWorker` with `OmniGenerationScheduler`.
+
+## Notes
+
+- Output audio is saved to `output.wav` by default. Use `--output` to change the path.
+- The model estimates duration from text automatically via `RuleDurationEstimator`.
+- Use `--stage-init-timeout` to increase the stage initialization timeout for first-time model downloads.
@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end OmniVoice TTS inference via vLLM-Omni.
+
+Supports:
+- Auto voice mode: text only → generated speech
+- Voice cloning mode: text + reference audio → cloned voice speech
+
+Usage:
+    # Auto voice
+    python end2end.py --model k2-fsa/OmniVoice --text "Hello world"
+
+    # Voice cloning
+    python end2end.py --model k2-fsa/OmniVoice --text "Hello" \
+        --ref-audio ref.wav --ref-text "reference transcription"
+"""
+
+import argparse
+import os
+
+import numpy as np
+import soundfile as sf
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def run_e2e():
+    parser = argparse.ArgumentParser(description="OmniVoice E2E TTS inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="k2-fsa/OmniVoice",
+        help="Model name or path (HuggingFace or local)",
+    )
+    parser.add_argument(
+        "--stage-config",
+        type=str,
+        default="vllm_omni/model_executor/stage_configs/omnivoice.yaml",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="Hello, this is a test of the OmniVoice text to speech system.",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio for voice cloning (WAV file)",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcription of reference audio",
+    )
+    parser.add_argument(
+        "--lang",
+        type=str,
+        default=None,
+        help="Language code (e.g., 'en', 'zh')",
+    )
+    parser.add_argument(
+        "--instruct",
+        type=str,
+        default=None,
+        help="Voice design instruction (e.g., 'female, low pitch, british accent')",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output.wav",
+        help="Output audio file path",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Stage initialization timeout in seconds",
+    )
+    args = parser.parse_args()
+
+    if not os.path.exists(args.stage_config):
+        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
+
+    print(f"Initializing OmniVoice with model={args.model}")
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_config,
+        trust_remote_code=True,
+        log_stats=True,
+    )
+
+    print("Model initialized. Preparing inputs...")
+
+    # Build prompt
+    mm_processor_kwargs = {}
+    multi_modal_data = {}
+
+    if args.ref_audio:
+        if not os.path.exists(args.ref_audio):
+            raise FileNotFoundError(f"Reference audio not found: {args.ref_audio}")
+
+        import librosa
+
+        audio_signal, sr = librosa.load(args.ref_audio, sr=None)
+        multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr)
+        mm_processor_kwargs["ref_text"] = args.ref_text or ""
+        mm_processor_kwargs["sample_rate"] = sr
+
+    if args.lang:
+        mm_processor_kwargs["lang"] = args.lang
+    if args.instruct:
+        mm_processor_kwargs["instruct"] = args.instruct
+
+    prompts = {"prompt": args.text}
+    if multi_modal_data:
+        prompts["multi_modal_data"] = multi_modal_data
+    if mm_processor_kwargs:
+        prompts["mm_processor_kwargs"] = mm_processor_kwargs
+
+    sampling_params_list = [OmniDiffusionSamplingParams()]
+
+    print(f"Generating speech for: {args.text}")
+
+    outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
+
+    print(f"Received {len(outputs)} outputs.")
+    for i, output in enumerate(outputs):
+        try:
+            ro = output.request_output
+            if ro is None:
+                print("No request_output found.")
+                continue
+
+            mm = getattr(ro, "multimodal_output", None)
+            if not mm and ro.outputs:
+                mm = getattr(ro.outputs[0], "multimodal_output", None)
+
+            if mm:
+                print(f"Multimodal output keys: {mm.keys()}")
+                if "audio" in mm:
+                    audio_out = mm["audio"]
+                    sr = mm.get("sr", 24000)
+                    if isinstance(audio_out, np.ndarray):
+                        audio_np = audio_out
+                    else:
+                        audio_np = audio_out.cpu().numpy().squeeze()
+                    out_path = args.output if i == 0 else f"output_{i}.wav"
+                    sf.write(out_path, audio_np, sr)
+                    print(f"Saved audio to {out_path} ({sr}Hz, {len(audio_np) / sr:.2f}s)")
+            else:
+                print("No multimodal output found.")
+        except Exception as e:
+            print(f"Error inspecting output: {e}")
+
+    omni.close()
+    print("Done.")
+
+
+if __name__ == "__main__":
+    run_e2e()
@@ -0,0 +1,131 @@
+# OmniVoice
+
+## Model Overview
+
+| Model | Description |
+|-------|-------------|
+| `k2-fsa/OmniVoice` | Zero-shot multilingual TTS (600+ languages) with diffusion language model (Qwen3-0.6B backbone) |
+
+> **Note:** Requires `transformers>=5.3.0` for voice cloning (HiggsAudioV2 tokenizer). Auto voice and voice design work with `transformers>=4.57.0`.
+
+## Launch the Server
+
+```bash
+vllm serve k2-fsa/OmniVoice \
+    --omni \
+    --port 8091 \
+    --trust-remote-code
+```
+
+Or use the convenience script:
+
+```bash
+./run_server.sh
+```
+
+## Send TTS Request
+
+### Using curl
+
+```bash
+# Basic TTS (auto voice)
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "response_format": "wav"
+    }' --output output.wav
+```
+
+### Using Python
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+### Using OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+response = client.audio.speech.create(
+    model="k2-fsa/OmniVoice",
+    voice="default",
+    input="Hello, how are you?",
+)
+
+response.stream_to_file("output.wav")
+```
+
+### Using the CLI Client
+
+```bash
+cd examples/online_serving/omnivoice
+
+# Basic TTS
+python speech_client.py --text "Hello, how are you?"
+
+# Specify language for improved quality
+python speech_client.py --text "Bonjour, comment allez-vous?" --language French
+```
+
+The CLI client supports:
+
+- `--api-base`: API base URL (default: `http://localhost:8091`)
+- `--model` (or `-m`): Model name (default: `k2-fsa/OmniVoice`)
+- `--text`: Text to synthesize (required)
+- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
+- `--language`: Language hint (default: Auto)
+- `--output` (or `-o`): Output file path (default: `omnivoice_output.wav`)
+
+## Inference Modes
+
+OmniVoice supports three inference modes. Currently, **auto voice** is supported through the online Speech API. Voice cloning and voice design are available via offline inference (see `examples/offline_inference/omnivoice/`).
+
+| Mode | Description | Online API | Offline |
+|------|-------------|:----------:|:-------:|
+| Auto Voice | Generate speech without reference | Yes | Yes |
+| Voice Clone | Clone from reference audio + transcript | - | Yes |
+| Voice Design | Control style via natural language instruction | - | Yes |
+
+## Architecture
+
+OmniVoice uses a single-stage diffusion pipeline:
+
+- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative masked unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text, then decodes to 24kHz waveform via HiggsAudioV2 RVQ quantizer + DAC acoustic decoder.
+
+## API Parameters
+
+OmniVoice uses the standard `/v1/audio/speech` endpoint. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full documentation.
+
+Key parameters:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text to synthesize |
+| `voice` | string | "default" | Voice name |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+
+## Troubleshooting
+
+1. **TTS model did not produce audio output**: Ensure the model is fully downloaded (`huggingface-cli download k2-fsa/OmniVoice`)
+2. **Connection refused**: Make sure the server is running on the correct port
+3. **Out of memory**: Reduce `--gpu-memory-utilization` (default stage config uses 0.5)
+4. **Slow first request**: The model performs warmup on first inference; subsequent requests are faster
@@ -0,0 +1,19 @@
+#!/bin/bash
+# Launch vLLM-Omni server for OmniVoice TTS
+#
+# Usage:
+#   ./run_server.sh
+#   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-k2-fsa/OmniVoice}"
+PORT="${PORT:-8091}"
+
+echo "Starting OmniVoice server with model: $MODEL"
+
+vllm serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --trust-remote-code \
+    --omni
@@ -0,0 +1,84 @@
+"""Client for OmniVoice TTS via /v1/audio/speech endpoint.
+
+Examples:
+    # Basic TTS (auto voice)
+    python speech_client.py --text "Hello, how are you?"
+
+    # Specify language
+    python speech_client.py --text "Bonjour, comment allez-vous?" --language French
+"""
+
+import argparse
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8091"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def run_tts(args) -> None:
+    """Generate speech via /v1/audio/speech API."""
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "voice": "default",
+        "response_format": args.response_format,
+    }
+
+    if args.language:
+        payload["language"] = args.language
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    if args.language:
+        print(f"Language: {args.language}")
+    print("Generating audio...")
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    with httpx.Client(timeout=300.0) as client:
+        response = client.post(api_url, json=payload, headers=headers)
+
+    if response.status_code != 200:
+        print(f"Error: {response.status_code}")
+        print(response.text)
+        return
+
+    try:
+        text = response.content.decode("utf-8")
+        if text.startswith('{"error"'):
+            print(f"Error: {text}")
+            return
+    except UnicodeDecodeError:
+        pass
+
+    output_path = args.output or "omnivoice_output.wav"
+    with open(output_path, "wb") as f:
+        f.write(response.content)
+    print(f"Audio saved to: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="OmniVoice TTS client")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
+    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
+    parser.add_argument("--model", "-m", default="k2-fsa/OmniVoice", help="Model name")
+    parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument("--language", default=None, help="Language hint (e.g., English, Chinese, French)")
+    parser.add_argument(
+        "--response-format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio format (default: wav)",
+    )
+    parser.add_argument("--output", "-o", default=None, help="Output file path")
+    args = parser.parse_args()
+    run_tts(args)
+
+
+if __name__ == "__main__":
+    main()
@@ -223,3 +223,4 @@ extend-ignore-identifiers-re = [
 ue = "ue"
 semantics = "semantics"
 fullset = "fullset"
+Vai = "Vai"
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E offline tests for OmniVoice TTS model with text input and audio output.
+
+Uses GPUGenerationWorker for both stages (iterative unmasking + DAC decoder).
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.utils import hardware_test
+
+MODEL = "k2-fsa/OmniVoice"
+
+
+def get_stage_config():
+    return str(
+        Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml"
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_omnivoice_text_to_audio() -> None:
+    """
+    Test OmniVoice text-to-audio generation via offline Omni runner.
+    Deploy Setting: omnivoice.yaml (enforce_eager=true)
+    Input Modal: text
+    Output Modal: audio
+    """
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni = Omni(
+        model=MODEL,
+        stage_configs_path=get_stage_config(),
+        trust_remote_code=True,
+        log_stats=True,
+    )
+
+    try:
+        prompts = {"prompt": "Hello, this is a test for text to audio."}
+
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        sampling_params_list = [OmniDiffusionSamplingParams()]
+
+        outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
+
+        assert len(outputs) > 0, "No outputs generated"
+
+        # Check final output has audio
+        final_output = outputs[-1]
+        ro = final_output.request_output
+        assert ro is not None, "No request_output"
+
+        mm = getattr(ro, "multimodal_output", None)
+        if not mm and ro.outputs:
+            mm = getattr(ro.outputs[0], "multimodal_output", None)
+
+        assert mm is not None, "No multimodal_output"
+        assert "audio" in mm, f"No 'audio' key in multimodal_output: {mm.keys()}"
+
+        audio = mm["audio"]
+        if isinstance(audio, np.ndarray):
+            audio_np = audio
+        else:
+            audio_np = audio.cpu().numpy().squeeze()
+
+        assert audio_np.size > 0, "Audio output is empty"
+        rms = np.sqrt(np.mean(audio_np**2))
+        assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely silence"
+
+        print(f"Generated audio: {len(audio_np) / 24000:.2f}s, rms={rms:.4f}")
+    finally:
+        omni.close()
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for OmniVoice TTS model via /v1/audio/speech endpoint.
+
+Tests verify that the OmniVoice model generates valid audio when
+accessed through the standard OpenAI-compatible speech API.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import httpx
+import pytest
+
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+MODEL = "k2-fsa/OmniVoice"
+
+STAGE_CONFIG = str(
+    Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml"
+)
+EXTRA_ARGS = [
+    "--trust-remote-code",
+    "--disable-log-stats",
+]
+TEST_PARAMS = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=STAGE_CONFIG,
+        server_args=EXTRA_ARGS,
+    )
+]
+
+MIN_AUDIO_BYTES = 5000
+
+
+def make_speech_request(
+    host: str,
+    port: int,
+    text: str,
+    timeout: float = 180.0,
+) -> httpx.Response:
+    """Make a request to the /v1/audio/speech endpoint for OmniVoice."""
+    url = f"http://{host}:{port}/v1/audio/speech"
+    payload = {"input": text}
+
+    with httpx.Client(timeout=timeout) as client:
+        return client.post(url, json=payload)
+
+
+def verify_wav_audio(content: bytes) -> bool:
+    """Verify that content is valid WAV audio data."""
+    if len(content) < 44:
+        return False
+    return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestOmniVoiceTTS:
+    """E2E tests for OmniVoice TTS model."""
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_speech_auto_voice(self, omni_server) -> None:
+        """Test auto voice TTS generation (text only, no reference audio)."""
+        response = make_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            text="Hello, this is a test of the OmniVoice text to speech system.",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert response.headers.get("content-type") == "audio/wav"
+        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
+        assert len(response.content) > MIN_AUDIO_BYTES, (
+            f"Audio too small ({len(response.content)} bytes), expected > {MIN_AUDIO_BYTES}"
+        )
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+OmniVoice TTS Pipeline for vLLM-Omni diffusion engine.
+
+Single-stage pipeline that runs the full text-to-speech flow:
+  text → tokenize → 32-step iterative unmasking → 8-codebook tokens → DAC decode → 24kHz audio
+
+Uses request-mode execution (all steps in one forward() call).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Iterable
+from typing import ClassVar
+
+import torch
+from tokenizers import Tokenizer as HFTokenizer
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.interface import SupportAudioOutput
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
+from vllm_omni.model_executor.models.omnivoice.duration import RuleDurationEstimator
+from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import OmniVoiceDecoder
+from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator
+
+logger = init_logger(__name__)
+
+
+def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig):
+    """Post-processing: convert audio tensor to numpy for WAV encoding."""
+
+    def post_process_func(audio: torch.Tensor, output_type: str = "np"):
+        if output_type == "pt":
+            return audio
+        return audio.cpu().float().numpy()
+
+    return post_process_func
+
+
+class OmniVoicePipeline(nn.Module, SupportAudioOutput):
+    """OmniVoice text-to-speech pipeline for the diffusion engine.
+
+    Wraps OmniVoiceGenerator (32-step iterative unmasking) and
+    OmniVoiceDecoder (HiggsAudioV2 RVQ + DAC) into a single forward() call.
+    """
+
+    support_audio_output: ClassVar[bool] = True
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        self.od_config = od_config
+        self.device = get_local_device()
+        self.model_path = od_config.model
+
+        # Resolve model path (HF hub ID → local cache)
+        if not os.path.isdir(self.model_path):
+            from huggingface_hub import snapshot_download
+
+            self.model_path = snapshot_download(self.model_path)
+
+        # Load OmniVoice config
+        config_path = os.path.join(self.model_path, "config.json")
+        with open(config_path) as f:
+            hf_config = json.load(f)
+        self.config = OmniVoiceConfig(**hf_config)
+
+        # Build generator and decoder
+        self.generator = OmniVoiceGenerator(self.config)
+        self.decoder = OmniVoiceDecoder(self.config)
+
+        # Tokenizer (low-level, avoids HF tokenizer extra_special_tokens issue)
+        tokenizer_path = os.path.join(self.model_path, "tokenizer.json")
+        self.tokenizer = HFTokenizer.from_file(tokenizer_path)
+
+        # Duration estimator
+        self.duration_estimator = RuleDurationEstimator()
+
+        # Generation parameters
+        self.num_step = self.config.num_step
+        self.guidance_scale = self.config.guidance_scale
+        self.t_shift = self.config.t_shift
+        self.layer_penalty_factor = self.config.layer_penalty_factor
+        self.position_temperature = self.config.position_temperature
+        self.class_temperature = self.config.class_temperature
+        self.sample_rate = self.config.sample_rate
+
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        """Generate speech audio from text.
+
+        Args:
+            req: Diffusion request containing text prompt(s).
+
+        Returns:
+            DiffusionOutput with audio tensor in .output
+        """
+        # Extract text from request
+        prompt = req.prompts[0] if req.prompts else ""
+        if isinstance(prompt, dict):
+            text = prompt.get("input", prompt.get("text", str(prompt)))
+        else:
+            text = str(prompt)
+
+        if not text:
+            return DiffusionOutput(error="Empty text prompt")
+
+        device = self.device
+        num_cb = self.config.num_audio_codebook
+        mask_id = self.config.audio_mask_id
+
+        # Estimate target duration
+        target_len = self.duration_estimator.estimate_duration(text, "Nice to meet you.", 25)
+        target_len = max(1, int(target_len))
+
+        # Tokenize with control tokens
+        style = "<|denoise|><|lang_start|>None<|lang_end|><|instruct_start|>None<|instruct_end|>"
+        full_prompt = f"{style}<|text_start|>{text}<|text_end|>"
+        encoding = self.tokenizer.encode(full_prompt)
+        text_tokens = torch.tensor(encoding.ids, dtype=torch.long, device=device)
+        text_len = text_tokens.shape[0]
+
+        # Build conditional + unconditional batches [2, 8, max_len]
+        text_ids = text_tokens.unsqueeze(0).repeat(num_cb, 1)
+        target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=device)
+        cond_ids = torch.cat([text_ids, target_ids], dim=1)
+        cond_len = cond_ids.shape[1]
+
+        uncond_ids = target_ids.clone()
+        uncond_len = target_len
+        max_len = max(cond_len, uncond_len)
+        if uncond_len < max_len:
+            pad = torch.full(
+                (num_cb, max_len - uncond_len),
+                mask_id,
+                dtype=torch.long,
+                device=device,
+            )
+            uncond_ids = torch.cat([uncond_ids, pad], dim=1)
+
+        batch_input_ids = torch.stack([cond_ids, uncond_ids])
+
+        batch_audio_mask = torch.zeros(2, max_len, dtype=torch.bool, device=device)
+        batch_audio_mask[0, text_len:cond_len] = True
+        batch_audio_mask[1, :uncond_len] = True
+
+        batch_attn_mask = torch.zeros(2, 1, max_len, max_len, dtype=torch.bool, device=device)
+        batch_attn_mask[0, :, :cond_len, :cond_len] = True
+        batch_attn_mask[1, :, :uncond_len, :uncond_len] = True
+
+        # Run 32-step iterative unmasking
+        tokens = self.generator(
+            input_ids=batch_input_ids,
+            audio_mask=batch_audio_mask,
+            attention_mask=batch_attn_mask,
+            target_lens=[target_len],
+            num_step=self.num_step,
+            guidance_scale=self.guidance_scale,
+            t_shift=self.t_shift,
+            layer_penalty_factor=self.layer_penalty_factor,
+            position_temperature=self.position_temperature,
+            class_temperature=self.class_temperature,
+        )
+
+        # Decode tokens to audio
+        audio = self.decoder(tokens)  # [1, 1, samples]
+
+        return DiffusionOutput(output=audio)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights from model directory (not from the iterator).
+
+        The diffusion model loader passes HF safetensors weights, but OmniVoice
+        has custom weight names (llm.* → generator.*, audio_tokenizer.* → decoder.*).
+        We load from model_path directly and return all param names to satisfy
+        the loader's "all weights initialized" check.
+        """
+        # Consume the iterator (required by the loader contract)
+        for _ in weights:
+            pass
+
+        device = self.device
+        self.generator.load_weights(self.model_path, device)
+        self.generator = self.generator.to(device).eval()
+        self.decoder.load_weights(self.model_path, device)
+        logger.info("OmniVoice pipeline loaded on %s", device)
+
+        # Return all parameter names to indicate they're initialized
+        return {name for name, _ in self.named_parameters()}
@@ -173,6 +173,16 @@ _DIFFUSION_MODELS = {
        "pipeline_hunyuan_video_1_5_i2v",
        "HunyuanVideo15I2VPipeline",
    ),
+    "OmniVoicePipeline": (
+        "omnivoice",
+        "pipeline_omnivoice",
+        "OmniVoicePipeline",
+    ),
+    "OmniVoice": (
+        "omnivoice",
+        "pipeline_omnivoice",
+        "OmniVoicePipeline",
+    ),
 }


@@ -358,6 +368,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
    "Flux2Pipeline": "get_flux2_post_process_func",
    "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
    "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
+    "OmniVoicePipeline": "get_omnivoice_post_process_func",
 }

 _DIFFUSION_PRE_PROCESS_FUNCS = {
@@ -18,6 +18,7 @@ def _register_omni_hf_configs() -> None:
        from transformers import AutoConfig

        from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
+        from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
            Qwen3TTSConfig,
        )
@@ -31,6 +32,7 @@ def _register_omni_hf_configs() -> None:
    for model_type, config_cls in [
        ("qwen3_tts", Qwen3TTSConfig),
        ("cosyvoice3", CosyVoice3Config),
+        ("omnivoice", OmniVoiceConfig),
        ("voxtral_tts", VoxtralTTSConfig),
    ]:
        try:
@@ -510,6 +510,13 @@ async def omni_init_app_state(
            stage_configs=diffusion_stage_configs,
        )

+        state.openai_serving_speech = OmniOpenAIServingSpeech.for_diffusion(
+            diffusion_engine=engine_client,
+            model_name=model_name,
+            stage_configs=diffusion_stage_configs,
+        )
+        state.openai_streaming_speech = None
+
        state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False)
        state.server_load_metrics = 0
        logger.info("Pure diffusion API server initialized for model: %s", model_name)
@@ -47,7 +47,10 @@ logger = init_logger(__name__)
 _VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"}
 _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"}
 _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"}
-_TTS_MODEL_STAGES: set[str] = _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES | _FISH_TTS_MODEL_STAGES
+_OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"}
+_TTS_MODEL_STAGES: set[str] = (
+    _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES | _FISH_TTS_MODEL_STAGES | _OMNIVOICE_TTS_MODEL_STAGES
+)
 _TTS_LANGUAGES: set[str] = {
    "Auto",
    "Chinese",
@@ -145,6 +148,27 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool:


 class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
+    _diffusion_mode: bool = False
+
+    @classmethod
+    def for_diffusion(
+        cls,
+        diffusion_engine: "Any",
+        model_name: str,
+        stage_configs: "list[Any] | None" = None,
+    ) -> "OmniOpenAIServingSpeech":
+        """Create a speech serving instance for pure diffusion TTS models.
+
+        Bypasses OpenAIServing.__init__ which requires a fully configured
+        engine client that pure diffusion engines don't provide.
+        """
+        instance = cls.__new__(cls)
+        instance._diffusion_mode = True
+        instance._diffusion_engine = diffusion_engine
+        instance._diffusion_model_name = model_name
+        instance._diffusion_stage_configs = stage_configs
+        return instance
+
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Initialize uploaded speakers storage
@@ -240,6 +264,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
            return "voxtral_tts"
        if model_stage in _FISH_TTS_MODEL_STAGES:
            return "fish_tts"
+        if model_stage in _OMNIVOICE_TTS_MODEL_STAGES:
+            return "omnivoice"
        return None

    def _compute_max_instructions_length(self) -> int:
@@ -1203,6 +1229,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                ref_audio_data = (wav_list, sr)
            prompt = self._build_fish_speech_prompt(request, ref_audio_data=ref_audio_data)
            tts_params = {}
+        elif self._tts_model_type == "omnivoice":
+            tts_params = {}
+            prompt = request.input  # Diffusion engine takes raw text
        elif self._is_tts:
            validation_error = self._validate_tts_request(request)
            if validation_error:
@@ -1324,6 +1353,79 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        audio_response: AudioResponse = self.create_audio(audio_obj)
        return audio_response.audio_data, audio_response.media_type

+    async def _create_diffusion_speech(
+        self,
+        request: OpenAICreateSpeechRequest,
+    ) -> Response:
+        """Handle speech generation for pure diffusion TTS models (e.g. OmniVoice)."""
+        from vllm_omni.outputs import OmniRequestOutput
+
+        try:
+            request_id = f"speech-{random_uuid()}"
+            prompt = request.input
+
+            logger.info(
+                "Diffusion TTS speech request %s: text=%r",
+                request_id,
+                prompt[:50] + "..." if len(prompt) > 50 else prompt,
+            )
+
+            generator = self._diffusion_engine.generate(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params_list=self._diffusion_engine.default_sampling_params_list,
+                output_modalities=["audio"],
+            )
+
+            final_output: OmniRequestOutput | None = None
+            async for res in generator:
+                final_output = res
+
+            if final_output is None:
+                raise ValueError("No output generated from the model.")
+
+            audio_output, audio_key = self._extract_audio_output(final_output)
+            if audio_key is None:
+                raise ValueError("TTS model did not produce audio output.")
+
+            audio_tensor = audio_output[audio_key]
+            sr_raw = audio_output.get("sr", 24000)
+            sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+            sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+
+            if isinstance(audio_tensor, list):
+                non_empty = [c for c in audio_tensor if c.numel() > 0]
+                audio_tensor = torch.cat(non_empty, dim=-1) if non_empty else np.zeros((0,), dtype=np.float32)
+            if hasattr(audio_tensor, "float"):
+                audio_tensor = audio_tensor.float().detach().cpu().numpy()
+            if audio_tensor.ndim > 1:
+                audio_tensor = audio_tensor.squeeze()
+
+            audio_obj = CreateAudio(
+                audio_tensor=audio_tensor,
+                sample_rate=sample_rate,
+                response_format=request.response_format or "wav",
+                speed=request.speed or 1.0,
+                stream_format=request.stream_format,
+                base64_encode=False,
+            )
+            audio_response: AudioResponse = self.create_audio(audio_obj)
+            return Response(content=audio_response.audio_data, media_type=audio_response.media_type)
+
+        except asyncio.CancelledError:
+            return self._diffusion_error_response("Client disconnected")
+        except ValueError as e:
+            return self._diffusion_error_response(str(e))
+        except Exception as e:
+            logger.exception("Diffusion speech generation failed: %s", e)
+            return self._diffusion_error_response(f"Speech generation failed: {e}")
+
+    @staticmethod
+    def _diffusion_error_response(message: str) -> Response:
+        """Create a JSON error response without depending on OpenAIServing."""
+        error_body = json.dumps({"error": {"message": message, "type": "server_error", "param": None, "code": 500}})
+        return Response(content=error_body, media_type="application/json", status_code=500)
+
    async def create_speech(
        self,
        request: OpenAICreateSpeechRequest,
@@ -1349,6 +1451,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        Each Code2Wav chunk is yielded as raw audio bytes as soon as it is decoded.
        For WAV format, a header with placeholder size values is emitted first.
        """
+        if self._diffusion_mode:
+            return await self._create_diffusion_speech(request)
+
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            logger.error("Error with model %s", error_check_ret)
@@ -1426,6 +1531,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        batch_request: BatchSpeechRequest,
    ) -> BatchSpeechResponse | ErrorResponse:
        """Generate speech for multiple items concurrently."""
+        if self._diffusion_mode:
+            raise ValueError("Batch speech is not supported in diffusion mode")
        if len(batch_request.items) > self._batch_max_items:
            raise ValueError(
                f"Batch contains {len(batch_request.items)} items, exceeding the maximum of {self._batch_max_items}."
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OmniVoice configuration for vLLM-Omni two-stage pipeline."""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class OmniVoiceConfig(PretrainedConfig):
+    """Configuration for OmniVoice model in vLLM-Omni.
+
+    This mirrors the HuggingFace OmniVoiceConfig but adds fields needed
+    for the two-stage serving pipeline.
+    """
+
+    model_type = "omnivoice"
+
+    def get_text_config(self, **kwargs):
+        """Return self so vLLM uses our top-level config (which has
+        num_attention_heads etc.) instead of trying to extract a sub-config."""
+        return self
+
+    def __init__(self, **kwargs):
+        # HF repos (e.g. k2-fsa/OmniVoice) may nest generation hyperparameters.
+        gen_cfg = kwargs.pop("generation_config", None)
+        if isinstance(gen_cfg, dict):
+            for k, v in gen_cfg.items():
+                kwargs.setdefault(k, v)
+
+        super().__init__(**kwargs)
+
+        # Audio codec params (prefer values set by PretrainedConfig from config.json)
+        self.audio_vocab_size = getattr(self, "audio_vocab_size", 1025)
+        self.audio_mask_id = getattr(self, "audio_mask_id", 1024)
+        self.num_audio_codebook = getattr(self, "num_audio_codebook", 8)
+        self.audio_codebook_weights = getattr(
+            self,
+            "audio_codebook_weights",
+            [8, 8, 6, 6, 4, 4, 2, 2],
+        )
+
+        # LLM backbone params (Qwen3-0.6B defaults from HF config)
+        llm_config = getattr(self, "llm_config", None) or {}
+        if isinstance(llm_config, PretrainedConfig):
+            llm_config = llm_config.to_dict()
+        elif not isinstance(llm_config, dict):
+            llm_config = {}
+        self.llm_hidden_size = llm_config.get("hidden_size", 1024)
+        self.llm_num_hidden_layers = llm_config.get("num_hidden_layers", 28)
+        self.llm_num_attention_heads = llm_config.get("num_attention_heads", 16)
+        self.llm_num_key_value_heads = llm_config.get("num_key_value_heads", 8)
+        self.llm_intermediate_size = llm_config.get("intermediate_size", 3072)
+        self.llm_vocab_size = llm_config.get("vocab_size", 151676)
+        self.llm_max_position_embeddings = llm_config.get("max_position_embeddings", 40960)
+        self.llm_rope_theta = llm_config.get("rope_theta", 1000000.0)
+        self.llm_rms_norm_eps = llm_config.get("rms_norm_eps", 1e-6)
+        self.llm_head_dim = llm_config.get("head_dim", self.llm_hidden_size // self.llm_num_attention_heads)
+
+        # Expose LLM params at top level for vLLM ModelConfig compatibility
+        # (vLLM expects num_attention_heads, hidden_size, etc. on the config)
+        self.num_attention_heads = self.llm_num_attention_heads
+        self.num_key_value_heads = self.llm_num_key_value_heads
+        self.num_hidden_layers = self.llm_num_hidden_layers
+        self.hidden_size = self.llm_hidden_size
+        self.head_dim = self.llm_head_dim
+        if not hasattr(self, "vocab_size"):
+            self.vocab_size = self.llm_vocab_size
+
+        # Generation params (defaults from OmniVoiceGenerationConfig)
+        self.num_step = getattr(self, "num_step", 32)
+        self.guidance_scale = getattr(self, "guidance_scale", 2.0)
+        self.t_shift = getattr(self, "t_shift", 0.1)
+        self.layer_penalty_factor = getattr(self, "layer_penalty_factor", 5.0)
+        self.position_temperature = getattr(self, "position_temperature", 5.0)
+        self.class_temperature = getattr(self, "class_temperature", 0.0)
+
+        # Audio output
+        self.sample_rate = getattr(self, "sample_rate", 24000)
+        self.frame_rate = getattr(self, "frame_rate", 25)
+
+        # Serving
+        self.speculative_config = None
@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+# Copyright    2026  Xiaomi Corp.        (authors:  Han Zhu)
+#
+# See ../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Text duration estimation for TTS generation.
+
+Provides ``RuleDurationEstimator``, which estimates audio duration from text
+using character phonetic weights across 600+ languages. Used by
+``OmniVoice.generate()`` to determine output length when no duration is specified.
+"""
+
+import bisect
+import unicodedata
+from functools import lru_cache
+
+
+class RuleDurationEstimator:
+    def __init__(self):
+        # ==========================================
+        # 1. Phonetic Weights Table
+        # ==========================================
+        # The weight represents the relative speaking time compared to
+        # a standard Latin letter.
+        # Benchmark: 1.0 = One Latin Character (~40-50ms)
+        self.weights = {
+            # --- Logographic (1 char = full syllable/word) ---
+            "cjk": 3.0,  # Chinese, Japanese Kanji, etc.
+            # --- Syllabic / Blocks
+            "hangul": 2.5,  # Korean Hangul
+            "kana": 2.2,  # Japanese Hiragana/Katakana
+            "ethiopic": 3.0,  # Amharic/Ge'ez
+            "yi": 3.0,  # Yi script
+            # --- Abugida (Consonant-Vowel complexes) ---
+            "indic": 1.8,  # Hindi, Bengali, Tamil, etc.
+            "thai_lao": 1.5,  # Thai, Lao
+            "khmer_myanmar": 1.8,  # Khmer, Myanmar
+            # --- Abjad (Consonant-heavy) ---
+            "arabic": 1.5,  # Arabic, Persian, Urdu
+            "hebrew": 1.5,  # Hebrew
+            # --- Alphabet (Segmental) ---
+            "latin": 1.0,  # English, Spanish, French, Vietnamese, etc. (Baseline)
+            "cyrillic": 1.0,  # Russian, Ukrainian
+            "greek": 1.0,  # Greek
+            "armenian": 1.0,  # Armenian
+            "georgian": 1.0,  # Georgian
+            # --- Symbols & Misc ---
+            "punctuation": 0.5,  # Pause capability
+            "space": 0.2,  # Word boundary/Breath (0.05 / 0.22)
+            "digit": 3.5,  # Numbers
+            "mark": 0.0,  # Diacritics/Accents (Silent modifiers)
+            "default": 1.0,  # Fallback for unknown scripts
+        }
+
+        # ==========================================
+        # 2. Unicode Range Mapping
+        # ==========================================
+        # Format: (End_Codepoint, Type_Key)
+        # Used for fast binary search (bisect).
+        self.ranges = [
+            (0x02AF, "latin"),  # Latin (Basic, Supplement, Ext, IPA)
+            (0x03FF, "greek"),  # Greek & Coptic
+            (0x052F, "cyrillic"),  # Cyrillic
+            (0x058F, "armenian"),  # Armenian
+            (0x05FF, "hebrew"),  # Hebrew
+            (0x077F, "arabic"),  # Arabic, Syriac, Arabic Supplement
+            (0x089F, "arabic"),  # Arabic Extended-B (+ Syriac Supp)
+            (0x08FF, "arabic"),  # Arabic Extended-A
+            (0x097F, "indic"),  # Devanagari
+            (0x09FF, "indic"),  # Bengali
+            (0x0A7F, "indic"),  # Gurmukhi
+            (0x0AFF, "indic"),  # Gujarati
+            (0x0B7F, "indic"),  # Oriya
+            (0x0BFF, "indic"),  # Tamil
+            (0x0C7F, "indic"),  # Telugu
+            (0x0CFF, "indic"),  # Kannada
+            (0x0D7F, "indic"),  # Malayalam
+            (0x0DFF, "indic"),  # Sinhala
+            (0x0EFF, "thai_lao"),  # Thai & Lao
+            (0x0FFF, "indic"),  # Tibetan (Abugida)
+            (0x109F, "khmer_myanmar"),  # Myanmar
+            (0x10FF, "georgian"),  # Georgian
+            (0x11FF, "hangul"),  # Hangul Jamo
+            (0x137F, "ethiopic"),  # Ethiopic
+            (0x139F, "ethiopic"),  # Ethiopic Supplement
+            (0x13FF, "default"),  # Cherokee
+            (0x167F, "default"),  # Canadian Aboriginal Syllabics
+            (0x169F, "default"),  # Ogham
+            (0x16FF, "default"),  # Runic
+            (0x171F, "default"),  # Tagalog (Baybayin)
+            (0x173F, "default"),  # Hanunoo
+            (0x175F, "default"),  # Buhid
+            (0x177F, "default"),  # Tagbanwa
+            (0x17FF, "khmer_myanmar"),  # Khmer
+            (0x18AF, "default"),  # Mongolian
+            (0x18FF, "default"),  # Canadian Aboriginal Syllabics Ext
+            (0x194F, "indic"),  # Limbu
+            (0x19DF, "indic"),  # Tai Le & New Tai Lue
+            (0x19FF, "khmer_myanmar"),  # Khmer Symbols
+            (0x1A1F, "indic"),  # Buginese
+            (0x1AAF, "indic"),  # Tai Than
+            (0x1B7F, "indic"),  # Balinese
+            (0x1BBF, "indic"),  # Sundanese
+            (0x1BFF, "indic"),  # Batak
+            (0x1C4F, "indic"),  # Lepcha
+            (0x1C7F, "indic"),  # Ol Chiki (Santali)
+            (0x1C8F, "cyrillic"),  # Cyrillic Extended-C
+            (0x1CBF, "georgian"),  # Georgian Extended
+            (0x1CCF, "indic"),  # Sundanese Supplement
+            (0x1CFF, "indic"),  # Vedic Extensions
+            (0x1D7F, "latin"),  # Phonetic Extensions
+            (0x1DBF, "latin"),  # Phonetic Extensions Supplement
+            (0x1DFF, "default"),  # Combining Diacritical Marks Supplement
+            (0x1EFF, "latin"),  # Latin Extended Additional (Vietnamese)
+            (0x309F, "kana"),  # Hiragana
+            (0x30FF, "kana"),  # Katakana
+            (0x312F, "cjk"),  # Bopomofo (Pinyin)
+            (0x318F, "hangul"),  # Hangul Compatibility Jamo
+            (0x9FFF, "cjk"),  # CJK Unified Ideographs (Main)
+            (0xA4CF, "yi"),  # Yi Syllables
+            (0xA4FF, "default"),  # Lisu
+            (0xA63F, "default"),  # Vai
+            (0xA69F, "cyrillic"),  # Cyrillic Extended-B
+            (0xA6FF, "default"),  # Bamum
+            (0xA7FF, "latin"),  # Latin Extended-D
+            (0xA82F, "indic"),  # Syloti Nagri
+            (0xA87F, "default"),  # Phags-pa
+            (0xA8DF, "indic"),  # Saurashtra
+            (0xA8FF, "indic"),  # Devanagari Extended
+            (0xA92F, "indic"),  # Kayah Li
+            (0xA95F, "indic"),  # Rejang
+            (0xA97F, "hangul"),  # Hangul Jamo Extended-A
+            (0xA9DF, "indic"),  # Javanese
+            (0xA9FF, "khmer_myanmar"),  # Myanmar Extended-B
+            (0xAA5F, "indic"),  # Cham
+            (0xAA7F, "khmer_myanmar"),  # Myanmar Extended-A
+            (0xAADF, "indic"),  # Tai Viet
+            (0xAAFF, "indic"),  # Meetei Mayek Extensions
+            (0xAB2F, "ethiopic"),  # Ethiopic Extended-A
+            (0xAB6F, "latin"),  # Latin Extended-E
+            (0xABBF, "default"),  # Cherokee Supplement
+            (0xABFF, "indic"),  # Meetei Mayek
+            (0xD7AF, "hangul"),  # Hangul Syllables
+            (0xFAFF, "cjk"),  # CJK Compatibility
+            (0xFDFF, "arabic"),  # Arabic Presentation Forms-A
+            (0xFE6F, "default"),  # Variation Selectors
+            (0xFEFF, "arabic"),  # Arabic Presentation Forms-B
+            (0xFFEF, "latin"),  # Fullwidth Latin
+        ]
+        self.breakpoints = [r[0] for r in self.ranges]
+
+    @lru_cache(maxsize=4096)
+    def _get_char_weight(self, char):
+        """Determines the weight of a single character."""
+        code = ord(char)
+        if (65 <= code <= 90) or (97 <= code <= 122):
+            return self.weights["latin"]
+        if code == 32:
+            return self.weights["space"]
+
+        # Ignore arabic Tatweel
+        if code == 0x0640:
+            return self.weights["mark"]
+
+        category = unicodedata.category(char)
+
+        if category.startswith("M"):
+            return self.weights["mark"]
+
+        if category.startswith("P") or category.startswith("S"):
+            return self.weights["punctuation"]
+
+        if category.startswith("Z"):
+            return self.weights["space"]
+
+        if category.startswith("N"):
+            return self.weights["digit"]
+
+        # 3. Binary search for Unicode Block (此时区间里绝不会再混进标点符号)
+        idx = bisect.bisect_left(self.breakpoints, code)
+        if idx < len(self.ranges):
+            script_type = self.ranges[idx][1]
+            return self.weights.get(script_type, self.weights["default"])
+
+        # 4. Handle upper planes (CJK Ext B/C/D, Historic scripts)
+        if code > 0x20000:
+            return self.weights["cjk"]
+
+        return self.weights["default"]
+
+    def calculate_total_weight(self, text):
+        """Sums up the normalized weights for a string."""
+        return sum(self._get_char_weight(c) for c in text)
+
+    def estimate_duration(
+        self,
+        target_text: str,
+        ref_text: str,
+        ref_duration: float,
+        low_threshold: float | None = 50,
+        boost_strength: float = 3,
+    ) -> float:
+        """
+
+        Args:
+            target_text (str): The text for which we want to estimate the duration.
+            ref_text (str): The reference text that was used to measure
+                the ref_duration.
+            ref_duration (float): The actual duration it took
+                to speak the ref_text.
+            low_threshold (float): The minimum duration threshold below which the
+                estimation will be considered unreliable.
+            boost_strength (float): Controls the power-curve boost for short durations.
+                Higher values boost small durations more aggressively.
+                1 = no boost (linear), 2 = sqrt-like
+
+        Returns:
+            float: The estimated duration for the target_text based
+                on the ref_text and ref_duration.
+        """
+        if ref_duration <= 0 or not ref_text:
+            return 0.0
+
+        ref_weight = self.calculate_total_weight(ref_text)
+        if ref_weight == 0:
+            return 0.0
+
+        speed_factor = ref_weight / ref_duration
+        target_weight = self.calculate_total_weight(target_text)
+
+        estimated_duration = target_weight / speed_factor
+        if low_threshold is not None and estimated_duration < low_threshold:
+            alpha = 1.0 / boost_strength
+            return low_threshold * (estimated_duration / low_threshold) ** alpha
+        else:
+            return estimated_duration
+
+
+# ==========================================
+# Example Usage
+# ==========================================
+if __name__ == "__main__":
+    estimator = RuleDurationEstimator()
+
+    ref_txt = "Hello, world."
+    ref_dur = 1.5
+
+    test_cases = [
+        ("Hindi (With complex marks)", "नमस्ते दुनिया"),
+        ("Arabic (With vowels)", "مَرْحَبًا بِالْعَالَم"),
+        ("Vietnamese (Lots of diacritics)", "Chào thế giới"),
+        ("Chinese", "你好，世界！"),
+        ("Mixed Emoji", "Hello 🌍! This is fun 🎉"),
+    ]
+
+    print("--- Reference ---")
+    print(f"Reference Text: '{ref_txt}'")
+    print(f"Reference Duration: {ref_dur}s")
+    print("-" * 30)
+
+    for lang, txt in test_cases:
+        est_time = estimator.estimate_duration(txt, ref_txt, ref_dur)
+        weight = estimator.calculate_total_weight(txt)
+
+        print(f"[{lang}]")
+        print(f"Text: {txt}")
+        print(f"Total Weight: {weight:.2f}")
+        print(f"Estimated Duration: {est_time:.2f} s")
+        print("-" * 30)
@@ -0,0 +1,520 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+OmniVoice model for vLLM-Omni two-stage TTS pipeline.
+
+Stage 0 (Generator): Qwen3 backbone + iterative unmasking → 8-codebook tokens
+Stage 1 (Decoder): HiggsAudioV2 decoder → 24kHz waveform
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Mapping, Sequence
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers.feature_extraction_utils import BatchFeature
+from vllm.config import VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.logger import init_logger
+from vllm.multimodal.inputs import (
+    MultiModalDataDict,
+    MultiModalFieldConfig,
+    MultiModalKwargsItems,
+)
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    ProcessorInputs,
+    PromptIndexTargets,
+    PromptInsertion,
+    PromptUpdate,
+)
+from vllm.sequence import IntermediateTensors
+
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Multimodal processing
+# ---------------------------------------------------------------------------
+
+
+class OmniVoiceMultiModalProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config(OmniVoiceConfig)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": None}
+
+    def get_data_parser(self):
+        return MultiModalDataParser(
+            target_sr=self.ctx.get_hf_config().sample_rate,
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+
+class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]):
+    """Processes text + optional reference audio for OmniVoice.
+
+    For voice cloning: text + reference audio → tokenized reference
+    For auto voice: text only
+    """
+
+    def _ensure_cached_runtime_components(self, model_dir: str, config: OmniVoiceConfig) -> None:
+        cached_model_dir = getattr(self, "_cached_model_dir", None)
+        if cached_model_dir == model_dir:
+            return
+
+        from transformers import AutoTokenizer
+
+        self.text_tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+        # Audio tokenizer for encoding reference audio
+        audio_tokenizer_path = os.path.join(model_dir, "audio_tokenizer")
+        if os.path.isdir(audio_tokenizer_path):
+            try:
+                from transformers import (
+                    AutoFeatureExtractor,
+                    HiggsAudioV2TokenizerModel,
+                )
+            except ImportError as e:
+                raise ImportError(
+                    "OmniVoice voice cloning requires transformers with "
+                    "HiggsAudioV2TokenizerModel. Upgrade transformers or "
+                    "use text-only mode (no reference audio)."
+                ) from e
+
+            self.audio_tokenizer = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_path, device_map="cpu")
+            self.feature_extractor = AutoFeatureExtractor.from_pretrained(audio_tokenizer_path)
+            self.audio_tokenizer.eval()
+        else:
+            self.audio_tokenizer = None
+            self.feature_extractor = None
+            logger.warning(
+                "audio_tokenizer not found at %s, voice cloning disabled",
+                audio_tokenizer_path,
+            )
+
+        self._cached_model_dir = model_dir
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        config = self.info.ctx.get_hf_config()
+        model_dir = self.info.ctx.model_config.model
+        self._ensure_cached_runtime_components(model_dir, config)
+
+        audio = mm_data.get("audio", None)
+        if audio is None:
+            audio = mm_data.get("audios")
+            if audio is not None:
+                audio = audio[0], config.sample_rate
+
+        # Build text prompt with control tokens
+        lang = mm_kwargs.get("lang", None)
+        instruct = mm_kwargs.get("instruct", None)
+        denoise = mm_kwargs.get("denoise", True)
+        ref_text = mm_kwargs.get("ref_text", None)
+
+        # Construct the style + text portion
+        style_text = ""
+        if denoise:
+            style_text += "<|denoise|>"
+        lang_str = lang if lang else "None"
+        instruct_str = instruct if instruct else "None"
+        style_text += f"<|lang_start|>{lang_str}<|lang_end|>"
+        style_text += f"<|instruct_start|>{instruct_str}<|instruct_end|>"
+
+        # Combine ref_text and main text
+        if ref_text:
+            full_text = f"{ref_text} {prompt}"
+        else:
+            full_text = prompt
+
+        text_prompt = f"{style_text}<|text_start|>{full_text}<|text_end|>"
+        text_tokens = self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0)  # [N_text]
+
+        if audio is None:
+            # Text-only path (auto voice mode)
+            return BatchFeature(
+                {
+                    "input_ids": text_tokens,
+                    "input_len": [len(text_tokens)],
+                }
+            )
+
+        # Voice cloning: encode reference audio to tokens
+        audio_signal, sr = audio
+        if isinstance(audio_signal, np.ndarray):
+            audio_signal = torch.from_numpy(audio_signal).float()
+        if audio_signal.dim() == 1:
+            audio_signal = audio_signal.unsqueeze(0)
+
+        # Resample to tokenizer sample rate if needed
+        if self.feature_extractor is not None:
+            target_sr = self.feature_extractor.sampling_rate
+            if sr != target_sr:
+                import torchaudio
+
+                audio_signal = torchaudio.functional.resample(audio_signal, sr, target_sr)
+
+        # Encode reference audio to 8-codebook tokens
+        if self.audio_tokenizer is not None:
+            with torch.inference_mode():
+                ref_audio_tokens = self.audio_tokenizer.encode(audio_signal)  # [8, T_ref]
+                if ref_audio_tokens.dim() == 3:
+                    ref_audio_tokens = ref_audio_tokens.squeeze(0)  # [8, T_ref]
+        else:
+            raise RuntimeError(
+                "Audio tokenizer not available for voice cloning. Ensure audio_tokenizer/ exists in model directory."
+            )
+
+        ft = BatchFeature(
+            {
+                "input_ids": text_tokens,
+                "ref_audio_tokens": ref_audio_tokens,  # [8, T_ref]
+                "ref_audio_len": [ref_audio_tokens.shape[1]],
+            }
+        )
+        return ft
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return {
+            "ref_audio_tokens": MultiModalFieldConfig.batched("audio"),
+            "ref_audio_len": MultiModalFieldConfig.batched("audio"),
+        }
+
+    def _hf_processor_applies_updates(
+        self,
+        prompt_text: str,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> bool:
+        return False
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        def insertion_end(item_idx):
+            if "audio" in out_mm_kwargs and out_mm_kwargs["audio"]:
+                ref_len = out_mm_kwargs["audio"][0]["ref_audio_len"].data[0].item()
+                return [1] * ref_len
+            return []
+
+        return [
+            PromptInsertion(
+                modality="audio",
+                target=PromptIndexTargets.start(),
+                insertion=insertion_end,
+            ),
+        ]
+
+
+class OmniVoiceDummyInputsBuilder(BaseDummyInputsBuilder[OmniVoiceMultiModalProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return "Hello, this is a test of the OmniVoice system."
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> MultiModalDataDict:
+        num_audios = mm_counts.get("audio")
+        max_prompt_seconds = 10
+        prompt_sample_rate = 24000
+        target_audio_length = max_prompt_seconds * prompt_sample_rate
+
+        audio_overrides = mm_options.get("audio") if mm_options else None
+        mm_data = {
+            "audio": (
+                self._get_dummy_audios(
+                    length=target_audio_length,
+                    num_audios=num_audios,
+                    overrides=audio_overrides,
+                )[0],
+                24000,
+            ),
+        }
+        return mm_data
+
+    def get_dummy_processor_inputs(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> ProcessorInputs:
+        inputs = super().get_dummy_processor_inputs(seq_len, mm_counts, mm_options)
+        inputs.hf_processor_mm_kwargs = {"ref_text": "Testing voice cloning."}
+        return inputs
+
+
+# ---------------------------------------------------------------------------
+# Main model class
+# ---------------------------------------------------------------------------
+
+
+class OmniVoiceModel(
+    nn.Module,
+):
+    """OmniVoice model for vLLM-Omni two-stage pipeline.
+
+    Routes to generator (Stage 0) or decoder (Stage 1) based on model_stage.
+    """
+
+    requires_raw_input_tokens = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.config = vllm_config.model_config.hf_config
+        self.have_multimodal_outputs = True
+        self.model_stage = vllm_config.model_config.model_stage
+        self.model_dir = vllm_config.model_config.model
+
+        if self.model_stage == "omnivoice_generator":
+            from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import (
+                OmniVoiceGenerator,
+            )
+
+            self.generator = OmniVoiceGenerator(self.config)
+            self.model = self.generator
+        elif self.model_stage == "omnivoice_decoder":
+            from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import (
+                OmniVoiceDecoder,
+            )
+
+            self.decoder = OmniVoiceDecoder(self.config)
+            self.model = self.decoder
+        else:
+            raise ValueError(f"Unsupported model_stage: {self.model_stage}")
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal=None,
+    ) -> torch.Tensor:
+        if self.model_stage == "omnivoice_generator":
+            # Generator handles its own embedding in forward()
+            hidden = int(self.config.llm_hidden_size)
+            return torch.zeros((input_ids.shape[0], hidden), device=input_ids.device)
+        elif self.model_stage == "omnivoice_decoder":
+            hidden = int(self.config.llm_hidden_size)
+            return torch.zeros((input_ids.shape[0], hidden), device=input_ids.device)
+        else:
+            raise RuntimeError(f"embed_input_ids not valid for {self.model_stage}")
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        additional_information: dict[str, object] | None = None,
+        **kwargs: object,
+    ) -> OmniOutput:
+        if self.model_stage == "omnivoice_generator":
+            return self._forward_generator(input_ids, kwargs)
+        elif self.model_stage == "omnivoice_decoder":
+            return self._forward_decoder(input_ids, kwargs)
+        else:
+            raise ValueError(f"Unsupported model_stage: {self.model_stage}")
+
+    def _forward_generator(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutput:
+        """Run generator stage: text → 8-codebook audio tokens."""
+        runtime_info = kwargs.get("runtime_additional_information", [])
+
+        if not runtime_info:
+            # Profiling / dummy run — return a plain tensor (not OmniOutput)
+            # so the v1 model runner's _dummy_run can index into it.
+            return torch.zeros(
+                (input_ids.shape[0], self.config.llm_hidden_size),
+                device=input_ids.device,
+                dtype=torch.float32,
+            )
+
+        info = runtime_info[0]
+        device = input_ids.device
+        num_codebooks = self.config.num_audio_codebook
+        mask_id = self.config.audio_mask_id
+
+        # Extract text tokens from input_ids
+        text_tokens = input_ids  # [N_text]
+        text_len = text_tokens.shape[0]
+
+        # Estimate target length using RuleDurationEstimator
+        # (same formula as reference OmniVoice: weight * 25 / 14.1)
+        from vllm_omni.model_executor.models.omnivoice.duration import (
+            RuleDurationEstimator,
+        )
+
+        if not hasattr(self, "_duration_estimator"):
+            self._duration_estimator = RuleDurationEstimator()
+        raw_text = info.get("raw_text", "")
+        if raw_text:
+            target_len = self._duration_estimator.estimate_duration(raw_text, "Nice to meet you.", 25)
+            target_len = max(1, int(target_len))
+        else:
+            # Fallback: use character weight formula on text tokens
+            # approximate ~1.77 frames per text token (25/14.1)
+            target_len = max(int(text_len * 1.77), 25)
+
+        # Get reference audio tokens if available
+        ref_audio_tokens = info.get("ref_audio_tokens", None)
+
+        # Build input_ids tensor: [2*B, 8, S]
+        # B=1, conditional + unconditional
+
+        # Replicate text tokens across 8 codebooks
+        text_ids = text_tokens.unsqueeze(0).repeat(num_codebooks, 1)  # [8, N_text]
+
+        # Target: all MASK
+        target_ids = torch.full((num_codebooks, target_len), mask_id, dtype=torch.long, device=device)
+
+        # Conditional: [text] [ref_audio?] [target_mask]
+        if ref_audio_tokens is not None:
+            ref_tokens = ref_audio_tokens.to(device)  # [8, T_ref]
+            cond_ids = torch.cat([text_ids, ref_tokens, target_ids], dim=1)
+            cond_audio_start = text_ids.shape[1]
+        else:
+            cond_ids = torch.cat([text_ids, target_ids], dim=1)
+            cond_audio_start = text_ids.shape[1]
+
+        cond_len = cond_ids.shape[1]
+
+        # Unconditional: [target_mask only]
+        uncond_ids = target_ids.clone()
+        uncond_len = target_len
+
+        # Pad to same length
+        max_len = max(cond_len, uncond_len)
+        if cond_len < max_len:
+            pad = torch.full(
+                (num_codebooks, max_len - cond_len),
+                mask_id,
+                dtype=torch.long,
+                device=device,
+            )
+            cond_ids = torch.cat([cond_ids, pad], dim=1)
+        if uncond_len < max_len:
+            pad = torch.full(
+                (num_codebooks, max_len - uncond_len),
+                mask_id,
+                dtype=torch.long,
+                device=device,
+            )
+            uncond_ids = torch.cat([uncond_ids, pad], dim=1)
+
+        batch_input_ids = torch.stack([cond_ids, uncond_ids], dim=0)  # [2, 8, max_len]
+
+        # Audio mask: True for audio positions
+        batch_audio_mask = torch.zeros((2, max_len), dtype=torch.bool, device=device)
+        batch_audio_mask[0, cond_audio_start:cond_len] = True
+        batch_audio_mask[1, :uncond_len] = True
+
+        # Attention mask: [2, 1, S, S]
+        batch_attention_mask = torch.zeros((2, 1, max_len, max_len), dtype=torch.bool, device=device)
+        batch_attention_mask[0, :, :cond_len, :cond_len] = True
+        batch_attention_mask[1, :, :uncond_len, :uncond_len] = True
+
+        # Run iterative generation
+        tokens = self.generator(
+            input_ids=batch_input_ids,
+            audio_mask=batch_audio_mask,
+            attention_mask=batch_attention_mask,
+            target_lens=[target_len],
+            num_step=self.config.num_step,
+            guidance_scale=self.config.guidance_scale,
+            t_shift=self.config.t_shift,
+            layer_penalty_factor=self.config.layer_penalty_factor,
+            position_temperature=self.config.position_temperature,
+            class_temperature=self.config.class_temperature,
+        )  # [1, 8, target_len]
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"audio_tokens": tokens},
+        )
+
+    def _forward_decoder(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutput:
+        """Run decoder stage: 8-codebook tokens → audio waveform."""
+        runtime_info = kwargs.get("runtime_additional_information", [])
+
+        if not runtime_info:
+            # Profiling / dummy run — return plain tensor for v1 runner compat
+            return torch.zeros(
+                (input_ids.shape[0], self.config.llm_hidden_size),
+                device=input_ids.device,
+                dtype=torch.float32,
+            )
+
+        info = runtime_info[0]
+        audio_tokens = info.get("audio_tokens", None)
+
+        if audio_tokens is None:
+            raise RuntimeError("No audio_tokens received from generator stage")
+
+        if isinstance(audio_tokens, np.ndarray):
+            audio_tokens = torch.from_numpy(audio_tokens)
+
+        # audio_tokens: [B, 8, T]; buffer may be CPU — move to decoder weights
+        if audio_tokens.dim() == 2:
+            audio_tokens = audio_tokens.unsqueeze(0)  # Add batch dim
+
+        dec_device = next(self.decoder.parameters()).device
+        audio_tokens = audio_tokens.to(device=dec_device, dtype=torch.long)
+
+        tts_speech = self.decoder(audio_tokens)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "audio": tts_speech,
+                "sr": self.config.sample_rate,
+            },
+        )
+
+    def _resolve_model_dir(self) -> str:
+        """Resolve model directory to local path (handles HF hub IDs)."""
+        model_dir = self.model_dir
+        if os.path.isdir(model_dir):
+            return model_dir
+        # HF hub model ID — resolve to local cache
+        from huggingface_hub import snapshot_download
+
+        return snapshot_download(model_dir)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        try:
+            device = next(self.parameters()).device
+        except StopIteration:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        model_dir = self._resolve_model_dir()
+
+        if self.model_stage == "omnivoice_generator":
+            self.generator.load_weights(model_dir, device)
+        elif self.model_stage == "omnivoice_decoder":
+            self.decoder.load_weights(model_dir, device)
+        else:
+            raise ValueError(f"{self.model_stage} not supported!")
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+OmniVoice Decoder (Stage 1) - Audio token to waveform conversion.
+
+Implements the HiggsAudioV2 decode path using transformers' DacModel decoder
+and a custom RVQ quantizer, compatible with transformers 4.x.
+
+Decode path:
+  audio_codes [B, 8, T]
+    → RVQ codebook lookup + project_out → sum → [B, 1024, T]
+    → fc2 Linear(1024, 256) → [B, 256, T]
+    → DAC acoustic decoder (conv transpose upsampling) → [B, 1, T*960]
+    → 24kHz waveform (25fps × 960 samples/frame)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+
+import torch
+import torch.nn as nn
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
+
+logger = init_logger(__name__)
+
+
+class HiggsAudioVQLayer(nn.Module):
+    """Single VQ layer: codebook lookup + project_out."""
+
+    def __init__(self, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024):
+        super().__init__()
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+        self.project_out = nn.Linear(codebook_dim, hidden_size)
+
+    def decode(self, indices: torch.Tensor) -> torch.Tensor:
+        """indices: [B, T] → [B, hidden_size, T]"""
+        quantized = self.codebook(indices)  # [B, T, codebook_dim]
+        quantized = self.project_out(quantized)  # [B, T, hidden_size]
+        return quantized.permute(0, 2, 1)  # [B, hidden_size, T]
+
+
+class HiggsAudioRVQ(nn.Module):
+    """Residual Vector Quantizer with 8 codebook layers."""
+
+    def __init__(
+        self, num_quantizers: int = 8, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024
+    ):
+        super().__init__()
+        self.quantizers = nn.ModuleList(
+            [HiggsAudioVQLayer(codebook_size, codebook_dim, hidden_size) for _ in range(num_quantizers)]
+        )
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """codes: [num_quantizers, B, T] → [B, hidden_size, T]"""
+        result = torch.zeros(
+            codes.shape[1],
+            self.quantizers[0].project_out.out_features,
+            codes.shape[2],
+            device=codes.device,
+            dtype=torch.float32,
+        )
+        for i, quantizer in enumerate(self.quantizers):
+            result = result + quantizer.decode(codes[i])
+        return result
+
+
+class OmniVoiceDecoder(nn.Module):
+    """OmniVoice Stage 1: Token-to-audio decoder.
+
+    Uses DAC acoustic decoder from transformers + custom HiggsAudio RVQ
+    quantizer to convert 8-codebook tokens into 24kHz waveform.
+    """
+
+    def __init__(self, config: OmniVoiceConfig):
+        super().__init__()
+        self.config = config
+        self.sample_rate = config.sample_rate
+        self._loaded = False
+
+        # These are populated by load_weights
+        self.quantizer = None
+        self.fc2 = None
+        self.acoustic_decoder = None
+
+    @torch.inference_mode()
+    def forward(self, audio_codes: torch.Tensor) -> torch.Tensor:
+        """Decode audio tokens to waveform.
+
+        Args:
+            audio_codes: [B, 8, T] - 8-codebook audio token IDs
+
+        Returns:
+            waveform: [B, 1, audio_samples] at 24kHz
+        """
+        if not self._loaded:
+            raise RuntimeError("Decoder not loaded. Call load_weights() first.")
+
+        device = audio_codes.device
+
+        # Transpose: [B, 8, T] → [8, B, T]
+        codes = audio_codes.transpose(0, 1).long()
+
+        # RVQ decode: sum codebook embeddings → [B, 1024, T]
+        quantized = self.quantizer.decode(codes)
+
+        # Project: [B, 1024, T] → fc2 → [B, 256, T]
+        quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
+
+        # Acoustic decoder: [B, 256, T] → [B, 1, T*960]
+        audio = self.acoustic_decoder(quantized)
+
+        # Ensure [B, 1, samples]
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(1)
+
+        return audio.to(device)
+
+    def _adjust_output_padding(self, decoder: nn.Module):
+        """Adjust ConvTranspose1d output_padding (HiggsAudioV2 modification)."""
+        for module in decoder.modules():
+            if isinstance(module, nn.ConvTranspose1d):
+                stride = module.stride[0] if isinstance(module.stride, tuple) else module.stride
+                module.output_padding = (stride % 2,)
+
+    def load_weights(self, model_dir: str, device: torch.device) -> None:
+        """Load decoder components from audio_tokenizer/model.safetensors."""
+        from safetensors.torch import load_file
+        from transformers import DacConfig, DacModel
+
+        audio_tokenizer_path = os.path.join(model_dir, "audio_tokenizer")
+        config_path = os.path.join(audio_tokenizer_path, "config.json")
+        weights_path = os.path.join(audio_tokenizer_path, "model.safetensors")
+
+        if not os.path.exists(weights_path):
+            raise FileNotFoundError(f"Audio tokenizer weights not found at {weights_path}")
+
+        with open(config_path) as f:
+            tokenizer_config = json.load(f)
+
+        state_dict = load_file(weights_path, device=str(device))
+
+        # 1. Build RVQ quantizer
+        codebook_dim = tokenizer_config.get("codebook_dim", 64)
+        codebook_size = tokenizer_config.get("codebook_size", 1024)
+        # Hidden size = quantizer project_out output dim
+        hidden_size = state_dict["quantizer.quantizers.0.project_out.weight"].shape[0]
+        num_quantizers = sum(
+            1 for k in state_dict if k.startswith("quantizer.quantizers.") and k.endswith(".codebook.embed")
+        )
+
+        self.quantizer = HiggsAudioRVQ(
+            num_quantizers=num_quantizers,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            hidden_size=hidden_size,
+        ).to(device)
+
+        # Load quantizer weights
+        for i in range(num_quantizers):
+            prefix = f"quantizer.quantizers.{i}"
+            embed_key = f"{prefix}.codebook.embed"
+            if embed_key in state_dict:
+                self.quantizer.quantizers[i].codebook.weight.data.copy_(state_dict[embed_key])
+            proj_out_w = f"{prefix}.project_out.weight"
+            proj_out_b = f"{prefix}.project_out.bias"
+            if proj_out_w in state_dict:
+                self.quantizer.quantizers[i].project_out.weight.data.copy_(state_dict[proj_out_w])
+            if proj_out_b in state_dict:
+                self.quantizer.quantizers[i].project_out.bias.data.copy_(state_dict[proj_out_b])
+
+        # 2. Build fc2 projection
+        fc2_w = state_dict["fc2.weight"]
+        fc2_b = state_dict["fc2.bias"]
+        self.fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device)
+        self.fc2.weight.data.copy_(fc2_w)
+        self.fc2.bias.data.copy_(fc2_b)
+
+        # 3. Build DAC acoustic decoder
+        dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"])
+        dac_model = DacModel(dac_cfg)
+        self.acoustic_decoder = dac_model.decoder.to(device)
+
+        # Load acoustic decoder weights
+        loaded = 0
+        for name, param in self.acoustic_decoder.named_parameters():
+            higgs_name = f"acoustic_decoder.{name}"
+            if higgs_name in state_dict:
+                param.data.copy_(state_dict[higgs_name])
+                loaded += 1
+
+        # Apply HiggsAudioV2 output padding adjustment
+        self._adjust_output_padding(self.acoustic_decoder)
+
+        # Remove tanh if present (HiggsAudioV2 uses Identity instead)
+        if hasattr(self.acoustic_decoder, "tanh"):
+            self.acoustic_decoder.tanh = nn.Identity()
+
+        self.acoustic_decoder.eval()
+        self._loaded = True
+
+        logger.info(
+            "Loaded OmniVoice decoder: %d quantizers, fc2(%d→%d), acoustic decoder (%d weights)",
+            num_quantizers,
+            fc2_w.shape[1],
+            fc2_w.shape[0],
+            loaded,
+        )
@@ -0,0 +1,588 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+OmniVoice Generator (Stage 0) - Iterative unmasking with Qwen3 backbone.
+
+Generates 8-codebook audio tokens from text via 32-step non-autoregressive
+iterative masked prediction with classifier-free guidance.
+
+Uses vLLM-Omni's DiffusionAttention for optimized full (bidirectional) attention
+via FlashAttention/SageAttention/SDPA backends.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Unmasking schedule helpers
+# ---------------------------------------------------------------------------
+
+
+def _get_time_steps(
+    t_start: float,
+    t_end: float,
+    num_step: int,
+    t_shift: float,
+) -> torch.Tensor:
+    """Compute the unmasking schedule with time shift.
+
+    Returns cumulative proportions [0, ..., 1] of length num_step.
+    Formula: r_n = t_shift * (n/N) / (1 + (t_shift - 1) * (n/N))
+    """
+    steps = torch.linspace(t_start, t_end, num_step)
+    shifted = t_shift * steps / (1.0 + (t_shift - 1.0) * steps)
+    return shifted
+
+
+def _gumbel_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor:
+    """Add Gumbel noise for stochastic position selection."""
+    noise = -torch.log(-torch.log(torch.rand_like(logits).clamp(min=1e-8)))
+    return logits / max(temperature, 1e-8) + noise
+
+
+# ---------------------------------------------------------------------------
+# Qwen3-style transformer blocks using DiffusionAttention
+# ---------------------------------------------------------------------------
+
+
+class OmniVoiceRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return self.weight * x.to(self.weight.dtype)
+
+
+class OmniVoiceAttention(nn.Module):
+    """Qwen3-style GQA attention using DiffusionAttention backend."""
+
+    def __init__(self, config: OmniVoiceConfig):
+        super().__init__()
+        self.hidden_size = config.llm_hidden_size
+        self.num_heads = config.llm_num_attention_heads
+        self.num_kv_heads = config.llm_num_key_value_heads
+        self.head_dim = config.llm_head_dim
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        # Qwen3 uses per-head QK norm
+        self.q_norm = OmniVoiceRMSNorm(self.head_dim)
+        self.k_norm = OmniVoiceRMSNorm(self.head_dim)
+
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        cos: torch.Tensor | None = None,
+        sin: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        # Per-head QK norm (Qwen3)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Apply RoPE
+        if cos is not None and sin is not None:
+            q = _apply_rotary_pos_emb(q, cos, sin)
+            k = _apply_rotary_pos_emb(k, cos, sin)
+
+        # Expand KV heads for GQA (8 KV heads → 16 Q heads)
+        if self.num_kv_heads != self.num_heads:
+            repeat_factor = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(repeat_factor, dim=2)
+            v = v.repeat_interleave(repeat_factor, dim=2)
+
+        # Full bidirectional attention via SDPA with proper mask support
+        # Permute to (batch, heads, seq, head_dim) for SDPA
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+
+        # Convert [B, 1, S, S] bool mask to float mask for SDPA
+        sdpa_mask = None
+        if attention_mask is not None:
+            sdpa_mask = attention_mask.to(dtype=q.dtype)
+            sdpa_mask = sdpa_mask.masked_fill(~attention_mask, float("-inf"))
+            sdpa_mask = sdpa_mask.masked_fill(attention_mask, 0.0)
+
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=sdpa_mask,
+            scale=1.0 / math.sqrt(self.head_dim),
+        )
+
+        # Back to (batch, seq, heads * head_dim)
+        out = out.permute(0, 2, 1, 3).contiguous()
+        out = out.view(batch_size, seq_len, self.num_heads * self.head_dim)
+        return self.o_proj(out)
+
+
+class OmniVoiceMLP(nn.Module):
+    """Qwen3-style MLP with SwiGLU."""
+
+    def __init__(self, config: OmniVoiceConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.llm_hidden_size, config.llm_intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.llm_hidden_size, config.llm_intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.llm_intermediate_size, config.llm_hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class OmniVoiceTransformerBlock(nn.Module):
+    """Single Qwen3 transformer block with DiffusionAttention."""
+
+    def __init__(self, config: OmniVoiceConfig):
+        super().__init__()
+        self.input_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps)
+        self.self_attn = OmniVoiceAttention(config)
+        self.post_attention_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps)
+        self.mlp = OmniVoiceMLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        cos: torch.Tensor | None = None,
+        sin: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask, cos=cos, sin=sin)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# RoPE
+# ---------------------------------------------------------------------------
+
+
+def _precompute_rope(
+    head_dim: int,
+    max_seq_len: int,
+    theta: float = 1000000.0,
+    device: torch.device | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Precompute RoPE cos/sin tensors."""
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
+    t = torch.arange(max_seq_len, device=device, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)
+    cos = freqs.cos()
+    sin = freqs.sin()
+    return cos, sin
+
+
+def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply rotary position embedding. x shape: (B, S, H, D)."""
+    seq_len = x.shape[1]
+    cos = cos[:seq_len].unsqueeze(0).unsqueeze(2)  # (1, S, 1, D/2)
+    sin = sin[:seq_len].unsqueeze(0).unsqueeze(2)
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    rotated = torch.cat([-x2, x1], dim=-1)
+    return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Generator model
+# ---------------------------------------------------------------------------
+
+
+class OmniVoiceGenerator(nn.Module):
+    """OmniVoice Stage 0: Iterative unmasking generator.
+
+    Architecture:
+    - Text embedding (from Qwen3 vocab) + Audio embedding (8*1025 entries)
+    - 28-layer Qwen3 transformer with full bidirectional attention
+    - 8-codebook prediction head (single linear: hidden → 8*1025)
+    - 32-step iterative unmasking with classifier-free guidance
+
+    Optimizations:
+    - DiffusionAttention (FlashAttn/SageAttn/SDPA auto-selected)
+    - TeaCache / Cache-DiT compatible (hook-based, non-intrusive)
+    - regionally_compile() compatible for torch.compile on repeated blocks
+    - Sequence parallelism via SP hooks for multi-GPU
+    """
+
+    # For regionally_compile() support
+    _repeated_blocks = ["layers"]
+
+    def __init__(self, config: OmniVoiceConfig):
+        super().__init__()
+        self.config = config
+
+        # Text embedding (shared with LLM)
+        self.text_embedding = nn.Embedding(config.llm_vocab_size, config.llm_hidden_size)
+
+        # Audio embedding: 8 codebooks * 1025 tokens
+        self.audio_embeddings = nn.Embedding(
+            config.num_audio_codebook * config.audio_vocab_size,
+            config.llm_hidden_size,
+        )
+        self.register_buffer(
+            "codebook_layer_offsets",
+            torch.arange(config.num_audio_codebook) * config.audio_vocab_size,
+        )
+
+        # Transformer layers
+        self.layers = nn.ModuleList([OmniVoiceTransformerBlock(config) for _ in range(config.llm_num_hidden_layers)])
+        self.norm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps)
+
+        # Prediction head: hidden → 8 * 1025
+        self.audio_heads = nn.Linear(
+            config.llm_hidden_size,
+            config.num_audio_codebook * config.audio_vocab_size,
+            bias=False,
+        )
+
+        # Precompute RoPE
+        self._rope_cos = None
+        self._rope_sin = None
+
+    def _ensure_rope(self, seq_len: int, device: torch.device) -> None:
+        """Lazily compute RoPE cos/sin if needed."""
+        if self._rope_cos is None or self._rope_cos.shape[0] < seq_len:
+            max_len = max(seq_len, 4096)
+            self._rope_cos, self._rope_sin = _precompute_rope(
+                self.config.llm_head_dim,
+                max_len,
+                theta=self.config.llm_rope_theta,
+                device=device,
+            )
+
+    def _prepare_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Prepare mixed text+audio embeddings.
+
+        Args:
+            input_ids: [B, 8, S] - text tokens replicated across codebooks,
+                       audio positions have per-codebook token IDs
+            audio_mask: [B, S] - True for audio positions, False for text
+
+        Returns:
+            embeddings: [B, S, hidden_size]
+        """
+        # Text embeddings from first codebook row (all rows identical for text)
+        text_embeds = self.text_embedding(input_ids[:, 0, :])
+
+        # Audio embeddings: offset per codebook, then sum across codebooks
+        shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1)
+        audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1)
+
+        # Merge: audio where audio_mask=True, text elsewhere
+        return torch.where(audio_mask.unsqueeze(-1), audio_embeds, text_embeds)
+
+    def _transformer_forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Run through transformer layers.
+
+        Args:
+            inputs_embeds: [B, S, hidden_size]
+            attention_mask: [B, 1, S, S] or None
+
+        Returns:
+            hidden_states: [B, S, hidden_size]
+        """
+        device = inputs_embeds.device
+        seq_len = inputs_embeds.shape[1]
+        self._ensure_rope(seq_len, device)
+
+        hidden_states = inputs_embeds
+        cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype)
+        sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype)
+
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                cos=cos,
+                sin=sin,
+            )
+
+        return self.norm(hidden_states)
+
+    def _get_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Project hidden states to per-codebook logits.
+
+        Args:
+            hidden_states: [B, S, hidden_size]
+
+        Returns:
+            logits: [B, 8, S, 1025]
+        """
+        batch_size, seq_len, _ = hidden_states.shape
+        logits_flat = self.audio_heads(hidden_states)  # [B, S, 8*1025]
+        return logits_flat.view(
+            batch_size,
+            seq_len,
+            self.config.num_audio_codebook,
+            self.config.audio_vocab_size,
+        ).permute(0, 2, 1, 3)  # [B, 8, S, 1025]
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        audio_mask: torch.Tensor,
+        attention_mask: torch.Tensor,
+        target_lens: list[int],
+        num_step: int = 32,
+        guidance_scale: float = 2.0,
+        t_shift: float = 0.1,
+        layer_penalty_factor: float = 5.0,
+        position_temperature: float = 5.0,
+        class_temperature: float = 0.0,
+    ) -> torch.Tensor:
+        """Run the full 32-step iterative unmasking generation.
+
+        Args:
+            input_ids: [2*B, 8, S] - conditional (0:B) + unconditional (B:2B)
+            audio_mask: [2*B, S] - True for audio positions
+            attention_mask: [2*B, 1, S, S] - attention mask
+            target_lens: List of target audio lengths per batch item
+            num_step: Number of unmasking steps
+            guidance_scale: CFG scale
+            t_shift: Time shift for schedule
+            layer_penalty_factor: Penalty for later codebooks
+            position_temperature: Gumbel temperature for position selection
+            class_temperature: Temperature for token prediction (0=greedy)
+
+        Returns:
+            tokens: [B, 8, max_target_len] - generated audio tokens
+        """
+        B = len(target_lens)
+        device = input_ids.device
+        max_target_len = max(target_lens)
+        mask_id = self.config.audio_mask_id
+        num_codebooks = self.config.num_audio_codebook
+
+        # Initialize all target tokens as [MASK]
+        tokens = torch.full(
+            (B, num_codebooks, max_target_len),
+            mask_id,
+            dtype=torch.long,
+            device=device,
+        )
+
+        # Compute unmasking schedule
+        timesteps = _get_time_steps(0.0, 1.0, num_step + 1, t_shift).tolist()
+        schedules = []
+        for t_len in target_lens:
+            total_mask = t_len * num_codebooks
+            rem = total_mask
+            sched = []
+            for step in range(num_step):
+                num = (
+                    rem
+                    if step == num_step - 1
+                    else min(
+                        math.ceil(total_mask * (timesteps[step + 1] - timesteps[step])),
+                        rem,
+                    )
+                )
+                sched.append(int(num))
+                rem -= int(num)
+            schedules.append(sched)
+
+        layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1)
+
+        # Compute c_lens for extracting target region from full sequence
+        c_lens = []
+        for i in range(B):
+            # Conditional sequence length = number of non-padding positions
+            c_len = attention_mask[i, 0, 0].sum().item()
+            c_lens.append(int(c_len))
+
+        # Main iterative loop
+        for step in range(num_step):
+            # Prepare embeddings and run transformer
+            inputs_embeds = self._prepare_embeddings(input_ids, audio_mask)
+            hidden_states = self._transformer_forward(inputs_embeds, attention_mask)
+            batch_logits = self._get_logits(hidden_states).to(torch.float32)
+            # batch_logits: [2*B, 8, S, 1025]
+
+            for i in range(B):
+                k = schedules[i][step]
+                if k <= 0:
+                    continue
+
+                c_len = c_lens[i]
+                t_len = target_lens[i]
+
+                # Extract logits for target region
+                c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :]  # [1, 8, T, 1025]
+                u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :]  # [1, 8, T, 1025]
+
+                # Classifier-free guidance
+                if guidance_scale != 0:
+                    c_log_probs = F.log_softmax(c_logits, dim=-1)
+                    u_log_probs = F.log_softmax(u_logits, dim=-1)
+                    log_probs = torch.log_softmax(
+                        c_log_probs + guidance_scale * (c_log_probs - u_log_probs),
+                        dim=-1,
+                    )
+                else:
+                    log_probs = F.log_softmax(c_logits, dim=-1)
+
+                # Prevent predicting [MASK]
+                log_probs[..., mask_id] = -float("inf")
+
+                # Token prediction
+                if class_temperature > 0.0:
+                    pred_tokens = _gumbel_sample(log_probs, class_temperature).argmax(dim=-1)
+                else:
+                    pred_tokens = log_probs.argmax(dim=-1)  # [1, 8, T]
+
+                # Confidence scores
+                scores = log_probs.max(dim=-1)[0]  # [1, 8, T]
+
+                # Layer penalty (earlier codebooks get higher priority)
+                scores = scores - (layer_ids * layer_penalty_factor)
+
+                # Gumbel noise for position selection
+                if position_temperature > 0.0:
+                    scores = _gumbel_sample(scores, position_temperature)
+
+                # Mask out already unmasked positions
+                sample_tokens = tokens[i : i + 1, :, :t_len]
+                scores.masked_fill_(sample_tokens != mask_id, -float("inf"))
+
+                # Select top-k positions to unmask
+                _, topk_idx = torch.topk(scores.flatten(), k)
+                flat_tokens = sample_tokens.flatten().clone()
+                flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx]
+                sample_tokens.copy_(flat_tokens.view_as(sample_tokens))
+
+                # Update tokens and batch inputs for next iteration
+                tokens[i : i + 1, :, :t_len] = sample_tokens
+                input_ids = input_ids.clone()
+                input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0)
+                input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0)
+
+        return tokens
+
+    def load_weights(self, model_dir: str, device: torch.device) -> None:
+        """Load weights from HuggingFace OmniVoice model.safetensors.
+
+        The HF checkpoint contains:
+        - llm.* -> Qwen3 transformer weights
+        - audio_embeddings.* -> audio embedding table
+        - audio_heads.* -> prediction head
+        """
+        import os
+
+        from safetensors.torch import load_file
+
+        weights_path = os.path.join(model_dir, "model.safetensors")
+        if not os.path.exists(weights_path):
+            raise FileNotFoundError(f"Model weights not found at {weights_path}")
+
+        state_dict = load_file(weights_path, device=str(device))
+
+        # Map HF weight names to our module names
+        loaded_keys = set()
+
+        # 1. Text embedding: llm.embed_tokens.weight -> text_embedding.weight
+        text_emb_key = "llm.embed_tokens.weight"
+        if text_emb_key in state_dict:
+            self.text_embedding.weight.data.copy_(state_dict[text_emb_key])
+            loaded_keys.add(text_emb_key)
+
+        # 2. Audio embeddings
+        for key in ["audio_embeddings.weight"]:
+            if key in state_dict:
+                self.audio_embeddings.weight.data.copy_(state_dict[key])
+                loaded_keys.add(key)
+
+        # 3. Audio heads
+        for key in ["audio_heads.weight"]:
+            if key in state_dict:
+                self.audio_heads.weight.data.copy_(state_dict[key])
+                loaded_keys.add(key)
+
+        # 4. Transformer layers: llm.layers.N.* -> layers.N.*
+        for key, value in state_dict.items():
+            if key.startswith("llm.layers."):
+                # llm.layers.0.self_attn.q_proj.weight -> layers.0.self_attn.q_proj.weight
+                our_key = key.replace("llm.layers.", "layers.")
+                parts = our_key.split(".")
+                module = self
+                try:
+                    for part in parts[:-1]:
+                        if part.isdigit():
+                            module = module[int(part)]
+                        else:
+                            module = getattr(module, part)
+                    param_name = parts[-1]
+                    param = getattr(module, param_name)
+                    if isinstance(param, nn.Parameter):
+                        param.data.copy_(value)
+                    elif isinstance(param, torch.Tensor):
+                        param.copy_(value)
+                    loaded_keys.add(key)
+                except (AttributeError, IndexError, KeyError) as e:
+                    logger.warning("Failed to load weight %s: %s", key, e)
+
+        # 5. Final norm: llm.norm.weight -> norm.weight
+        norm_key = "llm.norm.weight"
+        if norm_key in state_dict:
+            self.norm.weight.data.copy_(state_dict[norm_key])
+            loaded_keys.add(norm_key)
+
+        unloaded = set(state_dict.keys()) - loaded_keys
+        # Filter out audio_tokenizer weights (loaded in decoder stage)
+        unloaded = {k for k in unloaded if not k.startswith("audio_tokenizer.")}
+        if unloaded:
+            logger.info(
+                "Generator: %d/%d weights loaded, %d skipped (decoder weights)",
+                len(loaded_keys),
+                len(state_dict),
+                len(unloaded),
+            )
+        else:
+            logger.info("Generator: all %d weights loaded", len(loaded_keys))
@@ -57,6 +57,11 @@ _OMNI_MODELS = {
        "cosyvoice3",
        "CosyVoice3Model",
    ),
+    "OmniVoiceModel": (
+        "omnivoice",
+        "omnivoice",
+        "OmniVoiceModel",
+    ),
    "MammothModa2Qwen2ForCausalLM": (
        "mammoth_moda2",
        "mammoth_moda2",
@@ -0,0 +1,20 @@
+# Stage config for OmniVoice TTS via diffusion engine
+# Single stage: text → 32-step iterative unmasking → 8-codebook tokens → DAC decode → 24kHz audio
+
+stage_args:
+  - stage_id: 0
+    stage_type: diffusion
+    is_comprehension: true
+    runtime:
+      devices: 0
+    engine_args:
+      model_stage: dit
+      model_class_name: "OmniVoicePipeline"
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: audio
+      distributed_executor_backend: "mp"
+      dtype: "float32"
+    final_output: true
+    final_output_type: audio
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inter-stage processor for OmniVoice: Generator → Decoder."""
+
+from typing import Any
+
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+
+def tokens2audio(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    requires_multimodal_data: bool = True,
+):
+    """Build stage-1 (decoder) inputs from stage-0 (generator) outputs.
+
+    Takes the 8-codebook audio tokens from the generator and packages
+    them for the HiggsAudioV2 decoder.
+    """
+    source_stage_id = engine_input_source[0]
+    source_outputs = stage_list[source_stage_id].engine_outputs
+
+    if not isinstance(prompt, list):
+        prompt = [prompt]
+
+    source_output = source_outputs[0]
+    output = source_output.outputs[0]
+
+    multi_modal_data = output.multimodal_output
+    if multi_modal_data is None:
+        raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
+
+    # Pass audio_tokens from generator to decoder
+    engine_input = OmniTokensPrompt(
+        prompt_token_ids=output.token_ids,
+        additional_information=multi_modal_data,
+    )
+    return [engine_input]