Compare commits

..
Author SHA1 Message Date
hsliuandClaude Opus 4.7 3bda7ebd8b build: bump docker base images and install docs to v0.22.0
Update Dockerfiles, quickstart, and installation guides to reference
vLLM v0.22.0 (CUDA 12.9, ROCm 7.2.2) instead of v0.20.0/v0.21.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-09 06:13:06 +08:00
Yueqian Lin 969c9b5a61 docs: add Higgs-Audio v2 to supported models for v0.22.0
Higgs-Audio v2 (HiggsAudioV2ForConditionalGeneration, #3762) is included
in the v0.22.0 release but was missing from the supported models table.
Higgs-Audio v3 (#4169) landed after the v0.22.0 cut and is intentionally
not listed here.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-08 08:26:36 -07:00
Hongsheng Liu e592d29f4c docs: update README and supported models for v0.22.0 release
- Add v0.22.0 to Latest News section in README with Cosmos3 recipe link
- Add new models to supported_models.md: MiniCPM-o 4.5, OmniVoice, VoxCPM2,
  MOSS-TTS Nano, MOSS-TTS Delay, MOSS-TTS Realtime, and Lance
- Remove Higgs-Audio v2/v3 (not in 0.22 release)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-06-08 16:43:49 +08:00
Yueqian Lin 82cce633a3 [Perf][Higgs-Audio-V3] Turn on Stage-0 prefix caching by default (#4199)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-07 03:09:47 +00:00
Yueqian Lin fce0f51089 [Perf][Higgs-Audio-V3] LRU cache for voice-clone ref-audio encode (#4200)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-07 03:09:16 +00:00
Yuekai ZhangandYueqian Lin 82e74ce972 [TTS][New Model] support bosonai/higgs-audio-v3-tts-4b (#4169)
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
2026-06-06 15:59:12 +00:00
28 changed files with 4441 additions and 9 deletions

No files matched your search

+23
View File
@@ -399,6 +399,29 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "TTS · Higgs-Audio-V3 Test"
commands:
- |
timeout 20m bash -c "
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_v3.py -m 'core_model' --run-level 'core_model'
"
agents:
queue: "gpu_4_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"
commands:
- |
+1
View File
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*Latest News* 🔥
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with veRL-Omni.
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
@@ -45,7 +45,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
uv pip install vllm==0.22.0 --torch-backend=auto
```
The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -70,8 +70,8 @@ git checkout v0.22.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
+1 -1
View File
@@ -22,7 +22,7 @@ source .venv/bin/activate
uv pip install vllm==0.22.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
+8
View File
@@ -50,6 +50,8 @@ th {
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
@@ -62,6 +64,10 @@ th {
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
@@ -69,11 +75,13 @@ th {
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
@@ -0,0 +1,204 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline higgs-audio v3 inference example.
Runs Stage 0 (Qwen3 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 (plain TTS):
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
--texts "Hello world." \\
--output-dir results/higgs_v3_wavs
Example (voice clone):
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
--texts "Hello world." \\
--ref-audio path/to/reference.wav \\
--ref-text "Transcript of the reference clip." \\
--output-dir results/higgs_v3_wavs
"""
from __future__ import annotations
import os
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
import re
import time
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from vllm_omni import Omni
SAMPLE_RATE = 24_000
DEFAULT_TEXTS = (
"Hello world.",
"The quick brown fox jumps over the lazy dog.",
"Today is a beautiful day for a walk in the park.",
)
def parse_args():
import argparse
parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
parser.add_argument(
"--model",
type=str,
default="bosonai/higgs-audio-v3-tts-4b",
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/higgs_v3_wavs",
help="Directory to write per-prompt WAV files.",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override the deploy config path.",
)
parser.add_argument(
"--max-new-tokens",
type=int,
default=2048,
help="Cap on Stage-0 codec frames per request.",
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Transcript of the reference audio. Optional but improves fidelity.",
)
return parser.parse_args()
def _slugify(text: str) -> str:
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:
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()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
apply_delay_pattern,
encode_reference_audio,
)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
# Load and encode reference audio once (if voice cloning)
ref_codes_delayed: torch.Tensor | 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)
ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
print(f"Reference : {args.ref_audio}")
print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
if args.ref_text:
print(f"Ref text : {args.ref_text}")
print(f"Model : {args.model}")
print(f"Prompts : {len(args.texts)}")
print(f"Output dir : {output_dir}")
print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
total_elapsed = 0.0
total_dur = 0.0
for text in args.texts:
if ref_codes_delayed is not None:
prompt_ids = adapter.build_prompt(
text,
num_ref_tokens=int(ref_codes_delayed.shape[0]),
reference_text=args.ref_text,
)
prompt = {
"prompt_token_ids": prompt_ids,
"additional_information": {
"audio_input_ids": ref_codes_delayed.to(torch.long),
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
},
}
else:
prompt_ids = adapter.build_prompt(text)
prompt = {"prompt_token_ids": prompt_ids}
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__":
main()
@@ -0,0 +1,40 @@
# Higgs-Audio V3 Online Serving
## Start the server
```bash
# Default: GPU 0, port 8095
./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
# Custom GPU / port
PORT=8096 GPUS=0,1 ./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
```
## Plain text TTS
```bash
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_batch \
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
```
## Voice clone
```bash
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_clone \
--ref-audio path/to/reference.wav \
--ref-text "transcript of the reference clip" \
--prompts "Text to synthesize in the cloned voice."
```
## curl example
```bash
curl -X POST http://localhost:8095/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"model": "higgs_audio_v3", "input": "Hello world."}' \
--output hello.wav
```
@@ -0,0 +1,115 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Batch client for the higgs-audio v3 online server.
Sends prompts to ``/v1/audio/speech`` and saves the returned WAV files.
Usage (plain text -> speech):
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_batch \
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
Usage (voice clone):
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--output-dir /tmp/higgs_v3_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.",
"Today is a beautiful day for a walk in the park.",
"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:8095")
parser.add_argument("--model", default="higgs_audio_v3")
parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_v3_batch"))
parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
parser.add_argument("--max-new-tokens", type=int, default=2048)
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 (WAV/FLAC/MP3 path). Pair with --ref-text.",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Transcript of the reference clip. Optional but improves fidelity.",
)
args = parser.parse_args()
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
if args.ref_text:
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())
@@ -0,0 +1,32 @@
#!/bin/bash
# Launch vLLM-Omni server for higgs-audio v3.
#
# Supports plain text TTS and voice cloning via /v1/audio/speech.
#
# Usage:
# ./run_server.sh # default port 8095, GPU 0
# PORT=8096 GPUS=0,1 ./run_server.sh
# MODEL=/path/to/local/checkpoint ./run_server.sh
set -e
MODEL="${MODEL:-bosonai/higgs-audio-v3-tts-4b}"
PORT="${PORT:-8095}"
GPUS="${GPUS:-0}"
GPU_UTIL="${GPU_UTIL:-0.6}"
echo "Starting higgs-audio v3 server"
echo " MODEL=$MODEL"
echo " PORT=$PORT"
echo " CUDA_VISIBLE_DEVICES=$GPUS"
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_multimodal_qwen3.yaml \
--host 0.0.0.0 \
--port "$PORT" \
--gpu-memory-utilization "$GPU_UTIL" \
--trust-remote-code \
--omni
+110
View File
@@ -0,0 +1,110 @@
# Higgs-Audio V3 TTS
> Multilingual text-to-speech with voice cloning on 1×H100
## Summary
- Vendor: Boson AI
- Model: `bosonai/higgs-audio-v3-tts-4b`
- Task: Text-to-speech synthesis with optional voice cloning (100+ languages)
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API; offline batch inference
- Maintainer: @yuekaiz
## When to use this recipe
Use this recipe to serve `bosonai/higgs-audio-v3-tts-4b` for high-quality
multilingual TTS. The model generates 24 kHz speech, supports zero-shot voice
cloning from a reference clip, and handles 100+ languages with inline control
tokens for emotion, style, and prosody. The architecture is a ~4B Qwen3
backbone with fused multi-codebook embedding/head (8 codebooks × 1026 vocab,
MusicGen-style delay pattern).
## References
- Model card: [bosonai/higgs-audio-v3-tts-4b](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b)
- Offline example: [`examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py`](../../examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py)
- Online example: [`examples/online_serving/text_to_speech/higgs_audio_v3/`](../../examples/online_serving/text_to_speech/higgs_audio_v3/)
- Benchmark results: see Performance section below
## Hardware Support
## GPU
### 1×H100 80GB
#### Environment
- OS: Linux
- Python: 3.12+
- CUDA: 12.x
- vLLM version: 0.22.0
- vLLM-Omni version or commit: `36e048fd` (branch `higgs-v3`)
#### Command
**Online serving:**
```bash
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
--host 0.0.0.0 --port 8095 \
--trust-remote-code --omni
```
The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is
loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`).
**Offline batch inference:**
```bash
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
--output-dir results/higgs_v3_wavs
```
**Offline voice clone:**
```bash
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
--texts "Text to synthesize in the cloned voice." \
--ref-audio path/to/reference.wav \
--ref-text "Transcript of the reference clip." \
--output-dir results/higgs_v3_clone
```
#### Verification
Basic TTS via curl:
```bash
curl -X POST http://localhost:8095/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "bosonai/higgs-audio-v3-tts-4b",
"input": "Hello, how are you?"
}' --output hello.wav
```
Voice clone via Python client:
```bash
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
--base-url http://localhost:8095 \
--model bosonai/higgs-audio-v3-tts-4b \
--ref-audio path/to/reference.wav \
--ref-text "Transcript of the reference." \
--prompts "Text to clone."
```
#### Notes
- Memory usage: Stage 0 (talker, ~4B) uses ~60% GPU memory; Stage 1 (codec decoder) uses ~25%.
- Key flags: `--trust-remote-code` and `--omni` are required.
- Output: 24 kHz mono WAV.
- Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity.
- Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`).
- `max_num_seqs=16` for both stages.
- `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
- Known limitations:
- CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph.
- Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture).
- Async streaming (chunk-based) not yet implemented; sync-only pipeline.
@@ -0,0 +1,388 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end online tests for higgs-audio v3 against /v1/audio/speech.
Mirrors the higgs_audio_v2 test layout. Covers the plain-text-in / audio-out
happy path, a small concurrent burst (Stage 0 prefix caching has to coexist
with batching), the voice-clone path via ``ref_audio`` + ``ref_text``, the
upstream ``references[]`` cookbook alias added by
``normalize_references_alias`` in ``protocol/audio.py``, and the few
validator rejections that should remain 4xx.
"""
from __future__ import annotations
import os
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 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-v3-tts-4b"
STAGE_CONFIG = get_deploy_config_path("higgs_multimodal_qwen3.yaml")
SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
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_v3_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; conservative floor catches truncated /
# silence-only outputs without flagging short legitimate clips.
_MIN_AUDIO_BYTES = 20_000
# Reuse the shared TTS reference clip (clean ~5 s 24 kHz mono human speech)
# vendored under tests/assets/qwen3_tts/. Keeps a single WAV across TTS
# tests rather than duplicating asset bytes.
_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 TestHiggsAudioV3OnlineHappyPath:
"""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 - 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_with_max_new_tokens(self, omni_server, openai_client) -> None:
"""``max_new_tokens`` is one of the few extra fields the v3 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 per-slot audio state and Stage-0 PC under batching."""
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,
)
@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 talker2code2wav_async_chunk + the connector's codec_streaming path.
The deploy YAML pins ``async_chunk: true`` and ``codec_streaming: true`` so chunks
flow Stage 0 -> Stage 1 per AR loop. The Stage-1 codec honors
``meta.left_context_size`` and ``meta.right_holdback_size`` so per-chunk codec
windows stitch into a coherent PCM stream. The byte-count gate is the same as
the sync paths; per-chunk audio content is verified offline against Whisper.
NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
check still catches catastrophic codec failure (silence, white noise, sample
scramble all give HNR << 0) while allowing for the sliding-window codec's
slightly-noisier-than-sync output. The default 1.0 dB threshold has only
~0.16 dB margin over measured single-request output (1.16 dB on L4), so
flake risk is too high there.
"""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "The quick brown fox jumps over the lazy dog.",
"stream": True,
"response_format": "pcm",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
"min_hnr_db": 0.0,
}
)
@pytest.mark.core_model
@pytest.mark.tts
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None:
"""Three concurrent streaming requests - guards per-request frame cursors
in ``talker2code2wav_async_chunk`` and per-slot delay-pattern state under batched AR.
NOTE: ``min_hnr_db=-2.0`` is looser than the single-request streaming test
because batched AR adds per-request codec quality variance (DAC boundary
artifacts compounded across the 3-way batch); -2.0 dB gives ~1.6 dB margin
over the measured worst-of-3 on L4 (-0.42 dB) while still well above the
catastrophic-failure region.
"""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "She sells seashells by the seashore.",
"stream": True,
"response_format": "pcm",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
"min_hnr_db": -2.0,
},
request_num=3,
)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
class TestHiggsAudioV3OnlineInlineControlTokens:
"""Inline control-token surface from the bosonai cookbook.
The model exposes four categories of ``<|category:value|>`` tokens
(emotion, style, prosody, sfx). Delivery tokens (emotion / style /
prosody speed-pitch-expressive) go at the START of input; positional
tokens (pause, sfx) go inline. SFX tokens must pair with their written
onomatopoeia.
These tests check the *serving surface* - the validator accepts the
payload, the engine produces audio, and the WAV has non-trivial size.
They do not assert audio quality (out of scope for CI).
"""
@pytest.mark.core_model
@pytest.mark.tts
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None:
"""Delivery tokens (emotion + expressive_high) at the start of input."""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": ("<|emotion:amusement|><|prosody:expressive_high|>Wait, that was actually hilarious."),
"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_inline_style_whispering(self, omni_server, openai_client) -> None:
"""Style token at the start - ``whispering``."""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "<|style:whispering|>It is just between you and me, alright?",
"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_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None:
"""Two prosody tokens at the start - slow speed plus low pitch."""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": (
"<|prosody:speed_slow|><|prosody:pitch_low|>The radar shows a storm approaching from the east."
),
"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_inline_pause_mid_text(self, omni_server, openai_client) -> None:
"""Positional pause token placed inline between two clauses."""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hold on a moment <|prosody:pause|> let me think about it.",
"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_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None:
"""SFX token paired with its written onomatopoeia (``<|sfx:laughter|>Hehe``)."""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": (
"<|emotion:amusement|>I cannot believe that just happened. "
"<|sfx:laughter|>Hehe, I am still recovering from it."
),
"stream": False,
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
class TestHiggsAudioV3OnlineVoiceClone:
"""Voice clone via the two payload shapes the model serves."""
@pytest.mark.core_model
@pytest.mark.tts
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> None:
"""Canonical vllm-omni voice clone via ``ref_audio`` + ``ref_text``."""
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.core_model
@pytest.mark.tts
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_voice_clone_references_alias(self, omni_server, openai_client) -> None:
"""BosonAI cookbook payload: ``references=[{audio_path, text}]``.
``normalize_references_alias`` (``protocol/audio.py``) is supposed to
translate the cookbook field into ``ref_audio`` / ``ref_text``; without
it the request would silently fall through to zero-shot synthesis.
"""
openai_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hello world.",
"references": [{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT}],
"stream": False,
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
class TestHiggsAudioV3OnlineValidatorRejections:
"""Out-of-scope shapes must come back as 4xx.
These cases use ``send_audio_speech_http_request`` (raw HTTP POST against
``/v1/audio/speech``) rather than the OpenAI SDK helper because:
- We need the failure-path matcher; the SDK helper asserts success.
- The cookbook ``references`` field is not in the OpenAI SDK schema, so
the SDK strips it before the request leaves the client and the server
sees a plain TTS request that returns 200. Raw HTTP keeps the JSON body
intact so the validator actually fires.
"""
@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:
"""Empty text input - validator should reject before reaching the engine."""
openai_client.send_audio_speech_http_request(
{
"json": {
"model": omni_server.model,
"input": "",
"response_format": "wav",
},
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"err_code": (400, 422),
"err_message": ("input", "empty"),
}
)
@pytest.mark.core_model
@pytest.mark.tts
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_rejects_multi_reference_payload(self, omni_server, openai_client) -> None:
"""``references[]`` with more than one entry - multi-shot voice clone is not supported."""
openai_client.send_audio_speech_http_request(
{
"json": {
"model": omni_server.model,
"input": "Hello world.",
"references": [
{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT},
{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT},
],
"response_format": "wav",
},
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"err_code": (400, 422),
"err_message": "references",
}
)
@pytest.mark.core_model
@pytest.mark.tts
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_rejects_conflicting_ref_audio_and_references(self, omni_server, openai_client) -> None:
"""``ref_audio`` and ``references`` cannot both be set to different values."""
openai_client.send_audio_speech_http_request(
{
"json": {
"model": omni_server.model,
"input": "Hello world.",
"ref_audio": _REF_AUDIO_URL,
"references": [{"audio_path": "https://example.com/other.wav"}],
"response_format": "wav",
},
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"err_code": (400, 422),
"err_message": "mutually exclusive",
}
)
Whitespace-only changes.
@@ -0,0 +1,910 @@
# SPDX-License-Identifier: Apache-2.0
"""Deterministic unit tests for higgs-audio v3.
These tests verify AC-1 (config), AC-3 (prompt), AC-4 (fused modules),
AC-5 (delay pattern), AC-7 (stage processor), and AC-10 (registry)
without requiring the actual checkpoint or GPU.
"""
import pytest
import torch
# ---- AC-1: Configuration ----
class TestHiggsAudioV3Config:
def test_default_config_loads(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
config = HiggsAudioV3Config()
assert config.num_codebooks == 8
assert config.codebook_size == 1026
assert config.audio_stream_bos_id == 1024
assert config.audio_stream_eos_id == 1025
assert config.sample_rate == 24000
assert config.frame_rate == 25
assert config.num_real_codes == 1024
assert config.tie_modality_embeddings is True
def test_custom_audio_encoder_config(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
config = HiggsAudioV3Config(
audio_encoder_config={
"encoder_type": "discrete",
"num_codebooks": 4,
"vocab_size": 512,
"tie_word_embeddings": False,
}
)
assert config.num_codebooks == 4
assert config.codebook_size == 512
assert config.tie_modality_embeddings is False
def test_invalid_num_codebooks_rejected(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
with pytest.raises(ValueError, match="num_codebooks must be > 0"):
HiggsAudioV3Config(audio_encoder_config={"num_codebooks": 0, "vocab_size": 1026})
def test_negative_num_codebooks_rejected(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
with pytest.raises(ValueError, match="num_codebooks must be > 0"):
HiggsAudioV3Config(audio_encoder_config={"num_codebooks": -1, "vocab_size": 1026})
def test_special_token_ids_initially_none(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
config = HiggsAudioV3Config()
assert config.tts_token_id is None
assert config.text_token_id is None
assert config.audio_continuation_id is None
def test_hidden_size_from_text_config(self):
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
config = HiggsAudioV3Config(text_config={"model_type": "qwen3", "hidden_size": 2560})
assert config.hidden_size == 2560
assert config.audio_hidden_size == 2560
# ---- AC-4: Fused Multi-Codebook Modules ----
class TestFusedModules:
def test_embedding_shape(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsFusedMultiTextEmbedding,
)
embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=256)
assert embed.weight.shape == (8 * 1026, 256)
codes = torch.randint(0, 1024, (5, 8))
out = embed(codes)
assert out.shape == (5, 256)
def test_embedding_offset_indexing(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsFusedMultiTextEmbedding,
)
embed = HiggsFusedMultiTextEmbedding(num_codebooks=2, vocab_size=4, hidden_size=3)
# Set weights to identity-like pattern so we can verify offsets
with torch.no_grad():
embed.weight.zero_()
# Codebook 0, vocab [0..3] -> row 0..3
# Codebook 1, vocab [0..3] -> row 4..7
for i in range(8):
embed.weight[i, i % 3] = float(i)
# codes[0] = [1, 2] -> embed[1] + embed[4+2=6]
codes = torch.tensor([[1, 2]])
out = embed(codes)
expected = embed.weight[1] + embed.weight[6]
assert torch.allclose(out[0], expected)
def test_head_shape(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsFusedMultiTextHead,
)
head = HiggsFusedMultiTextHead(num_codebooks=8, vocab_size=1026, hidden_size=256)
assert head.weight.shape == (8 * 1026, 256)
hidden = torch.randn(3, 256)
out = head.generate(hidden)
assert out.shape == (3, 8, 1026)
def test_tying(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsFusedMultiTextEmbedding,
HiggsFusedMultiTextHead,
)
embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=64)
head = HiggsFusedMultiTextHead(num_codebooks=8, vocab_size=1026, hidden_size=64)
head.weight = embed.weight
assert head.weight is embed.weight
# Modification to one should reflect in the other
with torch.no_grad():
embed.weight[0, 0] = 42.0
assert head.weight[0, 0] == 42.0
def test_codes_out_of_range(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsFusedMultiTextEmbedding,
)
embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
# Code 1026 + offset for codebook 7 = 1026 + 7*1026 = 8208
# But weight only has 8*1026 = 8208 rows, so index 8208 is out of bounds
codes = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 1026]]) # Last cb: 1026 is OOB
with pytest.raises(IndexError):
embed(codes)
# ---- AC-5: Delay Pattern Behavior ----
class TestDelayPatternBehavior:
"""Test the delay pattern masking logic extracted from the talker."""
def test_boc_eoc_ids(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
BOC_ID,
EOC_ID,
)
assert BOC_ID == 1024
assert EOC_ID == 1025
def test_delay_phase_boc_masking(self):
"""During delay phase, codebooks beyond delay_count are forced to BOC."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
num_codebooks = 8
# Simulate delay phase at step 3 (delay_count=2, so CBs 3-7 should be BOC)
codes = torch.randint(0, 1024, (num_codebooks,))
delay_count = 2
next_cb = delay_count + 1
if next_cb < num_codebooks:
codes[next_cb:] = BOC_ID
# CBs 0-2 should have original codes, CBs 3-7 should be BOC
assert all(codes[i] != BOC_ID or i >= 3 for i in range(num_codebooks))
assert all(codes[i] == BOC_ID for i in range(3, num_codebooks))
def test_cb0_eoc_triggers_rampdown(self):
"""EOC on codebook 0 starts ramp-down; EOC on other codebooks does not."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
EOC_ID,
)
num_codebooks = 8
codes = torch.randint(0, 1024, (num_codebooks,))
# Simulate: cb0 emits EOC
codes[0] = EOC_ID
assert int(codes[0].item()) == EOC_ID
# Ramp-down should start with countdown = N-2 = 6
eoc_countdown = num_codebooks - 2
assert eoc_countdown == 6
# If cb3 emits EOC but cb0 doesn't, no ramp-down
codes2 = torch.randint(0, 1024, (num_codebooks,))
codes2[3] = EOC_ID
assert int(codes2[0].item()) != EOC_ID # cb0 is not EOC
def test_rampdown_termination(self):
"""After N-2 steps of ramp-down, generation_done becomes True."""
num_codebooks = 8
eoc_countdown = num_codebooks - 2 # = 6
generation_done = False
for _ in range(6):
eoc_countdown -= 1
if eoc_countdown <= 0:
generation_done = True
assert generation_done is True
def test_all_boc_seed(self):
"""First audio step should seed last_codes with all-BOC."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
seeded = torch.full((8,), BOC_ID, dtype=torch.long)
assert seeded.shape == (8,)
assert all(seeded[i] == BOC_ID for i in range(8))
# ---- AC-5: Real Sampler Method Tests ----
class TestSamplerMethods:
"""Test the actual _sample_audio_codes and _apply_delay_pattern_masking methods."""
def _make_minimal_talker(self):
"""Create a minimal talker-like object with sampler/masking methods."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
class FakeTalker:
num_codebooks = 8
codebook_size = 1026
_audio_state = {}
t = FakeTalker()
t._sample_audio_codes = mod.HiggsAudioV3TalkerForConditionalGeneration._sample_audio_codes.__get__(t)
t._apply_delay_pattern_masking = (
mod.HiggsAudioV3TalkerForConditionalGeneration._apply_delay_pattern_masking.__get__(t)
)
return t
def test_sample_respects_mask(self):
"""Tokens masked to -inf must never be sampled."""
t = self._make_minimal_talker()
logits = torch.full((10, 1026), float("-inf"))
# Only token 500 is allowed for each row
logits[:, 500] = 0.0
result = t._sample_audio_codes(logits)
assert result.shape == (10,)
assert (result == 500).all(), f"Expected all 500, got {result.tolist()}"
def test_sample_all_masked_falls_back_to_argmax(self):
"""All-masked row should fall back to argmax (least-negative logit)."""
t = self._make_minimal_talker()
logits = torch.full((2, 1026), float("-inf"))
# Row 0: all masked
# Row 1: only token 42 allowed
logits[1, 42] = 0.0
result = t._sample_audio_codes(logits)
assert result.shape == (2,)
assert result[1].item() == 42
def test_delay_masking_forces_boc_during_delay(self):
"""During delay phase, codebooks beyond delay_count must have only BOC allowed."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
t = self._make_minimal_talker()
# Simulate delay_count=2 for batch row 0
t._audio_state[0] = {"num_delay": 2, "num_remaining_delays": None}
cb_logits = torch.zeros(1, 8, 1026) # [1 audio row, 8 codebooks, 1026 vocab]
t._apply_delay_pattern_masking(cb_logits, [0])
# CBs 3-7 should have only BOC allowed (everything else -inf)
for q in range(3, 8):
row = cb_logits[0, q]
assert row[BOC_ID].item() == 0.0 # BOC kept
# All non-BOC should be -inf
non_boc = torch.cat([row[:BOC_ID], row[BOC_ID + 1 :]])
assert (non_boc == float("-inf")).all(), f"CB{q} has non-inf non-BOC values"
def test_delay_masking_disallows_boc_for_active_codebooks(self):
"""Active codebooks during delay should have BOC disallowed."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
t = self._make_minimal_talker()
t._audio_state[0] = {"num_delay": 3, "num_remaining_delays": None}
cb_logits = torch.zeros(1, 8, 1026)
t._apply_delay_pattern_masking(cb_logits, [0])
# CBs 0-3 are active: BOC should be -inf
for q in range(4):
assert cb_logits[0, q, BOC_ID].item() == float("-inf")
def test_delay_masking_only_cb0_allows_eoc(self):
"""Only codebook 0 should allow EOC during normal generation."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
t = self._make_minimal_talker()
t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": None}
cb_logits = torch.zeros(1, 8, 1026)
t._apply_delay_pattern_masking(cb_logits, [0])
# CB0 should keep EOC
assert cb_logits[0, 0, EOC_ID].item() != float("-inf")
# CB1-7 should have EOC masked
for q in range(1, 8):
assert cb_logits[0, q, EOC_ID].item() == float("-inf"), f"CB{q} EOC not masked"
def test_rampdown_masking_locks_to_eoc(self):
"""During ramp-down, locked codebooks should only allow EOC."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
BOC_ID,
EOC_ID,
)
t = self._make_minimal_talker()
# Ramp-down with 4 remaining delays: lock CBs 0-3 to EOC
t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": 4}
cb_logits = torch.zeros(1, 8, 1026)
t._apply_delay_pattern_masking(cb_logits, [0])
# CBs 0-3 locked: only EOC allowed
for q in range(4):
row = cb_logits[0, q]
assert row[EOC_ID].item() == 0.0
non_eoc = torch.cat([row[:EOC_ID], row[EOC_ID + 1 :]])
assert (non_eoc == float("-inf")).all(), f"Locked CB{q} has non-inf non-EOC"
# CBs 4-7 active: BOC and EOC disallowed
for q in range(4, 8):
assert cb_logits[0, q, BOC_ID].item() == float("-inf")
assert cb_logits[0, q, EOC_ID].item() == float("-inf")
# ---- AC-6: Feedback Method Tests ----
class TestFeedbackMethods:
def test_postprocess_emits_audio_codes(self):
"""postprocess() should return codes from _last_audio_codes."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
class FakeTalker:
_last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]])
_postprocess_cursor = 0
t = FakeTalker()
t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
result = t.postprocess(torch.zeros(1, 64))
assert "codes" in result
assert "audio" in result["codes"]
assert result["codes"]["audio"].shape == (1, 8)
def test_postprocess_skips_negative_rows(self):
"""postprocess() should skip rows with -1 (no audio)."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
class FakeTalker:
_last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]])
_postprocess_cursor = 0
t = FakeTalker()
t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
result = t.postprocess(torch.zeros(1, 64))
assert result == {}
def test_postprocess_advances_cursor(self):
"""postprocess() should advance cursor by 1 per call."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
class FakeTalker:
_last_audio_codes = torch.tensor(
[[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]]
)
_postprocess_cursor = 0
t = FakeTalker()
t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
r1 = t.postprocess(torch.zeros(1, 64))
assert "codes" in r1
r2 = t.postprocess(torch.zeros(1, 64))
assert r2 == {} # Second row is all -1
# ---- AC-6: Audio Feedback Method Tests ----
class TestAudioFeedback:
"""Test _apply_audio_feedback with minimal fake talker."""
def _make_feedback_talker(self):
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
torch.nn.init.ones_(embed.weight)
class FakeTalker:
num_codebooks = 8
codebook_size = 1026
_audio_continuation_id = 99999 # fake audio token
_last_step_query_start_loc = None
_audio_state = {}
multimodal_embedding = embed
model = type("M", (), {"embed_tokens": lambda self, ids: torch.zeros(ids.shape[0], 16)})()
t = FakeTalker()
t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
return t
def test_text_positions_unchanged(self):
"""Non-audio positions should not be modified."""
t = self._make_feedback_talker()
# No audio state → all positions unchanged
input_ids = torch.tensor([1, 2, 3, 4])
hidden = torch.randn(4, 16)
result = t._apply_audio_feedback(hidden, input_ids)
assert torch.equal(result, hidden)
def test_audio_position_replaced_with_state(self):
"""Audio position with state should have its embedding replaced."""
t = self._make_feedback_talker()
audio_id = t._audio_continuation_id
t._audio_state[1] = {"last_codes": torch.zeros(8, dtype=torch.long)}
input_ids = torch.tensor([1, audio_id, 3])
hidden = torch.zeros(3, 16)
result = t._apply_audio_feedback(hidden, input_ids)
# Position 1 (audio) should be non-zero (embedding of all-0 codes)
assert result[1].abs().sum() > 0
# Positions 0 and 2 should remain zero
assert result[0].abs().sum() == 0
assert result[2].abs().sum() == 0
def test_no_state_no_replacement(self):
"""Audio position without state should not be replaced."""
t = self._make_feedback_talker()
audio_id = t._audio_continuation_id
# No state for row 0
input_ids = torch.tensor([audio_id])
hidden = torch.zeros(1, 16)
result = t._apply_audio_feedback(hidden, input_ids)
assert torch.equal(result, hidden)
# ---- AC-8: Codec Strictness ----
class TestCodecStrictness:
def test_bundled_missing_quantizer_key_raises(self):
"""Bundled codec load must fail when a quantizer codebook key is missing."""
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
config = HiggsAudioV3Config()
c2w = HiggsAudioV3Code2Wav(config=config)
# State with only 7 quantizers (missing quantizer.quantizers.7.*)
codec_state = {}
for i in range(7):
codec_state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64)
codec_state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64)
codec_state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024)
codec_state["fc2.weight"] = torch.randn(256, 1024)
codec_state["fc2.bias"] = torch.randn(256)
with pytest.raises(KeyError, match="quantizer 7"):
c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
def test_bundled_missing_fc2_raises(self):
"""Bundled codec load must fail when fc2 keys are missing."""
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
config = HiggsAudioV3Config()
c2w = HiggsAudioV3Code2Wav(config=config)
codec_state = {}
for i in range(8):
codec_state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64)
codec_state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64)
codec_state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024)
# No fc2 keys
with pytest.raises(KeyError, match="fc2"):
c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
def test_bundled_unknown_decoder_key_rejected(self, monkeypatch):
"""Unknown decoder-side key should be rejected as unconsumed."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
class FakeDAC(torch.nn.Module):
def __init__(self):
super().__init__()
self.dummy = torch.nn.Parameter(torch.zeros(1))
monkeypatch.setattr(c2w_mod, "build_boson_dac_decoder", lambda device: FakeDAC().to(device))
monkeypatch.setattr(c2w_mod, "build_higgs_audio_acoustic_decoder", lambda cfg, device: FakeDAC().to(device))
config = HiggsAudioV3Config()
c2w = HiggsAudioV3Code2Wav(config=config)
codec_state = _build_full_codec_state()
codec_state["unknown_decoder_module.weight"] = torch.randn(10)
with pytest.raises(RuntimeError, match="unexpected decoder-side keys"):
c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
def test_bundled_encoder_side_keys_accepted(self, monkeypatch):
"""Known encoder-side keys should not trigger unconsumed-key rejection."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
HiggsAudioV3Code2Wav,
)
class FakeDAC(torch.nn.Module):
def __init__(self):
super().__init__()
self.dummy = torch.nn.Parameter(torch.zeros(1))
monkeypatch.setattr(c2w_mod, "build_boson_dac_decoder", lambda device: FakeDAC().to(device))
monkeypatch.setattr(c2w_mod, "build_higgs_audio_acoustic_decoder", lambda cfg, device: FakeDAC().to(device))
config = HiggsAudioV3Config()
c2w = HiggsAudioV3Code2Wav(config=config)
codec_state = _build_full_codec_state()
codec_state["acoustic_encoder.block.0.conv_t1.weight"] = torch.randn(10)
codec_state["semantic_model.encoder.layers.0.weight"] = torch.randn(10)
codec_state["fc.weight"] = torch.randn(10)
# Should NOT raise — encoder-side keys are allowed
c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
assert c2w._loaded
def _build_minimal_codec_state() -> dict[str, torch.Tensor]:
"""Build a minimal codec state dict (quantizer + fc2, no DAC)."""
state: dict[str, torch.Tensor] = {}
for i in range(8):
state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64)
state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64)
state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024)
state["fc2.weight"] = torch.randn(256, 1024)
state["fc2.bias"] = torch.randn(256)
return state
def _build_full_codec_state() -> dict[str, torch.Tensor]:
"""Build a full codec state dict including acoustic_decoder keys.
The acoustic_decoder keys here match what a FakeDAC module expects
(just 'dummy'). Real DAC has many more keys.
"""
state = _build_minimal_codec_state()
state["acoustic_decoder.dummy"] = torch.zeros(1)
return state
# ---- AC-2: Loader Required-Key Validation ----
class TestLoaderRequiredKeys:
"""Test exact required-key validation in load_weights."""
def test_build_required_keys_count(self):
"""Required key set should have 1 text_emb + 36*11 layers + 1 norm + 1 modality = 399."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsAudioV3TalkerForConditionalGeneration,
)
keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(36)
assert len(keys) == 1 + 36 * 11 + 1 + 1 # text + layers + norm + modality = 399
def test_build_required_keys_contains_norm(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsAudioV3TalkerForConditionalGeneration,
)
keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(2)
assert "body.norm.weight" in keys
def test_build_required_keys_contains_layer_subkeys(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsAudioV3TalkerForConditionalGeneration,
)
keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(2)
assert "body.layers.0.self_attn.k_proj.weight" in keys
assert "body.layers.1.mlp.gate_proj.weight" in keys
def test_build_required_keys_contains_modality_embedding(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsAudioV3TalkerForConditionalGeneration,
)
keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1)
assert "tied.embedding.modality_embeddings.0.embedding.weight" in keys
def test_build_required_keys_contains_text_embedding(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
HiggsAudioV3TalkerForConditionalGeneration,
)
keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1)
assert "tied.embedding.text_embedding.weight" in keys
class TestLoadWeightsEnforcement:
"""Test load_weights() itself with synthetic key sets via a minimal fake talker."""
def _make_fake_talker_and_keys(self):
"""Build a minimal talker-like object and the full 1-layer key set."""
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
# Minimal fake talker with just enough to run load_weights
talker = object.__new__(mod.HiggsAudioV3TalkerForConditionalGeneration)
torch.nn.Module.__init__(talker)
# Attach minimal attributes load_weights needs
talker.multimodal_embedding = mod.HiggsFusedMultiTextEmbedding(8, 1026, 16)
talker.modality_head = mod.HiggsFusedMultiTextHead(8, 1026, 16)
talker.tie_modality = True
talker.modality_head.weight = talker.multimodal_embedding.weight
# Fake backbone config with 1 layer
class FakeConfig:
num_hidden_layers = 1
tie_word_embeddings = True
talker._backbone_config = FakeConfig()
# Minimal model and lm_head so _BackboneWrapper can be constructed
talker.model = torch.nn.Module()
talker.lm_head = torch.nn.Module()
# Fake vllm_config for resolve_special_tokens (no-op)
class FakeModelConfig:
model = None
class FakeVllmConfig:
model_config = FakeModelConfig()
talker.vllm_config = FakeVllmConfig()
talker.config = type("C", (), {"resolve_special_tokens": lambda self, p: None})()
talker._resolved_tokens = True
talker._audio_continuation_id = None
talker._eos_token_id = None
# Build the full required key set for 1 layer
required = mod.HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1)
# Build synthetic weights: tiny tensors for every required key
weights = {}
for key in required:
if "modality_embeddings.0.embedding" in key:
weights[key] = torch.randn(8 * 1026, 16)
else:
weights[key] = torch.randn(4) # tiny placeholder
# Monkeypatch _BackboneWrapper.load_weights to accept and return keys
def fake_backbone_load(self_wrapper, ws):
consumed = set()
for name, _ in ws:
consumed.add(name)
return consumed
talker._fake_backbone_load = fake_backbone_load
return talker, weights, mod, fake_backbone_load
def test_load_weights_succeeds_with_all_keys(self, monkeypatch):
"""load_weights succeeds when all required keys are present."""
talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
result = talker.load_weights(iter(weights.items()))
assert isinstance(result, set)
assert "multimodal_embedding.weight" in result
def test_load_weights_fails_missing_norm(self, monkeypatch):
"""load_weights raises when body.norm.weight is missing."""
talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
del weights["body.norm.weight"]
with pytest.raises(RuntimeError, match="body.norm.weight"):
talker.load_weights(iter(weights.items()))
def test_load_weights_fails_missing_layer_subkey(self, monkeypatch):
"""load_weights raises when a specific layer subkey is missing."""
talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
del weights["body.layers.0.self_attn.k_proj.weight"]
with pytest.raises(RuntimeError, match="body.layers.0.self_attn.k_proj.weight"):
talker.load_weights(iter(weights.items()))
def test_load_weights_fails_missing_modality_embedding(self, monkeypatch):
"""load_weights raises when modality embedding is missing."""
talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
del weights["tied.embedding.modality_embeddings.0.embedding.weight"]
with pytest.raises(RuntimeError, match="modality_embeddings"):
talker.load_weights(iter(weights.items()))
def test_load_weights_fails_unknown_prefix(self, monkeypatch):
"""load_weights raises on unknown non-codec checkpoint prefix."""
talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
weights["unknown.prefix.weight"] = torch.randn(4)
with pytest.raises(ValueError, match="Unexpected checkpoint key"):
talker.load_weights(iter(weights.items()))
# ---- AC-7: Stage Input Processor ----
class TestStageInputProcessor:
def test_revert_delay_pattern(self):
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
_revert_delay_pattern,
)
# 8 codebooks, 3 real frames -> delayed shape [8, 3+8-1=10]
Q, T = 8, 3
delayed = torch.full((Q, T + Q - 1), 1024) # Fill with BOC
for i in range(Q):
for t in range(T):
delayed[i, i + t] = i * 100 + t # Real codes at shifted positions
result = _revert_delay_pattern(delayed)
assert result.shape == (Q, T)
for i in range(Q):
for t in range(T):
assert result[i, t].item() == i * 100 + t
def test_revert_delay_pattern_rejects_wrong_codebooks(self):
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
_revert_delay_pattern,
)
# 7 codebooks should be rejected
codes = torch.zeros(7, 20)
with pytest.raises(ValueError, match="Expected exactly 8 codebook rows"):
_revert_delay_pattern(codes)
def test_revert_delay_pattern_rejects_too_few_frames(self):
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
_revert_delay_pattern,
)
# 8 codebooks but only 5 frames (need at least 8)
codes = torch.zeros(8, 5)
with pytest.raises(ValueError, match="Not enough frames"):
_revert_delay_pattern(codes)
def test_filter_real_code_frames(self):
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
_filter_real_code_frames,
)
# 8 codebooks, 4 frames
codes = torch.tensor(
[
[100, 200, 1024, 300], # cb0: frame 2 has BOC
[101, 201, 1024, 301],
[102, 202, 1024, 302],
[103, 203, 1024, 303],
[104, 204, 1024, 304],
[105, 205, 1024, 305],
[106, 206, 1024, 306],
[107, 207, 1024, 307],
]
)
result = _filter_real_code_frames(codes)
# Frame 2 (column 2) has BOC in all codebooks -> filtered out
assert result.shape == (8, 3)
assert result[0, 0].item() == 100
assert result[0, 1].item() == 200
assert result[0, 2].item() == 300
# ---- AC-10: Registry ----
class TestRegistry:
def test_talker_registered(self):
from vllm_omni.model_executor.models.registry import _OMNI_MODELS
assert "HiggsMultimodalQwen3ForConditionalGeneration" in _OMNI_MODELS
assert "HiggsAudioV3TalkerForConditionalGeneration" in _OMNI_MODELS
def test_code2wav_registered(self):
from vllm_omni.model_executor.models.registry import _OMNI_MODELS
assert "HiggsAudioV3Code2WavForConditionalGeneration" in _OMNI_MODELS
def test_pipeline_registered(self):
from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
assert "higgs_multimodal_qwen3" in _OMNI_PIPELINES
def test_deploy_yaml_exists(self):
import os
# __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root
repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
yaml_path = os.path.join(repo_root, "vllm_omni", "deploy", "higgs_multimodal_qwen3.yaml")
assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
# ---- AC-3: Prompt Builder ----
class TestPromptBuilder:
def _make_mock_tokenizer(self):
"""Create a mock tokenizer with the required special tokens."""
class MockTokenizer:
def __init__(self):
self._added_vocab = {
"<|tts|>": 151700,
"<|text|>": 151701,
"<|audio|>": 151702,
"<|ref_audio|>": 151703,
"<|ref_text|>": 151704,
}
def get_added_vocab(self):
return self._added_vocab
def encode(self, text, add_special_tokens=True):
# Simple word-level tokenization for testing
return list(range(100, 100 + len(text.split())))
return MockTokenizer()
def test_plain_tts_prompt(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
)
tok = self._make_mock_tokenizer()
adapter = HiggsAudioV3TokenizerAdapter(tok)
ids = adapter.build_prompt("Hello world")
assert ids[0] == 151700 # <|tts|>
assert ids[1] == 151701 # <|text|>
assert ids[-1] == 151702 # <|audio|>
assert len(ids) == 2 + 2 + 1 # tts + text + 2 word tokens + audio
def test_empty_text_rejected(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
)
tok = self._make_mock_tokenizer()
adapter = HiggsAudioV3TokenizerAdapter(tok)
with pytest.raises(ValueError, match="non-empty"):
adapter.build_prompt("")
def test_whitespace_only_rejected(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
)
tok = self._make_mock_tokenizer()
adapter = HiggsAudioV3TokenizerAdapter(tok)
with pytest.raises(ValueError, match="non-empty"):
adapter.build_prompt(" ")
def test_missing_specials_rejected(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
)
class BadTokenizer:
def get_added_vocab(self):
return {"<|tts|>": 1} # Missing <|text|> and <|audio|>
with pytest.raises(ValueError, match="missing"):
HiggsAudioV3TokenizerAdapter(BadTokenizer())
def test_no_voice_clone_tokens_in_plain_tts(self):
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
)
tok = self._make_mock_tokenizer()
adapter = HiggsAudioV3TokenizerAdapter(tok)
ids = adapter.build_prompt("Hello")
# Should not contain ref_audio or ref_text token IDs
assert 151703 not in ids # <|ref_audio|>
assert 151704 not in ids # <|ref_text|>
+4
View File
@@ -149,6 +149,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
"vllm_omni.model_executor.models.higgs_audio_v2.pipeline",
"HIGGS_AUDIO_V2_PIPELINE",
),
"higgs_multimodal_qwen3": (
"vllm_omni.model_executor.models.higgs_audio_v3.pipeline",
"HIGGS_AUDIO_V3_PIPELINE",
),
"dynin_omni": (
"vllm_omni.model_executor.models.dynin_omni.pipeline",
"DYNIN_OMNI_PIPELINE",
+11 -5
View File
@@ -391,27 +391,33 @@ class OmniTensorPrefixCache:
self,
query_start_loc: torch.Tensor,
input_batch: InputBatch,
multimodal_outputs: dict,
multimodal_outputs: dict | None,
num_scheduled_tokens: dict[str, int],
):
"""Get the merged multimodal states if hidden state prefix caching is enabled."""
combined_multimodal_outputs = {}
# Talkers that produce multimodal outputs only at decode (e.g. Higgs
# Audio v3's audio codes) have ``multimodal_outputs is None`` at the
# initial prefill step. Treat it as an empty mapping so the merger
# short-circuits cleanly and the cache merge still runs once codes
# start arriving.
mm_outputs = multimodal_outputs if multimodal_outputs is not None else {}
# First get the prefix cached tensors that are present in the mm data
for mm_key in self.mm_cache_keys:
if mm_key in multimodal_outputs:
if mm_key in mm_outputs:
combined_multimodal_outputs[mm_key] = self._get_merged_tensors(
query_start_loc=query_start_loc,
input_batch=input_batch,
cache=self.mm_outputs_cache[mm_key],
hidden_states=multimodal_outputs[mm_key],
hidden_states=mm_outputs[mm_key],
num_scheduled_tokens=num_scheduled_tokens,
)
# Then, get everything else (passthrough data); first, convert to CPU
# tensors similarly to the non prefix cached path, and then populate
# the subdicts mapping request IDs -> payload objects
passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys
passthrough_mm_data = {k: v for k, v in multimodal_outputs.items() if k in passthrough_keys}
passthrough_keys = set(mm_outputs.keys()) - self.mm_cache_keys
passthrough_mm_data = {k: v for k, v in mm_outputs.items() if k in passthrough_keys}
mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data)
for mm_key, mm_val in mm_cpu.items():
+2
View File
@@ -63,6 +63,7 @@ class OmniPayloadMeta(TypedDict, total=False):
stream_finished: torch.Tensor
req_id: list[str]
left_context_size: int
right_holdback_size: int
override_keys: list[tuple[str, str]]
num_processed_tokens: int
next_stage_prompt_len: int
@@ -150,6 +151,7 @@ class MetaStruct(_StructBase):
stream_finished: torch.Tensor | None = None
req_id: list[str] | None = None
left_context_size: int | None = None
right_holdback_size: int | None = None
override_keys: list[tuple[str, str]] | None = None
num_processed_tokens: int | None = None
next_stage_prompt_len: int | None = None
@@ -0,0 +1,74 @@
# Higgs-audio v3 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory.
#
# Plain TTS only (no voice cloning). async_chunk streams codec chunks from
# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's
# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size
# so the per-chunk codec windows stitch into a coherent PCM stream.
async_chunk: true
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
# Stream PCM bytes back to the client as Stage 1 emits each chunk.
codec_streaming: true
# Streaming chunk shape for talker2code2wav_async_chunk:
# codec_chunk_frames steady-state new frames per emit
# codec_left_context_frames sliding past context to avoid leading-edge artifacts
# codec_right_holdback_frames trailing frames held back so the codec
# has future context for emit-region tail samples
# initial_codec_chunk_frames smaller first chunk to minimize TTFA
codec_chunk_frames: 25
codec_left_context_frames: 25
codec_right_holdback_frames: 4
initial_codec_chunk_frames: 1
connector_get_sleep_s: 0.01
connector_get_max_wait_first_chunk: 5000
connector_get_max_wait: 500
stages:
- stage_id: 0
max_num_seqs: 16
gpu_memory_utilization: 0.6
enforce_eager: true
trust_remote_code: true
# Stage-0 prefix caching is on by default. The talker's HS opt-out and
# deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py)
# keep per-step bookkeeping negligible; any shared text prefix across
# concurrent requests is a net throughput win. Stage 1 (codec decoder)
# has no KV cache and stays disabled.
enable_prefix_caching: true
async_scheduling: false
max_num_batched_tokens: 4096
max_model_len: 8192
devices: "0"
output_connectors:
to_stage_1: connector_of_shared_memory
default_sampling_params:
temperature: 1.0
top_p: 0.95
top_k: 50
max_tokens: 2048
seed: 42
repetition_penalty: 1.0
- stage_id: 1
max_num_seqs: 16
gpu_memory_utilization: 0.25
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
async_scheduling: false
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
@@ -115,6 +115,54 @@ class OpenAICreateSpeechRequest(BaseModel):
raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
return v
@model_validator(mode="before")
@classmethod
def normalize_references_alias(cls, data):
"""Map the BosonAI Higgs Audio v3 cookbook ``references`` array onto ``ref_audio`` / ``ref_text``.
Upstream Higgs Audio v3 examples and the boson.ai cookbook send voice
clones as::
{"references": [{"audio_path": str, "text": str}]}
vllm-omni's existing convention is ``ref_audio`` + ``ref_text``.
Without this alias pydantic silently drops the unknown ``references``
field and the request falls through to zero-shot synthesis with no
error, which is hard to debug from the client side. The normalizer
translates a single reference and rejects multi-reference payloads
(vllm-omni does not yet support multi-shot voice clone) or conflicts
with the explicit ``ref_audio`` / ``ref_text`` fields.
"""
if not isinstance(data, dict):
return data
refs = data.get("references")
if refs is None:
return data
if not isinstance(refs, list) or len(refs) == 0:
raise ValueError("'references' must be a non-empty array of {audio_path, text} objects")
if len(refs) > 1:
raise ValueError("'references' only supports a single reference; multi-shot voice clone is not supported")
first = refs[0]
if not isinstance(first, dict):
raise ValueError("'references[0]' must be an object with 'audio_path' and optional 'text'")
audio_path = first.get("audio_path")
if not isinstance(audio_path, str) or not audio_path:
raise ValueError(
"'references[0].audio_path' is required and must be a string (URL, data: URI, or file path)"
)
existing_ref_audio = data.get("ref_audio")
if existing_ref_audio and existing_ref_audio != audio_path:
raise ValueError("'references' and 'ref_audio' are mutually exclusive")
data["ref_audio"] = audio_path
text = first.get("text")
if isinstance(text, str) and text.strip():
existing_ref_text = data.get("ref_text")
if existing_ref_text and existing_ref_text != text:
raise ValueError("'references[0].text' and 'ref_text' conflict; supply only one")
data["ref_text"] = text
data.pop("references", None)
return data
@model_validator(mode="before")
@classmethod
def reject_higgs_audio_v2_unsupported_aliases(cls, data):
@@ -76,6 +76,7 @@ _MING_TTS_MODEL_STAGES = {"ming_tts"}
_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
_MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
_HIGGS_AUDIO_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
_GLM_TTS_MODEL_STAGES = {"glm_tts"}
_TTS_MODEL_STAGES: set[str] = (
_VOXTRAL_TTS_MODEL_STAGES
@@ -84,6 +85,7 @@ _TTS_MODEL_STAGES: set[str] = (
| _COSYVOICE3_TTS_MODEL_STAGES
| _OMNIVOICE_TTS_MODEL_STAGES
| _HIGGS_AUDIO_V2_TTS_MODEL_STAGES
| _HIGGS_AUDIO_V3_TTS_MODEL_STAGES
| _COVO_AUDIO_MODEL_STAGES
| _VOXCPM2_TTS_MODEL_STAGES
| _MING_TTS_MODEL_STAGES
@@ -98,6 +100,7 @@ _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
"cosyvoice3",
"voxcpm2",
"higgs_audio_v2",
"higgs_audio_v3",
}
_TTS_LANGUAGES: set[str] = {
"Auto",
@@ -594,6 +597,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
return "moss_tts"
if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
return "higgs_audio_v2"
if model_stage in _HIGGS_AUDIO_V3_TTS_MODEL_STAGES:
return "higgs_audio_v3"
if model_stage in _GLM_TTS_MODEL_STAGES:
return "glm_tts"
return None
@@ -934,6 +939,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
"moss_tts_nano",
"glm_tts",
"higgs_audio_v2",
"higgs_audio_v3",
):
label = {
"cosyvoice3": "CosyVoice3",
@@ -941,6 +947,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
"omnivoice": "OmniVoice",
"moss_tts_nano": "MOSS-TTS-Nano",
"higgs_audio_v2": "Higgs-Audio V2",
"higgs_audio_v3": "Higgs-Audio V3",
"glm_tts": "GLM-TTS",
}.get(self._tts_model_type, self._tts_model_type)
return (
@@ -1896,6 +1903,90 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
self._higgs_audio_v2_processor = processor
return processor
# ---- higgs-audio v3 ----
def _validate_higgs_audio_v3_request(self, request: OpenAICreateSpeechRequest) -> str | None:
"""Validate higgs_audio_v3 request parameters."""
if not request.input or not request.input.strip():
return "higgs_audio_v3: input text cannot be empty"
if request.ref_audio is not None and not request.ref_text:
# Voice clone ref_text is optional for v3 (improves fidelity but not required)
pass
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}"
return None
async def _build_higgs_audio_v3_params(self, request: OpenAICreateSpeechRequest):
"""Build prompt_token_ids for higgs_audio_v3.
Plain-text path: builds ``[tts, text, tokens, audio]``.
Voice-clone path: encodes reference audio, applies delay pattern,
builds ``[tts, (ref_text, tokens,) ref_audio, -100xN, text, tokens, audio]``.
"""
adapter = await self._resolve_higgs_audio_v3_adapter()
if request.ref_audio is None:
prompt_ids = adapter.build_prompt(request.input)
return tokens_input(prompt_token_ids=prompt_ids)
# Voice clone
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
apply_delay_pattern,
encode_reference_audio,
)
from vllm_omni.model_executor.models.higgs_audio_v3.ref_audio_cache import (
cached_encode_reference_audio,
)
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
wav = np.asarray(wav_list, dtype=np.float32)
# Cache the codec encode by content hash so identical ref_audio
# payloads skip the encode and reuse the same [num_frames,
# num_codebooks] codes tensor across requests.
ref_codes_raw = await asyncio.to_thread(cached_encode_reference_audio, wav, int(sr), encode_reference_audio)
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
prompt_ids = adapter.build_prompt(
request.input,
num_ref_tokens=int(ref_codes_delayed.shape[0]),
reference_text=request.ref_text or None,
)
prompt = tokens_input(prompt_token_ids=prompt_ids)
import torch
prompt["additional_information"] = {
"audio_input_ids": ref_codes_delayed.to(torch.long),
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
}
return prompt
async def _resolve_higgs_audio_v3_adapter(self):
"""Lazy-load the tokenizer adapter for higgs_audio_v3."""
cached = getattr(self, "_higgs_audio_v3_adapter", None)
if cached is not None:
return cached
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
HiggsAudioV3TokenizerAdapter,
)
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:
model_path = getattr(self.engine_client, "model", None)
if model_path is None:
raise RuntimeError("higgs_audio_v3 serving could not resolve model path")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
self._higgs_audio_v3_adapter = adapter
return adapter
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():
@@ -2796,6 +2887,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
additional["voice_name"] = voice_lower
additional["voice_created_at"] = self._voice_created_at(voice_lower)
tts_params = {}
elif self._tts_model_type == "higgs_audio_v3":
err = self._apply_uploaded_speaker(request)
if err:
raise ValueError(err)
validation_error = self._validate_higgs_audio_v3_request(request)
if validation_error:
raise ValueError(validation_error)
prompt = await self._build_higgs_audio_v3_params(request)
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.
@@ -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 v3 (two-stage TTS)."""
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
__all__ = ["HiggsAudioV3Config"]
@@ -0,0 +1,182 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Configuration class for higgs-audio v3 (HiggsMultimodalQwen3) in vllm-omni.
``HiggsAudioV3Config.from_pretrained(model_path)`` returns a config with
``tts_token_id``, ``text_token_id``, ``audio_continuation_id``, and
``eos_token_id`` already resolved from the checkpoint tokenizer. If the
tokenizer is unavailable or missing required specials, the load raises.
"""
from __future__ import annotations
import logging
import os
from typing import Any
import transformers
from transformers import AutoConfig, PretrainedConfig
__all__ = ["HiggsAudioV3Config"]
logger = logging.getLogger(__name__)
_QWEN3_ROPE_THETA = 1_000_000
_REQUIRED_SPECIALS = ("<|tts|>", "<|text|>", "<|audio|>")
def _build_text_config(raw: Any) -> PretrainedConfig:
if isinstance(raw, PretrainedConfig):
return raw
cfg = dict(raw or {})
model_type = cfg.get("model_type", "qwen3")
if model_type == "qwen3" and cfg.get("rope_theta") is None:
cfg["rope_theta"] = _QWEN3_ROPE_THETA
try:
cfg_cls = transformers.CONFIG_MAPPING[model_type]
except KeyError as exc:
raise ValueError(f"Unknown text backbone model_type {model_type!r}") from exc
return cfg_cls(**cfg)
def _resolve_model_dir(pretrained_model_name_or_path: str) -> str | None:
"""Resolve a model name/path to a local directory containing tokenizer files."""
if os.path.isdir(pretrained_model_name_or_path):
return pretrained_model_name_or_path
try:
from huggingface_hub import try_to_load_from_cache
cached = try_to_load_from_cache(repo_id=pretrained_model_name_or_path, filename="tokenizer_config.json")
if isinstance(cached, str) and os.path.isfile(cached):
return os.path.dirname(cached)
except Exception:
pass
try:
from huggingface_hub.constants import HF_HUB_CACHE
safe = pretrained_model_name_or_path.replace("/", "--")
snapshots = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots")
if os.path.isdir(snapshots):
for rev in os.listdir(snapshots):
candidate = os.path.join(snapshots, rev)
if os.path.isfile(os.path.join(candidate, "tokenizer_config.json")):
return candidate
except Exception:
pass
return None
class HiggsAudioV3Config(PretrainedConfig):
"""Typed config for higgs-audio v3 (HiggsMultimodalQwen3).
``from_pretrained()`` automatically resolves ``<|tts|>``, ``<|text|>``,
``<|audio|>`` and ``eos_token_id`` from the checkpoint tokenizer.
"""
model_type: str = "higgs_multimodal_qwen3"
is_composition = True
def __init__(
self,
audio_encoder_config: dict[str, Any] | None = None,
text_config: dict[str, Any] | PretrainedConfig | None = None,
audio_token_id: int = -100,
mel_per_sample: int = 8,
num_codebooks: int = 8,
codebook_size: int = 1026,
audio_stream_bos_id: int = 1024,
audio_stream_eos_id: int = 1025,
sample_rate: int = 24000,
frame_rate: int = 25,
tts_token_id: int | None = None,
text_token_id: int | None = None,
audio_continuation_id: int | None = None,
**kwargs: Any,
) -> None:
self.audio_token_id = audio_token_id
self.mel_per_sample = mel_per_sample
if audio_encoder_config is None:
audio_encoder_config = {
"encoder_type": "discrete",
"num_codebooks": num_codebooks,
"vocab_size": codebook_size,
"tie_word_embeddings": True,
}
self.audio_encoder_config = audio_encoder_config
self.num_codebooks = int(audio_encoder_config.get("num_codebooks", num_codebooks))
if self.num_codebooks <= 0:
raise ValueError(f"num_codebooks must be > 0, got {self.num_codebooks}")
self.codebook_size = int(audio_encoder_config.get("vocab_size", codebook_size))
self.tie_modality_embeddings = bool(audio_encoder_config.get("tie_word_embeddings", True))
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.text_config = _build_text_config(text_config)
self.audio_hidden_size = int(audio_encoder_config.get("out_dim", self.text_config.hidden_size))
self.tts_token_id = tts_token_id
self.text_token_id = text_token_id
self.audio_continuation_id = audio_continuation_id
super().__init__(**kwargs)
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs: Any) -> HiggsAudioV3Config:
"""Load config and resolve special token IDs from the checkpoint tokenizer.
Passes the original ``pretrained_model_name_or_path`` (local dir or
HF repo ID) directly to ``AutoTokenizer.from_pretrained()`` so it can
handle cache hits, downloads, and local paths uniformly. Raises if the
tokenizer is missing required specials.
"""
config = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
config.resolve_special_tokens(pretrained_model_name_or_path)
return config
def get_text_config(self, decoder: bool = False) -> PretrainedConfig:
del decoder
return self.text_config
@property
def num_real_codes(self) -> int:
return self.audio_stream_bos_id
@property
def hidden_size(self) -> int:
return self.text_config.hidden_size
def resolve_special_tokens(self, model_path: str) -> None:
"""Resolve <|tts|>, <|text|>, <|audio|> and eos from the HF tokenizer.
Raises ``ValueError`` if any of the 3 required specials is missing
from the tokenizer's added vocabulary.
"""
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
vocab = dict(tokenizer.get_added_vocab())
missing = [t for t in _REQUIRED_SPECIALS if t not in vocab]
if missing:
raise ValueError(
f"Tokenizer at {model_path} is missing required Higgs TTS v3 "
f"special tokens: {missing}. Available added tokens: "
f"{sorted(vocab.keys())[:20]}"
)
self.tts_token_id = vocab["<|tts|>"]
self.text_token_id = vocab["<|text|>"]
self.audio_continuation_id = vocab["<|audio|>"]
if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
self.eos_token_id = int(tokenizer.eos_token_id)
AutoConfig.register("higgs_multimodal_qwen3", HiggsAudioV3Config)
@@ -0,0 +1,524 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage 1 (codec decoder) for higgs-audio v3.
Reuses higgs-audio-v2's RVQ + DAC codec decoder but loads weights from the
v3 checkpoint's bundled codec (``tied.embedding.modality_embeddings.0.model.*``
prefix) rather than from a standalone tokenizer repo.
"""
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.higgs_audio_decoder import (
HiggsAudioRVQ,
build_boson_dac_decoder,
build_higgs_audio_acoustic_decoder,
load_higgs_audio_codec,
)
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput
__all__ = [
"HiggsAudioV3Code2Wav",
"HiggsAudioV3Code2WavForConditionalGeneration",
]
logger = init_logger(__name__)
# Prefix under which codec weights are stored in the v3 checkpoint.
_CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model."
class HiggsAudioV3Code2Wav(nn.Module):
"""Stage-1 codec decoder for higgs-audio v3.
The codec architecture is identical to v2 (HiggsAudioRVQ + fc2 +
BosonDacDecoder); only the weight loading path differs. V3 bundles
codec weights inside the main checkpoint under a known prefix, while
v2 uses a standalone tokenizer repo.
For backward compatibility and simplicity, we delegate to the v2
``load_higgs_audio_codec`` function when the audio tokenizer repo
is available. When not (or as primary path), we extract codec weights
from the v3 safetensors via the engine's weight iterator.
"""
input_modalities = "audio"
def __init__(
self,
config: HiggsAudioV3Config | None = None,
*,
vllm_config: VllmConfig | None = None,
prefix: str = "",
):
super().__init__()
if vllm_config is not None:
hf_config = vllm_config.model_config.hf_config
if isinstance(hf_config, HiggsAudioV3Config):
self.config = hf_config
else:
self.config = HiggsAudioV3Config(**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("HiggsAudioV3Code2Wav: provide either `config` or `vllm_config`.")
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)
self.hop_length: int = 960
# Engine-runner hooks
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
# Codec modules (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
# Do NOT eagerly load from standalone tokenizer here.
# load_weights() will try bundled V3 codec first, then fall back.
def _load_from_tokenizer_repo(self) -> None:
"""Load codec from the standalone higgs-audio-v2-tokenizer repo."""
tokenizer_id = "bosonai/higgs-audio-v2-tokenizer"
# Try to resolve from HF cache
from huggingface_hub import try_to_load_from_cache
cached = try_to_load_from_cache(repo_id=tokenizer_id, filename="config.json")
if isinstance(cached, str) and os.path.isfile(cached):
tokenizer_dir = os.path.dirname(cached)
else:
from huggingface_hub.constants import HF_HUB_CACHE
safe = tokenizer_id.replace("/", "--")
snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots")
if os.path.isdir(snapshots_dir):
tokenizer_dir = None
for rev in os.listdir(snapshots_dir):
candidate = os.path.join(snapshots_dir, rev)
if os.path.isfile(os.path.join(candidate, "config.json")):
tokenizer_dir = candidate
break
if tokenizer_dir is None:
raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}")
else:
raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
quantizer, fc2, acoustic_decoder, _cfg = load_higgs_audio_codec(tokenizer_dir, device)
self.quantizer = quantizer
self.fc2 = fc2
self.acoustic_decoder = acoustic_decoder
self._loaded = True
logger.info("Loaded HiggsAudioV3Code2Wav from standalone tokenizer repo.")
# ------------------------------------------------------------------ engine hooks
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
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):
if self._loaded:
if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)):
for _ in weights_or_model_dir:
pass
return {name for name, _ in self.named_parameters()}
if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)):
# Engine path: try to extract V3 bundled codec weights first
codec_state: dict[str, torch.Tensor] = {}
for name, tensor in weights_or_model_dir:
if name.startswith(_CODEC_PREFIX):
stripped = name[len(_CODEC_PREFIX) :]
codec_state[stripped] = tensor
if codec_state:
# Bundled keys found — load strictly, no fallback on error.
self._load_from_bundled_state(codec_state, device)
logger.info(
"Loaded HiggsAudioV3Code2Wav from bundled V3 checkpoint (%d codec keys consumed).",
len(codec_state),
)
return {name for name, _ in self.named_parameters()}
# No bundled keys — fallback to standalone tokenizer repo
if not self._loaded:
try:
self._load_from_tokenizer_repo()
except (FileNotFoundError, OSError) as exc2:
raise RuntimeError(
"HiggsAudioV3Code2Wav: could not load codec weights. "
"Neither V3 bundled codec keys (tied.embedding."
"modality_embeddings.0.model.*) were found in the "
"checkpoint, nor is bosonai/higgs-audio-v2-tokenizer "
f"cached in HF_HOME. Last error: {exc2}"
) from exc2
return {name for name, _ in self.named_parameters()}
# File-based path
if not self._loaded:
self._load_from_tokenizer_repo()
return None
def _load_from_bundled_state(
self,
codec_state: dict[str, torch.Tensor],
device: torch.device | None = None,
) -> None:
"""Load codec from V3 checkpoint's bundled weights.
Keys arrive with the ``tied.embedding.modality_embeddings.0.model.``
prefix already stripped. We remap to the same layout that
``load_higgs_audio_codec`` expects and build the modules.
"""
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import (
_remap_boson_model_pth_state_dict,
)
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# The bundled codec may use boson-ai key naming. Try remapping.
needs_remap = any(k.startswith("quantizer.vq.layers.") for k in codec_state)
if needs_remap:
codec_state = _remap_boson_model_pth_state_dict(codec_state)
# Discover parameters from the state dict
proj_key = "quantizer.quantizers.0.project_out.weight"
if proj_key not in codec_state:
raise KeyError(
f"Bundled codec state dict missing '{proj_key}'. "
f"Available keys (first 10): {list(codec_state.keys())[:10]}"
)
hidden_size = codec_state[proj_key].shape[0]
codebook_dim = 64
codebook_size = 1024
required_quantizers = self.num_codebooks
# Build RVQ — require exactly num_codebooks quantizers
quantizer = HiggsAudioRVQ(
num_quantizers=required_quantizers,
codebook_size=codebook_size,
codebook_dim=codebook_dim,
hidden_size=hidden_size,
).to(device)
for i in range(required_quantizers):
prefix = f"quantizer.quantizers.{i}"
# Require codebook weight
embed_key = None
for key_suffix in [".codebook.embed", ".codebook.weight"]:
candidate = f"{prefix}{key_suffix}"
if candidate in codec_state:
embed_key = candidate
break
if embed_key is None:
raise KeyError(
f"Bundled codec missing codebook weight for quantizer {i}. "
f"Expected {prefix}.codebook.embed or {prefix}.codebook.weight"
)
quantizer.quantizers[i].codebook.weight.data.copy_(codec_state[embed_key].to(device))
# Require projector weights
proj_w = f"{prefix}.project_out.weight"
proj_b = f"{prefix}.project_out.bias"
if proj_w not in codec_state:
raise KeyError(f"Bundled codec missing {proj_w}")
if proj_b not in codec_state:
raise KeyError(f"Bundled codec missing {proj_b}")
quantizer.quantizers[i].project_out.weight.data.copy_(codec_state[proj_w].to(device))
quantizer.quantizers[i].project_out.bias.data.copy_(codec_state[proj_b].to(device))
# Build fc2
fc2_w = codec_state.get("fc2.weight")
fc2_b = codec_state.get("fc2.bias")
if fc2_w is None or fc2_b is None:
raise KeyError(
"Bundled codec state dict missing fc2.weight and/or fc2.bias. "
f"Available keys (first 10): {list(codec_state.keys())[:10]}"
)
fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device)
fc2.weight.data.copy_(fc2_w.to(device))
fc2.bias.data.copy_(fc2_b.to(device))
# Build acoustic decoder — detect layout from key names.
# V3 checkpoint uses OmniVoice layout (block.N/conv1/conv2/snake1),
# not boson-ai standalone layout (model.0/model.1).
ad_keys = {
k[len("acoustic_decoder.") :]: v for k, v in codec_state.items() if k.startswith("acoustic_decoder.")
}
is_boson_layout = any(k.startswith("model.") for k in ad_keys)
if is_boson_layout:
acoustic_decoder = build_boson_dac_decoder(device)
else:
# OmniVoice layout — need the v2 tokenizer config for DacConfig
tokenizer_config = {
"acoustic_model_config": {
"codebook_dim": 8,
"codebook_size": 1024,
"decoder_hidden_size": 1024,
"downsampling_ratios": [8, 5, 4, 2, 3],
"encoder_hidden_size": 64,
"hidden_size": 256,
"hop_length": 960,
"model_type": "dac",
"n_codebooks": 9,
"sampling_rate": 16000,
"upsampling_ratios": [8, 5, 4, 2, 3],
}
}
acoustic_decoder = build_higgs_audio_acoustic_decoder(tokenizer_config, device)
if not ad_keys:
raise KeyError(
"Bundled codec state dict has no acoustic_decoder.* keys. "
f"Available keys (first 10): {list(codec_state.keys())[:10]}"
)
load_report = acoustic_decoder.load_state_dict(ad_keys, strict=False)
if load_report.missing_keys:
raise RuntimeError(
f"Bundled DAC decoder missing {len(load_report.missing_keys)} keys: {load_report.missing_keys[:5]}..."
)
if load_report.unexpected_keys:
raise RuntimeError(
f"Bundled DAC decoder has {len(load_report.unexpected_keys)} unexpected keys: "
f"{load_report.unexpected_keys[:5]}..."
)
acoustic_decoder.eval()
# Verify required decoder-side codec keys were consumed. The V3
# checkpoint also bundles encoder-side keys (acoustic_encoder.*,
# encoder_semantic.*, decoder_semantic.*, semantic_model.*, fc.*, fc1.*)
# that the DAC decoder does not need — those are allowed but ignored.
consumed: set[str] = set()
for i in range(required_quantizers):
prefix = f"quantizer.quantizers.{i}"
for suffix in [".codebook.embed", ".codebook.weight"]:
if f"{prefix}{suffix}" in codec_state:
consumed.add(f"{prefix}{suffix}")
consumed.add(f"{prefix}.project_out.weight")
consumed.add(f"{prefix}.project_out.bias")
consumed.add("fc2.weight")
consumed.add("fc2.bias")
for k in codec_state:
if k.startswith("acoustic_decoder."):
consumed.add(k)
# Known unused prefixes: encoder-side modules and quantizer training
# artifacts (cluster_size, embed_avg, inited, project_in) that are
# bundled in the checkpoint but not needed for inference decoding.
_KNOWN_UNUSED_PREFIXES = (
"acoustic_encoder.",
"encoder_semantic.",
"decoder_semantic.",
"semantic_model.",
"fc.",
"fc1.",
)
# Quantizer training-only suffixes (per-codebook)
_QUANTIZER_TRAINING_SUFFIXES = (
".codebook.cluster_size",
".codebook.embed_avg",
".codebook.inited",
".project_in.weight",
".project_in.bias",
)
unconsumed = set()
for k in codec_state:
if k in consumed:
continue
if any(k.startswith(p) for p in _KNOWN_UNUSED_PREFIXES):
continue
if any(k.endswith(s) for s in _QUANTIZER_TRAINING_SUFFIXES):
continue
unconsumed.add(k)
if unconsumed:
raise RuntimeError(
f"Bundled codec has {len(unconsumed)} unexpected decoder-side keys: {sorted(unconsumed)[:10]}..."
)
self.quantizer = quantizer
self.fc2 = fc2
self.acoustic_decoder = acoustic_decoder
self._loaded = True
# ------------------------------------------------------------------ decode
@torch.inference_mode()
def decode_codes(self, audio_codes: torch.Tensor) -> torch.Tensor:
"""Decode [B, num_codebooks=8, T] codes to PCM [B, 1, T*960]."""
if not self._loaded:
raise RuntimeError("HiggsAudioV3Code2Wav not loaded.")
codes = self._validate_codes(audio_codes)
rvq_codes = codes.transpose(0, 1).long()
quantized = self.quantizer.decode(rvq_codes)
quantized = quantized.to(dtype=self.fc2.weight.dtype)
quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
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(codes.device)
@torch.inference_mode()
def forward_chunk(
self,
audio_codes: torch.Tensor,
*,
left_context_size: int = 0,
right_holdback_size: int = 0,
hop_length: int | None = None,
) -> torch.Tensor:
# ``left_context_size`` frames trimmed off the front (already-emitted
# prefix). ``right_holdback_size`` frames trimmed off the end
# (deferred to the next chunk so the codec has future context for
# the emit region's trailing samples).
hop = int(hop_length) if hop_length is not None else self.hop_length
pcm = self.decode_codes(audio_codes)
right_trim = right_holdback_size * hop
left_trim = left_context_size * hop
if right_trim > 0:
if pcm.shape[-1] <= right_trim:
return pcm[..., :0]
pcm = pcm[..., :-right_trim]
if left_trim == 0:
return pcm
if pcm.shape[-1] <= left_trim:
return pcm[..., :0]
return pcm[..., left_trim:]
# ------------------------------------------------------------------ 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:
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)
right_holdback_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"])
if "right_holdback_size" in meta:
right_holdback_size[i] = int(meta["right_holdback_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(
"HiggsAudioV3Code2Wav: flat code length %d not divisible by %d",
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],
right_holdback_size=right_holdback_size[i],
hop_length=self.hop_length,
)
except ValueError as exc:
logger.warning("HiggsAudioV3Code2Wav: 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:
if audio_codes.ndim != 3:
raise ValueError(
f"audio_codes must have shape [B, {self.num_codebooks}, T]; got {tuple(audio_codes.shape)}"
)
if int(audio_codes.shape[1]) != self.num_codebooks:
raise ValueError(f"dim 1 must equal num_codebooks={self.num_codebooks}; got {int(audio_codes.shape[1])}")
# ``.item()`` triggers a GPU->CPU sync and is illegal inside CUDA-graph
# capture (the dummy run before warmup hits this path). Skip the
# range check while a graph is being captured; the per-step validator
# still fires at real inference time.
if audio_codes.numel() > 0 and not torch.cuda.is_current_stream_capturing():
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(
f"audio_codes out of range: min={min_val}, max={max_val}; expected [0, {self.num_real_codes - 1}]"
)
return audio_codes
@staticmethod
def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
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]
HiggsAudioV3Code2WavForConditionalGeneration = HiggsAudioV3Code2Wav
@@ -0,0 +1,958 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage-0 talker for higgs-audio v3 (Qwen3 backbone, fused multi-codebook).
Architecture:
- Backbone: Qwen3 (~4B, 36 layers, 2560 hidden, GQA 32/8). No DualFFN.
- Fused multi-codebook embedding: [N*V, D] weight, offset lookup, sum across N
- Fused multi-codebook head: same weight (tied), reshape to [L, N, V]
- MusicGen-style delay pattern [0,1,...,7] with BOC/EOC
- Audio feedback: replace continuation-token embedding with fused codebook embed
Weight loading maps from the HF checkpoint's prefixes:
tied.embedding.text_embedding. -> model.embed_tokens.
body.layers. -> model.layers.
body.norm. -> model.norm.
tied.head.text_head. -> lm_head.
tied.embedding.modality_embeddings.0.embedding. -> multimodal_embedding.
tied.embedding.modality_embeddings.0.model.* -> skipped (codec for code2wav)
tied.head.modality_heads.0.* -> skipped when tied
"""
from __future__ import annotations
import copy
from collections.abc import Iterable
from typing import Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.model_executor.models.qwen3 import Qwen3Model
from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
HiggsAudioV3Config,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput
__all__ = ["HiggsAudioV3TalkerForConditionalGeneration"]
logger = init_logger(__name__)
# Delay pattern constants
BOC_ID = 1024 # beginning of codebook
EOC_ID = 1025 # end of codebook
# Checkpoint prefix mapping
_BACKBONE_PREFIX_MAP = {
"tied.embedding.text_embedding.": "model.embed_tokens.",
"body.layers.": "model.layers.",
"body.norm.": "model.norm.",
"tied.head.text_head.": "lm_head.",
}
_MODALITY_EMBEDDING_PREFIX = "tied.embedding.modality_embeddings.0.embedding."
_MODALITY_HEAD_PREFIX = "tied.head.modality_heads.0."
_CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model."
class HiggsFusedMultiTextEmbedding(nn.Module):
"""Fused multi-codebook embedding: [N*V, D] weight + offset lookup."""
def __init__(self, num_codebooks: int, vocab_size: int, hidden_size: int):
super().__init__()
self.weight = nn.Parameter(torch.empty(num_codebooks * vocab_size, hidden_size))
self.num_codebooks = num_codebooks
self.vocab_size = vocab_size
def forward(self, codes: torch.Tensor) -> torch.Tensor:
N = self.num_codebooks
V = self.vocab_size
offsets = torch.arange(N, device=codes.device, dtype=codes.dtype) * V
fused_ids = codes + offsets
return F.embedding(fused_ids, self.weight).sum(dim=-2)
class HiggsFusedMultiTextHead(nn.Module):
"""Fused multi-codebook head: [L, D] -> [L, N, V] via one linear."""
def __init__(self, num_codebooks: int, vocab_size: int, hidden_size: int):
super().__init__()
self.weight = nn.Parameter(torch.empty(num_codebooks * vocab_size, hidden_size))
self.num_codebooks = num_codebooks
self.vocab_size = vocab_size
def generate(self, hidden: torch.Tensor) -> torch.Tensor:
logits = F.linear(hidden, self.weight)
return logits.reshape(hidden.shape[0], self.num_codebooks, self.vocab_size)
class HiggsAudioV3TalkerForConditionalGeneration(nn.Module):
"""Stage-0 talker for higgs-audio v3.
Wraps Qwen3Model backbone + fused multi-codebook modules for TTS generation
with MusicGen-style delay pattern sampling and audio feedback embedding.
"""
# Tell the AR runner to call model.sample() instead of the stock sampler.
prefer_model_sampler: bool = True
# Tell the runner to call postprocess() to emit per-step audio codes.
have_multimodal_outputs: bool = True
has_postprocess: bool = True
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
hf_config = vllm_config.model_config.hf_config
if isinstance(hf_config, HiggsAudioV3Config):
self.config = hf_config
else:
self.config = HiggsAudioV3Config(**hf_config.to_dict())
self.vllm_config = vllm_config
self.num_codebooks = int(self.config.num_codebooks)
self.codebook_size = int(self.config.codebook_size)
hidden_size = int(self.config.audio_hidden_size)
self.tie_modality = self.config.tie_modality_embeddings
# Fused multi-codebook modules
self.multimodal_embedding = HiggsFusedMultiTextEmbedding(self.num_codebooks, self.codebook_size, hidden_size)
self.modality_head = HiggsFusedMultiTextHead(self.num_codebooks, self.codebook_size, hidden_size)
if self.tie_modality:
self.modality_head.weight = self.multimodal_embedding.weight
# Qwen3 backbone
self._backbone_config = self.config.text_config
backbone_vllm_config = copy.copy(vllm_config)
backbone_model_config = copy.copy(vllm_config.model_config)
backbone_model_config.hf_config = self._backbone_config
backbone_vllm_config.model_config = backbone_model_config
self.model = Qwen3Model(
vllm_config=backbone_vllm_config,
prefix=f"{prefix}.model" if prefix else "model",
)
if self._backbone_config.tie_word_embeddings:
self.lm_head = self.model.embed_tokens
else:
self.lm_head = ParallelLMHead(
self._backbone_config.vocab_size,
self._backbone_config.hidden_size,
prefix=f"{prefix}.lm_head" if prefix else "lm_head",
)
self.logits_processor = LogitsProcessor(self._backbone_config.vocab_size)
# Audio continuation token ID — resolved lazily from tokenizer.
# This is the <|audio|> token that serves as the LM-level continuation
# marker during audio generation (equivalent to v2's audio_token_id).
self._audio_continuation_id: int | None = None
self._eos_token_id: int | None = None
self._resolved_tokens = False
# Per-request audio state keyed by batch row index.
# Reset per slot via _slot_output_len tracking (same pattern as v2).
self._audio_state: dict[int, dict[str, Any]] = {}
self._slot_output_len: dict[int, int] = {}
self._last_logits_hidden: torch.Tensor | None = None
self._last_step_input_ids: torch.Tensor | None = None
self._last_step_query_start_loc: torch.Tensor | None = None
self._last_first_audio_after_start: torch.Tensor | None = None
self._last_audio_codes: torch.Tensor | None = None
self._postprocess_cursor: int = 0
# Pre-allocated decode-step audio feedback buffers (CUDA-graph safe).
# Populated by sample(), read by forward() via torch.where (no dict).
max_bs = 64 # safe upper bound; will grow if needed
self._decode_last_codes = torch.zeros(max_bs, self.num_codebooks, dtype=torch.long)
self._decode_has_codes = torch.zeros(max_bs, dtype=torch.bool)
# PrefixCache opt-outs (mirror qwen3_tts pattern):
# 1. The talker only consumes the last token's hidden state, so the
# runner can skip the per-step full hidden-state GPU->CPU merge
# that PrefixCache otherwise does.
# 2. Per-step ``codes.audio`` rows stay GPU-resident; defer the CPU
# write of the prefix-cache mm-output copy to request finish so
# the per-step bookkeeping does not block batching. Stage 0 can
# then set ``enable_prefix_caching: true`` without the regression
# observed in qwen3_tts (#3665).
self.requires_full_prefix_cached_hidden_states = False
self.deferred_prefix_cache_mm_keys = {"codes.audio"}
def _resolve_token_ids(self) -> None:
"""Resolve <|audio|> and eos token IDs.
Prefers config's pre-resolved IDs (from ``resolve_special_tokens()``),
falls back to loading the HF tokenizer directly.
"""
if self._resolved_tokens:
return
self._resolved_tokens = True
# Try config first (populated by resolve_special_tokens or from_pretrained)
cfg_audio = getattr(self.config, "audio_continuation_id", None)
cfg_eos = getattr(self.config, "eos_token_id", None)
if cfg_audio is not None:
self._audio_continuation_id = int(cfg_audio)
if cfg_eos is not None:
self._eos_token_id = int(cfg_eos)
if self._audio_continuation_id is not None:
logger.info(
"Resolved v3 token IDs from config: audio_continuation=%s, eos=%s",
self._audio_continuation_id,
self._eos_token_id,
)
return
# Fallback: load tokenizer directly
model_path = getattr(self.vllm_config.model_config, "model", None)
if model_path is None:
return
try:
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
vocab = dict(tokenizer.get_added_vocab())
if "<|audio|>" in vocab:
self._audio_continuation_id = vocab["<|audio|>"]
if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
self._eos_token_id = int(tokenizer.eos_token_id)
logger.info(
"Resolved v3 token IDs from tokenizer: audio_continuation=%s, eos=%s",
self._audio_continuation_id,
self._eos_token_id,
)
except Exception as exc:
logger.warning("Failed to resolve token IDs from tokenizer: %s", exc)
# ------------------------------------------------------------------ forward
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: Any | None = None,
inputs_embeds: torch.Tensor | None = None,
**kwargs: Any,
) -> torch.Tensor:
if inputs_embeds is None:
# Mask -100 placeholders to 0 before embedding. Use torch.where
# (no Python data-dependent branch) so this is CUDA-graph safe.
safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids)
hidden_states = self.model.embed_tokens(safe_ids)
else:
hidden_states = inputs_embeds
if input_ids is not None:
self._last_step_input_ids = input_ids
# Stash query_start_loc for _apply_audio_mode_bias
try:
from vllm.forward_context import get_forward_context
attn_metadata = get_forward_context().attn_metadata
if isinstance(attn_metadata, dict) and attn_metadata:
attn = next(iter(attn_metadata.values()))
else:
attn = attn_metadata
qsl = getattr(attn, "query_start_loc", None)
if isinstance(qsl, torch.Tensor):
self._last_step_query_start_loc = qsl.detach().clone()
else:
self._last_step_query_start_loc = None
except Exception:
self._last_step_query_start_loc = None
# Prefill-only operations: ref audio substitution and audio feedback
# require Python dict/list ops that break CUDA graph capture.
# Detect prefill (sequence length > batch size heuristic) vs decode.
is_prefill = input_ids is not None and inputs_embeds is None and int(input_ids.numel()) > 1
if is_prefill:
# Voice clone: replace -100 placeholder positions with ref audio embeddings
info_dicts = kwargs.get("model_intermediate_buffer")
if info_dicts is None:
info_dicts = kwargs.get("runtime_additional_information")
hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, info_dicts)
# Audio feedback at decode: replace continuation token embeddings
if input_ids is not None and inputs_embeds is None:
hidden_states = self._apply_audio_feedback(hidden_states, input_ids)
residual: torch.Tensor | None = None
for layer in self.model.layers:
hidden_states, residual = layer(positions, hidden_states, residual)
norm_out = self.model.norm(hidden_states, residual)
if isinstance(norm_out, tuple):
norm_out = norm_out[0]
return norm_out
def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor:
self._last_logits_hidden = hidden_states
return self.logits_processor(self.lm_head, hidden_states, sampling_metadata)
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
safe_ids = input_ids
if input_ids is not None and (input_ids < 0).any():
safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids)
text_embed = self.model.embed_tokens(safe_ids)
return self._apply_audio_feedback(text_embed, input_ids)
# ------------------------------------------------------------------ ref audio substitution
def _apply_ref_audio_substitution(
self,
hidden_states: torch.Tensor,
input_ids: torch.Tensor,
info_dicts: list[dict[str, Any]] | None,
) -> torch.Tensor:
"""Replace -100 placeholder positions with fused multi-codebook embeddings
of the delay-pattern-encoded reference audio codes.
Called at prefill to inject voice clone reference. ``info_dicts`` is a
list of per-request dicts from ``model_intermediate_buffer``, each
containing ``audio_input_ids`` ([T, N] delayed codes) and
``audio_input_ids_mask`` ([T] bool mask).
"""
if not info_dicts:
return hidden_states
PLACEHOLDER = -100
flat_ids = input_ids.reshape(-1)
placeholder_mask = flat_ids == PLACEHOLDER
if not placeholder_mask.any():
return hidden_states
# Use query_start_loc to map placeholders to per-request spans
q_start = self._last_step_query_start_loc
if not isinstance(q_start, torch.Tensor) or q_start.numel() < 2:
# Fallback: single-request batch
q_start_list = [0, int(flat_ids.numel())]
else:
q_start_list = q_start.detach().to("cpu").tolist()
new_hidden: torch.Tensor | None = None
num_requests = min(len(info_dicts), len(q_start_list) - 1)
for i in range(num_requests):
info = info_dicts[i]
if not isinstance(info, dict):
continue
codes = info.get("audio_input_ids")
mask = info.get("audio_input_ids_mask")
# Handle msgspec serialization (may be list-wrapped)
if isinstance(codes, list):
codes = codes[0] if codes else None
if isinstance(mask, list):
mask = mask[0] if mask else None
if not isinstance(codes, torch.Tensor):
continue
# codes shape: [T, num_codebooks] delayed reference codes
if codes.ndim == 3:
codes = codes[0]
if codes.ndim != 2:
continue
if isinstance(mask, torch.Tensor):
if mask.ndim == 2:
mask = mask[0]
codes = codes[mask.to(dtype=torch.bool)]
if codes.numel() == 0:
continue
# Find placeholder positions in this request's span
s = int(q_start_list[i])
e = int(q_start_list[i + 1])
if e - s <= 1:
continue # Decode step, skip
span_mask = placeholder_mask[s:e]
placeholders = span_mask.nonzero(as_tuple=True)[0]
n_codes = int(codes.shape[0])
if int(placeholders.numel()) < n_codes:
continue # Mismatch
# Embed delayed codes via fused multi-codebook embedding
target = placeholders[:n_codes] + s
codes_device = codes.to(device=hidden_states.device, dtype=torch.long)
embeds = self.multimodal_embedding(codes_device) # [n_codes, hidden]
if new_hidden is None:
new_hidden = hidden_states.clone()
flat_hidden = new_hidden.reshape(-1, new_hidden.shape[-1])
flat_hidden[target] = embeds.to(new_hidden.dtype)
return new_hidden if new_hidden is not None else hidden_states
# ------------------------------------------------------------------ audio feedback
def _apply_audio_feedback(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
"""Replace decode-step embeddings with audio feedback from pre-allocated buffers.
CUDA-graph safe: reads from pre-allocated _decode_last_codes and
_decode_has_codes tensors using torch.where (no Python dict lookup).
The buffers are populated by sample() after each step.
For decode steps (1 token per request): position i maps to row i.
For prefill: audio feedback is not needed (ref audio substitution
handles the prefill path separately).
"""
bs = hidden_states.shape[0]
# Ensure buffers are on the right device and large enough
if self._decode_last_codes.device != hidden_states.device:
self._decode_last_codes = self._decode_last_codes.to(hidden_states.device)
self._decode_has_codes = self._decode_has_codes.to(hidden_states.device)
if bs > self._decode_last_codes.shape[0]:
new_size = max(bs, self._decode_last_codes.shape[0] * 2)
self._decode_last_codes = torch.zeros(
new_size, self.num_codebooks, dtype=torch.long, device=hidden_states.device
)
self._decode_has_codes = torch.zeros(new_size, dtype=torch.bool, device=hidden_states.device)
# Compute audio embeddings from last_codes for ALL rows (graph-safe)
codes_slice = self._decode_last_codes[:bs] # [bs, N]
has_codes = self._decode_has_codes[:bs].unsqueeze(-1) # [bs, 1]
audio_embeds = self.multimodal_embedding(codes_slice) # [bs, D]
audio_embeds = audio_embeds.to(dtype=hidden_states.dtype)
# Select: where has_codes, use audio embed; else keep text embed
return torch.where(has_codes, audio_embeds, hidden_states)
# ------------------------------------------------------------------ sampling
def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> Any:
"""Model-owned sampler with delay-pattern audio dispatch.
Mirrors v2's pattern: bias LM logits to force audio continuation,
sample multi-codebook codes via the fused head, apply delay pattern,
and accumulate per-request state.
"""
self._resolve_token_ids()
sampler = getattr(self, "_stock_sampler", None)
if sampler is None:
from vllm.v1.sample.sampler import Sampler
sampler = Sampler()
self._stock_sampler = sampler
audio_id = self._audio_continuation_id
num_codebooks = self.num_codebooks
# Slot reuse detection: check output_token_ids length for each row.
# If a slot's output length decreased (fresh request reusing a finished
# slot), drop stale _audio_state for that row.
output_ids = getattr(sampling_metadata, "output_token_ids", None)
if output_ids is not None:
for i in range(len(output_ids)):
cur_len = len(output_ids[i]) if output_ids[i] else 0
prev_len = self._slot_output_len.get(i, 0)
if cur_len < prev_len:
# Fresh request reusing this slot — clear stale state
self._audio_state.pop(i, None)
self._slot_output_len[i] = 0
self._slot_output_len[i] = cur_len
# Bias LM logits for audio continuation
self._apply_audio_mode_bias(logits, sampling_metadata)
sampler_output = sampler(logits=logits, sampling_metadata=sampling_metadata)
hidden = self._last_logits_hidden
self._last_logits_hidden = None
if hidden is None or audio_id is None:
self._last_audio_codes = None
return sampler_output
sampled = getattr(sampler_output, "sampled_token_ids", None)
if sampled is None:
self._last_audio_codes = None
return sampler_output
sampled_flat = sampled.reshape(-1)
if int(sampled_flat.numel()) != int(hidden.shape[0]):
self._last_audio_codes = None
return sampler_output
is_audio = sampled_flat == audio_id
# Handle first-after-<|audio|> transition: initialize state but do NOT
# skip sampling. Unlike v2 (which seeds an all-BOC frame and skips the
# first audio step), v3/sglang samples codebook 0 on the very first
# step with delay_count=0 (codebooks 1-7 masked to BOC).
first_after_start = self._last_first_audio_after_start
self._last_first_audio_after_start = None
if isinstance(first_after_start, torch.Tensor) and first_after_start.numel() == is_audio.shape[0]:
first_after_start = first_after_start.to(is_audio.device)
init_rows = first_after_start & is_audio
if bool(init_rows.any()):
boc_frame = torch.full((num_codebooks,), BOC_ID, dtype=torch.long, device=hidden.device)
for bi in torch.nonzero(init_rows, as_tuple=False).reshape(-1).tolist():
bi = int(bi)
self._slot_output_len[bi] = 0
self._audio_state[bi] = {
"num_delay": 0,
"num_remaining_delays": None,
"audio_out_ids": None,
"last_codes": boc_frame.clone(),
"should_terminate": False,
}
# Update CUDA-graph-safe feedback buffers
self._decode_last_codes[bi] = boc_frame.to(self._decode_last_codes.device)
self._decode_has_codes[bi] = True
# Do NOT remove init_rows from is_audio — let them be sampled
if not bool(is_audio.any()):
self._last_audio_codes = None
return sampler_output
audio_row_indices = torch.nonzero(is_audio, as_tuple=False).reshape(-1).tolist()
# Per-codebook logits at audio positions
cb_logits = self._audio_codebook_logits(hidden, is_audio)
# Apply delay pattern masking BEFORE sampling
self._apply_delay_pattern_masking(cb_logits, audio_row_indices)
# Sample per-codebook
cb_logits_2d = cb_logits.reshape(-1, cb_logits.shape[-1])
codes_2d = self._sample_audio_codes(cb_logits_2d)
codes_flat = codes_2d.view(cb_logits.shape[0], cb_logits.shape[1]).to(torch.long)
# Update delay pattern state
eos_stream = EOC_ID
bos = BOC_ID
new_codes_flat: list[torch.Tensor] = []
for local_i, batch_i in enumerate(audio_row_indices):
state = self._audio_state.setdefault(
int(batch_i),
{
"num_delay": 0,
"num_remaining_delays": None,
"audio_out_ids": None,
"last_codes": torch.full((num_codebooks,), bos, dtype=torch.long, device=hidden.device),
"should_terminate": False,
},
)
num_delay: int = state["num_delay"]
num_remaining_delays: int | None = state["num_remaining_delays"]
this_codes = codes_flat[local_i].clone()
# Leading delay-pattern BOS pad
if num_delay + 1 < num_codebooks:
this_codes[num_delay + 1 :] = bos
num_delay += 1
# Trailing eos ramp-down
if num_remaining_delays is not None:
this_codes[: num_codebooks - num_remaining_delays] = eos_stream
num_remaining_delays -= 1
else:
eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1)
if eos_positions.numel() > 0:
last_eos_idx = int(eos_positions[-1].item())
this_codes[: last_eos_idx + 1] = eos_stream
num_remaining_delays = num_codebooks - last_eos_idx - 1
if num_remaining_delays is not None and num_remaining_delays <= 0:
# Ramp-down complete — terminate
state["num_delay"] = 0
state["num_remaining_delays"] = None
state["should_terminate"] = True
new_codes_flat.append(torch.full_like(this_codes, -1))
continue
state["num_delay"] = num_delay
state["num_remaining_delays"] = num_remaining_delays
state["last_codes"] = this_codes.clone()
# Update CUDA-graph-safe feedback buffers
self._decode_last_codes[int(batch_i)] = this_codes.to(self._decode_last_codes.device)
self._decode_has_codes[int(batch_i)] = True
if state["audio_out_ids"] is None:
state["audio_out_ids"] = this_codes.unsqueeze(-1).clone()
else:
state["audio_out_ids"] = torch.cat([state["audio_out_ids"], this_codes.unsqueeze(-1)], dim=-1)
new_codes_flat.append(this_codes)
# Build full codes tensor [batch_size, num_codebooks]
# -1 marks "no audio code at this position"
codes_full = torch.full(
(int(sampled_flat.numel()), num_codebooks),
-1,
dtype=torch.long,
device=hidden.device,
)
if new_codes_flat:
stacked = torch.stack(new_codes_flat, dim=0).to(hidden.device)
codes_full[is_audio] = stacked
self._last_audio_codes = codes_full
self._postprocess_cursor = 0
return sampler_output
# ------------------------------------------------------------------ postprocess
def postprocess(
self,
hidden_states_slice: torch.Tensor,
multimodal_outputs: Any = None,
**req_infos: Any,
) -> dict[str, Any]:
"""Publish per-request audio codes into model_intermediate_buffer.
Called once per request in batch order. Indexes _last_audio_codes
by a running cursor (one row per request per step).
"""
_ = multimodal_outputs
codes_full = self._last_audio_codes
if codes_full is None:
return {}
cursor = int(self._postprocess_cursor)
if cursor >= int(codes_full.shape[0]):
self._postprocess_cursor = 0
return {}
slice_codes = codes_full[cursor : cursor + 1]
self._postprocess_cursor = cursor + 1
# Drop placeholder rows (-1)
audio_rows = slice_codes[:, 0] >= 0
if not bool(audio_rows.any()):
return {}
new_codes = slice_codes[audio_rows].to(torch.int32)
return {"codes": {"audio": new_codes}}
# ------------------------------------------------------------------ helpers
def _audio_codebook_logits(self, hidden_states: torch.Tensor, audio_mask: torch.Tensor) -> torch.Tensor:
mask = audio_mask.reshape(-1).to(hidden_states.device)
hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1])
if not mask.any():
return torch.empty(
(0, self.num_codebooks, self.codebook_size),
dtype=hidden_states.dtype,
device=hidden_states.device,
)
return self.modality_head.generate(hidden_flat[mask])
def _apply_delay_pattern_masking(self, cb_logits: torch.Tensor, audio_row_indices: list[int]) -> None:
"""Mask per-codebook logits according to delay pattern state, in-place.
During delay phase: codebooks beyond delay_count only allow BOC.
During ramp-down: locked codebooks only allow EOC.
Normal generation: BOC disallowed; only cb0 allows EOC.
"""
bos_pre = BOC_ID
eos_pre = EOC_ID
num_codebooks = self.num_codebooks
for local_i, batch_i in enumerate(audio_row_indices):
state = self._audio_state.get(int(batch_i))
num_delay = int(state["num_delay"]) if state else 0
num_rem = state.get("num_remaining_delays") if state else None
if num_rem is not None:
lock_until = num_codebooks - int(num_rem)
for q in range(num_codebooks):
row = cb_logits[local_i, q]
if q < lock_until:
mask = torch.full_like(row, float("-inf"))
mask[eos_pre] = row[eos_pre]
cb_logits[local_i, q] = mask
else:
cb_logits[local_i, q, bos_pre] = float("-inf")
cb_logits[local_i, q, eos_pre] = float("-inf")
else:
for q in range(num_codebooks):
row = cb_logits[local_i, q]
if q > num_delay:
mask = torch.full_like(row, float("-inf"))
mask[bos_pre] = row[bos_pre]
cb_logits[local_i, q] = mask
else:
cb_logits[local_i, q, bos_pre] = float("-inf")
if q != 0:
cb_logits[local_i, q, eos_pre] = float("-inf")
def _sample_audio_codes(self, logits_2d: torch.Tensor) -> torch.Tensor:
"""Replicate upstream sampling: temperature → top-k → top-p → multinomial."""
x = logits_2d.float()
top_k = 50
top_p = 0.95
if 0 < top_k < x.shape[-1]:
kth = x.topk(top_k, dim=-1).values[..., -1:]
x = x.masked_fill(x < kth, float("-inf"))
if 0.0 < top_p < 1.0:
sorted_logits, sorted_idx = x.sort(dim=-1, descending=True)
cumprobs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
sorted_mask = cumprobs > top_p
sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
sorted_mask[..., 0] = False
mask = torch.zeros_like(x, dtype=torch.bool)
mask.scatter_(-1, sorted_idx, sorted_mask)
x = x.masked_fill(mask, float("-inf"))
# Detect all-masked rows BEFORE softmax (softmax of all-inf yields NaN).
has_finite = torch.isfinite(x).any(dim=-1)
all_masked = ~has_finite
if all_masked.any():
# For all-masked rows, fall back to argmax (picks least-negative).
fallback = x.argmax(dim=-1)
if has_finite.any():
probs = x[has_finite].softmax(dim=-1)
sampled = torch.multinomial(probs, num_samples=1).squeeze(-1)
result = fallback.clone()
result[has_finite] = sampled
return result
return fallback
probs = x.softmax(dim=-1)
return torch.multinomial(probs, num_samples=1).squeeze(-1)
def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> None:
"""Detect <|audio|> transition, force continuation, force eos at ramp-down.
Mirrors v2's _apply_audio_mode_bias: walks per-request to find the
previous token, and if it was <|audio|> or the continuation token,
forces the next emit to <|audio|>. On ramp-down completion, forces eos.
"""
if logits is None or logits.ndim != 2:
return
audio_id = self._audio_continuation_id
eos_id = self._eos_token_id
if audio_id is None:
return
num_rows = int(logits.shape[0])
prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None)
output_ids = getattr(sampling_metadata, "output_token_ids", None)
# Fallback prev-token source from stashed input_ids
stash_ids = self._last_step_input_ids
stash_tail: list[int] | None = None
if isinstance(stash_ids, torch.Tensor) and stash_ids.numel() > 0:
q_start = self._last_step_query_start_loc
if isinstance(q_start, torch.Tensor) and int(q_start.numel()) == num_rows + 1:
q_start_cpu = q_start.detach().to("cpu").tolist()
tail_idx = [max(0, int(q_start_cpu[i + 1]) - 1) for i in range(num_rows)]
flat_ids = stash_ids.detach().to("cpu").tolist()
stash_tail = [int(flat_ids[idx]) if idx < len(flat_ids) else -1 for idx in tail_idx]
elif int(stash_ids.numel()) >= num_rows:
stash_tail = stash_ids[-num_rows:].detach().to("cpu").tolist()
first_after_start = torch.zeros(num_rows, dtype=torch.bool, device=logits.device)
for i in range(num_rows):
prev: int | None = None
if output_ids is not None and i < len(output_ids):
hist = output_ids[i]
if hist:
prev = int(hist[-1])
if prev is None and prompt_ids is not None:
try:
p_i = prompt_ids[i]
if hasattr(p_i, "tolist"):
p_i = p_i.tolist()
if p_i:
prev = int(p_i[-1])
except (IndexError, TypeError):
prev = None
if prev is None and stash_tail is not None and i < len(stash_tail):
prev = int(stash_tail[i])
if prev is None:
continue
# Only bias if previous token was <|audio|> (the continuation token)
if prev != audio_id:
continue
# Check if this is the FIRST step after <|audio|> appears
# (i.e., transitioning from prompt to audio generation)
audio_state = self._audio_state.get(i)
if audio_state is None or audio_state.get("should_terminate"):
# No state yet, or stale state from a finished prior request
# reusing this slot — treat as first audio step
self._audio_state.pop(i, None)
# Clear CUDA-graph feedback buffers for this slot
self._decode_has_codes[i] = False
first_after_start[i] = True
# Check for ramp-down termination
should_terminate = bool(isinstance(audio_state, dict) and audio_state.get("should_terminate"))
if should_terminate and eos_id is not None and 0 <= eos_id < int(logits.shape[-1]):
row = logits[i]
mask = torch.full_like(row, float("-inf"))
mask[eos_id] = row[eos_id]
logits[i].copy_(mask)
audio_state["should_terminate"] = False
continue
# Force audio continuation token
row = logits[i]
mask = torch.full_like(row, float("-inf"))
if 0 <= audio_id < row.shape[-1]:
mask[audio_id] = row[audio_id]
logits[i].copy_(mask)
self._last_first_audio_after_start = first_after_start
# ------------------------------------------------------------------ omni output
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
return model_outputs
hidden = model_outputs
info_dicts = kwargs.get("model_intermediate_buffer")
if info_dicts is None:
info_dicts = kwargs.get("runtime_additional_information")
if info_dicts is None:
info_dicts = []
audio_codes_list: list[torch.Tensor] = []
any_nonempty = False
for info in info_dicts:
ac: torch.Tensor | None = None
if isinstance(info, dict):
codes_field = info.get("codes")
if isinstance(codes_field, dict):
ac = codes_field.get("audio")
else:
ac = info.get("audio_codes")
if isinstance(ac, torch.Tensor) and ac.numel() > 0:
audio_codes_list.append(ac)
any_nonempty = True
else:
audio_codes_list.append(torch.empty(0, dtype=torch.long))
if any_nonempty:
return OmniOutput(
text_hidden_states=hidden,
multimodal_outputs={"codes": {"audio": audio_codes_list}},
)
return OmniOutput(text_hidden_states=hidden, multimodal_outputs=None)
# ------------------------------------------------------------------ weight loading
# Per-layer suffixes from the actual V3 checkpoint (results/higgs_v3_checkpoint_analysis.txt)
_V3_LAYER_SUFFIXES = (
"input_layernorm.weight",
"mlp.down_proj.weight",
"mlp.gate_proj.weight",
"mlp.up_proj.weight",
"post_attention_layernorm.weight",
"self_attn.k_norm.weight",
"self_attn.k_proj.weight",
"self_attn.o_proj.weight",
"self_attn.q_norm.weight",
"self_attn.q_proj.weight",
"self_attn.v_proj.weight",
)
@classmethod
def _build_required_keys(cls, num_layers: int) -> set[str]:
"""Build the exact set of required V3 checkpoint keys."""
keys = {
"tied.embedding.text_embedding.weight",
"body.norm.weight",
f"{_MODALITY_EMBEDDING_PREFIX}weight",
}
for i in range(num_layers):
for suffix in cls._V3_LAYER_SUFFIXES:
keys.add(f"body.layers.{i}.{suffix}")
return keys
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
backbone_weights: list[tuple[str, torch.Tensor]] = []
loaded_params: set[str] = set()
own_params = dict(self.named_parameters())
seen_checkpoint_keys: set[str] = set()
for name, tensor in weights:
seen_checkpoint_keys.add(name)
mapped = self._map_weight_name(name)
if mapped is None:
continue
if mapped.startswith("model.") or mapped.startswith("lm_head."):
backbone_weights.append((mapped, tensor))
elif mapped in own_params:
param = own_params[mapped]
if param.shape != tensor.shape:
raise ValueError(
f"Shape mismatch for {mapped}: expected {param.shape}, "
f"got {tensor.shape} (checkpoint key: {name})"
)
param.data.copy_(tensor.to(param.dtype))
loaded_params.add(mapped)
if backbone_weights:
backbone_module = _BackboneWrapper(self.model, self.lm_head, self._backbone_config)
loaded = backbone_module.load_weights(iter(backbone_weights))
loaded_params.update(loaded)
# Resolve special token IDs from the tokenizer
model_path = getattr(self.vllm_config.model_config, "model", None)
if model_path:
self.config.resolve_special_tokens(model_path)
self._resolve_token_ids()
# Verify every required checkpoint key was seen.
num_layers = int(self._backbone_config.num_hidden_layers)
required = self._build_required_keys(num_layers)
missing = required - seen_checkpoint_keys
if missing:
raise RuntimeError(
f"HiggsAudioV3Talker: {len(missing)} required checkpoint keys missing: {sorted(missing)[:5]}..."
)
logger.info(
"HiggsAudioV3Talker: loaded %d params, modality_embedding=%s, tied=%s",
len(loaded_params),
tuple(self.multimodal_embedding.weight.shape),
self.tie_modality,
)
return loaded_params
def _map_weight_name(self, name: str) -> str | None:
if name.startswith(_CODEC_PREFIX):
return None
if name.startswith(_MODALITY_HEAD_PREFIX):
if self.tie_modality:
return None
return name.replace(_MODALITY_HEAD_PREFIX, "modality_head.")
if name.startswith(_MODALITY_EMBEDDING_PREFIX):
return name.replace(_MODALITY_EMBEDDING_PREFIX, "multimodal_embedding.")
for ckpt_prefix, model_prefix in _BACKBONE_PREFIX_MAP.items():
if name.startswith(ckpt_prefix):
return name.replace(ckpt_prefix, model_prefix, 1)
# Reject unexpected non-codec Higgs checkpoint keys
raise ValueError(
f"Unexpected checkpoint key with no known mapping: {name!r}. "
f"Known prefixes: {list(_BACKBONE_PREFIX_MAP.keys())}, "
f"{_MODALITY_EMBEDDING_PREFIX!r}, {_MODALITY_HEAD_PREFIX!r}, {_CODEC_PREFIX!r}"
)
class _BackboneWrapper(nn.Module):
"""Wrapper to use AutoWeightsLoader for Qwen3 backbone."""
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
def __init__(self, model, lm_head, config):
super().__init__()
self.model = model
self.lm_head = lm_head
self.config = config
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
from vllm.model_executor.models.utils import AutoWeightsLoader
skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None
loader = AutoWeightsLoader(self, skip_prefixes=skip)
return loader.load_weights(weights)
@@ -0,0 +1,135 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Prompt builder for higgs-audio v3 TTS.
Prompt formats:
Zero-shot:
<|tts|> <|text|> {text tokens} <|audio|>
Voice clone (no ref text):
<|tts|> <|ref_audio|> [-100]×N <|text|> {text tokens} <|audio|>
Voice clone (with ref text):
<|tts|> <|ref_text|> {ref text tokens} <|ref_audio|> [-100]×N <|text|> {text tokens} <|audio|>
-100 placeholders are replaced at prefill time with fused multi-codebook
embeddings of the delay-pattern-encoded reference audio codes.
"""
from __future__ import annotations
from typing import Any
import numpy as np
import torch
__all__ = [
"HiggsAudioV3TokenizerAdapter",
"AUDIO_PLACEHOLDER_ID",
"encode_reference_audio",
"apply_delay_pattern",
]
AUDIO_PLACEHOLDER_ID = -100
_REQUIRED_SPECIALS: tuple[str, ...] = (
"<|tts|>",
"<|text|>",
"<|audio|>",
)
# Codec stream special token IDs (inside the codebook vocab, not the LM vocab)
BOC_ID = 1024
EOC_ID = 1025
def apply_delay_pattern(codes_tn: torch.Tensor) -> torch.Tensor:
"""Apply MusicGen-style delay pattern to raw codes.
Input: [T, N] raw codes (T frames, N codebooks).
Output: [T + N - 1, N] delayed codes with BOC/EOC padding.
Codebook c is delayed by c positions: rows 0..c-1 get BOC,
rows c..c+T-1 get real codes, rows c+T..T+N-2 get EOC.
"""
T, N = codes_tn.shape
L = T + N - 1
delayed = torch.full((L, N), EOC_ID, dtype=codes_tn.dtype)
for c in range(N):
delayed[:c, c] = BOC_ID
delayed[c : c + T, c] = codes_tn[:, c]
return delayed
def encode_reference_audio(
wav: np.ndarray,
sr: int,
) -> torch.Tensor:
"""Encode a reference audio clip to codec codes [T, num_codebooks].
Uses the same HiggsAudioV2TokenizerModel as v2 (same codec).
Returns raw codes before delay pattern application.
"""
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
_encode_ref_audio_codes,
)
# _encode_ref_audio_codes returns [num_codebooks, T_raw]
codes_qt = _encode_ref_audio_codes(wav, sr)
# Transpose to [T_raw, num_codebooks]
return codes_qt.transpose(0, 1).contiguous()
class HiggsAudioV3TokenizerAdapter:
"""Wraps the HF tokenizer and builds TTS prompts."""
def __init__(self, tokenizer: Any) -> None:
self._tok = tokenizer
vocab = dict(tokenizer.get_added_vocab())
missing = [t for t in _REQUIRED_SPECIALS if t not in vocab]
if missing:
raise ValueError(f"Tokenizer is missing Higgs TTS v3 specials: {missing}")
self.tts_id: int = vocab["<|tts|>"]
self.text_id: int = vocab["<|text|>"]
self.audio_id: int = vocab["<|audio|>"]
self.ref_audio_id: int | None = vocab.get("<|ref_audio|>")
self.ref_text_id: int | None = vocab.get("<|ref_text|>")
@property
def tokenizer(self) -> Any:
return self._tok
def build_prompt(
self,
text: str,
*,
num_ref_tokens: int = 0,
reference_text: str | None = None,
) -> list[int]:
"""Build a TTS prompt.
Args:
text: Target text to synthesize.
num_ref_tokens: Number of delay-pattern reference code rows.
0 means zero-shot (no voice clone).
reference_text: Optional transcript of the reference audio.
"""
if not text or not text.strip():
raise ValueError("Text input must be non-empty for TTS")
if num_ref_tokens < 0:
raise ValueError(f"num_ref_tokens must be >= 0, got {num_ref_tokens}")
ids: list[int] = [self.tts_id]
# Voice clone: optional ref text + ref audio placeholders
if reference_text and num_ref_tokens > 0 and self.ref_text_id is not None:
ids.append(self.ref_text_id)
ids.extend(self._tok.encode(reference_text, add_special_tokens=False))
if num_ref_tokens > 0:
if self.ref_audio_id is None:
raise ValueError("Tokenizer missing <|ref_audio|> for voice clone")
ids.append(self.ref_audio_id)
ids.extend([AUDIO_PLACEHOLDER_ID] * num_ref_tokens)
ids.append(self.text_id)
ids.extend(self._tok.encode(text, add_special_tokens=False))
ids.append(self.audio_id)
return ids
@@ -0,0 +1,60 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""higgs-audio v3 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM).
Two delivery modes are wired here:
* ``sync_process_input_func`` runs when the deploy YAML has
``async_chunk: false``. The orchestrator collects the entire Stage-0 emit,
reverts the delay pattern once, and hands a single payload to Stage 1.
* ``async_chunk_process_next_stage_input_func`` runs when the deploy YAML
has ``async_chunk: true``. Stage 0 dispatches per AR step; the streaming
adapter buffers raw delay-pattern rows, slides a window with left context
and right holdback, and emits codec-ready frames per chunk. Stage 1
trims the overlap on both ends so the client sees a coherent PCM stream.
"""
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
StagePipelineConfig,
)
_PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v3"
HIGGS_AUDIO_V3_PIPELINE = PipelineConfig(
model_type="higgs_multimodal_qwen3",
model_arch="HiggsAudioV3TalkerForConditionalGeneration",
hf_architectures=("HiggsMultimodalQwen3ForConditionalGeneration",),
stages=(
StagePipelineConfig(
stage_id=0,
model_stage="higgs_audio_v3",
execution_type=StageExecutionType.LLM_AR,
input_sources=(),
owns_tokenizer=True,
engine_output_type="latent",
# stop_token_ids: the model-owned sampler forces eos at ramp-down
# completion. Safety stops from the actual V3 checkpoint:
# 151643 = <|endoftext|> (eos_token_id from config.json)
# 151671 = <|audio_end|> (audio generation end marker)
sampling_constraints={
"detokenize": False,
"stop_token_ids": [151643, 151671],
},
async_chunk_process_next_stage_input_func=f"{_PROC}.talker2code2wav_async_chunk",
),
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="HiggsAudioV3Code2WavForConditionalGeneration",
sync_process_input_func=f"{_PROC}.talker2code2wav",
sampling_constraints={"detokenize": True},
),
),
)
@@ -0,0 +1,111 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Bounded LRU cache for Higgs Audio v3 reference-audio codec encoding.
Voice-clone requests that re-send the same ref_audio are otherwise re-encoded
through the DAC encoder on every call. The cache keys by a stable content
hash (size + sample rate + head/mid/tail sentinel bytes) so the same audio
served via different URLs or repeated base64 payloads also hits warm.
Tensors are cloned on get/put so callers and the cache cannot alias storage.
"""
from __future__ import annotations
import hashlib
import threading
from collections import OrderedDict
from collections.abc import Callable
import numpy as np
import torch
__all__ = ["cached_encode_reference_audio", "cache_stats", "clear_cache"]
_DEFAULT_MAX_ENTRIES = 128
_SENTINEL_BYTES = 4096
_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
_cache_lock = threading.Lock()
_hits = 0
_misses = 0
def _content_hash(wav: np.ndarray, sr: int) -> str:
"""Stable content hash. Head/mid/tail sentinel keeps cost O(1)."""
raw = wav.tobytes()
n = len(raw)
h = hashlib.sha256()
h.update(n.to_bytes(8, "little"))
h.update(int(sr).to_bytes(4, "little"))
if n == 0:
return h.hexdigest()
if n <= _SENTINEL_BYTES * 3:
h.update(raw)
return h.hexdigest()
h.update(raw[:_SENTINEL_BYTES])
mid = (n - _SENTINEL_BYTES) // 2
h.update(raw[mid : mid + _SENTINEL_BYTES])
h.update(raw[-_SENTINEL_BYTES:])
return h.hexdigest()
def cached_encode_reference_audio(
wav: np.ndarray,
sr: int,
encode_fn: Callable[[np.ndarray, int], torch.Tensor],
max_entries: int = _DEFAULT_MAX_ENTRIES,
) -> torch.Tensor:
"""LRU-cached wrapper around ``encode_reference_audio``.
``encode_fn`` is the actual codec encoder; passed in so this module
stays decoupled from the tokenizer import (avoids circular imports
during stage-input-processor discovery).
Returns a CPU tensor clone — never the cache's interior storage.
"""
global _hits, _misses
key = _content_hash(wav, sr)
with _cache_lock:
cached = _cache.get(key)
if cached is not None:
_hits += 1
_cache.move_to_end(key)
return cached.clone()
# Miss: encode outside the lock so concurrent misses on different
# refs don't serialize on this single mutex.
codes = encode_fn(wav, sr)
codes_cpu = codes.detach().to("cpu").contiguous()
with _cache_lock:
if key in _cache:
# Concurrent miss winner already stored it; refresh LRU position
# and return that copy to avoid divergent tensors per request.
_cache.move_to_end(key)
_misses += 1 # we still paid the encode cost
return _cache[key].clone()
_cache[key] = codes_cpu
if len(_cache) > max_entries:
_cache.popitem(last=False)
_misses += 1
return codes_cpu.clone()
def cache_stats() -> dict:
"""Return (hits, misses, size, hit_rate) — for /metrics-style reporting."""
with _cache_lock:
total = _hits + _misses
return {
"hits": _hits,
"misses": _misses,
"size": len(_cache),
"hit_rate": (_hits / total) if total > 0 else 0.0,
}
def clear_cache() -> None:
"""Drop all cached entries. Used by tests."""
global _hits, _misses
with _cache_lock:
_cache.clear()
_hits = 0
_misses = 0
@@ -119,6 +119,22 @@ _OMNI_MODELS = {
"higgs_audio_v2_code2wav",
"HiggsAudioV2Code2WavForConditionalGeneration",
),
## higgs-audio v3
"HiggsMultimodalQwen3ForConditionalGeneration": (
"higgs_audio_v3",
"higgs_audio_v3_talker",
"HiggsAudioV3TalkerForConditionalGeneration",
),
"HiggsAudioV3TalkerForConditionalGeneration": (
"higgs_audio_v3",
"higgs_audio_v3_talker",
"HiggsAudioV3TalkerForConditionalGeneration",
),
"HiggsAudioV3Code2WavForConditionalGeneration": (
"higgs_audio_v3",
"higgs_audio_v3_code2wav",
"HiggsAudioV3Code2WavForConditionalGeneration",
),
## mimo_audio
"MiMoAudioModel": (
"mimo_audio",
@@ -0,0 +1,372 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage-input processor for higgs-audio v3: Talker -> Code2Wav.
Two adapters:
* ``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 talker step. It accumulates raw delay-pattern rows per request and
flushes a sliding-window chunk to Stage 1 once enough new frames are
available, holding back ``codec_right_holdback_frames`` trailing frames
so the codec has future context for the emit region's tail samples. Stage 1
trims ``left_context_size * hop`` samples off the front and
``right_holdback_size * hop`` off the end before emitting PCM.
Key difference from v2: BOC/EOC filtering happens AFTER delay pattern reversal
to avoid corrupting valid tail content.
"""
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__)
_NUM_CODEBOOKS = 8
_NUM_REAL_CODES = 1024 # codes in [0, 1023] are real
# Defaults for async_chunk streaming when the deploy YAML connector config
# does not override them. Mirrors the qwen3_tts canonical setting at 25 fps
# codec rate (25 frames = ~1.0 s of audio).
_DEFAULT_CODEC_CHUNK_FRAMES = 25
_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 25
_DEFAULT_CODEC_RIGHT_HOLDBACK_FRAMES = 4
def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor:
"""Reverse the MusicGen-style delay pattern.
Input shape: [num_codebooks, seq_len + num_codebooks - 1].
Output shape: [num_codebooks, seq_len].
For each codebook i, extract delayed[i, i : i + seq_len] to remove
the i leading BOC pads and Q-1-i trailing EOC 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 q != _NUM_CODEBOOKS:
raise ValueError(f"Expected exactly {_NUM_CODEBOOKS} codebook rows, got {q}. Input shape: [{q}, {t}]")
if t < q:
raise ValueError(f"Not enough frames to revert delay pattern: T={t} < Q={q}")
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 _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor:
"""Keep only frames where ALL codebook values are in [0, 1023].
Input shape: [num_codebooks, num_frames].
Called AFTER delay pattern reversal.
"""
if audio_codes_qt.numel() == 0:
return audio_codes_qt
# Transpose to [num_frames, num_codebooks] for per-frame filtering
frames = audio_codes_qt.t()
valid = (frames >= 0).all(dim=1) & (frames < _NUM_REAL_CODES).all(dim=1)
return frames[valid].t().contiguous()
def talker2code2wav(
source_outputs: list[Any],
prompt: Any = None,
_requires_multimodal_data: bool = False,
) -> list[Any]:
"""Sync: collect all talker codes, revert delay pattern, filter, pass to code2wav."""
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:
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:
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)}")
# Validate codebook count
if audio_codes.shape[1] != _NUM_CODEBOOKS:
raise ValueError(
f"Expected {_NUM_CODEBOOKS} codebooks per frame, "
f"got {audio_codes.shape[1]}. Audio codes shape: {tuple(audio_codes.shape)}"
)
# Transpose to [Q, T] for delay pattern reversal
codes_qt = audio_codes.transpose(0, 1).contiguous().cpu()
# Step 1: Revert delay pattern
codes_qt = _revert_delay_pattern(codes_qt)
# Step 2: Replace out-of-range codes (BOC=1024, EOC=1025, -1) with 0.
# Must use torch.where, NOT clamp: clamp(max=1023) turns 1025→1023
# which is a valid codec code and decodes to audio artifacts.
# Matches sglang's: torch.where(codes >= codec_vocab, 0, codes)
codes_qt = torch.where(
(codes_qt >= _NUM_REAL_CODES) | (codes_qt < 0),
torch.zeros_like(codes_qt),
codes_qt,
)
# Step 3: Trim the last frame. After de-delay, the final frame
# contains residual ramp-down codes (EOC→0 substituted) that
# decode to a brief noise artifact at the end of the audio.
if codes_qt.shape[-1] >= 2:
codes_qt = codes_qt[:, :-1]
if codes_qt.numel() == 0:
code2wav_inputs.append(
OmniTokensPrompt(
prompt_token_ids=[],
multi_modal_data=None,
mm_processor_kwargs=None,
additional_information=None,
)
)
continue
# 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
# ---------------------------------------------------------------------------
# Async-chunk streaming adapter (sliding window + holdback).
# ---------------------------------------------------------------------------
def _extract_last_step_row(pooling_output: OmniPayload) -> torch.Tensor | None:
"""Return the talker's emit row for this AR step as a flat ``[Q]`` tensor.
The talker postprocess emits ``codes.audio`` as ``[1, Q]`` (most often) or
``[N, Q]`` for multi-active stages. In the async-chunk path
``pooling_output`` is already request-scoped so the typical shape is
``[1, Q]``. Returns ``None`` if there is no row to append at this step.
Rows may carry BOC=1024 in trailing codebooks during the ramp-up phase
(AR steps 0..Q-2). Those are KEPT verbatim — the de-delay window
operation downstream slices them away by construction.
"""
codes = pooling_output.get("codes", {}) if isinstance(pooling_output, dict) else {}
audio_codes = codes.get("audio") if isinstance(codes, dict) else None
if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
return None
if audio_codes.ndim == 2:
row = audio_codes[-1]
elif audio_codes.ndim == 1:
row = audio_codes
else:
raise ValueError(f"unexpected audio_codes shape for higgs_audio_v3 async_chunk: {tuple(audio_codes.shape)}")
if row.numel() != _NUM_CODEBOOKS:
raise ValueError(
f"talker emit row has {row.numel()} codebooks; expected {_NUM_CODEBOOKS} for higgs_audio_v3 async_chunk."
)
return row.to(torch.long).reshape(-1)
def talker2code2wav_async_chunk(
transfer_manager: Any,
pooling_output: OmniPayload | None,
request: Any,
is_finished: bool = False,
) -> OmniPayloadStruct | None:
"""Streaming adapter: per-step accumulator + sliding-window de-delay flush.
Algorithm:
* Append each step's raw delay-pattern row to a per-request buffer.
* Track ``emitted_frames``: how many de-delayed frames' worth of audio
Stage 1 has already emitted for this request.
* With ``N`` accumulated rows we can recover ``N - Q + 1`` de-delayed
frames (the de-delay drops Q-1 leading BOC pads).
* Hold back the trailing ``H = codec_right_holdback_frames`` de-delayed
frames so the codec has future context for the emit region's last
samples. The emit boundary therefore lives at
``target_emit = max(0, de_delayed_total - H)`` for non-final calls
and ``de_delayed_total`` for the final flush.
* When ``target_emit - emitted_frames >= chunk_size`` (or the request
finishes), build a sliding window spanning
``[emitted_frames - L, target_emit + H]`` de-delayed frames; that
corresponds to AR rows
``[emitted_frames - L, target_emit + H + Q - 1]``. Apply de-delay,
replace out-of-range codes with 0, and emit codebook-major flat.
* Stage 1 honors ``meta.left_context_size = L`` (trim front) and
``meta.right_holdback_size = H`` (trim end). Net new audio per call
is ``(target_emit - emitted_frames) * hop_length`` samples.
"""
request_id = request.external_req_id
finished = bool(is_finished or request.is_finished())
# Per-request rolling state. ``code_prompt_token_ids`` is the
# framework-provided buffer; ``higgs_v3_emitted_frames`` tracks how many
# de-delayed frames we have already flushed for this request.
emitted_frames = getattr(transfer_manager, "higgs_v3_emitted_frames", None)
if emitted_frames is None:
emitted_frames = {}
transfer_manager.higgs_v3_emitted_frames = emitted_frames
if isinstance(pooling_output, dict):
row = _extract_last_step_row(pooling_output)
if row is not None:
transfer_manager.code_prompt_token_ids[request_id].append(row.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))
right_holdback_size_config = int(cfg.get("codec_right_holdback_frames", _DEFAULT_CODEC_RIGHT_HOLDBACK_FRAMES))
configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0)
if (
chunk_size <= 0
or left_context_size_config < 0
or right_holdback_size_config < 0
or configured_initial_chunk_size < 0
):
raise ValueError(
f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
f"codec_left_context_frames={left_context_size_config}, "
f"codec_right_holdback_frames={right_holdback_size_config}, "
f"initial_codec_chunk_frames={configured_initial_chunk_size}"
)
n_rows = len(transfer_manager.code_prompt_token_ids[request_id])
de_delayed_total = max(0, n_rows - (_NUM_CODEBOOKS - 1))
emitted = int(emitted_frames.get(request_id, 0))
# Hold back the trailing H frames from the emit boundary unless the
# request is finishing - the final chunk needs no future context.
H = 0 if finished else right_holdback_size_config
target_emit = max(0, de_delayed_total - H)
pending = target_emit - emitted
# First-chunk fast path for TTFA.
if emitted == 0 and 0 < configured_initial_chunk_size < chunk_size:
target_chunk = configured_initial_chunk_size
else:
target_chunk = chunk_size
if not finished and pending < target_chunk:
return None
if pending <= 0:
if finished:
emitted_frames.pop(request_id, None)
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
)
return None
actual_chunk = pending
desired_left_context = min(left_context_size_config, emitted)
# Window in AR-row space: span the de-delayed frames
# [emitted - L, target_emit + H) which require AR rows
# [emitted - L, target_emit + H + Q - 1).
window_row_start = emitted - desired_left_context
window_row_end_exclusive = target_emit + H + (_NUM_CODEBOOKS - 1)
if window_row_end_exclusive > n_rows:
# Defensive: pending >= 0 was checked above; this should be unreachable
# for non-final calls. On the final call we cap to n_rows.
if not finished:
return None
window_row_end_exclusive = n_rows
window_rows = transfer_manager.code_prompt_token_ids[request_id][window_row_start:window_row_end_exclusive]
if len(window_rows) < _NUM_CODEBOOKS:
if finished:
emitted_frames.pop(request_id, None)
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
)
return None
codes_qt = torch.tensor(window_rows, dtype=torch.long).t().contiguous()
de_delayed = _revert_delay_pattern(codes_qt)
# Replace BOC=1024/EOC=1025 (and any negative pads) with 0; matches the
# sync-path substitution. Clamp would turn 1025 into 1023 which is a
# VALID codec code and decodes to audible artifacts.
de_delayed = torch.where(
(de_delayed >= _NUM_REAL_CODES) | (de_delayed < 0),
torch.zeros_like(de_delayed),
de_delayed,
)
# On the FINAL chunk only: trim the trailing residual frame (the last
# de-delayed frame still carries EOC-substituted codes from ramp-down
# which decode to a brief noise artifact). Mirrors the sync-path trim.
if finished and window_row_end_exclusive == n_rows and de_delayed.shape[-1] >= 2:
de_delayed = de_delayed[:, :-1]
actual_chunk = max(actual_chunk - 1, 0)
codec_codes = de_delayed.reshape(-1)
left_context_emitted = desired_left_context
right_holdback_emitted = H
emitted_frames[request_id] = emitted + actual_chunk
if finished:
emitted_frames.pop(request_id, None)
meta = MetaStruct(
left_context_size=left_context_emitted,
right_holdback_size=right_holdback_emitted,
finished=torch.tensor(finished, dtype=torch.bool),
)
return OmniPayloadStruct(
codes=CodesStruct(audio=codec_codes),
meta=meta,
)