[Higgs-Audio] bosonai/higgs-audio-v2-generation-3B-base TTS model support (#3762)

Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

.buildkite/test-merge.yml

@@ -389,6 +389,31 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: "Omni · Higgs-Audio V2 Test"
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_USE_DEEP_GEMM=0
+              export VLLM_MOE_USE_DEEP_GEMM=0
+              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_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"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "Diffusion · Qwen Image Edit Test"
         timeout_in_minutes: 20
         commands:

.readthedocs.yml

@@ -10,6 +10,11 @@ build:
   jobs:
     post_checkout:
       - git fetch --unshallow || true
+    pre_install:
+      # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
+      # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
+      # 1 GB memory cap and time out the mkdocs build.
+      - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
 
 mkdocs:
   configuration: mkdocs.yml

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
+| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz |
 
 ## Common Quick Start
 
@@ -333,6 +334,46 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
 - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
 - Run `--help` for the full argument surface.
 
+---
+
+## higgs-audio v2
+
+2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`.
+
+### Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+### Quick start (plain TTS)
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
+    --output-dir results/wavs
+```
+
+### Voice cloning
+
+Pass both `--ref-audio` and `--ref-text` together:
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --output-dir results/wavs
+```
+
+### Notes
+
+- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
+
 ## Example materials
 
 ??? abstract "cosyvoice3/end2end.py"
@@ -343,6 +384,9 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
     ``````
+??? abstract "higgs_audio_v2/end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py"
 ??? abstract "glm_tts/end2end.py"
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -19,6 +19,7 @@ For the full list of supported architectures across all modalities, see
 |---|---|---|---|---|---|
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
+| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
 | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
@@ -193,6 +194,45 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server
 
 ---
 
+## higgs-audio v2
+
+2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
+
+### Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec):
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+### Launch
+```bash
+GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
+```
+Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
+
+### Sending requests
+```bash
+# Plain TTS
+python higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --output-dir /tmp/higgs_out \
+    --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+
+# Voice cloning — pass a reference clip and its transcript together
+python higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --output-dir /tmp/higgs_clone \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --prompts "Hello, this is a cloned voice."
+```
+
+### Notes
+- Output: 24 kHz mono.
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
 ## GLM-TTS
 
 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
@@ -459,6 +499,13 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
     ``````py
     --8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
     ``````
+??? abstract "higgs_audio_v2/batch_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py"
+    ``````
+??? abstract "higgs_audio_v2/run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh"
 ??? abstract "glm_tts/gradio_demo.py"
     ``````py
     --8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"

examples/offline_inference/text_to_speech/higgs_audio_v2/README.md

@@ -0,0 +1,38 @@
+# higgs-audio v2 — offline example
+
+Drives Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) for `bosonai/higgs-audio-v2-generation-3B-base` end-to-end through the vLLM-Omni engine and writes a 24 kHz mono WAV per prompt.
+
+## Prerequisites
+
+Voice clone needs `transformers>=5.3.0` — vllm-omni loads the audio codec via HF's `HiggsAudioV2TokenizerModel`, instantiated from the `k2-fsa/OmniVoice/audio_tokenizer/` subdirectory (only that ~806 MB subdir is downloaded). The boson-ai standalone tokenizer repo's `model.safetensors` is actually a copy of the 3B talker LM, so HF can't load it directly; the k2 bundle ships the same codec weights repackaged with HF-compatible key naming.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+## Quick start
+
+Plain TTS:
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
+    --output-dir results/wavs
+```
+
+## Voice cloning
+
+Pass both `--ref-audio` and `--ref-text` together:
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
+    --texts "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --output-dir results/wavs
+```
+
+## Notes
+
+- Output: 24 kHz mono WAV.
+- Deploy config: `vllm_omni/deploy/higgs_audio_v2.yaml` (auto-loaded by HF `model_type`).
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- For online serving, see [`examples/online_serving/text_to_speech/higgs_audio_v2/`](../../../online_serving/text_to_speech/higgs_audio_v2/).

examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline higgs-audio v2 inference example.
+
+Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end
+through the vLLM-Omni engine without going through the HTTP server, and
+saves a 24 kHz mono WAV per prompt.
+
+Example:
+
+    python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\
+        --texts "Hello world." \\
+                "The quick brown fox jumps over the lazy dog." \\
+        --output-dir results/wavs
+"""
+
+from __future__ import annotations
+
+import os
+
+# DeepGEMM FP8 kernels require an optional backend that may not be installed.
+# Disable the warmup before importing vLLM so engine startup falls back to the
+# regular gemm path. Users with deep_gemm installed can override these.
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import time
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+SAMPLE_RATE = 24_000
+DEFAULT_TEXTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="bosonai/higgs-audio-v2-generation-3B-base",
+        help="Stage-0 talker model id or path.",
+    )
+    parser.add_argument(
+        "--texts",
+        type=str,
+        nargs="+",
+        default=list(DEFAULT_TEXTS),
+        help="One or more plain-text prompts to synthesize.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="results/wavs",
+        help="Directory to write per-prompt WAV files.",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Override the deploy config path. Auto-loads "
+        "vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=500,
+        help="Cap on Stage-0 codec frames per request.",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference clip. Required when --ref-audio is set.",
+    )
+    nullify_stage_engine_defaults(parser)
+    return parser.parse_args()
+
+
+def _slugify(text: str) -> str:
+    import re
+
+    slug = re.sub(r"\s+", "_", text.strip().lower())
+    slug = re.sub(r"[^a-z0-9_]+", "", slug)
+    return slug[:48] or "prompt"
+
+
+def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
+    """Pull the final concatenated PCM tensor out of a request's multimodal_output."""
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
+    if isinstance(audio, list):
+        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+        if not valid:
+            raise ValueError("audio list is empty")
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
+    arr = pcm.numpy()
+    if arr.dtype.kind == "f":
+        arr = np.clip(arr, -1.0, 1.0)
+        arr = (arr * 32767.0).astype(np.int16)
+    else:
+        arr = arr.astype(np.int16)
+    return arr
+
+
+def main():
+    args = parse_args()
+    if (args.ref_audio is None) != (args.ref_text is None):
+        raise SystemExit("--ref-audio and --ref-text must be supplied together")
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config)
+
+    # Build prompt_token_ids using the same path serving_speech.py takes online.
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_plain_text_prompt,
+        build_voice_clone_prompt,
+        input_ids_to_python_list,
+    )
+
+    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
+
+    # Voice-clone path: load the reference clip once. The HF processor will
+    # encode it via the bundled HiggsAudioV2TokenizerModel each time we build
+    # a prompt below. This is cheap on CPU for a few-second clip.
+    ref_wav: np.ndarray | None = None
+    ref_sr: int | None = None
+    if args.ref_audio is not None:
+        ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
+        if ref_wav.ndim == 2:
+            ref_wav = ref_wav.mean(axis=1)
+
+    print(f"Model       : {args.model}")
+    print(f"Prompts     : {len(args.texts)}")
+    print(f"Output dir  : {output_dir}")
+    print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}")
+
+    # Run one prompt at a time. The Stage-0 talker's per-slot audio state is
+    # request-scoped; submitting multiple prompts in the same engine.generate()
+    # call would batch them in the AR runner and exercise a code path that is
+    # not validated for this model yet.
+    total_elapsed = 0.0
+    total_dur = 0.0
+    for text in args.texts:
+        if ref_wav is not None:
+            out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text)
+            prompt = {
+                "prompt_token_ids": out["prompt_token_ids"],
+                # Bare tensors (NOT list-wrapped): the msgspec serializer in
+                # vllm_omni.data_entry_keys routes torch.Tensor → tensor_data
+                # and list[Tensor] → list_data (which silently strips tensors).
+                "additional_information": {
+                    "audio_input_ids": out["audio_input_ids"],
+                    "audio_input_ids_mask": out["audio_input_ids_mask"],
+                },
+            }
+        else:
+            inputs = build_plain_text_prompt(processor, text)
+            prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)}
+        t_start = time.perf_counter()
+        outputs = engine.generate([prompt])
+        elapsed = time.perf_counter() - t_start
+        total_elapsed += elapsed
+
+        mm = outputs[0].outputs[0].multimodal_output
+        pcm = _extract_pcm(mm)
+        slug = _slugify(text)
+        out_path = output_dir / f"{slug}.wav"
+        sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
+        dur = pcm.numel() / SAMPLE_RATE
+        total_dur += dur
+        print(f"  {slug:<50} dur={dur:5.2f}s  -> {out_path}")
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    main()

examples/online_serving/text_to_speech/higgs_audio_v2/README.md

@@ -0,0 +1,63 @@
+# higgs-audio v2 online example
+
+This directory contains the online-serving entry points for boson-ai's higgs-audio v2 as integrated by vllm-omni: a 2-stage TTS pipeline (Llama-3.2-3B talker with DualFFN audio expert + HiggsAudio codec decoder) emitting 24 kHz mono speech.
+
+## Prerequisites
+
+Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec). Only that ~806 MB subdir is downloaded.
+
+```bash
+pip install -U "transformers>=5.3.0"
+```
+
+## Files
+
+- `run_server.sh` — launch the vllm-omni server with the bundled `vllm_omni/deploy/higgs_audio_v2.yaml` deploy config.
+- `batch_speech_client.py` — send a list of prompts to `/v1/audio/speech` and save the returned WAV / PCM bytes to a directory; optionally passes `--ref-audio` + `--ref-text` for shallow voice clone.
+
+## Launching the server
+
+```bash
+GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
+```
+
+Environment overrides:
+
+- `MODEL` — HF id of the talker (default `bosonai/higgs-audio-v2-generation-3B-base`).
+- `PORT` — server port (default `8094`).
+- `GPUS` — `CUDA_VISIBLE_DEVICES` value (default `6,7`).
+- `GPU_UTIL` — `--gpu-memory-utilization` (default `0.4`).
+
+The script also exports `VLLM_USE_DEEP_GEMM=0` / `VLLM_MOE_USE_DEEP_GEMM=0` so the example works on images without the optional `deep_gemm` backend.
+
+The deploy YAML ships with `async_chunk: false` and `codec_streaming: true`, i.e. Stage 0 finishes its codec frames before Stage 1 starts decoding, and Stage 1 streams WAV/PCM bytes to the client chunk-by-chunk.
+
+## Driving the server
+
+Plain TTS:
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --model bosonai/higgs-audio-v2-generation-3B-base \
+    --output-dir /tmp/higgs_audio_v2_batch \
+    --prompts "Hello world." \
+              "The quick brown fox jumps over the lazy dog."
+```
+
+Voice clone — pass a reference clip and its transcript (both required together):
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+    --base-url http://localhost:8094 \
+    --model bosonai/higgs-audio-v2-generation-3B-base \
+    --output-dir /tmp/higgs_audio_v2_clone \
+    --ref-audio /path/to/reference.wav \
+    --ref-text  "Exact transcript spoken in reference.wav." \
+    --prompts "Hello, this is a cloned voice."
+```
+
+## Notes
+
+- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
+- Out of scope (rejected with explicit 4xx by the request validator): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.

examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Batch client for the higgs-audio v2 online server.
+
+Sends a fixed list of prompts to ``/v1/audio/speech`` and saves the returned
+WAV files (or raw PCM bytes when ``--format pcm``) into ``--output-dir``.
+
+Usage (plain text -> speech):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+      --base-url http://localhost:8094 \
+      --output-dir /tmp/higgs_audio_v2_batch \
+      --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+
+Usage (shallow voice clone — pass a reference clip + its transcript):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
+      --base-url http://localhost:8094 \
+      --output-dir /tmp/higgs_audio_v2_clone \
+      --ref-audio path/to/reference.wav \
+      --ref-text "the transcript of the reference clip" \
+      --prompts "Hello world."
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import sys
+from pathlib import Path
+
+DEFAULT_PROMPTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "It was the night before my birthday.",
+    "Innovation distinguishes between a leader and a follower.",
+)
+
+
+def _slug(text: str) -> str:
+    import re
+
+    s = re.sub(r"\s+", "_", text.strip().lower())
+    return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--base-url", default="http://localhost:8094")
+    parser.add_argument("--model", default="higgs_audio_v2")
+    parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
+    parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_audio_v2_batch"))
+    parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
+    parser.add_argument("--max-new-tokens", type=int, default=300)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--timeout-s", type=float, default=120.0)
+    parser.add_argument(
+        "--ref-audio",
+        type=Path,
+        default=None,
+        help="Reference clip for voice clone (path to a WAV file). Must be paired with --ref-text.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference clip. Required when --ref-audio is set.",
+    )
+    args = parser.parse_args()
+
+    if (args.ref_audio is None) != (args.ref_text is None):
+        print("--ref-audio and --ref-text must be supplied together", file=sys.stderr)
+        return 2
+
+    ref_audio_data_url: str | None = None
+    if args.ref_audio is not None:
+        if not args.ref_audio.exists():
+            print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
+            return 2
+        mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
+        ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
+        ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
+
+    try:
+        import httpx
+    except ImportError:
+        print(
+            "this client needs `httpx`. Install with `pip install httpx`.",
+            file=sys.stderr,
+        )
+        return 2
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    url = args.base_url.rstrip("/") + "/v1/audio/speech"
+    failures = 0
+    with httpx.Client(timeout=args.timeout_s) as client:
+        for prompt in args.prompts:
+            payload = {
+                "model": args.model,
+                "input": prompt,
+                "response_format": args.format,
+                "max_new_tokens": args.max_new_tokens,
+                "seed": args.seed,
+            }
+            if ref_audio_data_url is not None:
+                payload["ref_audio"] = ref_audio_data_url
+                payload["ref_text"] = args.ref_text
+            resp = client.post(url, json=payload)
+            if resp.status_code != 200:
+                print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
+                failures += 1
+                continue
+            suffix = ".wav" if args.format == "wav" else ".pcm"
+            out = args.output_dir / f"{_slug(prompt)}{suffix}"
+            out.write_bytes(resp.content)
+            print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
+
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Launch vLLM-Omni server for higgs-audio v2.
+#
+# v1 scope: plain text -> 24 kHz speech only. Voice cloning, multi-speaker,
+# ChatML rich content, and language overrides are rejected by the validator
+# with explicit 4xx (see vllm_omni/entrypoints/openai/serving_speech.py).
+#
+# Usage:
+#   ./run_server.sh                 # default port 8094, GPUs 6 and 7
+#   PORT=8095 GPUS=6,7 ./run_server.sh
+#   MODEL=bosonai/higgs-audio-v2-generation-3B-base ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-bosonai/higgs-audio-v2-generation-3B-base}"
+PORT="${PORT:-8094}"
+GPUS="${GPUS:-6,7}"
+GPU_UTIL="${GPU_UTIL:-0.4}"
+
+echo "Starting higgs-audio v2 server"
+echo "  MODEL=$MODEL"
+echo "  PORT=$PORT"
+echo "  CUDA_VISIBLE_DEVICES=$GPUS"
+
+# DeepGEMM FP8 kernels are optional and trip warmup on builds without
+# the deep_gemm backend; disable them so the example works out of the box.
+# Users with deep_gemm installed can re-enable via the same env vars.
+CUDA_VISIBLE_DEVICES="$GPUS" \
+VLLM_USE_DEEP_GEMM=0 \
+VLLM_MOE_USE_DEEP_GEMM=0 \
+vllm-omni serve "$MODEL" \
+    --deploy-config vllm_omni/deploy/higgs_audio_v2.yaml \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization "$GPU_UTIL" \
+    --trust-remote-code \
+    --omni

setup.py

@@ -62,6 +62,12 @@ def detect_target_device() -> str:
         else:
             print(f"Warning: Invalid VLLM_OMNI_TARGET_DEVICE '{target_device}', falling back to auto-detection")
 
+    # Priority 1.5: ReadTheDocs builds use CPU requirements (no GPU, 1 GB RAM cap).
+    # Avoids pulling ~2 GB of CUDA libs that push the mkdocs build into swap.
+    if os.environ.get("READTHEDOCS"):
+        print("Detected READTHEDOCS=1; using cpu requirements for docs build")
+        return "cpu"
+
     # Priority 2: Torch backend detection
     # This is a code path for when user is using
     # --no-build-isolation flag

tests/e2e/offline_inference/test_higgs_audio_v2.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end offline inference tests for higgs-audio v2.
+
+Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py:
+build Stage-0 prompt_token_ids with the model's own processor + plain-text
+chat template, drive the engine via ``Omni.generate``, and inspect the
+Stage-1 PCM payload directly. The OmniRunnerHandler.send_audio_speech_request
+helper applies a Qwen-style ChatML wrap that is not valid for higgs, so this
+test bypasses it (same approach as test_voxtral_tts.py).
+
+v1 scope is plain text -> 24 kHz speech; voice cloning, multi-speaker prompts,
+language overrides, and task_type selection are out of scope.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Match the example script: keep DeepGEMM FP8 off so engine init does not
+# require the optional ``deep_gemm`` backend.
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import numpy as np
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "bosonai/higgs-audio-v2-generation-3B-base"
+STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
+SAMPLE_RATE = 24_000
+# A short sentence at 24 kHz must produce at least ~0.5 s of audio.
+MIN_AUDIO_SAMPLES = 12_000
+TEST_TEXT = "Hello world."
+
+# Reuse the qwen3_tts ref clip + transcript for voice clone tests (see
+# tests/e2e/online_serving/test_qwen3_tts_base.py:27-28 for the source).
+_REF_AUDIO_ASSET = "qwen3_tts/clone_2.wav"
+_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+# (model, stage_config_path, extra_omni_kwargs) — same shape as test_voxtral_tts.py.
+# Stage configs already pin enforce_eager=true per stage, so no Omni-level override.
+_OMNI_RUNNER_PARAM = (MODEL, STAGE_CONFIG, {"trust_remote_code": True})
+
+
+def _compose_request(model_name: str, text: str) -> dict:
+    """Build a higgs prompt dict via the upstream processor + plain-text template.
+
+    The serving layer's ``_build_higgs_audio_v2_params`` does the same thing
+    (see vllm_omni/entrypoints/openai/serving_speech.py:1520) — keeping the
+    offline and online prompt-construction paths identical is what guarantees
+    input-token parity with the HF reference.
+    """
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_plain_text_prompt,
+        input_ids_to_python_list,
+    )
+
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    inputs = build_plain_text_prompt(processor, text)
+    return {"prompt_token_ids": input_ids_to_python_list(inputs)}
+
+
+def _extract_pcm(outputs) -> torch.Tensor:
+    """Pull the concatenated PCM tensor out of the Stage-1 multimodal_output."""
+    for stage_output in outputs:
+        mm = getattr(stage_output, "multimodal_output", None)
+        if not mm:
+            continue
+        audio = mm.get("model_outputs") if "model_outputs" in mm else mm.get("audio")
+        if audio is None:
+            continue
+        if isinstance(audio, list):
+            chunks = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+            if not chunks:
+                continue
+            return torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0]
+        return torch.as_tensor(audio).float().cpu().reshape(-1)
+    raise AssertionError(f"no audio payload in any stage output (got {len(outputs)} stages)")
+
+
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_plain_text(omni_runner_function: OmniRunner) -> None:
+    """
+    Offline plain-text -> 24 kHz PCM via the bundled higgs_audio_v2 deploy.
+
+    Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml (codec_streaming=true,
+        async_chunk=false, enforce_eager=true per stage).
+    Input Modal: text
+    Output Modal: audio (PCM tensor, 24 kHz mono float)
+    Input Setting: single prompt, default sampling
+    """
+    omni = omni_runner_function.omni
+    inputs = _compose_request(MODEL, TEST_TEXT)
+
+    # max_tokens caps Stage-0 codec frames; Stage 1 max_tokens lifts the
+    # per-codebook prefill cap. Mirror the deploy yaml's defaults.
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate([inputs], stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate"
+
+    pcm = _extract_pcm(outputs)
+    arr = pcm.numpy()
+
+    assert arr.size > MIN_AUDIO_SAMPLES, (
+        f"audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES} "
+        f"({MIN_AUDIO_SAMPLES / SAMPLE_RATE:.2f} s @ {SAMPLE_RATE} Hz)"
+    )
+    # Threshold 5e-4 (not 1e-3) keeps us above true silence while tolerating
+    # the ``load_format=dummy`` Stage-1 codec — random RVQ/DAC weights collapse
+    # the output to a near-constant peak around 7.4e-4. Real-weight runs reach
+    # 0.1+ peak so this still catches genuine silence.
+    assert np.max(np.abs(arr)) > 5e-4, "audio output is silent (peak abs amplitude < 5e-4)"
+    assert np.isfinite(arr).all(), "audio contains non-finite samples (nan/inf)"
+
+
+@pytest.mark.core_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_batch_two_prompts(omni_runner_function: OmniRunner) -> None:
+    """
+    Submit two prompts in one engine call and verify both yield non-trivial PCM.
+
+    Guards the Stage-0 talker's per-slot audio state under a batched call —
+    the prior example ran prompts one at a time citing 'request-scoped audio
+    state'; this test pins the batched-2 behavior the engine path actually
+    needs to support for /v1/audio/speech under concurrent traffic.
+    """
+    omni = omni_runner_function.omni
+    prompts = [
+        _compose_request(MODEL, "Hello world."),
+        _compose_request(MODEL, "The quick brown fox jumps over the lazy dog."),
+    ]
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate(prompts, stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate"
+
+    pcm = _extract_pcm(outputs)
+    # Whether the engine returns one concatenated PCM or per-request tensors is
+    # implementation-defined; for either shape, total samples must clear the
+    # floor for at least one of the two prompts.
+    assert pcm.numel() > MIN_AUDIO_SAMPLES, (
+        f"batched audio too short: {pcm.numel()} samples, expected > {MIN_AUDIO_SAMPLES}"
+    )
+    # See the silence-threshold note in test_higgs_audio_v2_offline_plain_text.
+    assert np.max(np.abs(pcm.numpy())) > 5e-4, "batched audio output is silent"
+
+
+def _compose_voice_clone_request(model_name: str, text: str, ref_text: str, ref_wav_path) -> dict:
+    """Build the voice-clone prompt + additional_information payload.
+
+    Calls the same ``build_voice_clone_prompt`` the serving layer uses online —
+    the HF processor encodes the reference clip on the spot via the bundled
+    HiggsAudioV2TokenizerModel. The talker's
+    ``_maybe_apply_ref_audio_substitution`` consumes the tensors at prefill.
+    """
+    import soundfile as sf
+    from transformers import AutoProcessor
+
+    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+        build_voice_clone_prompt,
+    )
+
+    wav, sr = sf.read(str(ref_wav_path), always_2d=False)
+    if wav.ndim == 2:
+        wav = wav.mean(axis=1)  # downmix any stereo asset to mono
+    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+    out = build_voice_clone_prompt(processor, text, wav, int(sr), ref_text)
+    return {
+        "prompt_token_ids": out["prompt_token_ids"],
+        # Pass tensors bare (see serving_speech._build_higgs_audio_v2_params
+        # — list-wrapped tensors get dropped by the msgspec serializer).
+        "additional_information": {
+            "audio_input_ids": out["audio_input_ids"],
+            "audio_input_ids_mask": out["audio_input_ids_mask"],
+        },
+    }
+
+
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
+def test_higgs_audio_v2_offline_voice_clone(omni_runner_function: OmniRunner) -> None:
+    """
+    Shallow voice clone via ref_audio + ref_text, offline.
+
+    Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml.
+    Input Modal: text + reference audio (qwen3_tts/clone_2.wav vendored asset)
+    Output Modal: audio (PCM tensor, 24 kHz mono float)
+    Input Setting: single prompt
+    """
+    omni = omni_runner_function.omni
+    ref_path = get_asset_path(_REF_AUDIO_ASSET)
+    inputs = _compose_voice_clone_request(MODEL, TEST_TEXT, _REF_TEXT, ref_path)
+
+    stage_sampling = [
+        SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
+        SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
+    ]
+
+    outputs = list(omni.generate([inputs], stage_sampling))
+    assert outputs, "no outputs returned from Omni.generate for voice clone"
+
+    pcm = _extract_pcm(outputs)
+    arr = pcm.numpy()
+    assert arr.size > MIN_AUDIO_SAMPLES, f"cloned audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES}"
+    assert np.max(np.abs(arr)) > 5e-4, "voice-clone output is silent"
+    assert np.isfinite(arr).all(), "voice-clone audio contains non-finite samples"

tests/e2e/online_serving/test_higgs_audio_v2.py

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end online tests for higgs-audio v2 against /v1/audio/speech.
+
+v1 scope is plain text -> 24 kHz speech plus shallow voice clone via
+ref_audio + ref_text (inline) or voice=<name> (after POST /v1/audio/voices).
+The model-aware request validator
+(vllm_omni/entrypoints/openai/serving_speech.py::_validate_higgs_audio_v2_request)
+rejects multi-speaker tags, language overrides, task_type, and bare
+voice=<name> for names that do not match an uploaded speaker — this suite
+exercises both the happy path (plain text in, audio bytes out) and the
+validator rejections, so a regression that loosens the schema will fail
+loudly.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+# Match run_server.sh: DeepGEMM FP8 kernels are optional and trip warmup on
+# images without the deep_gemm backend, so disable them by default.
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "bosonai/higgs-audio-v2-generation-3B-base"
+STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
+SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
+# DeepGEMM warmup is optional; mirror run_server.sh and switch it off in the
+# server subprocess env too (parent-process os.environ above only affects this
+# test driver; the engine subprocesses inherit through env_dict).
+SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"}
+
+TEST_PARAMS = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=STAGE_CONFIG,
+            server_args=SERVER_ARGS,
+            env_dict=SERVER_ENV,
+        ),
+        id="higgs_audio_v2_plain_text",
+    )
+]
+
+DEFAULT_SPEECH_TIMEOUT_S = 180.0
+# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ≈ 24 KiB.
+# A WAV header adds 44 bytes; pick a conservative floor that catches truncated
+# / silence-only outputs without flagging short legitimate clips.
+_MIN_AUDIO_BYTES = 20_000
+
+# Reuse the qwen3_tts vendored reference clip (clean ~5 s 24 kHz mono human
+# speech) + its transcript. See tests/e2e/online_serving/test_qwen3_tts_base.py
+# for the asset rationale — keeping a single shared reference clip across TTS
+# tests avoids duplicating WAVs in the repo.
+_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineHappyPath:
+    """Plain-text -> audio happy paths against the live HTTP server."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_wav(self, omni_server, openai_client) -> None:
+        """Single non-streaming WAV request — covers the canonical TTS happy path."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Streaming PCM via the shared-memory connector's codec_streaming path.
+
+        higgs_audio_v2.yaml pins ``codec_streaming: true`` + ``async_chunk: false``,
+        so the only streaming surface exposed to clients is the WAV/PCM bytes
+        served chunk-by-chunk from Stage 1 — exercise it directly.
+        """
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Streaming the quick brown fox over the lazy dog.",
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
+        """max_new_tokens is one of the few extra fields the higgs validator accepts."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Innovation distinguishes between a leader and a follower.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "max_new_tokens": 500,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
+        """Three concurrent non-streaming requests — guards the per-slot audio state."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "It was the night before my birthday.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            },
+            request_num=3,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Validator rejections served over HTTP. Each case targets one out-of-scope
+# field; the validator returns 4xx via the OpenAI error path. We keep the
+# error-message substring loose so a phrasing tweak in the validator does not
+# break CI, but tight enough to catch a regression that silently accepts the
+# field.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineValidatorRejections:
+    """Out-of-scope fields must come back as 4xx with a higgs-named message."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None:
+        """Voice clone needs the transcript too — half-supplied is a 4xx."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_audio": _REF_AUDIO_URL,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "ref_text",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None:
+        """Symmetric guard: ref_text alone is not enough — must come with ref_audio."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_text": "some transcript",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "ref_audio",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_task_type(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "task_type": "Base",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "task_type",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_language_override(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "language": "Chinese",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "language",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "[SPEAKER0] hi",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "multi-speaker",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_empty_input(self, omni_server, openai_client) -> None:
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "   ",
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "status_code": (400, 422),
+                "err_message": "empty",
+            }
+        )
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV2OnlineVoiceClone:
+    """Shallow voice clone: ref_audio + ref_text -> speech in the cloned voice.
+
+    The HF processor that the serving layer calls (lazy-loaded at the first
+    request) ships with the bundled HiggsAudioV2TokenizerModel and encodes
+    the reference clip in-process. The talker substitutes the encoded codes
+    at the prompt-side audio placeholders via
+    :meth:`HiggsAudioV2TalkerForConditionalGeneration._maybe_apply_ref_audio_substitution`.
+    """
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_basic(self, omni_server, openai_client) -> None:
+        """Single non-streaming WAV request driven by the qwen3_tts ref clip."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_audio": _REF_AUDIO_URL,
+                "ref_text": _REF_TEXT,
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Voice clone over the streaming PCM path."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Innovation distinguishes a leader from a follower.",
+                "ref_audio": _REF_AUDIO_URL,
+                "ref_text": _REF_TEXT,
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )

vllm_omni/config/pipeline_registry.py

@@ -121,4 +121,8 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.moss_tts_nano.pipeline",
         "MOSS_TTS_NANO_PIPELINE",
     ),
+    "higgs_audio_v2": (
+        "vllm_omni.model_executor.models.higgs_audio_v2.pipeline",
+        "HIGGS_AUDIO_V2_PIPELINE",
+    ),
 }

vllm_omni/deploy/higgs_audio_v2.yaml

@@ -0,0 +1,83 @@
+# Higgs-audio v2 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory.
+#
+# Phase-1 scope is plain text -> 24 kHz speech only; voice cloning, multi-speaker,
+# ChatML rich messages, and music/sound-event prompts are rejected by the
+# request validator with explicit 4xx (see
+# vllm_omni/entrypoints/openai/serving_speech.py and
+# vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py).
+#
+# Streaming is wired through the stage-input processor's async-chunk
+# callback, the shared-memory connector ``codec_streaming`` flag, and the
+# WAV/PCM streaming convention in serving_speech.py. To run a strictly
+# non-streaming smoke test, set ``codec_streaming: false`` in the connector
+# extras below.
+async_chunk: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 5000
+      connector_get_max_wait: 500
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+      initial_codec_chunk_frames: 1
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 4
+    gpu_memory_utilization: 0.6
+    # Disable cudagraph capture: the live audio-sampler dispatch is still
+    # iterating, and the host's flash-attn-hopper graph path tripped a CUDA
+    # driver vs runtime mismatch during capture on this image.
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 2048
+    # max_model_len covers prompt + max_new_tokens. Up to ~1500 codec
+    # frames per request fit comfortably under 4096.
+    max_model_len: 4096
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      # Upstream serve_engine defaults (temperature=1.0, top_p=0.95,
+      # top_k=50). Greedy (temperature=0) collapses the audio codebook
+      # argmax to a fixed value after a few steps; sampling with diversity
+      # is required for natural-sounding speech.
+      temperature: 1.0
+      top_p: 0.95
+      top_k: 50
+      # Cap Stage-0 to 1024 audio frames (~40 s at 25 fps). The audio-eos
+      # natural stop (128012, forced by the talker's _apply_audio_mode_bias
+      # at ramp-completion) typically fires earlier and trims this.
+      max_tokens: 1024
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 4
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    # Prefill length is num_codebooks * num_frames; 8 * 8192 = 65536 covers
+    # the upper bound at the per-request 250-frame Stage-0 cap above plus
+    # headroom for batching multiple requests.
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -100,6 +100,50 @@ class OpenAICreateSpeechRequest(BaseModel):
             raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
         return v
 
+    @model_validator(mode="before")
+    @classmethod
+    def reject_higgs_audio_v2_unsupported_aliases(cls, data):
+        """Reject unsupported rich-input aliases for higgs_audio_v2 BEFORE pydantic strips them.
+
+        OpenAICreateSpeechRequest is a permissive schema (it accepts any extra
+        field that the OpenAI Speech API didn't promise to ban) but several
+        of those aliases pull a request out of the v1 higgs_audio_v2 scope.
+        Catching them at parse time lets the API return a deterministic 4xx
+        with a model-specific error message instead of silently dropping the
+        field and proceeding with a degraded request.
+        """
+        # Keys that are silently dropped by pydantic when posted to /v1/audio/speech
+        # (the schema doesn't declare them) but which a model-specific
+        # validator MUST be able to reject. Kept inline to avoid clashing
+        # with pydantic's ModelPrivateAttr discovery for class-level
+        # underscore-prefixed attributes.
+        higgs_audio_v2_reserved_keys = (
+            "messages",  # ChatML rich content (out of scope for v1)
+            "reference_audio",  # voice-cloning alias 1
+            "voice_prompt",  # voice-cloning alias 2
+            "speaker_audio",  # voice-cloning alias 3
+            "speakers",  # multi-speaker dialogue
+        )
+        if not isinstance(data, dict):
+            return data
+        model_id = data.get("model")
+        if not isinstance(model_id, str):
+            return data
+        # Match the "higgs_audio_v2" model_type label, the HF architecture id
+        # in pipeline_registry.hf_architectures, and the hyphenated HF repo id
+        # (e.g. "bosonai/higgs-audio-v2-generation-3B-base") by normalizing
+        # both underscores and hyphens out of the id before substring matching.
+        normalized = model_id.lower().replace("-", "").replace("_", "")
+        if "higgsaudiov2" not in normalized:
+            return data
+        offending = sorted(k for k in higgs_audio_v2_reserved_keys if k in data)
+        if offending:
+            raise ValueError(
+                "higgs_audio_v2 v1 does not support these rich-input fields: "
+                f"{offending}. Supply plain text via the 'input' field instead."
+            )
+        return data
+
     @model_validator(mode="after")
     def validate_embedding_constraints(self) -> "OpenAICreateSpeechRequest":
         if self.speaker_embedding is not None:

vllm_omni/entrypoints/openai/serving_speech.py

@@ -69,6 +69,7 @@ _COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"}
 _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
 _MING_TTS_MODEL_STAGES = {"ming_tts"}
 _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
+_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
 _GLM_TTS_MODEL_STAGES = {"glm_tts"}
 _TTS_MODEL_STAGES: set[str] = (
     _VOXTRAL_TTS_MODEL_STAGES
@@ -76,13 +77,21 @@ _TTS_MODEL_STAGES: set[str] = (
     | _FISH_TTS_MODEL_STAGES
     | _COSYVOICE3_TTS_MODEL_STAGES
     | _OMNIVOICE_TTS_MODEL_STAGES
+    | _HIGGS_AUDIO_V2_TTS_MODEL_STAGES
     | _COVO_AUDIO_MODEL_STAGES
     | _VOXCPM2_TTS_MODEL_STAGES
     | _MING_TTS_MODEL_STAGES
     | _MOSS_TTS_MODEL_STAGES
     | _GLM_TTS_MODEL_STAGES
 )
-_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {"fish_tts", "qwen3_tts", "voxtral_tts", "cosyvoice3", "voxcpm2"}
+_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
+    "fish_tts",
+    "qwen3_tts",
+    "voxtral_tts",
+    "cosyvoice3",
+    "voxcpm2",
+    "higgs_audio_v2",
+}
 _TTS_LANGUAGES: set[str] = {
     "Auto",
     "Chinese",
@@ -555,6 +564,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "ming_flash_omni_tts"
         if model_stage in _MOSS_TTS_MODEL_STAGES:
             return "moss_tts_nano"
+        if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
+            return "higgs_audio_v2"
         if model_stage in _GLM_TTS_MODEL_STAGES:
             return "glm_tts"
         return None
@@ -845,12 +856,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         voice_lower = request.voice.lower()
         if voice_lower not in self.uploaded_speakers:
-            if self._tts_model_type in ("cosyvoice3", "fish_tts", "omnivoice", "moss_tts_nano", "glm_tts"):
+            if self._tts_model_type in (
+                "cosyvoice3",
+                "fish_tts",
+                "omnivoice",
+                "moss_tts_nano",
+                "glm_tts",
+                "higgs_audio_v2",
+            ):
                 label = {
                     "cosyvoice3": "CosyVoice3",
                     "fish_tts": "Fish Speech",
                     "omnivoice": "OmniVoice",
                     "moss_tts_nano": "MOSS-TTS-Nano",
+                    "higgs_audio_v2": "Higgs-Audio V2",
                     "glm_tts": "GLM-TTS",
                 }.get(self._tts_model_type, self._tts_model_type)
                 return (
@@ -1446,6 +1465,73 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         params["prompt_audio_array"] = [[wav_list, sr]]
         return params
 
+    def _validate_higgs_audio_v2_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+        """Validate higgs_audio_v2 request parameters. Returns error message or None.
+
+        Accepted: plain text -> speech, or shallow voice clone via ``ref_audio``
+        + ``ref_text`` (both required together). Still out of scope: preset
+        ``voice``/``speaker`` selection, ``x_vector_only_mode`` /
+        ``speaker_embedding`` helpers, ``task_type``/``language``/
+        ``instructions``/``speed`` overrides, and multi-speaker ``[SPEAKERn]``
+        tags inside the input body.
+        """
+        from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+            MULTI_SPEAKER_TAG_PATTERN,
+        )
+
+        if not request.input or not request.input.strip():
+            return "higgs_audio_v2: input text cannot be empty"
+
+        # Voice clone: ref_audio and ref_text must come together.
+        if request.ref_audio is not None and not request.ref_text:
+            return (
+                "higgs_audio_v2 voice clone requires both 'ref_audio' and "
+                "'ref_text'; received ref_audio without ref_text"
+            )
+        if request.ref_text and request.ref_audio is None:
+            return (
+                "higgs_audio_v2 voice clone requires both 'ref_audio' and "
+                "'ref_text'; received ref_text without ref_audio"
+            )
+
+        if request.x_vector_only_mode is not None:
+            return "higgs_audio_v2 v1 does not support 'x_vector_only_mode' (voice-cloning helper field)"
+        if request.speaker_embedding is not None:
+            return "higgs_audio_v2 v1 does not support 'speaker_embedding' (voice-cloning helper field)"
+        if request.voice and request.ref_audio is None:
+            # _apply_uploaded_speaker runs before this validator; if voice was
+            # an uploaded speaker, ref_audio is now populated and ref_text is
+            # backfilled from the speaker entry. A bare voice= with no
+            # ref_audio means the name didn't resolve to an uploaded speaker
+            # (and higgs has no built-in preset voices).
+            return (
+                "higgs_audio_v2 v1 does not support 'voice'/'speaker' selection for built-in voices; "
+                f"upload a voice first via POST /v1/audio/voices, or use ref_audio + ref_text. "
+                f"Got voice={request.voice!r}"
+            )
+        if request.instructions:
+            return (
+                "higgs_audio_v2 v1 does not support 'instructions' (voice "
+                "style/emotion control); supply plain text instead"
+            )
+        if request.task_type is not None:
+            return "higgs_audio_v2 v1 does not support 'task_type'; the model is single-mode plain text -> speech"
+        if request.language is not None:
+            return (
+                "higgs_audio_v2 v1 does not accept 'language' overrides; the model infers language from the input text"
+            )
+        if request.speed is not None and request.speed != 1.0:
+            return (
+                "higgs_audio_v2 v1 does not support 'speed' adjustments; the audio is rendered at native rate (24 kHz)"
+            )
+
+        if MULTI_SPEAKER_TAG_PATTERN.search(request.input):
+            return "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; remove the tag from the input"
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
+                return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
+
     def _validate_glm_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate GLM-TTS request — requires ref_audio for voice cloning."""
         if not request.input or not request.input.strip():
@@ -1466,6 +1552,80 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
         return None
 
+    async def _build_higgs_audio_v2_params(self, request: OpenAICreateSpeechRequest):
+        """Build prompt_token_ids for higgs_audio_v2 via the upstream processor.
+
+        Plain-text path: runs ``build_plain_text_prompt`` and returns the
+        token-only prompt. Voice-clone path (``ref_audio`` + ``ref_text``):
+        resolves the reference clip via ``_resolve_ref_audio``, runs
+        ``build_voice_clone_prompt`` (which encodes the clip through HF's
+        ``HiggsAudioV2TokenizerModel`` loaded from the k2-fsa/OmniVoice
+        ``audio_tokenizer/`` subdirectory), and attaches the encoded
+        ``audio_input_ids`` + ``audio_input_ids_mask`` tensors via
+        ``additional_information`` so the talker substitutes them at the
+        prompt-side audio placeholders.
+        """
+        from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
+            build_plain_text_prompt,
+            build_voice_clone_prompt,
+            input_ids_to_python_list,
+        )
+
+        processor = await self._resolve_higgs_audio_v2_processor()
+
+        if request.ref_audio is None:
+            inputs = await asyncio.to_thread(build_plain_text_prompt, processor, request.input)
+            prompt_token_ids = input_ids_to_python_list(inputs)
+            return tokens_input(prompt_token_ids=prompt_token_ids)
+
+        wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
+        wav = np.asarray(wav_list, dtype=np.float32)
+        out = await asyncio.to_thread(
+            build_voice_clone_prompt,
+            processor,
+            request.input,
+            wav,
+            int(sr),
+            request.ref_text or "",
+        )
+        prompt = tokens_input(prompt_token_ids=out["prompt_token_ids"])
+        # Pass tensors at the top level of additional_information (NOT list-
+        # wrapped). ``vllm_omni.data_entry_keys.serialize_payload`` routes
+        # bare ``torch.Tensor`` values through ``_serialize_tensor``; a list
+        # containing tensors would fall into the ``list_data`` field which
+        # msgspec cannot serialize and the tensors would be dropped over the
+        # process boundary (silent voice-clone failure).
+        prompt["additional_information"] = {
+            "audio_input_ids": out["audio_input_ids"],
+            "audio_input_ids_mask": out["audio_input_ids_mask"],
+        }
+        return prompt
+
+    async def _resolve_higgs_audio_v2_processor(self):
+        """Lazy-load the AutoProcessor for higgs_audio_v2 (once per serving instance)."""
+        cached = getattr(self, "_higgs_audio_v2_processor", None)
+        if cached is not None:
+            return cached
+
+        from transformers import AutoProcessor
+
+        model_path = None
+        for stage in self.engine_client.stage_configs:
+            model_path = getattr(getattr(stage, "engine_args", None), "model", None)
+            if model_path:
+                break
+        if model_path is None:
+            # Fallback: the orchestrator stores the served model id on the engine
+            # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args
+            # may not surface ``model`` when the deploy yaml doesn't set it per
+            # stage (the CLI-passed model id is the single source of truth).
+            model_path = getattr(self.engine_client, "model", None)
+        if model_path is None:
+            raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs")
+        processor = AutoProcessor.from_pretrained(model_path)
+        self._higgs_audio_v2_processor = processor
+        return processor
+
     def _validate_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate Fish Speech request parameters. Returns error message or None."""
         if not request.input or not request.input.strip():
@@ -2226,6 +2386,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         elif self._tts_model_type == "covo_audio":
             prompt = self._build_covo_audio_prompt(request)
             tts_params = {}
+        elif self._tts_model_type == "higgs_audio_v2":
+            # Explicit higgs_audio_v2 branch: do NOT fall through to the
+            # Qwen-style generic placeholder path; build prompt_token_ids
+            # directly via the upstream processor + the project's plain-text
+            # helper.
+            #
+            # Resolve uploaded voices first so a voice=<name> request (after
+            # POST /v1/audio/voices) populates ref_audio + ref_text from the
+            # cached speaker entry, mirroring the cosyvoice3 / voxcpm2 flow.
+            err = self._apply_uploaded_speaker(request)
+            if err:
+                raise ValueError(err)
+            validation_error = self._validate_higgs_audio_v2_request(request)
+            if validation_error:
+                raise ValueError(validation_error)
+            prompt = await self._build_higgs_audio_v2_params(request)
+            if request.voice:
+                voice_lower = request.voice.lower()
+                if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
+                    additional = prompt.setdefault("additional_information", {})
+                    additional["voice_name"] = voice_lower
+                    additional["voice_created_at"] = self._voice_created_at(voice_lower)
+            tts_params = {}
         elif self._tts_model_type == "voxcpm2":
             # voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the
             # raw waveform tuple for prefill-length accounting, not a base64 data URL.
@@ -2338,6 +2521,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             model_type = "ming_flash_omni_tts"
         elif self._tts_model_type == "moss_tts_nano":
             model_type = "moss_tts_nano"
+        elif self._tts_model_type == "higgs_audio_v2":
+            model_type = "higgs_audio_v2"
         elif self._tts_model_type == "glm_tts":
             model_type = "glm_tts"
         elif self._is_tts:

vllm_omni/model_executor/models/higgs_audio_v2/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vllm-omni integration for boson-ai's higgs-audio v2 (two-stage TTS)."""
+
+from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import (
+    HiggsAudioV2Config,
+)
+
+__all__ = ["HiggsAudioV2Config"]

vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration class for higgs-audio v2 in vllm-omni.
+
+Mirrors `bosonai/higgs-audio-v2-generation-3B-base/config.json` and exposes the
+Stage-1 codec decoder reference (sample rate, codebook layout, stream-special
+token ids) so the rest of the integration can read a single typed object.
+
+The text-LM backbone constants match `Llama-3.2-3B` (hidden=3072, 28 layers,
+GQA 24Q/8KV, head_dim=128, vocab=128256, max_position_embeddings=2048, RoPE
+``llama3`` scaling with factor=32). The audio side adds 8 codebooks of size
+1026 (real codes ``[0..1023]``, stream specials ``audio_stream_bos_id=1024``
+and ``audio_stream_eos_id=1025``) plus the DualFFN audio expert.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import PretrainedConfig
+
+__all__ = ["HiggsAudioV2Config"]
+
+
+class HiggsAudioV2Config(PretrainedConfig):
+    """Typed wrapper around the upstream `higgs_audio_v2` config.
+
+    The field set here is a strict subset of the upstream
+    ``transformers.models.higgs_audio_v2.HiggsAudioV2Config``: only the
+    knobs that vllm-omni reads directly are surfaced. Anything extra coming
+    from the HF config dict is preserved via ``PretrainedConfig`` so
+    ``AutoConfig.from_pretrained(...)`` round-trips correctly.
+    """
+
+    model_type: str = "higgs_audio_v2"
+    keys_to_ignore_at_inference = ("past_key_values",)
+
+    # Architecture identifier persisted in config.json and used by both
+    # vllm_omni.config.pipeline_registry and vllm_omni.model_executor.models.registry.
+    DEFAULT_ARCHITECTURES = ("HiggsAudioV2ForConditionalGeneration",)
+
+    def __init__(
+        self,
+        # --- Text LM (Llama-3.2-3B) ---
+        vocab_size: int = 128256,
+        hidden_size: int = 3072,
+        intermediate_size: int = 8192,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 24,
+        num_key_value_heads: int = 8,
+        head_dim: int = 128,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 2048,
+        rms_norm_eps: float = 1e-5,
+        rope_parameters: dict[str, Any] | None = None,
+        initializer_range: float = 0.02,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        mlp_bias: bool = False,
+        tie_word_embeddings: bool = False,
+        pretraining_tp: int = 1,
+        # --- Audio side ---
+        num_codebooks: int = 8,
+        codebook_size: int = 1026,
+        audio_token_id: int = 128016,
+        audio_bos_token_id: int = 128013,
+        audio_eos_token_id: int = 128012,
+        audio_delay_token_id: int = 128014,
+        audio_stream_bos_id: int = 1024,
+        audio_stream_eos_id: int = 1025,
+        sample_rate: int = 24000,
+        frame_rate: int = 25,
+        # --- Standard LM specials ---
+        bos_token_id: int = 1,
+        eos_token_id: int = 128009,
+        pad_token_id: int = 128001,
+        # --- DualFFN ---
+        # Upstream `HiggsAudioV2DecoderLayer` always carries a parallel `audio_mlp`
+        # expert; this flag is kept for forward-compat in case a future config
+        # disables it (no-op for the boson-ai checkpoint).
+        use_audio_dual_ffn: bool = True,
+        # --- Stage-1 tokenizer / codec decoder ---
+        # Path to the audio-tokenizer subdir (relative to the model directory).
+        # Empty string means the model_dir itself contains config.json +
+        # model.safetensors (the layout used by the standalone
+        # ``bosonai/higgs-audio-v2-tokenizer`` HF repo). Non-empty means the
+        # tokenizer files live in ``<model_dir>/<subdir>/`` (the layout used by
+        # the 3B Stage-0 checkpoint which bundles its tokenizer under
+        # ``audio_tokenizer/``).
+        audio_tokenizer_subdir: str = "",
+        # HF repo id (or local path) of a STANDALONE audio tokenizer. When set,
+        # Stage-1 resolves this via ``huggingface_hub.snapshot_download`` and
+        # loads the codec from there, ignoring ``audio_tokenizer_subdir``. This
+        # matches the boson-ai release layout where the codec lives in a
+        # separate repo (``bosonai/higgs-audio-v2-tokenizer``).
+        audio_tokenizer_id: str | None = "bosonai/higgs-audio-v2-tokenizer",
+        # --- Misc ---
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        if rope_parameters is None:
+            rope_parameters = {
+                "factor": 32.0,
+                "high_freq_factor": 0.5,
+                "low_freq_factor": 0.125,
+                "original_max_position_embeddings": 1024,
+                "rope_theta": 500000.0,
+                "rope_type": "llama3",
+            }
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_parameters = dict(rope_parameters)
+        self.initializer_range = initializer_range
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.pretraining_tp = pretraining_tp
+
+        self.num_codebooks = num_codebooks
+        self.codebook_size = codebook_size
+        self.audio_token_id = audio_token_id
+        self.audio_bos_token_id = audio_bos_token_id
+        self.audio_eos_token_id = audio_eos_token_id
+        self.audio_delay_token_id = audio_delay_token_id
+        self.audio_stream_bos_id = audio_stream_bos_id
+        self.audio_stream_eos_id = audio_stream_eos_id
+        self.sample_rate = sample_rate
+        self.frame_rate = frame_rate
+
+        self.use_audio_dual_ffn = use_audio_dual_ffn
+        self.audio_tokenizer_subdir = audio_tokenizer_subdir
+        self.audio_tokenizer_id = audio_tokenizer_id
+        self.use_cache = use_cache
+
+    # Codes [0..codebook_size-1] include 2 stream specials at indices
+    # `audio_stream_bos_id` and `audio_stream_eos_id`. Real codes ready for the
+    # codec decoder are the strict prefix [0..audio_stream_bos_id).
+    @property
+    def num_real_codes(self) -> int:
+        return self.audio_stream_bos_id  # i.e. 1024 for the boson-ai checkpoint
+
+    # The full vocab size of each per-codebook output head.
+    @property
+    def codebook_output_size(self) -> int:
+        return self.codebook_size  # i.e. 1026 for the boson-ai checkpoint
+
+    # Canonical MusicGen-style delay pattern used by
+    # `HiggsAudioV2DelayPatternLogitsProcessor`: codebook k starts emitting
+    # real codes only after k frames. Kept as a function so subclasses can
+    # override if upstream ever publishes a different convention.
+    def default_delay_pattern(self) -> list[int]:
+        return list(range(self.num_codebooks))

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py

@@ -0,0 +1,471 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HiggsAudio codec decoder kernel for higgs-audio v2.
+
+This module hosts the parameter-side building blocks for the higgs-audio-v2
+audio tokenizer's decoder path:
+
+  audio_codes [B, 8, T]
+    -> RVQ codebook lookup + project_out -> sum -> [B, hidden_size, T]
+    -> fc2 Linear(hidden_size, 256) -> [B, 256, T]
+    -> DAC acoustic decoder (conv-transpose upsampling) -> [B, 1, T*960]
+    -> 24 kHz waveform (25 fps x 960 samples/frame)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "HiggsAudioVQLayer",
+    "HiggsAudioRVQ",
+    "adjust_conv_transpose_output_padding",
+    "build_higgs_audio_acoustic_decoder",
+    "BosonDacDecoder",
+    "build_boson_dac_decoder",
+    "load_higgs_audio_codec",
+    "_remap_boson_model_pth_state_dict",  # exported for unit-testing the mapper
+]
+
+
+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)
+        quantized = self.project_out(quantized)
+        return quantized.permute(0, 2, 1)
+
+
+class HiggsAudioRVQ(nn.Module):
+    """Residual Vector Quantizer with ``num_quantizers`` 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
+
+
+@torch.jit.script
+def _snake(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    return x.reshape(shape)
+
+
+class _Snake1d(nn.Module):
+    """Snake activation (per-channel learnable alpha). Mirrors the boson-ai DAC
+    Snake1d implementation byte-for-byte so the boson ``model.pth`` state-dict
+    loads directly into ``self.alpha``.
+    """
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return _snake(x, self.alpha)
+
+
+def _wn_conv1d(*args: Any, **kwargs: Any) -> nn.Module:
+    from torch.nn.utils import weight_norm
+
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def _wn_conv_transpose1d(*args: Any, **kwargs: Any) -> nn.Module:
+    from torch.nn.utils import weight_norm
+
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class _BosonResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            _Snake1d(dim),
+            _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            _Snake1d(dim),
+            _wn_conv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class _BosonDecoderBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, stride: int):
+        super().__init__()
+        import math
+
+        self.block = nn.Sequential(
+            _Snake1d(input_dim),
+            _wn_conv_transpose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+            _BosonResidualUnit(output_dim, dilation=1),
+            _BosonResidualUnit(output_dim, dilation=3),
+            _BosonResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.block(x)
+
+
+class BosonDacDecoder(nn.Module):
+    """Vendored copy of ``boson_multimodal``'s standalone DAC decoder, with the
+    exact ``decoder_2.model.<i>.{...}`` key layout the boson-ai
+    ``bosonai/higgs-audio-v2-tokenizer`` checkpoint expects.
+
+    See ``boson_multimodal/audio_processing/descriptaudiocodec/dac/model/dac.py``
+    (``class Decoder``) in the upstream package for the reference implementation
+    we mirror.
+
+    Architecture for the v2 24 kHz tokenizer:
+      input_channel=256, channels=1024, rates=[8, 5, 4, 2, 3], d_out=1
+    Cumulative upsample factor is ``prod(rates) = 960`` (25 fps -> 24 kHz).
+    """
+
+    def __init__(
+        self,
+        input_channel: int = 256,
+        channels: int = 1024,
+        rates: tuple[int, ...] = (8, 5, 4, 2, 3),
+        d_out: int = 1,
+    ):
+        super().__init__()
+        layers: list[nn.Module] = [_wn_conv1d(input_channel, channels, kernel_size=7, padding=3)]
+        last_dim = channels
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            layers.append(_BosonDecoderBlock(input_dim, output_dim, stride))
+            last_dim = output_dim
+        layers.append(_Snake1d(last_dim))
+        layers.append(_wn_conv1d(last_dim, d_out, kernel_size=7, padding=3))
+        self.model = nn.Sequential(*layers)
+        # Cache hop length (= prod(rates)) so callers don't need to recompute it.
+        hop = 1
+        for s in rates:
+            hop *= int(s)
+        self._hop_length = hop
+
+    @property
+    def hop_length(self) -> int:
+        return self._hop_length
+
+    # Expose conv1 with a ``.in_channels`` attribute so the existing fc2
+    # fallback shape inference (acoustic_decoder.conv1.in_channels) works
+    # unchanged for the boson decoder too.
+    @property
+    def conv1(self) -> nn.Module:
+        return self.model[0]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+def build_boson_dac_decoder(device: torch.device) -> BosonDacDecoder:
+    """Construct the vendored boson DAC decoder on ``device``.
+
+    The architectural constants (input_channel=256, channels=1024,
+    rates=[8,5,4,2,3], d_out=1) are baked in because they're invariant for the
+    v2 24 kHz tokenizer ``bosonai/higgs-audio-v2-tokenizer`` checkpoint we
+    target. Other channel counts are not supported until we encounter a
+    different boson checkpoint in the wild.
+    """
+    return BosonDacDecoder().to(device)
+
+
+def adjust_conv_transpose_output_padding(decoder: nn.Module) -> None:
+    """Set ConvTranspose1d output_padding = stride % 2 (HiggsAudioV2 modification).
+
+    The vanilla DAC decoder ships with the default output_padding (0); the
+    boson-ai checkpoint expects ``stride % 2`` instead. This is a no-op for
+    even strides and adds a single sample for odd strides.
+    """
+    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 build_higgs_audio_acoustic_decoder(
+    tokenizer_config: dict[str, Any],
+    device: torch.device,
+) -> nn.Module:
+    """Build the DAC acoustic decoder used by the HiggsAudioV2 tokenizer.
+
+    Returns the decoder sub-module of ``transformers.DacModel`` with the
+    HiggsAudioV2 output-padding fix already applied. The tanh activation is
+    replaced with ``Identity`` so the network matches the boson-ai checkpoint.
+    Weights are NOT loaded here; the caller is responsible for copying them in.
+    """
+    from transformers import DacConfig, DacModel
+
+    dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"])
+    dac_model = DacModel(dac_cfg)
+    decoder = dac_model.decoder.to(device)
+    adjust_conv_transpose_output_padding(decoder)
+    if hasattr(decoder, "tanh"):
+        decoder.tanh = nn.Identity()
+    return decoder
+
+
+def _load_higgs_audio_state_dict(audio_tokenizer_dir: str, device: torch.device) -> dict[str, torch.Tensor]:
+    """Load the codec state dict from either layout used in the wild.
+
+    Tries layouts in order:
+    1. ``<dir>/model.safetensors`` (OmniVoice-bundled layout used by
+       ``k2-fsa/OmniVoice/audio_tokenizer/``).
+    2. ``<dir>/model.pth`` (boson-ai standalone ``bosonai/higgs-audio-v2-tokenizer``
+       layout). The state-dict keys differ structurally
+       (``quantizer.vq.layers.<i>._codebook.embed`` etc.) and are remapped to
+       the OmniVoice-style names this kernel expects via
+       :func:`_remap_boson_model_pth_state_dict`.
+    """
+    safetensors_path = os.path.join(audio_tokenizer_dir, "model.safetensors")
+    pth_path = os.path.join(audio_tokenizer_dir, "model.pth")
+    # Preference order: try model.pth first when present (boson-ai standalone
+    # ships the codec there; its model.safetensors is a different artefact).
+    # Fall back to model.safetensors (the OmniVoice-bundled layout).
+    sd: dict[str, torch.Tensor]
+    if os.path.exists(pth_path):
+        sd = torch.load(pth_path, map_location=device, weights_only=False)
+    elif os.path.exists(safetensors_path):
+        from safetensors.torch import load_file
+
+        sd = load_file(safetensors_path, device=str(device))
+    else:
+        raise FileNotFoundError(f"Audio tokenizer weights not found at {pth_path} or {safetensors_path}")
+    # Run the boson-layout remap if needed; for the OmniVoice-bundled layout
+    # it is a no-op (no ``quantizer.vq.layers.*`` keys to rewrite).
+    if any(k.startswith("quantizer.vq.layers.") for k in sd):
+        sd = _remap_boson_model_pth_state_dict(sd)
+    return sd
+
+
+def _remap_boson_model_pth_state_dict(sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Translate boson-ai's standalone ``model.pth`` keys into OmniVoice-style names
+    that the shared kernel's RVQ + fc2 + DAC sites consume.
+
+    Best-effort: only the RVQ-side keys map cleanly. The boson-ai decoder
+    uses ``decoder_2.model.<i>.weight_g/weight_v`` (DAC with weight-norm +
+    Snake activations) rather than the OmniVoice DAC layout; that side
+    requires either vendoring the upstream decoder module or rewriting the
+    DAC builder to consume weight-normed tensors. This function returns the
+    mapped RVQ keys plus any acoustic_decoder.* / fc2.* / fc.* tensors that
+    happen to share the OmniVoice names, and leaves the decoder-2 keys
+    untouched (the caller will see them as MISSING when copying into the
+    OmniVoice DAC decoder).
+
+    Mapping (RVQ side only):
+        quantizer.vq.layers.<i>._codebook.embed -> quantizer.quantizers.<i>.codebook.embed
+        quantizer.vq.layers.<i>.project_out.weight -> quantizer.quantizers.<i>.project_out.weight
+        quantizer.vq.layers.<i>.project_out.bias   -> quantizer.quantizers.<i>.project_out.bias
+    """
+    if not isinstance(sd, dict):
+        raise TypeError(f"expected a state dict, got {type(sd)!r}")
+    remapped: dict[str, torch.Tensor] = {}
+    for key, tensor in sd.items():
+        if not isinstance(tensor, torch.Tensor):
+            continue
+        # Quantizer rewrite: vq.layers.<i>._codebook.embed -> quantizers.<i>.codebook.embed
+        if key.startswith("quantizer.vq.layers."):
+            parts = key.split(".")
+            # parts[3] is the layer index, parts[4]+ is the tail
+            if len(parts) >= 5:
+                idx = parts[3]
+                tail = ".".join(parts[4:])
+                # The actual codebook centroids the decode path uses live at
+                # ``_codebook.embed`` (a [codebook_size, codebook_dim] buffer).
+                # ``_codebook.embed_avg`` is the un-normalised EMA accumulator
+                # used during training (~cluster_size x bigger); MATCHING IT
+                # HERE produces a 4-5x amplified codebook and unintelligible
+                # PCM downstream. Use an EXACT-match check, not startswith.
+                if tail == "_codebook.embed":
+                    new_key = f"quantizer.quantizers.{idx}.codebook.embed"
+                    remapped[new_key] = tensor
+                    continue
+                if tail.startswith("project_out."):
+                    new_key = f"quantizer.quantizers.{idx}.{tail}"
+                    remapped[new_key] = tensor
+                    continue
+                # project_in / _codebook.cluster_size / _codebook.embed_avg / inited
+                # are encoder-side or training-only state that the decode path
+                # does not need; drop them.
+                continue
+        # boson decoder + fc_post2 -> our vendored kernel names.
+        if key.startswith("decoder_2."):
+            # decoder_2.model.<i>.{...} -> acoustic_decoder.model.<i>.{...}
+            new_key = "acoustic_decoder." + key[len("decoder_2.") :]
+            remapped[new_key] = tensor
+            continue
+        if key.startswith("fc_post2."):
+            # fc_post2.{weight,bias} -> fc2.{weight,bias}
+            new_key = "fc2." + key[len("fc_post2.") :]
+            remapped[new_key] = tensor
+            continue
+        # Anything else (acoustic_decoder.*, fc2.*, fc.*, fc_post*.*, decoder_2.*,
+        # decoder_semantic.*, semantic_model.*) passes through unchanged so
+        # the caller's lookup is unambiguous.
+        remapped[key] = tensor
+    return remapped
+
+
+def load_higgs_audio_codec(
+    audio_tokenizer_dir: str,
+    device: torch.device,
+) -> tuple[HiggsAudioRVQ, nn.Linear, nn.Module, dict[str, Any]]:
+    """Load the HiggsAudioV2 RVQ + fc2 + DAC decoder from a checkpoint folder.
+
+    Accepts both layouts: ``<dir>/model.safetensors`` (OmniVoice-bundled) and
+    ``<dir>/model.pth`` (boson-ai standalone). The standalone path remaps
+    quantizer keys before consumption; structural decoder differences
+    (boson uses Snake + weight-norm) still leave some DAC parameters missing
+    when loading the standalone path -- those entries log warnings but the
+    RVQ side completes successfully.
+
+    Args:
+        audio_tokenizer_dir: Path to a directory containing ``config.json``
+            and EITHER ``model.safetensors`` (OmniVoice layout) OR
+            ``model.pth`` (boson-ai standalone layout).
+        device: Device to place the loaded modules and state dict on.
+
+    Returns:
+        (quantizer, fc2, acoustic_decoder, tokenizer_config)
+        - quantizer: ``HiggsAudioRVQ`` with ``num_quantizers`` discovered from
+            the state dict (defaults to 8 for boson-ai checkpoints).
+        - fc2: ``nn.Linear`` projecting RVQ output (1024) into the DAC's
+            hidden dimension (typically 256). May be uninitialized when loaded
+            from the boson-ai standalone layout (no ``fc2.*`` keys present).
+        - acoustic_decoder: DAC decoder with HiggsAudioV2 output-padding fix
+            and tanh-replaced-by-Identity. Fully initialized for the OmniVoice
+            layout; partially initialized for boson-ai standalone (missing
+            ``decoder_2.*`` -> ``acoustic_decoder.*`` mapping is not yet
+            implemented; full upstream-decoder vendoring is the next step).
+        - tokenizer_config: The loaded ``config.json`` dict; useful for
+            callers that need ``sample_rate`` or other tokenizer metadata.
+    """
+    config_path = os.path.join(audio_tokenizer_dir, "config.json")
+
+    with open(config_path) as f:
+        tokenizer_config: dict[str, Any] = json.load(f)
+
+    state_dict = _load_higgs_audio_state_dict(audio_tokenizer_dir, device)
+
+    codebook_dim = tokenizer_config.get("codebook_dim", 64)
+    codebook_size = tokenizer_config.get("codebook_size", 1024)
+    # Discover hidden_size and num_quantizers from the (possibly remapped) state dict.
+    if "quantizer.quantizers.0.project_out.weight" not in state_dict:
+        raise KeyError(
+            "Codec state dict is missing 'quantizer.quantizers.0.project_out.weight'. "
+            "If you loaded a boson-ai standalone tokenizer, ensure the model.pth "
+            "remap fired (see _remap_boson_model_pth_state_dict)."
+        )
+    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")
+    )
+
+    quantizer = HiggsAudioRVQ(
+        num_quantizers=num_quantizers,
+        codebook_size=codebook_size,
+        codebook_dim=codebook_dim,
+        hidden_size=hidden_size,
+    ).to(device)
+    for i in range(num_quantizers):
+        prefix = f"quantizer.quantizers.{i}"
+        embed_key = f"{prefix}.codebook.embed"
+        if embed_key in state_dict:
+            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:
+            quantizer.quantizers[i].project_out.weight.data.copy_(state_dict[proj_out_w])
+        if proj_out_b in state_dict:
+            quantizer.quantizers[i].project_out.bias.data.copy_(state_dict[proj_out_b])
+
+    # Decide which acoustic decoder we need. A boson-layout state dict (after
+    # _remap_boson_model_pth_state_dict) carries ``acoustic_decoder.model.0.weight_v``
+    # (the WNConv first conv) which the OmniVoice DacModel doesn't have; the
+    # OmniVoice layout carries ``acoustic_decoder.block.0.*`` instead.
+    is_boson_layout = "acoustic_decoder.model.0.weight_v" in state_dict
+    if is_boson_layout:
+        acoustic_decoder = build_boson_dac_decoder(device)
+        boson_sd = {
+            k[len("acoustic_decoder.") :]: v for k, v in state_dict.items() if k.startswith("acoustic_decoder.")
+        }
+        load_report = acoustic_decoder.load_state_dict(boson_sd, strict=False)
+        if load_report.missing_keys:
+            raise RuntimeError(
+                "Boson DAC decoder is missing keys after load: "
+                f"{load_report.missing_keys[:5]}... ({len(load_report.missing_keys)} total)"
+            )
+        # ``unexpected_keys`` is informational only.
+    else:
+        acoustic_decoder = build_higgs_audio_acoustic_decoder(tokenizer_config, device)
+        for name, param in acoustic_decoder.named_parameters():
+            higgs_name = f"acoustic_decoder.{name}"
+            if higgs_name in state_dict:
+                param.data.copy_(state_dict[higgs_name])
+    acoustic_decoder.eval()
+
+    if "fc2.weight" in state_dict and "fc2.bias" in state_dict:
+        fc2_w = state_dict["fc2.weight"]
+        fc2_b = state_dict["fc2.bias"]
+        fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device)
+        fc2.weight.data.copy_(fc2_w)
+        fc2.bias.data.copy_(fc2_b)
+    else:
+        # Last-resort fallback: shape fc2 from the decoder's first conv input
+        # channels so the matmul produces compatible shapes. Weights stay
+        # uninitialised — PCM quality will be noise on this path.
+        first_conv = getattr(acoustic_decoder, "conv1", None)
+        if first_conv is None or not hasattr(first_conv, "in_channels"):
+            raise RuntimeError("acoustic_decoder has no conv1.in_channels; cannot size fc2 fallback")
+        fc2 = nn.Linear(hidden_size, int(first_conv.in_channels)).to(device)
+
+    return quantizer, fc2, acoustic_decoder, tokenizer_config

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py

@@ -0,0 +1,413 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage 1 (codec decoder) for higgs-audio v2.
+
+Two surfaces exposed:
+
+1. **Direct decode API** (kept for offline / unit tests):
+   ``decode_codes(audio_codes: [B, num_codebooks=8, T])`` -> ``[B, 1, T*hop]`` PCM.
+   ``forward_chunk(audio_codes, *, left_context_size, hop_length=960)``
+   trims overlap from streamed PCM. ``forward(audio_codes)`` retains the same
+   single-tensor signature for backward compatibility with code that imports
+   the class directly.
+
+2. **vLLM stage runtime API** (used by the engine):
+   ``__init__(*, vllm_config, prefix)``, ``embed_input_ids``,
+   ``compute_logits=None``, plus a runtime ``forward(input_ids=..., positions=...,
+   runtime_additional_information=...)`` that takes flat codebook-major
+   ``input_ids`` (``[Q * num_frames]``) per request and returns an
+   :class:`vllm_omni.model_executor.models.output_templates.OmniOutput`.
+
+The kernel is the HiggsAudio codec helper in
+``vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py``; both
+surfaces share the same RVQ + DAC weights.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import (
+    HiggsAudioV2Config,
+)
+from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import (
+    HiggsAudioRVQ,
+    load_higgs_audio_codec,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+__all__ = [
+    "HiggsAudioV2Code2Wav",
+    "HiggsAudioV2Code2WavForConditionalGeneration",
+]
+
+logger = init_logger(__name__)
+
+
+class HiggsAudioV2Code2Wav(nn.Module):
+    """Stage-1 codec decoder for higgs-audio v2.
+
+    Constructor accepts either the vLLM-runtime signature (``*, vllm_config,
+    prefix``) or the direct config-object form (``HiggsAudioV2Code2Wav(config)``)
+    so both engine-side and unit-test callers work. When the runtime form is
+    used we read the higgs_audio_v2 config out of ``vllm_config.model_config.hf_config``
+    and the model_path off ``vllm_config.model_config.model`` so the codec can
+    be loaded directly from disk.
+    """
+
+    input_modalities = "audio"
+
+    def __init__(
+        self,
+        config: HiggsAudioV2Config | None = None,
+        *,
+        vllm_config: VllmConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        # Resolve config from either positional (unit-test) or kwarg (engine) form.
+        if vllm_config is not None:
+            hf_config = vllm_config.model_config.hf_config
+            if isinstance(hf_config, HiggsAudioV2Config):
+                self.config = hf_config
+            else:
+                self.config = HiggsAudioV2Config(**hf_config.to_dict())
+            self._model_path: str | None = vllm_config.model_config.model
+            self.vllm_config: VllmConfig | None = vllm_config
+        else:
+            if config is None:
+                raise TypeError(
+                    "HiggsAudioV2Code2Wav: provide either positional `config` "
+                    "(HiggsAudioV2Config) or keyword `vllm_config` (VllmConfig)."
+                )
+            self.config = config
+            self._model_path = None
+            self.vllm_config = None
+
+        self.sample_rate: int = int(self.config.sample_rate)
+        self.num_codebooks: int = int(self.config.num_codebooks)
+        self.num_real_codes: int = int(self.config.num_real_codes)
+        # Each codec frame upsamples to 960 24 kHz samples (= 25 fps * 960 = 24000).
+        self.hop_length: int = 960
+
+        # Engine-runner hooks (Stage 1 has no token sampling).
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        # Populated by load_weights().
+        self.quantizer: HiggsAudioRVQ | None = None
+        self.fc2: nn.Linear | None = None
+        self.acoustic_decoder: nn.Module | None = None
+        self._loaded: bool = False
+
+        # When constructed via the engine path, eagerly load codec weights from
+        # the model directory. The engine's subsequent ``load_weights(weights)``
+        # call (with the Stage-0 safetensors iterator) is then a cheap no-op.
+        if self._model_path is not None:
+            try:
+                self.load_codec_from_disk(self._model_path)
+            except FileNotFoundError as exc:
+                logger.warning("HiggsAudioV2Code2Wav: eager codec load deferred (%s)", exc)
+
+    # ------------------------------------------------------------- engine hooks
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        """Stage 1 ignores embeddings; vLLM's runner still needs a stable shape."""
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
+        return None
+
+    # ------------------------------------------------------------------ load
+    def load_weights(self, weights_or_model_dir, device: torch.device | None = None):
+        """Dual-signature load:
+
+        - ``load_weights(weights_iterator)`` -- the vLLM engine path. The iterator
+          yields ``(name, tensor)`` pairs from the Stage-0 safetensors which DO NOT
+          contain codec weights, so we ignore them and eagerly load the codec from
+          ``self._model_path``. Returns the set of all our params as "initialized"
+          so the engine's loader passes its sanity check.
+        - ``load_weights(model_dir: str)`` -- the legacy file-based call used by
+          offline unit tests and ``end2end.py --mode stage1_only``. Returns ``None``.
+        """
+        if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)):
+            # Engine call: consume the iterator (we don't need it) and eagerly
+            # load the codec from disk.
+            try:
+                for _ in weights_or_model_dir:
+                    pass
+            except TypeError:
+                pass
+            if self._model_path is None:
+                raise RuntimeError(
+                    "HiggsAudioV2Code2Wav.load_weights called via the engine path but "
+                    "no model directory is known. Construct with vllm_config so the "
+                    "constructor caches vllm_config.model_config.model."
+                )
+            self.load_codec_from_disk(self._model_path, device=device)
+            return {name for name, _ in self.named_parameters()}
+
+        # Legacy file-based call.
+        self.load_codec_from_disk(weights_or_model_dir, device=device)
+        return None
+
+    def load_codec_from_disk(self, model_dir: str, device: torch.device | None = None) -> None:
+        """Load codec weights from disk (RVQ + fc2 + DAC decoder).
+
+        Resolution order for the codec directory:
+        1. If :attr:`HiggsAudioV2Config.audio_tokenizer_id` is set, treat it as
+           either a local path (if it exists on disk) or an HF repo id that we
+           resolve via ``huggingface_hub.snapshot_download``. This is the
+           default for the boson-ai release (codec in a separate repo).
+        2. Otherwise fall back to ``<model_dir>/<audio_tokenizer_subdir>``.
+        """
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        audio_tokenizer_path = self._resolve_audio_tokenizer_path(model_dir)
+        quantizer, fc2, acoustic_decoder, _tokenizer_config = load_higgs_audio_codec(audio_tokenizer_path, device)
+        if len(quantizer.quantizers) != self.num_codebooks:
+            raise ValueError(
+                f"checkpoint has {len(quantizer.quantizers)} quantizers but config.num_codebooks={self.num_codebooks}"
+            )
+        self.quantizer = quantizer
+        self.fc2 = fc2
+        self.acoustic_decoder = acoustic_decoder
+        self._loaded = True
+        logger.info(
+            "Loaded HiggsAudioV2Code2Wav: %d quantizers, fc2(%d->%d), sample_rate=%d",
+            len(self.quantizer.quantizers),
+            self.fc2.in_features,
+            self.fc2.out_features,
+            self.sample_rate,
+        )
+
+    def _resolve_audio_tokenizer_path(self, model_dir: str) -> str:
+        """Return a local filesystem path containing ``config.json`` + codec
+        weights. See :meth:`load_codec_from_disk` for the resolution rules."""
+        tokenizer_id = getattr(self.config, "audio_tokenizer_id", None)
+        if tokenizer_id:
+            if os.path.isdir(tokenizer_id):
+                return tokenizer_id
+            return self._resolve_hf_id_to_local(tokenizer_id)
+        subdir = self.config.audio_tokenizer_subdir or ""
+        if os.path.isdir(model_dir):
+            return os.path.join(model_dir, subdir) if subdir else model_dir
+        local = self._resolve_hf_id_to_local(model_dir)
+        return os.path.join(local, subdir) if subdir else local
+
+    @staticmethod
+    def _resolve_hf_id_to_local(repo_id: str) -> str:
+        """Resolve a HF repo id to a local snapshot directory.
+
+        Prefers a read-only cache lookup (no write locks) so a quota-constrained
+        Lustre cache can still serve already-downloaded weights. Falls back to
+        ``snapshot_download`` only when nothing is cached locally.
+        """
+        from huggingface_hub import try_to_load_from_cache
+
+        cached = try_to_load_from_cache(repo_id=repo_id, filename="config.json")
+        if isinstance(cached, str) and os.path.isfile(cached):
+            return os.path.dirname(cached)
+
+        from huggingface_hub.constants import HF_HUB_CACHE
+
+        safe = repo_id.replace("/", "--")
+        snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots")
+        if os.path.isdir(snapshots_dir):
+            for rev in os.listdir(snapshots_dir):
+                candidate = os.path.join(snapshots_dir, rev)
+                if os.path.isfile(os.path.join(candidate, "config.json")):
+                    return candidate
+
+        from huggingface_hub import snapshot_download
+
+        return snapshot_download(repo_id=repo_id)
+
+    # ------------------------------------------------------ direct decode API
+    @torch.inference_mode()
+    def decode_codes(self, audio_codes: torch.Tensor) -> torch.Tensor:
+        """Decode a ``[B, num_codebooks=8, T]`` code tensor to PCM ``[B, 1, T*hop]``."""
+        if not self._loaded:
+            raise RuntimeError("HiggsAudioV2Code2Wav not loaded. Call load_weights() first.")
+
+        codes = self._validate_codes(audio_codes)
+        device = codes.device
+
+        rvq_codes = codes.transpose(0, 1).long()  # [num_codebooks, B, T]
+        quantized = self.quantizer.decode(rvq_codes)  # [B, hidden, T]
+        # Match fc2's parameter dtype — vLLM may auto-cast the codec to the
+        # stage's model dtype (e.g. bf16), while the quantizer's RVQ codebooks
+        # often stay in float32 because they're embedded lookups. A defensive
+        # cast keeps the matmul homogeneous.
+        quantized = quantized.to(dtype=self.fc2.weight.dtype)
+        quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
+        # Acoustic decoder may have parameters in yet another dtype — match
+        # its first parameter so the linear/conv ops stay consistent.
+        first_param = next(self.acoustic_decoder.parameters(), None)
+        if first_param is not None and quantized.dtype != first_param.dtype:
+            quantized = quantized.to(dtype=first_param.dtype)
+        audio = self.acoustic_decoder(quantized)
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(1)
+        return audio.to(device)
+
+    @torch.inference_mode()
+    def forward_chunk(
+        self,
+        audio_codes: torch.Tensor,
+        *,
+        left_context_size: int = 0,
+        hop_length: int | None = None,
+    ) -> torch.Tensor:
+        """Chunked decode that trims ``left_context_size * hop_length`` samples
+        off the leading edge of the upsampled PCM. Mirrors the qwen3_tts
+        ``talker2code2wav_async_chunk`` overlap contract.
+        """
+        if left_context_size < 0:
+            raise ValueError(f"left_context_size must be >= 0; got {left_context_size}")
+        hop = int(hop_length) if hop_length is not None else self.hop_length
+        if hop <= 0:
+            raise ValueError(f"hop_length must be > 0; got {hop}")
+        pcm = self.decode_codes(audio_codes)
+        if left_context_size == 0:
+            return pcm
+        trim = left_context_size * hop
+        if pcm.shape[-1] <= trim:
+            return pcm[..., :0]
+        return pcm[..., trim:]
+
+    # ---------------------------------------------------- vLLM runtime forward
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | OmniOutput:
+        """Dual-mode forward.
+
+        - When ``input_ids`` is a 3-D ``[B, num_codebooks, T]`` tensor we treat
+          this as a direct decode call (legacy API) and return raw PCM.
+        - When ``input_ids`` is a flat 1-D / 2-D tensor we treat it as the
+          engine runtime payload (codebook-major flat per request) and return
+          an :class:`OmniOutput` carrying multimodal audio.
+        """
+        # Legacy direct-decode signature: caller passed a 3-D code tensor as input_ids.
+        if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 3:
+            return self.decode_codes(input_ids)
+
+        sr_val = int(self.sample_rate)
+        sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
+        empty = torch.zeros((0,), dtype=torch.float32)
+
+        if input_ids is None or input_ids.numel() == 0:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]},
+            )
+
+        ids = input_ids.reshape(-1).to(dtype=torch.long)
+        request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
+
+        left_context_size = [0] * len(request_ids_list)
+        if runtime_additional_information is not None:
+            for i, info in enumerate(runtime_additional_information):
+                if i >= len(left_context_size):
+                    break
+                meta = info.get("meta", {}) if isinstance(info, dict) else {}
+                if "left_context_size" in meta:
+                    left_context_size[i] = int(meta["left_context_size"])
+
+        wavs: list[torch.Tensor] = []
+        for i, req_ids in enumerate(request_ids_list):
+            n = int(req_ids.numel())
+            if n == 0:
+                wavs.append(empty)
+                continue
+            if n % self.num_codebooks != 0:
+                logger.warning(
+                    "HiggsAudioV2Code2Wav: flat code length %d not divisible by num_codebooks=%d; dropping",
+                    n,
+                    self.num_codebooks,
+                )
+                wavs.append(empty)
+                continue
+            frames = n // self.num_codebooks
+            codes_qf = req_ids.reshape(self.num_codebooks, frames)
+            codes_bqf = codes_qf.unsqueeze(0)
+            try:
+                pcm = self.forward_chunk(
+                    codes_bqf,
+                    left_context_size=left_context_size[i],
+                    hop_length=self.hop_length,
+                )
+            except ValueError as exc:
+                logger.warning("HiggsAudioV2Code2Wav: decode skipped (%s)", exc)
+                wavs.append(empty)
+                continue
+            wavs.append(pcm.squeeze(0).squeeze(0).to(torch.float32).cpu())
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "model_outputs": wavs,
+                "sr": [sr_tensor] * len(wavs),
+            },
+        )
+
+    # --------------------------------------------------------------- helpers
+    def _validate_codes(self, audio_codes: torch.Tensor) -> torch.Tensor:
+        """Ensure shape and value range; reject stream specials with ValueError."""
+        if not isinstance(audio_codes, torch.Tensor):
+            raise TypeError(f"audio_codes must be a torch.Tensor, got {type(audio_codes)!r}")
+        if audio_codes.ndim != 3:
+            raise ValueError(
+                f"audio_codes must have shape [B, num_codebooks={self.num_codebooks}, T]; "
+                f"got shape {tuple(audio_codes.shape)}"
+            )
+        if int(audio_codes.shape[1]) != self.num_codebooks:
+            raise ValueError(
+                f"audio_codes second dim must equal num_codebooks={self.num_codebooks}; got {int(audio_codes.shape[1])}"
+            )
+        if audio_codes.numel() > 0:
+            max_val = int(audio_codes.max().item())
+            min_val = int(audio_codes.min().item())
+            if max_val >= self.num_real_codes or min_val < 0:
+                raise ValueError(
+                    "audio_codes contains stream-special or out-of-range IDs: "
+                    f"min={min_val}, max={max_val}; real code range is "
+                    f"[0, {self.num_real_codes - 1}]. Filter audio_stream_bos_id="
+                    f"{self.config.audio_stream_bos_id} and audio_stream_eos_id="
+                    f"{self.config.audio_stream_eos_id} (and anything above) at "
+                    "Stage 0 before sending codes to the codec decoder."
+                )
+        return audio_codes
+
+    @staticmethod
+    def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
+        """Split a concatenated flat-codes tensor into per-request segments."""
+        if seq_token_counts is not None and len(seq_token_counts) > 1:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + int(count))
+            n = int(ids.numel())
+            return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))]
+        return [ids]
+
+
+# Engine-side architecture identifier alias (mirrors Qwen3TTSCode2Wav usage).
+HiggsAudioV2Code2WavForConditionalGeneration = HiggsAudioV2Code2Wav

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py

@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TTS prompt builder + scope validators for higgs-audio v2.
+
+vllm-omni's higgs path supports two request shapes:
+
+* **Plain text -> 24 kHz speech.** :func:`build_plain_text_prompt` runs the
+  upstream HF processor with a bare ``"Generate audio following instruction."``
+  system prompt; the emitted ``input_ids`` are byte-identical to upstream.
+* **Voice clone (shallow).** :func:`build_voice_clone_prompt` runs the
+  upstream HF processor with both a target text turn and a reference
+  ``(ref_audio, ref_text)`` ICL turn; the processor returns ``input_ids``
+  pre-expanded with audio placeholders plus an ``audio_input_ids`` tensor
+  carrying the encoded reference codes (HF tokenizer encodes audio on the
+  spot — no encoder vendored in vllm-omni).
+
+Still out of scope and rejected with explicit 4xx: multi-speaker
+``[SPEAKERn]`` dialogue, ``profile:`` text-only speaker descriptions, the
+``ref_audio_in_system_message`` system-block variant, and chunked long-form
+generation.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from typing import Any
+
+import numpy as np
+import torch
+
+__all__ = [
+    "UnsupportedInputError",
+    "MULTI_SPEAKER_TAG_PATTERN",
+    "validate_plain_text_input",
+    "build_plain_text_conversation",
+    "build_plain_text_prompt",
+    "build_voice_clone_conversation",
+    "build_voice_clone_prompt",
+    "input_ids_to_python_list",
+]
+
+
+class UnsupportedInputError(ValueError):
+    """Raised when a request asks for an out-of-scope higgs_audio_v2 feature."""
+
+
+# Matches the upstream multi-speaker SPEAKERn tag, e.g. [SPEAKER0], [SPEAKER12].
+MULTI_SPEAKER_TAG_PATTERN = re.compile(r"\[SPEAKER\d+\]", re.IGNORECASE)
+
+
+def validate_plain_text_input(text: str) -> None:
+    """Reject multi-speaker tags inside the user text body.
+
+    Phase-1 explicitly does NOT support multi-speaker dialogue. Catching the
+    pattern here means the rejection happens at the tokenizer boundary and is
+    visible to both offline (`pipeline.py`) and online (`serving_speech.py`)
+    code paths.
+    """
+    if not isinstance(text, str):
+        raise UnsupportedInputError(f"higgs_audio_v2 expects plain text input; got {type(text).__name__}")
+    if MULTI_SPEAKER_TAG_PATTERN.search(text):
+        raise UnsupportedInputError(
+            "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; received text contains a speaker tag"
+        )
+
+
+def build_plain_text_conversation(text: str) -> list[dict[str, Any]]:
+    """Build the canonical single-speaker plain-text conversation.
+
+    Uses the bare system prompt ``"Generate audio following instruction."``
+    that matches the upstream HF reference's input formatting; this exact
+    wording is required for input-token parity with the upstream processor.
+    """
+    validate_plain_text_input(text)
+    system_prompt = "Generate audio following instruction."
+    return [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": system_prompt}],
+        },
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": text}],
+        },
+    ]
+
+
+def build_plain_text_prompt(
+    processor: Any,
+    text: str,
+    *,
+    sampling_rate: int = 24000,
+    return_tensors: str | None = "pt",
+) -> dict[str, Any]:
+    """Run the upstream processor's chat template on a plain-text input.
+
+    Returns the processor output dict (``input_ids`` plus any auxiliary tensors
+    such as ``attention_mask``). The serving layer passes ``input_ids`` to
+    Stage 0 as ``prompt_token_ids`` after a ``.tolist()``.
+
+    Using the upstream processor verbatim (no system-prompt rewriting) is
+    what preserves input-token parity with the HF reference.
+    """
+    conversation = build_plain_text_conversation(text)
+    inputs = processor.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        sampling_rate=sampling_rate,
+        return_tensors=return_tensors,
+    )
+    if "input_ids" not in inputs:
+        raise RuntimeError(f"HiggsAudioV2 processor returned no input_ids; got keys {list(inputs.keys())!r}")
+    return inputs
+
+
+def input_ids_to_python_list(inputs: dict[str, Any]) -> list[int]:
+    """Convenience: pull a flat ``list[int]`` of token IDs from a processor output."""
+    ids = inputs["input_ids"]
+    if isinstance(ids, torch.Tensor):
+        if ids.ndim == 2 and int(ids.shape[0]) != 1:
+            raise ValueError(f"expected batch=1 prompt; got input_ids shape {tuple(ids.shape)}")
+        return ids.reshape(-1).tolist()
+    return list(ids)
+
+
+_AUDIO_OUT_TOKEN = "<|AUDIO_OUT|>"
+_AUDIO_OUT_BOS_TOKEN = "<|audio_out_bos|>"
+_AUDIO_EOS_TOKEN = "<|audio_eos|>"
+_AUDIO_DELAY_TOKEN = "<|reserved_special_token_6|>"
+_AUDIO_STREAM_BOS_ID = 1024
+_AUDIO_STREAM_EOS_ID = 1025
+
+
+def _build_delay_pattern(codes: torch.Tensor) -> torch.Tensor:
+    """Apply the upstream delay-pattern wrap to a ``[num_codebooks, T]`` code tensor.
+
+    Mirrors ``HiggsAudioV2Processor.build_delay_pattern``: prepend BOS, append
+    EOS, then arrange in a triangular pattern that stretches the sequence by
+    ``num_codebooks - 1`` frames so codebook k starts emitting real codes at
+    frame k.
+    """
+    num_codebooks, seq_len = codes.shape
+    bos = codes.new_full((num_codebooks, 1), _AUDIO_STREAM_BOS_ID)
+    eos = codes.new_full((num_codebooks, 1), _AUDIO_STREAM_EOS_ID)
+    wrapped = torch.cat([bos, codes, eos], dim=1)
+    wrapped_len = wrapped.shape[1]
+    new_seq_len = wrapped_len + num_codebooks - 1
+
+    output = torch.ones((1, num_codebooks, new_seq_len), dtype=codes.dtype, device=codes.device)
+    bos_mask = torch.tril(output, -1) > 0
+    eos_mask = torch.triu(output, wrapped_len) > 0
+    data_mask = ~(bos_mask | eos_mask)
+    output[bos_mask] = _AUDIO_STREAM_BOS_ID
+    output[data_mask] = wrapped.reshape(-1)
+    output[eos_mask] = _AUDIO_STREAM_EOS_ID
+    return output[0]
+
+
+_ENCODER_CACHE: Any | None = None
+_K2_OMNIVOICE_REPO = "k2-fsa/OmniVoice"
+_K2_OMNIVOICE_SUBDIR = "audio_tokenizer"
+
+
+def _load_audio_tokenizer():
+    """Load the higgs-audio v2 codec via HF ``HiggsAudioV2TokenizerModel``.
+
+    The boson-ai standalone tokenizer repo (``bosonai/higgs-audio-v2-tokenizer``)
+    ships a ``model.safetensors`` that is actually a copy of the 3B talker LM,
+    not the codec — pointing HF's ``from_pretrained`` at it yields an
+    encoder/decoder with randomly initialized structural weights, producing
+    noise codes at encode time. The k2-fsa OmniVoice repo bundles the same
+    codec weights repackaged under
+    ``audio_tokenizer/{config.json,model.safetensors,preprocessor_config.json}``
+    with key naming aligned to ``HiggsAudioV2TokenizerModel``'s class layout
+    (``acoustic_encoder.*`` / ``acoustic_decoder.*`` / ``quantizer.quantizers.*``
+    / ``fc/fc2`` / ``semantic_model.*``), so HF can load it directly.
+
+    Only the ``audio_tokenizer/`` subdirectory (~806 MB) is downloaded.
+    """
+    from huggingface_hub import snapshot_download
+    from transformers import HiggsAudioV2TokenizerModel
+
+    repo_path = snapshot_download(
+        _K2_OMNIVOICE_REPO,
+        allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"],
+    )
+    audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device)
+    return model.eval()
+
+
+def _encode_ref_audio_codes(
+    wav: np.ndarray,
+    sr: int,
+) -> torch.Tensor:
+    """Encode a single ref clip to codec codes via HF ``HiggsAudioV2TokenizerModel``.
+
+    Returns shape ``[num_codebooks, T_raw]`` (before BOS/EOS + delay-pattern wrap).
+    """
+    import torchaudio
+
+    global _ENCODER_CACHE
+    if _ENCODER_CACHE is None:
+        try:
+            _ENCODER_CACHE = _load_audio_tokenizer()
+        except ImportError as exc:
+            raise RuntimeError(
+                "higgs_audio_v2 voice clone needs `transformers>=5.3.0` "
+                "(which ships the HiggsAudioV2TokenizerModel class). "
+                "Install via `pip install -U 'transformers>=5.3.0'`."
+            ) from exc
+
+    target_sr = int(getattr(_ENCODER_CACHE.config, "sample_rate", 24000))
+    wav_t = torch.as_tensor(wav, dtype=torch.float32).reshape(-1)
+    if sr and sr != target_sr:
+        wav_t = torchaudio.functional.resample(wav_t, sr, target_sr)
+    # HF encode() expects (batch=1, channels=1, num_samples).
+    input_values = wav_t.unsqueeze(0).unsqueeze(0).to(_ENCODER_CACHE.device)
+    with torch.inference_mode():
+        codes = _ENCODER_CACHE.encode(input_values, return_dict=False)
+    # codes shape: (batch=1, num_quantizers, codes_length) -> (num_quantizers, T).
+    if codes.dim() == 3:
+        codes = codes.squeeze(0)
+    return codes.detach().to("cpu").long()
+
+
+def build_voice_clone_conversation(
+    text: str,
+    ref_text: str,
+) -> list[dict[str, Any]]:
+    """ChatML conversation for voice clone (HF jinja-template-compatible).
+
+    The assistant turn is a list with a single ``{"type": "audio"}`` content
+    block — no ``"audio"``/``"url"``/``"path"`` key, so ``apply_chat_template``
+    won't try to extract audio data (that path would collide with our explicit
+    audio encoder). The template renders the assistant block as the literal
+    ``<|audio_out_bos|><|AUDIO_OUT|><|audio_eos|>``; we expand the single
+    ``<|AUDIO_OUT|>`` token to ``N × audio_token + (num_codebooks-1) × delay_token``
+    in :func:`build_voice_clone_prompt` to match the reference clip's frame count.
+    """
+    validate_plain_text_input(text)
+    validate_plain_text_input(ref_text)
+    return [
+        {"role": "system", "content": "Generate audio following instruction."},
+        {"role": "user", "content": ref_text},
+        {"role": "assistant", "content": [{"type": "audio"}]},
+        {"role": "user", "content": text},
+    ]
+
+
+def build_voice_clone_prompt(
+    processor: Any,
+    text: str,
+    ref_audio_wav: np.ndarray | torch.Tensor,
+    ref_audio_sr: int,
+    ref_text: str,
+    *,
+    return_tensors: str | None = "pt",
+) -> dict[str, Any]:
+    """Build a voice-clone prompt using the upstream boson audio tokenizer.
+
+    Returns a dict carrying:
+      - ``prompt_token_ids``: ``list[int]`` — input_ids with the ref-audio
+        ``<|AUDIO_OUT|>`` + ``<|reserved_special_token_6|>`` placeholders
+        already expanded to match the encoded reference clip's frame count.
+      - ``audio_input_ids``: ``Tensor[T_frames, num_codebooks]`` — encoded
+        reference codes with BOS/EOS + delay pattern, transposed to ``[T, Q]``
+        for parity with the HF processor output shape.
+      - ``audio_input_ids_mask``: ``Tensor[T_frames]`` — all-True bool mask.
+    """
+    if isinstance(ref_audio_wav, torch.Tensor):
+        wav = ref_audio_wav.detach().to("cpu").float().reshape(-1).numpy()
+    else:
+        wav = np.asarray(ref_audio_wav, dtype=np.float32).reshape(-1)
+
+    # 1. Encode via the upstream boson tokenizer (real weights from model.pth).
+    codes_qt = _encode_ref_audio_codes(wav, ref_audio_sr)  # [num_codebooks, T_raw]
+    if codes_qt.ndim == 3:
+        codes_qt = codes_qt[0]
+    num_codebooks = int(codes_qt.shape[0])
+
+    # 2. Apply BOS/EOS + delay-pattern wrap.
+    audio_input_ids = _build_delay_pattern(codes_qt)  # [num_codebooks, T_full]
+    T_full = int(audio_input_ids.shape[1])
+
+    # 3. Render the chat template, then expand the single <|AUDIO_OUT|> marker
+    #    in the assistant turn into the full placeholder block. After expansion
+    #    the text contains exactly T_full audio-mask positions, matching the
+    #    audio_input_ids time axis.
+    conversation = build_voice_clone_conversation(text, ref_text)
+    rendered = processor.tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        tokenize=False,
+    )
+    n_delay = num_codebooks - 1
+    n_audio = T_full - n_delay
+    if n_audio < 0:
+        raise RuntimeError(f"ref clip too short ({T_full} frames) for delay pattern with num_codebooks={num_codebooks}")
+    placeholders = _AUDIO_OUT_TOKEN * n_audio + _AUDIO_DELAY_TOKEN * n_delay
+    expanded_assistant = f"{_AUDIO_OUT_BOS_TOKEN}{placeholders}{_AUDIO_EOS_TOKEN}"
+    if _AUDIO_OUT_TOKEN not in rendered:
+        raise RuntimeError(
+            f"Voice-clone chat-template render is missing the assistant "
+            f"audio placeholder marker. conversation={conversation!r}"
+        )
+    rendered = rendered.replace(_AUDIO_OUT_TOKEN, expanded_assistant, 1)
+
+    # 4. Tokenize the rendered prompt. ``apply_chat_template`` already emits
+    #    ``<|begin_of_text|>``, so disable add_special_tokens.
+    encoded = processor.tokenizer(
+        rendered,
+        add_special_tokens=False,
+        return_tensors=return_tensors,
+    )
+    prompt_token_ids = encoded["input_ids"]
+    if isinstance(prompt_token_ids, torch.Tensor):
+        prompt_token_ids = prompt_token_ids.reshape(-1).tolist()
+
+    # 5. Transpose codes to [T_full, num_codebooks] for HF parity.
+    audio_input_ids_t = audio_input_ids.transpose(0, 1).contiguous().to(torch.long)
+    audio_input_ids_mask = torch.ones(T_full, dtype=torch.bool)
+
+    return {
+        "prompt_token_ids": prompt_token_ids,
+        "audio_input_ids": audio_input_ids_t,
+        "audio_input_ids_mask": audio_input_ids_mask,
+    }

vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""higgs-audio v2 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM)."""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v2"
+
+HIGGS_AUDIO_V2_PIPELINE = PipelineConfig(
+    model_type="higgs_audio_v2",
+    model_arch="HiggsAudioV2ForConditionalGeneration",
+    hf_architectures=("HiggsAudioV2ForConditionalGeneration",),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="higgs_audio_v2",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            sampling_constraints={
+                "detokenize": False,
+                # Two stop signals for Stage-0:
+                #   * 128009 = standard LM eos_token_id (sequence-level stop)
+                #   * 128012 = audio_eos_token_id (forced by the talker's
+                #     _apply_audio_mode_bias at audio-ramp completion; matches
+                #     upstream HiggsAudioModel._sample_audio_tokens override).
+                "stop_token_ids": [128009, 128012],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="HiggsAudioV2Code2WavForConditionalGeneration",
+            sync_process_input_func=f"{_PROC}.talker2code2wav",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)

vllm_omni/model_executor/models/registry.py

@@ -102,6 +102,22 @@ _OMNI_MODELS = {
         "qwen3_tts_code2wav",
         "Qwen3TTSCode2Wav",
     ),
+    ## higgs-audio v2
+    "HiggsAudioV2ForConditionalGeneration": (
+        "higgs_audio_v2",
+        "higgs_audio_v2_talker",
+        "HiggsAudioV2TalkerForConditionalGeneration",
+    ),
+    "HiggsAudioV2TalkerForConditionalGeneration": (
+        "higgs_audio_v2",
+        "higgs_audio_v2_talker",
+        "HiggsAudioV2TalkerForConditionalGeneration",
+    ),
+    "HiggsAudioV2Code2WavForConditionalGeneration": (
+        "higgs_audio_v2",
+        "higgs_audio_v2_code2wav",
+        "HiggsAudioV2Code2WavForConditionalGeneration",
+    ),
     ## mimo_audio
     "MiMoAudioModel": (
         "mimo_audio",

vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py

@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage-input processor for higgs-audio v2: Talker -> Code2Wav.
+
+Mirrors the qwen3_tts contract:
+
+- ``talker2code2wav`` is the sync (non-async-chunk) adapter that runs once,
+  after Stage 0 has finished, collecting every audio frame into a single
+  ``OmniTokensPrompt`` for Stage 1.
+- ``talker2code2wav_async_chunk`` is the streaming adapter invoked once per
+  emitted codec frame; it accumulates frames in a per-request buffer and
+  flushes an ``OmniPayloadStruct`` once a chunk is ready (or at end-of-stream).
+
+Stage 1 strictly rejects any code id >= ``audio_stream_bos_id`` (1024). This
+processor filters frames that still carry stream specials so the codec sees
+only real codes in [0, 1023], regardless of where Stage 0 emitted them.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.data_entry_keys import (
+    CodesStruct,
+    MetaStruct,
+    OmniPayload,
+    OmniPayloadStruct,
+)
+
+__all__ = ["talker2code2wav", "talker2code2wav_async_chunk"]
+
+logger = init_logger(__name__)
+
+# Default codec chunk-shape knobs for async_chunk mode.  These match the
+# canonical Qwen3-TTS defaults (25 frames = 1 s at 25 fps).  Real values
+# come from the connector config; these are only fallbacks.
+_DEFAULT_CODEC_CHUNK_FRAMES = 25
+_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 25
+
+# Boson-ai constants. Kept local (rather than imported from the model package)
+# so this module remains usable without circular imports during
+# stage-input-processor discovery.
+_NUM_CODEBOOKS = 8
+_AUDIO_STREAM_BOS_ID = 1024
+_AUDIO_STREAM_EOS_ID = 1025
+_NUM_REAL_CODES = _AUDIO_STREAM_BOS_ID  # codes in [0, 1023] are real
+
+
+def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor:
+    """Keep only frames whose codes are entirely in [0, 1023].
+
+    ``audio_codes`` has shape ``[num_frames, num_codebooks]`` after the
+    standard Stage-0 transpose. We drop frames containing stream specials
+    or anything padded outside the real-code range.
+    """
+    if audio_codes.numel() == 0:
+        return audio_codes
+    if audio_codes.ndim != 2:
+        raise ValueError(f"expected [num_frames, num_codebooks] audio_codes; got shape {tuple(audio_codes.shape)}")
+    valid = (audio_codes >= 0).all(dim=1) & (audio_codes < _NUM_REAL_CODES).all(dim=1)
+    return audio_codes[valid]
+
+
+def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor:
+    """Reverse the delay-pattern shift, mirroring upstream's
+    ``boson_multimodal.model.higgs_audio.utils.revert_delay_pattern``.
+
+    Input shape: ``[num_codebooks, seq_len + num_codebooks - 1]``.
+    Output shape: ``[num_codebooks, seq_len]``.
+
+    For each codebook ``i``, slice rows ``[i:i+1]`` from columns
+    ``[i : seq_len + i]`` to remove ``i`` leading BOS pads and ``Q-1-i``
+    trailing PAD entries.
+    """
+    if audio_codes_qt.ndim != 2:
+        raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}")
+    q, t = audio_codes_qt.shape
+    if t < q:
+        # Not enough frames to revert delay pattern; return as-is.
+        return audio_codes_qt
+    seq_len = t - q + 1
+    out_l = []
+    for i in range(q):
+        out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i])
+    return torch.cat(out_l, dim=0)
+
+
+def talker2code2wav(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Sync: collect all talker codes, then pass to code2wav at once."""
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    code2wav_inputs: list[OmniTokensPrompt] = []
+    for talker_output in source_outputs:
+        if not talker_output.finished:
+            continue
+        output = talker_output.outputs[0]
+        mm = output.multimodal_output
+        mm_codes = mm.get("codes", {})
+
+        audio_codes = mm_codes.get("audio")
+        if audio_codes is None or not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+            # Nothing to decode for this request; emit an empty payload so the
+            # downstream stage can still close the response.
+            code2wav_inputs.append(
+                OmniTokensPrompt(
+                    prompt_token_ids=[],
+                    multi_modal_data=None,
+                    mm_processor_kwargs=None,
+                    additional_information=None,
+                )
+            )
+            continue
+
+        audio_codes = audio_codes.to(torch.long)
+        if audio_codes.ndim == 1:
+            # Stage 0 emitted a flat sequence; reshape to [num_frames, num_codebooks].
+            if audio_codes.numel() % _NUM_CODEBOOKS != 0:
+                raise ValueError(
+                    f"flat audio_codes length {audio_codes.numel()} not divisible by num_codebooks={_NUM_CODEBOOKS}"
+                )
+            audio_codes = audio_codes.reshape(-1, _NUM_CODEBOOKS)
+
+        if audio_codes.ndim != 2:
+            raise ValueError(f"audio_codes must be 1D or 2D; got shape {tuple(audio_codes.shape)}")
+
+        # Audio_codes from Stage-0 are emitted in DELAY-PATTERN layout.
+        # Revert the shift to recover the canonical ``[Q, seq_len]`` form,
+        # then clip to the real-code range and trim the first/last frame
+        # (the codec doesn't consume the BOS/EOS edge frames). This mirrors
+        # the upstream ``HiggsAudioServeEngine.generate`` decode line:
+        #     vq_code = revert_delay_pattern(out).clip(0, real-1)[:, 1:-1]
+        codes_qt = audio_codes.transpose(0, 1).contiguous().cpu()
+        codes_qt = _revert_delay_pattern(codes_qt)
+        codes_qt = codes_qt.clamp_(min=0, max=_NUM_REAL_CODES - 1)
+        if codes_qt.shape[-1] >= 3:
+            codes_qt = codes_qt[:, 1:-1]
+        # Code2Wav expects codebook-major flat: [Q * num_frames].
+        codec_codes = codes_qt.reshape(-1).tolist()
+
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+                additional_information=None,
+            )
+        )
+    return code2wav_inputs
+
+
+def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:
+    audio_codes = pooling_output.get("codes", {}).get("audio")
+    if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
+        return None
+    if audio_codes.ndim == 2:
+        frame = audio_codes[-1]
+    elif audio_codes.ndim == 1:
+        frame = audio_codes
+    else:
+        raise ValueError(f"unexpected audio_codes shape for higgs_audio_v2 async_chunk: {tuple(audio_codes.shape)}")
+    if frame.numel() == 0:
+        return None
+    frame = frame.to(torch.long).reshape(-1)
+    # Filter frames that still carry stream specials.
+    if int(frame.min().item()) < 0 or int(frame.max().item()) >= _NUM_REAL_CODES:
+        return None
+    return frame
+
+
+def talker2code2wav_async_chunk(
+    transfer_manager: Any,
+    pooling_output: OmniPayload | None,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Streaming adapter: buffer frames per request and flush a chunk at a time.
+
+    Emits codes in codebook-major flat layout (``[Q * num_frames]``) plus a
+    ``MetaStruct(left_context_size, finished)`` so Stage 1 can stitch chunks
+    together with the conventional left-context overlap.
+    """
+    request_id = request.external_req_id
+    finished = bool(is_finished or request.is_finished())
+
+    if isinstance(pooling_output, dict):
+        frame = _extract_last_frame(pooling_output)
+        if frame is not None:
+            transfer_manager.code_prompt_token_ids[request_id].append(frame.cpu().tolist())
+    elif not finished:
+        return None
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    chunk_size = int(cfg.get("codec_chunk_frames", _DEFAULT_CODEC_CHUNK_FRAMES))
+    left_context_size_config = int(cfg.get("codec_left_context_frames", _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES))
+
+    if chunk_size <= 0 or left_context_size_config < 0:
+        raise ValueError(
+            f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
+            f"codec_left_context_frames={left_context_size_config}"
+        )
+
+    length = len(transfer_manager.code_prompt_token_ids[request_id])
+    if length <= 0:
+        if finished:
+            return OmniPayloadStruct(
+                codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+            )
+        return None
+
+    chunk_remainder = length % chunk_size
+    if not finished and chunk_remainder != 0:
+        # Wait until we have a clean chunk boundary.
+        return None
+
+    context_length = chunk_remainder if (finished and chunk_remainder > 0) else chunk_size
+    end_index = min(length, left_context_size_config + context_length)
+    left_context_size = max(0, end_index - context_length)
+    window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+
+    if not window_frames:
+        # Defensive: should not happen because length > 0, but keeps the contract.
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+            meta=MetaStruct(finished=torch.tensor(finished, dtype=torch.bool)),
+        )
+
+    num_codebooks = len(window_frames[0])
+    if num_codebooks != _NUM_CODEBOOKS:
+        raise ValueError(f"expected {_NUM_CODEBOOKS} codebooks per frame; got {num_codebooks}")
+    num_frames = len(window_frames)
+    flat = torch.tensor(
+        [window_frames[f][q] for q in range(num_codebooks) for f in range(num_frames)],
+        dtype=torch.long,
+    )
+
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=flat),
+        meta=MetaStruct(
+            left_context_size=left_context_size,
+            finished=torch.tensor(finished, dtype=torch.bool),
+        ),
+    )