Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3bda7ebd8b | ||
|
|
969c9b5a61 | ||
|
|
e592d29f4c | ||
|
|
82cce633a3 | ||
|
|
fce0f51089 | ||
|
|
82e74ce972 |
No files matched your search
@@ -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:
|
||||
- |
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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|>
|
||||
@@ -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",
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
Reference in new issue
Block a user