From 8f45e68bcd4b5beb44c45ddbb576e59ea9c7a741 Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Mon, 25 May 2026 23:16:07 +0800 Subject: [PATCH] [Higgs-Audio] bosonai/higgs-audio-v2-generation-3B-base TTS model support (#3762) Signed-off-by: Yuekai Zhang Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .buildkite/test-merge.yml | 25 + .readthedocs.yml | 5 + .../offline_inference/text_to_speech.md | 44 + .../examples/online_serving/text_to_speech.md | 47 + .../text_to_speech/higgs_audio_v2/README.md | 38 + .../text_to_speech/higgs_audio_v2/end2end.py | 205 ++ .../text_to_speech/higgs_audio_v2/README.md | 63 + .../higgs_audio_v2/batch_speech_client.py | 122 ++ .../higgs_audio_v2/run_server.sh | 37 + setup.py | 6 + .../offline_inference/test_higgs_audio_v2.py | 231 +++ .../e2e/online_serving/test_higgs_audio_v2.py | 299 +++ vllm_omni/config/pipeline_registry.py | 4 + vllm_omni/deploy/higgs_audio_v2.yaml | 83 + .../entrypoints/openai/protocol/audio.py | 44 + .../entrypoints/openai/serving_speech.py | 189 +- .../models/higgs_audio_v2/__init__.py | 9 + .../configuration_higgs_audio_v2.py | 169 ++ .../higgs_audio_v2/higgs_audio_decoder.py | 471 +++++ .../higgs_audio_v2/higgs_audio_v2_code2wav.py | 413 ++++ .../higgs_audio_v2/higgs_audio_v2_talker.py | 1684 +++++++++++++++++ .../higgs_audio_v2_tokenizer.py | 332 ++++ .../models/higgs_audio_v2/pipeline.py | 49 + vllm_omni/model_executor/models/registry.py | 16 + .../stage_input_processors/higgs_audio_v2.py | 252 +++ 25 files changed, 4835 insertions(+), 2 deletions(-) create mode 100644 examples/offline_inference/text_to_speech/higgs_audio_v2/README.md create mode 100644 examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py create mode 100644 examples/online_serving/text_to_speech/higgs_audio_v2/README.md create mode 100644 examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py create mode 100755 examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh create mode 100644 tests/e2e/offline_inference/test_higgs_audio_v2.py create mode 100644 tests/e2e/online_serving/test_higgs_audio_v2.py create mode 100644 vllm_omni/deploy/higgs_audio_v2.yaml create mode 100644 vllm_omni/model_executor/models/higgs_audio_v2/__init__.py create mode 100644 vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py create mode 100644 vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py create mode 100644 vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py create mode 100644 vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py create mode 100644 vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py create mode 100644 vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py create mode 100644 vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 5b8d7e104..318b872b5 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -389,6 +389,31 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Omni · Higgs-Audio V2 Test" + timeout_in_minutes: 30 + commands: + - | + timeout 30m bash -c " + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_USE_DEEP_GEMM=0 + export VLLM_MOE_USE_DEEP_GEMM=0 + pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model' + " + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Diffusion · Qwen Image Edit Test" timeout_in_minutes: 20 commands: diff --git a/.readthedocs.yml b/.readthedocs.yml index 5a06c6634..2b5b81bbc 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -10,6 +10,11 @@ build: jobs: post_checkout: - git fetch --unshallow || true + pre_install: + # Pre-install CPU torch so the subsequent `pip install .[docs]` does not + # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's + # 1 GB memory cap and time out the mkdocs build. + - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch mkdocs: configuration: mkdocs.yml diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index 7937b1180..ad7777014 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -23,6 +23,7 @@ list of supported architectures across all modalities, see | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz | +| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz | ## Common Quick Start @@ -333,6 +334,46 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`. - Run `--help` for the full argument surface. +--- + +## higgs-audio v2 + +2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`. + +### Prerequisites + +Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded. + +```bash +pip install -U "transformers>=5.3.0" +``` + +### Quick start (plain TTS) + +```bash +python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \ + --texts "Hello world." "The quick brown fox jumps over the lazy dog." \ + --output-dir results/wavs +``` + +### Voice cloning + +Pass both `--ref-audio` and `--ref-text` together: + +```bash +python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \ + --texts "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Exact transcript spoken in reference.wav." \ + --output-dir results/wavs +``` + +### Notes + +- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`. +- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality. +- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`. + ## Example materials ??? abstract "cosyvoice3/end2end.py" @@ -343,6 +384,9 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T ``````py --8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py" `````` +??? abstract "higgs_audio_v2/end2end.py" + ``````py + --8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py" ??? abstract "glm_tts/end2end.py" ``````py --8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py" diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 37bfde7ad..3662dc107 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -19,6 +19,7 @@ For the full list of supported architectures across all modalities, see |---|---|---|---|---|---| | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | +| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | @@ -193,6 +194,45 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server --- +## higgs-audio v2 + +2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client. + +### Prerequisites + +Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec): + +```bash +pip install -U "transformers>=5.3.0" +``` + +### Launch +```bash +GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh +``` +Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`. + +### Sending requests +```bash +# Plain TTS +python higgs_audio_v2/batch_speech_client.py \ + --base-url http://localhost:8094 \ + --output-dir /tmp/higgs_out \ + --prompts "Hello world." "The quick brown fox jumps over the lazy dog." + +# Voice cloning — pass a reference clip and its transcript together +python higgs_audio_v2/batch_speech_client.py \ + --base-url http://localhost:8094 \ + --output-dir /tmp/higgs_clone \ + --ref-audio /path/to/reference.wav \ + --ref-text "Exact transcript spoken in reference.wav." \ + --prompts "Hello, this is a cloned voice." +``` + +### Notes +- Output: 24 kHz mono. +- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality. +- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`. ## GLM-TTS 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`. @@ -459,6 +499,13 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts ``````py --8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py" `````` +??? abstract "higgs_audio_v2/batch_speech_client.py" + ``````py + --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py" + `````` +??? abstract "higgs_audio_v2/run_server.sh" + ``````sh + --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh" ??? abstract "glm_tts/gradio_demo.py" ``````py --8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py" diff --git a/examples/offline_inference/text_to_speech/higgs_audio_v2/README.md b/examples/offline_inference/text_to_speech/higgs_audio_v2/README.md new file mode 100644 index 000000000..95a200d88 --- /dev/null +++ b/examples/offline_inference/text_to_speech/higgs_audio_v2/README.md @@ -0,0 +1,38 @@ +# higgs-audio v2 — offline example + +Drives Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) for `bosonai/higgs-audio-v2-generation-3B-base` end-to-end through the vLLM-Omni engine and writes a 24 kHz mono WAV per prompt. + +## Prerequisites + +Voice clone needs `transformers>=5.3.0` — vllm-omni loads the audio codec via HF's `HiggsAudioV2TokenizerModel`, instantiated from the `k2-fsa/OmniVoice/audio_tokenizer/` subdirectory (only that ~806 MB subdir is downloaded). The boson-ai standalone tokenizer repo's `model.safetensors` is actually a copy of the 3B talker LM, so HF can't load it directly; the k2 bundle ships the same codec weights repackaged with HF-compatible key naming. + +```bash +pip install -U "transformers>=5.3.0" +``` + +## Quick start + +Plain TTS: +```bash +python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \ + --texts "Hello world." "The quick brown fox jumps over the lazy dog." \ + --output-dir results/wavs +``` + +## Voice cloning + +Pass both `--ref-audio` and `--ref-text` together: +```bash +python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \ + --texts "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Exact transcript spoken in reference.wav." \ + --output-dir results/wavs +``` + +## Notes + +- Output: 24 kHz mono WAV. +- Deploy config: `vllm_omni/deploy/higgs_audio_v2.yaml` (auto-loaded by HF `model_type`). +- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality. +- For online serving, see [`examples/online_serving/text_to_speech/higgs_audio_v2/`](../../../online_serving/text_to_speech/higgs_audio_v2/). diff --git a/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py b/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py new file mode 100644 index 000000000..62a346621 --- /dev/null +++ b/examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline higgs-audio v2 inference example. + +Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end +through the vLLM-Omni engine without going through the HTTP server, and +saves a 24 kHz mono WAV per prompt. + +Example: + + python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\ + --texts "Hello world." \\ + "The quick brown fox jumps over the lazy dog." \\ + --output-dir results/wavs +""" + +from __future__ import annotations + +import os + +# DeepGEMM FP8 kernels require an optional backend that may not be installed. +# Disable the warmup before importing vLLM so engine startup falls back to the +# regular gemm path. Users with deep_gemm installed can override these. +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") +os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0") + +import time +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import Omni +from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults + +SAMPLE_RATE = 24_000 +DEFAULT_TEXTS = ( + "Hello world.", + "The quick brown fox jumps over the lazy dog.", +) + + +def parse_args(): + parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference") + parser.add_argument( + "--model", + type=str, + default="bosonai/higgs-audio-v2-generation-3B-base", + help="Stage-0 talker model id or path.", + ) + parser.add_argument( + "--texts", + type=str, + nargs="+", + default=list(DEFAULT_TEXTS), + help="One or more plain-text prompts to synthesize.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="results/wavs", + help="Directory to write per-prompt WAV files.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path. Auto-loads " + "vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=500, + help="Cap on Stage-0 codec frames per request.", + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcript of the reference clip. Required when --ref-audio is set.", + ) + nullify_stage_engine_defaults(parser) + return parser.parse_args() + + +def _slugify(text: str) -> str: + import re + + slug = re.sub(r"\s+", "_", text.strip().lower()) + slug = re.sub(r"[^a-z0-9_]+", "", slug) + return slug[:48] or "prompt" + + +def _extract_pcm(multimodal_output: dict) -> torch.Tensor: + """Pull the final concatenated PCM tensor out of a request's multimodal_output.""" + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}") + if isinstance(audio, list): + valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] + if not valid: + raise ValueError("audio list is empty") + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] + return torch.as_tensor(audio).float().cpu().reshape(-1) + + +def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray: + arr = pcm.numpy() + if arr.dtype.kind == "f": + arr = np.clip(arr, -1.0, 1.0) + arr = (arr * 32767.0).astype(np.int16) + else: + arr = arr.astype(np.int16) + return arr + + +def main(): + args = parse_args() + if (args.ref_audio is None) != (args.ref_text is None): + raise SystemExit("--ref-audio and --ref-text must be supplied together") + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + engine = Omni(model=args.model, deploy_config=args.deploy_config) + + # Build prompt_token_ids using the same path serving_speech.py takes online. + from transformers import AutoProcessor + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + build_plain_text_prompt, + build_voice_clone_prompt, + input_ids_to_python_list, + ) + + processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True) + + # Voice-clone path: load the reference clip once. The HF processor will + # encode it via the bundled HiggsAudioV2TokenizerModel each time we build + # a prompt below. This is cheap on CPU for a few-second clip. + ref_wav: np.ndarray | None = None + ref_sr: int | None = None + if args.ref_audio is not None: + ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False) + if ref_wav.ndim == 2: + ref_wav = ref_wav.mean(axis=1) + + print(f"Model : {args.model}") + print(f"Prompts : {len(args.texts)}") + print(f"Output dir : {output_dir}") + print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}") + + # Run one prompt at a time. The Stage-0 talker's per-slot audio state is + # request-scoped; submitting multiple prompts in the same engine.generate() + # call would batch them in the AR runner and exercise a code path that is + # not validated for this model yet. + total_elapsed = 0.0 + total_dur = 0.0 + for text in args.texts: + if ref_wav is not None: + out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text) + prompt = { + "prompt_token_ids": out["prompt_token_ids"], + # Bare tensors (NOT list-wrapped): the msgspec serializer in + # vllm_omni.data_entry_keys routes torch.Tensor → tensor_data + # and list[Tensor] → list_data (which silently strips tensors). + "additional_information": { + "audio_input_ids": out["audio_input_ids"], + "audio_input_ids_mask": out["audio_input_ids_mask"], + }, + } + else: + inputs = build_plain_text_prompt(processor, text) + prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)} + t_start = time.perf_counter() + outputs = engine.generate([prompt]) + elapsed = time.perf_counter() - t_start + total_elapsed += elapsed + + mm = outputs[0].outputs[0].multimodal_output + pcm = _extract_pcm(mm) + slug = _slugify(text) + out_path = output_dir / f"{slug}.wav" + sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16") + dur = pcm.numel() / SAMPLE_RATE + total_dur += dur + print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}") + + rtf = total_elapsed / total_dur if total_dur > 0 else float("inf") + print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}") + + +if __name__ == "__main__": + os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + main() diff --git a/examples/online_serving/text_to_speech/higgs_audio_v2/README.md b/examples/online_serving/text_to_speech/higgs_audio_v2/README.md new file mode 100644 index 000000000..7b5e272df --- /dev/null +++ b/examples/online_serving/text_to_speech/higgs_audio_v2/README.md @@ -0,0 +1,63 @@ +# higgs-audio v2 online example + +This directory contains the online-serving entry points for boson-ai's higgs-audio v2 as integrated by vllm-omni: a 2-stage TTS pipeline (Llama-3.2-3B talker with DualFFN audio expert + HiggsAudio codec decoder) emitting 24 kHz mono speech. + +## Prerequisites + +Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec). Only that ~806 MB subdir is downloaded. + +```bash +pip install -U "transformers>=5.3.0" +``` + +## Files + +- `run_server.sh` — launch the vllm-omni server with the bundled `vllm_omni/deploy/higgs_audio_v2.yaml` deploy config. +- `batch_speech_client.py` — send a list of prompts to `/v1/audio/speech` and save the returned WAV / PCM bytes to a directory; optionally passes `--ref-audio` + `--ref-text` for shallow voice clone. + +## Launching the server + +```bash +GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh +``` + +Environment overrides: + +- `MODEL` — HF id of the talker (default `bosonai/higgs-audio-v2-generation-3B-base`). +- `PORT` — server port (default `8094`). +- `GPUS` — `CUDA_VISIBLE_DEVICES` value (default `6,7`). +- `GPU_UTIL` — `--gpu-memory-utilization` (default `0.4`). + +The script also exports `VLLM_USE_DEEP_GEMM=0` / `VLLM_MOE_USE_DEEP_GEMM=0` so the example works on images without the optional `deep_gemm` backend. + +The deploy YAML ships with `async_chunk: false` and `codec_streaming: true`, i.e. Stage 0 finishes its codec frames before Stage 1 starts decoding, and Stage 1 streams WAV/PCM bytes to the client chunk-by-chunk. + +## Driving the server + +Plain TTS: + +```bash +python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \ + --base-url http://localhost:8094 \ + --model bosonai/higgs-audio-v2-generation-3B-base \ + --output-dir /tmp/higgs_audio_v2_batch \ + --prompts "Hello world." \ + "The quick brown fox jumps over the lazy dog." +``` + +Voice clone — pass a reference clip and its transcript (both required together): + +```bash +python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \ + --base-url http://localhost:8094 \ + --model bosonai/higgs-audio-v2-generation-3B-base \ + --output-dir /tmp/higgs_audio_v2_clone \ + --ref-audio /path/to/reference.wav \ + --ref-text "Exact transcript spoken in reference.wav." \ + --prompts "Hello, this is a cloned voice." +``` + +## Notes + +- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality. +- Out of scope (rejected with explicit 4xx by the request validator): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`. diff --git a/examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py b/examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py new file mode 100644 index 000000000..1e227b14d --- /dev/null +++ b/examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Batch client for the higgs-audio v2 online server. + +Sends a fixed list of prompts to ``/v1/audio/speech`` and saves the returned +WAV files (or raw PCM bytes when ``--format pcm``) into ``--output-dir``. + +Usage (plain text -> speech): + + python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \ + --base-url http://localhost:8094 \ + --output-dir /tmp/higgs_audio_v2_batch \ + --prompts "Hello world." "The quick brown fox jumps over the lazy dog." + +Usage (shallow voice clone — pass a reference clip + its transcript): + + python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \ + --base-url http://localhost:8094 \ + --output-dir /tmp/higgs_audio_v2_clone \ + --ref-audio path/to/reference.wav \ + --ref-text "the transcript of the reference clip" \ + --prompts "Hello world." +""" + +from __future__ import annotations + +import argparse +import base64 +import sys +from pathlib import Path + +DEFAULT_PROMPTS = ( + "Hello world.", + "The quick brown fox jumps over the lazy dog.", + "It was the night before my birthday.", + "Innovation distinguishes between a leader and a follower.", +) + + +def _slug(text: str) -> str: + import re + + s = re.sub(r"\s+", "_", text.strip().lower()) + return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--base-url", default="http://localhost:8094") + parser.add_argument("--model", default="higgs_audio_v2") + parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS)) + parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_audio_v2_batch")) + parser.add_argument("--format", choices=("wav", "pcm"), default="wav") + parser.add_argument("--max-new-tokens", type=int, default=300) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--timeout-s", type=float, default=120.0) + parser.add_argument( + "--ref-audio", + type=Path, + default=None, + help="Reference clip for voice clone (path to a WAV file). Must be paired with --ref-text.", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcript of the reference clip. Required when --ref-audio is set.", + ) + args = parser.parse_args() + + if (args.ref_audio is None) != (args.ref_text is None): + print("--ref-audio and --ref-text must be supplied together", file=sys.stderr) + return 2 + + ref_audio_data_url: str | None = None + if args.ref_audio is not None: + if not args.ref_audio.exists(): + print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr) + return 2 + mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg" + ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii") + ref_audio_data_url = f"data:{mime};base64,{ref_b64}" + + try: + import httpx + except ImportError: + print( + "this client needs `httpx`. Install with `pip install httpx`.", + file=sys.stderr, + ) + return 2 + + args.output_dir.mkdir(parents=True, exist_ok=True) + url = args.base_url.rstrip("/") + "/v1/audio/speech" + failures = 0 + with httpx.Client(timeout=args.timeout_s) as client: + for prompt in args.prompts: + payload = { + "model": args.model, + "input": prompt, + "response_format": args.format, + "max_new_tokens": args.max_new_tokens, + "seed": args.seed, + } + if ref_audio_data_url is not None: + payload["ref_audio"] = ref_audio_data_url + payload["ref_text"] = args.ref_text + resp = client.post(url, json=payload) + if resp.status_code != 200: + print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr) + failures += 1 + continue + suffix = ".wav" if args.format == "wav" else ".pcm" + out = args.output_dir / f"{_slug(prompt)}{suffix}" + out.write_bytes(resp.content) + print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)") + + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh b/examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh new file mode 100755 index 000000000..12db89558 --- /dev/null +++ b/examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Launch vLLM-Omni server for higgs-audio v2. +# +# v1 scope: plain text -> 24 kHz speech only. Voice cloning, multi-speaker, +# ChatML rich content, and language overrides are rejected by the validator +# with explicit 4xx (see vllm_omni/entrypoints/openai/serving_speech.py). +# +# Usage: +# ./run_server.sh # default port 8094, GPUs 6 and 7 +# PORT=8095 GPUS=6,7 ./run_server.sh +# MODEL=bosonai/higgs-audio-v2-generation-3B-base ./run_server.sh + +set -e + +MODEL="${MODEL:-bosonai/higgs-audio-v2-generation-3B-base}" +PORT="${PORT:-8094}" +GPUS="${GPUS:-6,7}" +GPU_UTIL="${GPU_UTIL:-0.4}" + +echo "Starting higgs-audio v2 server" +echo " MODEL=$MODEL" +echo " PORT=$PORT" +echo " CUDA_VISIBLE_DEVICES=$GPUS" + +# DeepGEMM FP8 kernels are optional and trip warmup on builds without +# the deep_gemm backend; disable them so the example works out of the box. +# Users with deep_gemm installed can re-enable via the same env vars. +CUDA_VISIBLE_DEVICES="$GPUS" \ +VLLM_USE_DEEP_GEMM=0 \ +VLLM_MOE_USE_DEEP_GEMM=0 \ +vllm-omni serve "$MODEL" \ + --deploy-config vllm_omni/deploy/higgs_audio_v2.yaml \ + --host 0.0.0.0 \ + --port "$PORT" \ + --gpu-memory-utilization "$GPU_UTIL" \ + --trust-remote-code \ + --omni diff --git a/setup.py b/setup.py index 057212d67..403d8a416 100644 --- a/setup.py +++ b/setup.py @@ -62,6 +62,12 @@ def detect_target_device() -> str: else: print(f"Warning: Invalid VLLM_OMNI_TARGET_DEVICE '{target_device}', falling back to auto-detection") + # Priority 1.5: ReadTheDocs builds use CPU requirements (no GPU, 1 GB RAM cap). + # Avoids pulling ~2 GB of CUDA libs that push the mkdocs build into swap. + if os.environ.get("READTHEDOCS"): + print("Detected READTHEDOCS=1; using cpu requirements for docs build") + return "cpu" + # Priority 2: Torch backend detection # This is a code path for when user is using # --no-build-isolation flag diff --git a/tests/e2e/offline_inference/test_higgs_audio_v2.py b/tests/e2e/offline_inference/test_higgs_audio_v2.py new file mode 100644 index 000000000..fdeccc924 --- /dev/null +++ b/tests/e2e/offline_inference/test_higgs_audio_v2.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +End-to-end offline inference tests for higgs-audio v2. + +Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py: +build Stage-0 prompt_token_ids with the model's own processor + plain-text +chat template, drive the engine via ``Omni.generate``, and inspect the +Stage-1 PCM payload directly. The OmniRunnerHandler.send_audio_speech_request +helper applies a Qwen-style ChatML wrap that is not valid for higgs, so this +test bypasses it (same approach as test_voxtral_tts.py). + +v1 scope is plain text -> 24 kHz speech; voice cloning, multi-speaker prompts, +language overrides, and task_type selection are out of scope. +""" + +from __future__ import annotations + +import os + +# Match the example script: keep DeepGEMM FP8 off so engine init does not +# require the optional ``deep_gemm`` backend. +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") +os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0") + +import numpy as np +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "bosonai/higgs-audio-v2-generation-3B-base" +STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml") +SAMPLE_RATE = 24_000 +# A short sentence at 24 kHz must produce at least ~0.5 s of audio. +MIN_AUDIO_SAMPLES = 12_000 +TEST_TEXT = "Hello world." + +# Reuse the qwen3_tts ref clip + transcript for voice clone tests (see +# tests/e2e/online_serving/test_qwen3_tts_base.py:27-28 for the source). +_REF_AUDIO_ASSET = "qwen3_tts/clone_2.wav" +_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." + +# (model, stage_config_path, extra_omni_kwargs) — same shape as test_voxtral_tts.py. +# Stage configs already pin enforce_eager=true per stage, so no Omni-level override. +_OMNI_RUNNER_PARAM = (MODEL, STAGE_CONFIG, {"trust_remote_code": True}) + + +def _compose_request(model_name: str, text: str) -> dict: + """Build a higgs prompt dict via the upstream processor + plain-text template. + + The serving layer's ``_build_higgs_audio_v2_params`` does the same thing + (see vllm_omni/entrypoints/openai/serving_speech.py:1520) — keeping the + offline and online prompt-construction paths identical is what guarantees + input-token parity with the HF reference. + """ + from transformers import AutoProcessor + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + build_plain_text_prompt, + input_ids_to_python_list, + ) + + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) + inputs = build_plain_text_prompt(processor, text) + return {"prompt_token_ids": input_ids_to_python_list(inputs)} + + +def _extract_pcm(outputs) -> torch.Tensor: + """Pull the concatenated PCM tensor out of the Stage-1 multimodal_output.""" + for stage_output in outputs: + mm = getattr(stage_output, "multimodal_output", None) + if not mm: + continue + audio = mm.get("model_outputs") if "model_outputs" in mm else mm.get("audio") + if audio is None: + continue + if isinstance(audio, list): + chunks = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] + if not chunks: + continue + return torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0] + return torch.as_tensor(audio).float().cpu().reshape(-1) + raise AssertionError(f"no audio payload in any stage output (got {len(outputs)} stages)") + + +@pytest.mark.core_model +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) +def test_higgs_audio_v2_offline_plain_text(omni_runner_function: OmniRunner) -> None: + """ + Offline plain-text -> 24 kHz PCM via the bundled higgs_audio_v2 deploy. + + Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml (codec_streaming=true, + async_chunk=false, enforce_eager=true per stage). + Input Modal: text + Output Modal: audio (PCM tensor, 24 kHz mono float) + Input Setting: single prompt, default sampling + """ + omni = omni_runner_function.omni + inputs = _compose_request(MODEL, TEST_TEXT) + + # max_tokens caps Stage-0 codec frames; Stage 1 max_tokens lifts the + # per-codebook prefill cap. Mirror the deploy yaml's defaults. + stage_sampling = [ + SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42), + SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42), + ] + + outputs = list(omni.generate([inputs], stage_sampling)) + assert outputs, "no outputs returned from Omni.generate" + + pcm = _extract_pcm(outputs) + arr = pcm.numpy() + + assert arr.size > MIN_AUDIO_SAMPLES, ( + f"audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES} " + f"({MIN_AUDIO_SAMPLES / SAMPLE_RATE:.2f} s @ {SAMPLE_RATE} Hz)" + ) + # Threshold 5e-4 (not 1e-3) keeps us above true silence while tolerating + # the ``load_format=dummy`` Stage-1 codec — random RVQ/DAC weights collapse + # the output to a near-constant peak around 7.4e-4. Real-weight runs reach + # 0.1+ peak so this still catches genuine silence. + assert np.max(np.abs(arr)) > 5e-4, "audio output is silent (peak abs amplitude < 5e-4)" + assert np.isfinite(arr).all(), "audio contains non-finite samples (nan/inf)" + + +@pytest.mark.core_model +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) +def test_higgs_audio_v2_offline_batch_two_prompts(omni_runner_function: OmniRunner) -> None: + """ + Submit two prompts in one engine call and verify both yield non-trivial PCM. + + Guards the Stage-0 talker's per-slot audio state under a batched call — + the prior example ran prompts one at a time citing 'request-scoped audio + state'; this test pins the batched-2 behavior the engine path actually + needs to support for /v1/audio/speech under concurrent traffic. + """ + omni = omni_runner_function.omni + prompts = [ + _compose_request(MODEL, "Hello world."), + _compose_request(MODEL, "The quick brown fox jumps over the lazy dog."), + ] + stage_sampling = [ + SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42), + SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42), + ] + + outputs = list(omni.generate(prompts, stage_sampling)) + assert outputs, "no outputs returned from Omni.generate" + + pcm = _extract_pcm(outputs) + # Whether the engine returns one concatenated PCM or per-request tensors is + # implementation-defined; for either shape, total samples must clear the + # floor for at least one of the two prompts. + assert pcm.numel() > MIN_AUDIO_SAMPLES, ( + f"batched audio too short: {pcm.numel()} samples, expected > {MIN_AUDIO_SAMPLES}" + ) + # See the silence-threshold note in test_higgs_audio_v2_offline_plain_text. + assert np.max(np.abs(pcm.numpy())) > 5e-4, "batched audio output is silent" + + +def _compose_voice_clone_request(model_name: str, text: str, ref_text: str, ref_wav_path) -> dict: + """Build the voice-clone prompt + additional_information payload. + + Calls the same ``build_voice_clone_prompt`` the serving layer uses online — + the HF processor encodes the reference clip on the spot via the bundled + HiggsAudioV2TokenizerModel. The talker's + ``_maybe_apply_ref_audio_substitution`` consumes the tensors at prefill. + """ + import soundfile as sf + from transformers import AutoProcessor + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + build_voice_clone_prompt, + ) + + wav, sr = sf.read(str(ref_wav_path), always_2d=False) + if wav.ndim == 2: + wav = wav.mean(axis=1) # downmix any stereo asset to mono + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) + out = build_voice_clone_prompt(processor, text, wav, int(sr), ref_text) + return { + "prompt_token_ids": out["prompt_token_ids"], + # Pass tensors bare (see serving_speech._build_higgs_audio_v2_params + # — list-wrapped tensors get dropped by the msgspec serializer). + "additional_information": { + "audio_input_ids": out["audio_input_ids"], + "audio_input_ids_mask": out["audio_input_ids_mask"], + }, + } + + +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) +def test_higgs_audio_v2_offline_voice_clone(omni_runner_function: OmniRunner) -> None: + """ + Shallow voice clone via ref_audio + ref_text, offline. + + Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml. + Input Modal: text + reference audio (qwen3_tts/clone_2.wav vendored asset) + Output Modal: audio (PCM tensor, 24 kHz mono float) + Input Setting: single prompt + """ + omni = omni_runner_function.omni + ref_path = get_asset_path(_REF_AUDIO_ASSET) + inputs = _compose_voice_clone_request(MODEL, TEST_TEXT, _REF_TEXT, ref_path) + + stage_sampling = [ + SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42), + SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42), + ] + + outputs = list(omni.generate([inputs], stage_sampling)) + assert outputs, "no outputs returned from Omni.generate for voice clone" + + pcm = _extract_pcm(outputs) + arr = pcm.numpy() + assert arr.size > MIN_AUDIO_SAMPLES, f"cloned audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES}" + assert np.max(np.abs(arr)) > 5e-4, "voice-clone output is silent" + assert np.isfinite(arr).all(), "voice-clone audio contains non-finite samples" diff --git a/tests/e2e/online_serving/test_higgs_audio_v2.py b/tests/e2e/online_serving/test_higgs_audio_v2.py new file mode 100644 index 000000000..92688ad0b --- /dev/null +++ b/tests/e2e/online_serving/test_higgs_audio_v2.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +End-to-end online tests for higgs-audio v2 against /v1/audio/speech. + +v1 scope is plain text -> 24 kHz speech plus shallow voice clone via +ref_audio + ref_text (inline) or voice= (after POST /v1/audio/voices). +The model-aware request validator +(vllm_omni/entrypoints/openai/serving_speech.py::_validate_higgs_audio_v2_request) +rejects multi-speaker tags, language overrides, task_type, and bare +voice= for names that do not match an uploaded speaker — this suite +exercises both the happy path (plain text in, audio bytes out) and the +validator rejections, so a regression that loosens the schema will fail +loudly. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") +# Match run_server.sh: DeepGEMM FP8 kernels are optional and trip warmup on +# images without the deep_gemm backend, so disable them by default. +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") +os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0") + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import load_test_audio_data_url +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "bosonai/higgs-audio-v2-generation-3B-base" +STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml") +SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"] +# DeepGEMM warmup is optional; mirror run_server.sh and switch it off in the +# server subprocess env too (parent-process os.environ above only affects this +# test driver; the engine subprocesses inherit through env_dict). +SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"} + +TEST_PARAMS = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + server_args=SERVER_ARGS, + env_dict=SERVER_ENV, + ), + id="higgs_audio_v2_plain_text", + ) +] + +DEFAULT_SPEECH_TIMEOUT_S = 180.0 +# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ≈ 24 KiB. +# A WAV header adds 44 bytes; pick a conservative floor that catches truncated +# / silence-only outputs without flagging short legitimate clips. +_MIN_AUDIO_BYTES = 20_000 + +# Reuse the qwen3_tts vendored reference clip (clean ~5 s 24 kHz mono human +# speech) + its transcript. See tests/e2e/online_serving/test_qwen3_tts_base.py +# for the asset rationale — keeping a single shared reference clip across TTS +# tests avoids duplicating WAVs in the repo. +_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestHiggsAudioV2OnlineHappyPath: + """Plain-text -> audio happy paths against the live HTTP server.""" + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_plain_text_wav(self, omni_server, openai_client) -> None: + """Single non-streaming WAV request — covers the canonical TTS happy path.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: + """Streaming PCM via the shared-memory connector's codec_streaming path. + + higgs_audio_v2.yaml pins ``codec_streaming: true`` + ``async_chunk: false``, + so the only streaming surface exposed to clients is the WAV/PCM bytes + served chunk-by-chunk from Stage 1 — exercise it directly. + """ + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Streaming the quick brown fox over the lazy dog.", + "stream": True, + "response_format": "pcm", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None: + """max_new_tokens is one of the few extra fields the higgs validator accepts.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Innovation distinguishes between a leader and a follower.", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "max_new_tokens": 500, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_concurrent_plain_text(self, omni_server, openai_client) -> None: + """Three concurrent non-streaming requests — guards the per-slot audio state.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "It was the night before my birthday.", + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + }, + request_num=3, + ) + + +# --------------------------------------------------------------------------- +# Validator rejections served over HTTP. Each case targets one out-of-scope +# field; the validator returns 4xx via the OpenAI error path. We keep the +# error-message substring loose so a phrasing tweak in the validator does not +# break CI, but tight enough to catch a regression that silently accepts the +# field. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestHiggsAudioV2OnlineValidatorRejections: + """Out-of-scope fields must come back as 4xx with a higgs-named message.""" + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None: + """Voice clone needs the transcript too — half-supplied is a 4xx.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "ref_audio": _REF_AUDIO_URL, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "status_code": (400, 422), + "err_message": "ref_text", + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None: + """Symmetric guard: ref_text alone is not enough — must come with ref_audio.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "ref_text": "some transcript", + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "status_code": (400, 422), + "err_message": "ref_audio", + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_task_type(self, omni_server, openai_client) -> None: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "task_type": "Base", + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "status_code": (400, 422), + "err_message": "task_type", + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_language_override(self, omni_server, openai_client) -> None: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "language": "Chinese", + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "status_code": (400, 422), + "err_message": "language", + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "[SPEAKER0] hi", + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "status_code": (400, 422), + "err_message": "multi-speaker", + } + ) + + @pytest.mark.core_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_rejects_empty_input(self, omni_server, openai_client) -> None: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": " ", + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "status_code": (400, 422), + "err_message": "empty", + } + ) + + +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +class TestHiggsAudioV2OnlineVoiceClone: + """Shallow voice clone: ref_audio + ref_text -> speech in the cloned voice. + + The HF processor that the serving layer calls (lazy-loaded at the first + request) ships with the bundled HiggsAudioV2TokenizerModel and encodes + the reference clip in-process. The talker substitutes the encoded codes + at the prompt-side audio placeholders via + :meth:`HiggsAudioV2TalkerForConditionalGeneration._maybe_apply_ref_audio_substitution`. + """ + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_voice_clone_basic(self, omni_server, openai_client) -> None: + """Single non-streaming WAV request driven by the qwen3_tts ref clip.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello world.", + "ref_audio": _REF_AUDIO_URL, + "ref_text": _REF_TEXT, + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "H100"}, num_cards=1) + def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None: + """Voice clone over the streaming PCM path.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Innovation distinguishes a leader from a follower.", + "ref_audio": _REF_AUDIO_URL, + "ref_text": _REF_TEXT, + "stream": True, + "response_format": "pcm", + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + ) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 1474e325c..4a8f48de6 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -121,4 +121,8 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.moss_tts_nano.pipeline", "MOSS_TTS_NANO_PIPELINE", ), + "higgs_audio_v2": ( + "vllm_omni.model_executor.models.higgs_audio_v2.pipeline", + "HIGGS_AUDIO_V2_PIPELINE", + ), } diff --git a/vllm_omni/deploy/higgs_audio_v2.yaml b/vllm_omni/deploy/higgs_audio_v2.yaml new file mode 100644 index 000000000..07e25d5c1 --- /dev/null +++ b/vllm_omni/deploy/higgs_audio_v2.yaml @@ -0,0 +1,83 @@ +# Higgs-audio v2 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory. +# +# Phase-1 scope is plain text -> 24 kHz speech only; voice cloning, multi-speaker, +# ChatML rich messages, and music/sound-event prompts are rejected by the +# request validator with explicit 4xx (see +# vllm_omni/entrypoints/openai/serving_speech.py and +# vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py). +# +# Streaming is wired through the stage-input processor's async-chunk +# callback, the shared-memory connector ``codec_streaming`` flag, and the +# WAV/PCM streaming convention in serving_speech.py. To run a strictly +# non-streaming smoke test, set ``codec_streaming: false`` in the connector +# extras below. +async_chunk: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 5000 + connector_get_max_wait: 500 + codec_chunk_frames: 25 + codec_left_context_frames: 25 + initial_codec_chunk_frames: 1 + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.6 + # Disable cudagraph capture: the live audio-sampler dispatch is still + # iterating, and the host's flash-attn-hopper graph path tripped a CUDA + # driver vs runtime mismatch during capture on this image. + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 2048 + # max_model_len covers prompt + max_new_tokens. Up to ~1500 codec + # frames per request fit comfortably under 4096. + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + # Upstream serve_engine defaults (temperature=1.0, top_p=0.95, + # top_k=50). Greedy (temperature=0) collapses the audio codebook + # argmax to a fixed value after a few steps; sampling with diversity + # is required for natural-sounding speech. + temperature: 1.0 + top_p: 0.95 + top_k: 50 + # Cap Stage-0 to 1024 audio frames (~40 s at 25 fps). The audio-eos + # natural stop (128012, forced by the talker's _apply_audio_mode_bias + # at ramp-completion) typically fires earlier and trims this. + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.25 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + # Prefill length is num_codebooks * num_frames; 8 * 8192 = 65536 covers + # the upper bound at the per-request 250-frame Stage-0 cap above plus + # headroom for batching multiple requests. + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 3ec2424c0..0325017b4 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -100,6 +100,50 @@ class OpenAICreateSpeechRequest(BaseModel): raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") return v + @model_validator(mode="before") + @classmethod + def reject_higgs_audio_v2_unsupported_aliases(cls, data): + """Reject unsupported rich-input aliases for higgs_audio_v2 BEFORE pydantic strips them. + + OpenAICreateSpeechRequest is a permissive schema (it accepts any extra + field that the OpenAI Speech API didn't promise to ban) but several + of those aliases pull a request out of the v1 higgs_audio_v2 scope. + Catching them at parse time lets the API return a deterministic 4xx + with a model-specific error message instead of silently dropping the + field and proceeding with a degraded request. + """ + # Keys that are silently dropped by pydantic when posted to /v1/audio/speech + # (the schema doesn't declare them) but which a model-specific + # validator MUST be able to reject. Kept inline to avoid clashing + # with pydantic's ModelPrivateAttr discovery for class-level + # underscore-prefixed attributes. + higgs_audio_v2_reserved_keys = ( + "messages", # ChatML rich content (out of scope for v1) + "reference_audio", # voice-cloning alias 1 + "voice_prompt", # voice-cloning alias 2 + "speaker_audio", # voice-cloning alias 3 + "speakers", # multi-speaker dialogue + ) + if not isinstance(data, dict): + return data + model_id = data.get("model") + if not isinstance(model_id, str): + return data + # Match the "higgs_audio_v2" model_type label, the HF architecture id + # in pipeline_registry.hf_architectures, and the hyphenated HF repo id + # (e.g. "bosonai/higgs-audio-v2-generation-3B-base") by normalizing + # both underscores and hyphens out of the id before substring matching. + normalized = model_id.lower().replace("-", "").replace("_", "") + if "higgsaudiov2" not in normalized: + return data + offending = sorted(k for k in higgs_audio_v2_reserved_keys if k in data) + if offending: + raise ValueError( + "higgs_audio_v2 v1 does not support these rich-input fields: " + f"{offending}. Supply plain text via the 'input' field instead." + ) + return data + @model_validator(mode="after") def validate_embedding_constraints(self) -> "OpenAICreateSpeechRequest": if self.speaker_embedding is not None: diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 7643ed625..4497f0409 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -69,6 +69,7 @@ _COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} _MING_TTS_MODEL_STAGES = {"ming_tts"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} +_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} _GLM_TTS_MODEL_STAGES = {"glm_tts"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES @@ -76,13 +77,21 @@ _TTS_MODEL_STAGES: set[str] = ( | _FISH_TTS_MODEL_STAGES | _COSYVOICE3_TTS_MODEL_STAGES | _OMNIVOICE_TTS_MODEL_STAGES + | _HIGGS_AUDIO_V2_TTS_MODEL_STAGES | _COVO_AUDIO_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES | _MING_TTS_MODEL_STAGES | _MOSS_TTS_MODEL_STAGES | _GLM_TTS_MODEL_STAGES ) -_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {"fish_tts", "qwen3_tts", "voxtral_tts", "cosyvoice3", "voxcpm2"} +_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = { + "fish_tts", + "qwen3_tts", + "voxtral_tts", + "cosyvoice3", + "voxcpm2", + "higgs_audio_v2", +} _TTS_LANGUAGES: set[str] = { "Auto", "Chinese", @@ -555,6 +564,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "ming_flash_omni_tts" if model_stage in _MOSS_TTS_MODEL_STAGES: return "moss_tts_nano" + if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES: + return "higgs_audio_v2" if model_stage in _GLM_TTS_MODEL_STAGES: return "glm_tts" return None @@ -845,12 +856,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): voice_lower = request.voice.lower() if voice_lower not in self.uploaded_speakers: - if self._tts_model_type in ("cosyvoice3", "fish_tts", "omnivoice", "moss_tts_nano", "glm_tts"): + if self._tts_model_type in ( + "cosyvoice3", + "fish_tts", + "omnivoice", + "moss_tts_nano", + "glm_tts", + "higgs_audio_v2", + ): label = { "cosyvoice3": "CosyVoice3", "fish_tts": "Fish Speech", "omnivoice": "OmniVoice", "moss_tts_nano": "MOSS-TTS-Nano", + "higgs_audio_v2": "Higgs-Audio V2", "glm_tts": "GLM-TTS", }.get(self._tts_model_type, self._tts_model_type) return ( @@ -1446,6 +1465,73 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): params["prompt_audio_array"] = [[wav_list, sr]] return params + def _validate_higgs_audio_v2_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate higgs_audio_v2 request parameters. Returns error message or None. + + Accepted: plain text -> speech, or shallow voice clone via ``ref_audio`` + + ``ref_text`` (both required together). Still out of scope: preset + ``voice``/``speaker`` selection, ``x_vector_only_mode`` / + ``speaker_embedding`` helpers, ``task_type``/``language``/ + ``instructions``/``speed`` overrides, and multi-speaker ``[SPEAKERn]`` + tags inside the input body. + """ + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + MULTI_SPEAKER_TAG_PATTERN, + ) + + if not request.input or not request.input.strip(): + return "higgs_audio_v2: input text cannot be empty" + + # Voice clone: ref_audio and ref_text must come together. + if request.ref_audio is not None and not request.ref_text: + return ( + "higgs_audio_v2 voice clone requires both 'ref_audio' and " + "'ref_text'; received ref_audio without ref_text" + ) + if request.ref_text and request.ref_audio is None: + return ( + "higgs_audio_v2 voice clone requires both 'ref_audio' and " + "'ref_text'; received ref_text without ref_audio" + ) + + if request.x_vector_only_mode is not None: + return "higgs_audio_v2 v1 does not support 'x_vector_only_mode' (voice-cloning helper field)" + if request.speaker_embedding is not None: + return "higgs_audio_v2 v1 does not support 'speaker_embedding' (voice-cloning helper field)" + if request.voice and request.ref_audio is None: + # _apply_uploaded_speaker runs before this validator; if voice was + # an uploaded speaker, ref_audio is now populated and ref_text is + # backfilled from the speaker entry. A bare voice= with no + # ref_audio means the name didn't resolve to an uploaded speaker + # (and higgs has no built-in preset voices). + return ( + "higgs_audio_v2 v1 does not support 'voice'/'speaker' selection for built-in voices; " + f"upload a voice first via POST /v1/audio/voices, or use ref_audio + ref_text. " + f"Got voice={request.voice!r}" + ) + if request.instructions: + return ( + "higgs_audio_v2 v1 does not support 'instructions' (voice " + "style/emotion control); supply plain text instead" + ) + if request.task_type is not None: + return "higgs_audio_v2 v1 does not support 'task_type'; the model is single-mode plain text -> speech" + if request.language is not None: + return ( + "higgs_audio_v2 v1 does not accept 'language' overrides; the model infers language from the input text" + ) + if request.speed is not None and request.speed != 1.0: + return ( + "higgs_audio_v2 v1 does not support 'speed' adjustments; the audio is rendered at native rate (24 kHz)" + ) + + if MULTI_SPEAKER_TAG_PATTERN.search(request.input): + return "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; remove the tag from the input" + + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + def _validate_glm_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate GLM-TTS request — requires ref_audio for voice cloning.""" if not request.input or not request.input.strip(): @@ -1466,6 +1552,80 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" return None + async def _build_higgs_audio_v2_params(self, request: OpenAICreateSpeechRequest): + """Build prompt_token_ids for higgs_audio_v2 via the upstream processor. + + Plain-text path: runs ``build_plain_text_prompt`` and returns the + token-only prompt. Voice-clone path (``ref_audio`` + ``ref_text``): + resolves the reference clip via ``_resolve_ref_audio``, runs + ``build_voice_clone_prompt`` (which encodes the clip through HF's + ``HiggsAudioV2TokenizerModel`` loaded from the k2-fsa/OmniVoice + ``audio_tokenizer/`` subdirectory), and attaches the encoded + ``audio_input_ids`` + ``audio_input_ids_mask`` tensors via + ``additional_information`` so the talker substitutes them at the + prompt-side audio placeholders. + """ + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + build_plain_text_prompt, + build_voice_clone_prompt, + input_ids_to_python_list, + ) + + processor = await self._resolve_higgs_audio_v2_processor() + + if request.ref_audio is None: + inputs = await asyncio.to_thread(build_plain_text_prompt, processor, request.input) + prompt_token_ids = input_ids_to_python_list(inputs) + return tokens_input(prompt_token_ids=prompt_token_ids) + + wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + wav = np.asarray(wav_list, dtype=np.float32) + out = await asyncio.to_thread( + build_voice_clone_prompt, + processor, + request.input, + wav, + int(sr), + request.ref_text or "", + ) + prompt = tokens_input(prompt_token_ids=out["prompt_token_ids"]) + # Pass tensors at the top level of additional_information (NOT list- + # wrapped). ``vllm_omni.data_entry_keys.serialize_payload`` routes + # bare ``torch.Tensor`` values through ``_serialize_tensor``; a list + # containing tensors would fall into the ``list_data`` field which + # msgspec cannot serialize and the tensors would be dropped over the + # process boundary (silent voice-clone failure). + prompt["additional_information"] = { + "audio_input_ids": out["audio_input_ids"], + "audio_input_ids_mask": out["audio_input_ids_mask"], + } + return prompt + + async def _resolve_higgs_audio_v2_processor(self): + """Lazy-load the AutoProcessor for higgs_audio_v2 (once per serving instance).""" + cached = getattr(self, "_higgs_audio_v2_processor", None) + if cached is not None: + return cached + + from transformers import AutoProcessor + + model_path = None + for stage in self.engine_client.stage_configs: + model_path = getattr(getattr(stage, "engine_args", None), "model", None) + if model_path: + break + if model_path is None: + # Fallback: the orchestrator stores the served model id on the engine + # itself (set by AsyncOmniEngine.__init__). Stage-level engine_args + # may not surface ``model`` when the deploy yaml doesn't set it per + # stage (the CLI-passed model id is the single source of truth). + model_path = getattr(self.engine_client, "model", None) + if model_path is None: + raise RuntimeError("higgs_audio_v2 serving could not resolve the model path from the engine stage configs") + processor = AutoProcessor.from_pretrained(model_path) + self._higgs_audio_v2_processor = processor + return processor + def _validate_fish_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Fish Speech request parameters. Returns error message or None.""" if not request.input or not request.input.strip(): @@ -2226,6 +2386,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): elif self._tts_model_type == "covo_audio": prompt = self._build_covo_audio_prompt(request) tts_params = {} + elif self._tts_model_type == "higgs_audio_v2": + # Explicit higgs_audio_v2 branch: do NOT fall through to the + # Qwen-style generic placeholder path; build prompt_token_ids + # directly via the upstream processor + the project's plain-text + # helper. + # + # Resolve uploaded voices first so a voice= request (after + # POST /v1/audio/voices) populates ref_audio + ref_text from the + # cached speaker entry, mirroring the cosyvoice3 / voxcpm2 flow. + err = self._apply_uploaded_speaker(request) + if err: + raise ValueError(err) + validation_error = self._validate_higgs_audio_v2_request(request) + if validation_error: + raise ValueError(validation_error) + prompt = await self._build_higgs_audio_v2_params(request) + if request.voice: + voice_lower = request.voice.lower() + if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: + additional = prompt.setdefault("additional_information", {}) + additional["voice_name"] = voice_lower + additional["voice_created_at"] = self._voice_created_at(voice_lower) + tts_params = {} elif self._tts_model_type == "voxcpm2": # voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the # raw waveform tuple for prefill-length accounting, not a base64 data URL. @@ -2338,6 +2521,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "ming_flash_omni_tts" elif self._tts_model_type == "moss_tts_nano": model_type = "moss_tts_nano" + elif self._tts_model_type == "higgs_audio_v2": + model_type = "higgs_audio_v2" elif self._tts_model_type == "glm_tts": model_type = "glm_tts" elif self._is_tts: diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py new file mode 100644 index 000000000..1a7c23eba --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vllm-omni integration for boson-ai's higgs-audio v2 (two-stage TTS).""" + +from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import ( + HiggsAudioV2Config, +) + +__all__ = ["HiggsAudioV2Config"] diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py new file mode 100644 index 000000000..dc5eb12b4 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration class for higgs-audio v2 in vllm-omni. + +Mirrors `bosonai/higgs-audio-v2-generation-3B-base/config.json` and exposes the +Stage-1 codec decoder reference (sample rate, codebook layout, stream-special +token ids) so the rest of the integration can read a single typed object. + +The text-LM backbone constants match `Llama-3.2-3B` (hidden=3072, 28 layers, +GQA 24Q/8KV, head_dim=128, vocab=128256, max_position_embeddings=2048, RoPE +``llama3`` scaling with factor=32). The audio side adds 8 codebooks of size +1026 (real codes ``[0..1023]``, stream specials ``audio_stream_bos_id=1024`` +and ``audio_stream_eos_id=1025``) plus the DualFFN audio expert. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import PretrainedConfig + +__all__ = ["HiggsAudioV2Config"] + + +class HiggsAudioV2Config(PretrainedConfig): + """Typed wrapper around the upstream `higgs_audio_v2` config. + + The field set here is a strict subset of the upstream + ``transformers.models.higgs_audio_v2.HiggsAudioV2Config``: only the + knobs that vllm-omni reads directly are surfaced. Anything extra coming + from the HF config dict is preserved via ``PretrainedConfig`` so + ``AutoConfig.from_pretrained(...)`` round-trips correctly. + """ + + model_type: str = "higgs_audio_v2" + keys_to_ignore_at_inference = ("past_key_values",) + + # Architecture identifier persisted in config.json and used by both + # vllm_omni.config.pipeline_registry and vllm_omni.model_executor.models.registry. + DEFAULT_ARCHITECTURES = ("HiggsAudioV2ForConditionalGeneration",) + + def __init__( + self, + # --- Text LM (Llama-3.2-3B) --- + vocab_size: int = 128256, + hidden_size: int = 3072, + intermediate_size: int = 8192, + num_hidden_layers: int = 28, + num_attention_heads: int = 24, + num_key_value_heads: int = 8, + head_dim: int = 128, + hidden_act: str = "silu", + max_position_embeddings: int = 2048, + rms_norm_eps: float = 1e-5, + rope_parameters: dict[str, Any] | None = None, + initializer_range: float = 0.02, + attention_bias: bool = False, + attention_dropout: float = 0.0, + mlp_bias: bool = False, + tie_word_embeddings: bool = False, + pretraining_tp: int = 1, + # --- Audio side --- + num_codebooks: int = 8, + codebook_size: int = 1026, + audio_token_id: int = 128016, + audio_bos_token_id: int = 128013, + audio_eos_token_id: int = 128012, + audio_delay_token_id: int = 128014, + audio_stream_bos_id: int = 1024, + audio_stream_eos_id: int = 1025, + sample_rate: int = 24000, + frame_rate: int = 25, + # --- Standard LM specials --- + bos_token_id: int = 1, + eos_token_id: int = 128009, + pad_token_id: int = 128001, + # --- DualFFN --- + # Upstream `HiggsAudioV2DecoderLayer` always carries a parallel `audio_mlp` + # expert; this flag is kept for forward-compat in case a future config + # disables it (no-op for the boson-ai checkpoint). + use_audio_dual_ffn: bool = True, + # --- Stage-1 tokenizer / codec decoder --- + # Path to the audio-tokenizer subdir (relative to the model directory). + # Empty string means the model_dir itself contains config.json + + # model.safetensors (the layout used by the standalone + # ``bosonai/higgs-audio-v2-tokenizer`` HF repo). Non-empty means the + # tokenizer files live in ``//`` (the layout used by + # the 3B Stage-0 checkpoint which bundles its tokenizer under + # ``audio_tokenizer/``). + audio_tokenizer_subdir: str = "", + # HF repo id (or local path) of a STANDALONE audio tokenizer. When set, + # Stage-1 resolves this via ``huggingface_hub.snapshot_download`` and + # loads the codec from there, ignoring ``audio_tokenizer_subdir``. This + # matches the boson-ai release layout where the codec lives in a + # separate repo (``bosonai/higgs-audio-v2-tokenizer``). + audio_tokenizer_id: str | None = "bosonai/higgs-audio-v2-tokenizer", + # --- Misc --- + use_cache: bool = True, + **kwargs: Any, + ) -> None: + if rope_parameters is None: + rope_parameters = { + "factor": 32.0, + "high_freq_factor": 0.5, + "low_freq_factor": 0.125, + "original_max_position_embeddings": 1024, + "rope_theta": 500000.0, + "rope_type": "llama3", + } + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.rope_parameters = dict(rope_parameters) + self.initializer_range = initializer_range + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.mlp_bias = mlp_bias + self.pretraining_tp = pretraining_tp + + self.num_codebooks = num_codebooks + self.codebook_size = codebook_size + self.audio_token_id = audio_token_id + self.audio_bos_token_id = audio_bos_token_id + self.audio_eos_token_id = audio_eos_token_id + self.audio_delay_token_id = audio_delay_token_id + self.audio_stream_bos_id = audio_stream_bos_id + self.audio_stream_eos_id = audio_stream_eos_id + self.sample_rate = sample_rate + self.frame_rate = frame_rate + + self.use_audio_dual_ffn = use_audio_dual_ffn + self.audio_tokenizer_subdir = audio_tokenizer_subdir + self.audio_tokenizer_id = audio_tokenizer_id + self.use_cache = use_cache + + # Codes [0..codebook_size-1] include 2 stream specials at indices + # `audio_stream_bos_id` and `audio_stream_eos_id`. Real codes ready for the + # codec decoder are the strict prefix [0..audio_stream_bos_id). + @property + def num_real_codes(self) -> int: + return self.audio_stream_bos_id # i.e. 1024 for the boson-ai checkpoint + + # The full vocab size of each per-codebook output head. + @property + def codebook_output_size(self) -> int: + return self.codebook_size # i.e. 1026 for the boson-ai checkpoint + + # Canonical MusicGen-style delay pattern used by + # `HiggsAudioV2DelayPatternLogitsProcessor`: codebook k starts emitting + # real codes only after k frames. Kept as a function so subclasses can + # override if upstream ever publishes a different convention. + def default_delay_pattern(self) -> list[int]: + return list(range(self.num_codebooks)) diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py new file mode 100644 index 000000000..74cdbb2db --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HiggsAudio codec decoder kernel for higgs-audio v2. + +This module hosts the parameter-side building blocks for the higgs-audio-v2 +audio tokenizer's decoder path: + + audio_codes [B, 8, T] + -> RVQ codebook lookup + project_out -> sum -> [B, hidden_size, T] + -> fc2 Linear(hidden_size, 256) -> [B, 256, T] + -> DAC acoustic decoder (conv-transpose upsampling) -> [B, 1, T*960] + -> 24 kHz waveform (25 fps x 960 samples/frame) +""" + +from __future__ import annotations + +import json +import os +from typing import Any + +import torch +import torch.nn as nn + +__all__ = [ + "HiggsAudioVQLayer", + "HiggsAudioRVQ", + "adjust_conv_transpose_output_padding", + "build_higgs_audio_acoustic_decoder", + "BosonDacDecoder", + "build_boson_dac_decoder", + "load_higgs_audio_codec", + "_remap_boson_model_pth_state_dict", # exported for unit-testing the mapper +] + + +class HiggsAudioVQLayer(nn.Module): + """Single VQ layer: codebook lookup + project_out.""" + + def __init__(self, codebook_size: int = 1024, codebook_dim: int = 64, hidden_size: int = 1024): + super().__init__() + self.codebook = nn.Embedding(codebook_size, codebook_dim) + self.project_out = nn.Linear(codebook_dim, hidden_size) + + def decode(self, indices: torch.Tensor) -> torch.Tensor: + """indices: [B, T] -> [B, hidden_size, T].""" + quantized = self.codebook(indices) + quantized = self.project_out(quantized) + return quantized.permute(0, 2, 1) + + +class HiggsAudioRVQ(nn.Module): + """Residual Vector Quantizer with ``num_quantizers`` codebook layers.""" + + def __init__( + self, + num_quantizers: int = 8, + codebook_size: int = 1024, + codebook_dim: int = 64, + hidden_size: int = 1024, + ): + super().__init__() + self.quantizers = nn.ModuleList( + [HiggsAudioVQLayer(codebook_size, codebook_dim, hidden_size) for _ in range(num_quantizers)] + ) + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """codes: [num_quantizers, B, T] -> [B, hidden_size, T].""" + result = torch.zeros( + codes.shape[1], + self.quantizers[0].project_out.out_features, + codes.shape[2], + device=codes.device, + dtype=torch.float32, + ) + for i, quantizer in enumerate(self.quantizers): + result = result + quantizer.decode(codes[i]) + return result + + +@torch.jit.script +def _snake(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor: + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + return x.reshape(shape) + + +class _Snake1d(nn.Module): + """Snake activation (per-channel learnable alpha). Mirrors the boson-ai DAC + Snake1d implementation byte-for-byte so the boson ``model.pth`` state-dict + loads directly into ``self.alpha``. + """ + + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return _snake(x, self.alpha) + + +def _wn_conv1d(*args: Any, **kwargs: Any) -> nn.Module: + from torch.nn.utils import weight_norm + + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def _wn_conv_transpose1d(*args: Any, **kwargs: Any) -> nn.Module: + from torch.nn.utils import weight_norm + + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +class _BosonResidualUnit(nn.Module): + def __init__(self, dim: int, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + _Snake1d(dim), + _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), + _Snake1d(dim), + _wn_conv1d(dim, dim, kernel_size=1), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.block(x) + pad = (x.shape[-1] - y.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + y + + +class _BosonDecoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int): + super().__init__() + import math + + self.block = nn.Sequential( + _Snake1d(input_dim), + _wn_conv_transpose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=stride % 2, + ), + _BosonResidualUnit(output_dim, dilation=1), + _BosonResidualUnit(output_dim, dilation=3), + _BosonResidualUnit(output_dim, dilation=9), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.block(x) + + +class BosonDacDecoder(nn.Module): + """Vendored copy of ``boson_multimodal``'s standalone DAC decoder, with the + exact ``decoder_2.model..{...}`` key layout the boson-ai + ``bosonai/higgs-audio-v2-tokenizer`` checkpoint expects. + + See ``boson_multimodal/audio_processing/descriptaudiocodec/dac/model/dac.py`` + (``class Decoder``) in the upstream package for the reference implementation + we mirror. + + Architecture for the v2 24 kHz tokenizer: + input_channel=256, channels=1024, rates=[8, 5, 4, 2, 3], d_out=1 + Cumulative upsample factor is ``prod(rates) = 960`` (25 fps -> 24 kHz). + """ + + def __init__( + self, + input_channel: int = 256, + channels: int = 1024, + rates: tuple[int, ...] = (8, 5, 4, 2, 3), + d_out: int = 1, + ): + super().__init__() + layers: list[nn.Module] = [_wn_conv1d(input_channel, channels, kernel_size=7, padding=3)] + last_dim = channels + for i, stride in enumerate(rates): + input_dim = channels // 2**i + output_dim = channels // 2 ** (i + 1) + layers.append(_BosonDecoderBlock(input_dim, output_dim, stride)) + last_dim = output_dim + layers.append(_Snake1d(last_dim)) + layers.append(_wn_conv1d(last_dim, d_out, kernel_size=7, padding=3)) + self.model = nn.Sequential(*layers) + # Cache hop length (= prod(rates)) so callers don't need to recompute it. + hop = 1 + for s in rates: + hop *= int(s) + self._hop_length = hop + + @property + def hop_length(self) -> int: + return self._hop_length + + # Expose conv1 with a ``.in_channels`` attribute so the existing fc2 + # fallback shape inference (acoustic_decoder.conv1.in_channels) works + # unchanged for the boson decoder too. + @property + def conv1(self) -> nn.Module: + return self.model[0] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.model(x) + + +def build_boson_dac_decoder(device: torch.device) -> BosonDacDecoder: + """Construct the vendored boson DAC decoder on ``device``. + + The architectural constants (input_channel=256, channels=1024, + rates=[8,5,4,2,3], d_out=1) are baked in because they're invariant for the + v2 24 kHz tokenizer ``bosonai/higgs-audio-v2-tokenizer`` checkpoint we + target. Other channel counts are not supported until we encounter a + different boson checkpoint in the wild. + """ + return BosonDacDecoder().to(device) + + +def adjust_conv_transpose_output_padding(decoder: nn.Module) -> None: + """Set ConvTranspose1d output_padding = stride % 2 (HiggsAudioV2 modification). + + The vanilla DAC decoder ships with the default output_padding (0); the + boson-ai checkpoint expects ``stride % 2`` instead. This is a no-op for + even strides and adds a single sample for odd strides. + """ + for module in decoder.modules(): + if isinstance(module, nn.ConvTranspose1d): + stride = module.stride[0] if isinstance(module.stride, tuple) else module.stride + module.output_padding = (stride % 2,) + + +def build_higgs_audio_acoustic_decoder( + tokenizer_config: dict[str, Any], + device: torch.device, +) -> nn.Module: + """Build the DAC acoustic decoder used by the HiggsAudioV2 tokenizer. + + Returns the decoder sub-module of ``transformers.DacModel`` with the + HiggsAudioV2 output-padding fix already applied. The tanh activation is + replaced with ``Identity`` so the network matches the boson-ai checkpoint. + Weights are NOT loaded here; the caller is responsible for copying them in. + """ + from transformers import DacConfig, DacModel + + dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"]) + dac_model = DacModel(dac_cfg) + decoder = dac_model.decoder.to(device) + adjust_conv_transpose_output_padding(decoder) + if hasattr(decoder, "tanh"): + decoder.tanh = nn.Identity() + return decoder + + +def _load_higgs_audio_state_dict(audio_tokenizer_dir: str, device: torch.device) -> dict[str, torch.Tensor]: + """Load the codec state dict from either layout used in the wild. + + Tries layouts in order: + 1. ``/model.safetensors`` (OmniVoice-bundled layout used by + ``k2-fsa/OmniVoice/audio_tokenizer/``). + 2. ``/model.pth`` (boson-ai standalone ``bosonai/higgs-audio-v2-tokenizer`` + layout). The state-dict keys differ structurally + (``quantizer.vq.layers.._codebook.embed`` etc.) and are remapped to + the OmniVoice-style names this kernel expects via + :func:`_remap_boson_model_pth_state_dict`. + """ + safetensors_path = os.path.join(audio_tokenizer_dir, "model.safetensors") + pth_path = os.path.join(audio_tokenizer_dir, "model.pth") + # Preference order: try model.pth first when present (boson-ai standalone + # ships the codec there; its model.safetensors is a different artefact). + # Fall back to model.safetensors (the OmniVoice-bundled layout). + sd: dict[str, torch.Tensor] + if os.path.exists(pth_path): + sd = torch.load(pth_path, map_location=device, weights_only=False) + elif os.path.exists(safetensors_path): + from safetensors.torch import load_file + + sd = load_file(safetensors_path, device=str(device)) + else: + raise FileNotFoundError(f"Audio tokenizer weights not found at {pth_path} or {safetensors_path}") + # Run the boson-layout remap if needed; for the OmniVoice-bundled layout + # it is a no-op (no ``quantizer.vq.layers.*`` keys to rewrite). + if any(k.startswith("quantizer.vq.layers.") for k in sd): + sd = _remap_boson_model_pth_state_dict(sd) + return sd + + +def _remap_boson_model_pth_state_dict(sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Translate boson-ai's standalone ``model.pth`` keys into OmniVoice-style names + that the shared kernel's RVQ + fc2 + DAC sites consume. + + Best-effort: only the RVQ-side keys map cleanly. The boson-ai decoder + uses ``decoder_2.model..weight_g/weight_v`` (DAC with weight-norm + + Snake activations) rather than the OmniVoice DAC layout; that side + requires either vendoring the upstream decoder module or rewriting the + DAC builder to consume weight-normed tensors. This function returns the + mapped RVQ keys plus any acoustic_decoder.* / fc2.* / fc.* tensors that + happen to share the OmniVoice names, and leaves the decoder-2 keys + untouched (the caller will see them as MISSING when copying into the + OmniVoice DAC decoder). + + Mapping (RVQ side only): + quantizer.vq.layers.._codebook.embed -> quantizer.quantizers..codebook.embed + quantizer.vq.layers..project_out.weight -> quantizer.quantizers..project_out.weight + quantizer.vq.layers..project_out.bias -> quantizer.quantizers..project_out.bias + """ + if not isinstance(sd, dict): + raise TypeError(f"expected a state dict, got {type(sd)!r}") + remapped: dict[str, torch.Tensor] = {} + for key, tensor in sd.items(): + if not isinstance(tensor, torch.Tensor): + continue + # Quantizer rewrite: vq.layers.._codebook.embed -> quantizers..codebook.embed + if key.startswith("quantizer.vq.layers."): + parts = key.split(".") + # parts[3] is the layer index, parts[4]+ is the tail + if len(parts) >= 5: + idx = parts[3] + tail = ".".join(parts[4:]) + # The actual codebook centroids the decode path uses live at + # ``_codebook.embed`` (a [codebook_size, codebook_dim] buffer). + # ``_codebook.embed_avg`` is the un-normalised EMA accumulator + # used during training (~cluster_size x bigger); MATCHING IT + # HERE produces a 4-5x amplified codebook and unintelligible + # PCM downstream. Use an EXACT-match check, not startswith. + if tail == "_codebook.embed": + new_key = f"quantizer.quantizers.{idx}.codebook.embed" + remapped[new_key] = tensor + continue + if tail.startswith("project_out."): + new_key = f"quantizer.quantizers.{idx}.{tail}" + remapped[new_key] = tensor + continue + # project_in / _codebook.cluster_size / _codebook.embed_avg / inited + # are encoder-side or training-only state that the decode path + # does not need; drop them. + continue + # boson decoder + fc_post2 -> our vendored kernel names. + if key.startswith("decoder_2."): + # decoder_2.model..{...} -> acoustic_decoder.model..{...} + new_key = "acoustic_decoder." + key[len("decoder_2.") :] + remapped[new_key] = tensor + continue + if key.startswith("fc_post2."): + # fc_post2.{weight,bias} -> fc2.{weight,bias} + new_key = "fc2." + key[len("fc_post2.") :] + remapped[new_key] = tensor + continue + # Anything else (acoustic_decoder.*, fc2.*, fc.*, fc_post*.*, decoder_2.*, + # decoder_semantic.*, semantic_model.*) passes through unchanged so + # the caller's lookup is unambiguous. + remapped[key] = tensor + return remapped + + +def load_higgs_audio_codec( + audio_tokenizer_dir: str, + device: torch.device, +) -> tuple[HiggsAudioRVQ, nn.Linear, nn.Module, dict[str, Any]]: + """Load the HiggsAudioV2 RVQ + fc2 + DAC decoder from a checkpoint folder. + + Accepts both layouts: ``/model.safetensors`` (OmniVoice-bundled) and + ``/model.pth`` (boson-ai standalone). The standalone path remaps + quantizer keys before consumption; structural decoder differences + (boson uses Snake + weight-norm) still leave some DAC parameters missing + when loading the standalone path -- those entries log warnings but the + RVQ side completes successfully. + + Args: + audio_tokenizer_dir: Path to a directory containing ``config.json`` + and EITHER ``model.safetensors`` (OmniVoice layout) OR + ``model.pth`` (boson-ai standalone layout). + device: Device to place the loaded modules and state dict on. + + Returns: + (quantizer, fc2, acoustic_decoder, tokenizer_config) + - quantizer: ``HiggsAudioRVQ`` with ``num_quantizers`` discovered from + the state dict (defaults to 8 for boson-ai checkpoints). + - fc2: ``nn.Linear`` projecting RVQ output (1024) into the DAC's + hidden dimension (typically 256). May be uninitialized when loaded + from the boson-ai standalone layout (no ``fc2.*`` keys present). + - acoustic_decoder: DAC decoder with HiggsAudioV2 output-padding fix + and tanh-replaced-by-Identity. Fully initialized for the OmniVoice + layout; partially initialized for boson-ai standalone (missing + ``decoder_2.*`` -> ``acoustic_decoder.*`` mapping is not yet + implemented; full upstream-decoder vendoring is the next step). + - tokenizer_config: The loaded ``config.json`` dict; useful for + callers that need ``sample_rate`` or other tokenizer metadata. + """ + config_path = os.path.join(audio_tokenizer_dir, "config.json") + + with open(config_path) as f: + tokenizer_config: dict[str, Any] = json.load(f) + + state_dict = _load_higgs_audio_state_dict(audio_tokenizer_dir, device) + + codebook_dim = tokenizer_config.get("codebook_dim", 64) + codebook_size = tokenizer_config.get("codebook_size", 1024) + # Discover hidden_size and num_quantizers from the (possibly remapped) state dict. + if "quantizer.quantizers.0.project_out.weight" not in state_dict: + raise KeyError( + "Codec state dict is missing 'quantizer.quantizers.0.project_out.weight'. " + "If you loaded a boson-ai standalone tokenizer, ensure the model.pth " + "remap fired (see _remap_boson_model_pth_state_dict)." + ) + hidden_size = state_dict["quantizer.quantizers.0.project_out.weight"].shape[0] + num_quantizers = sum( + 1 for k in state_dict if k.startswith("quantizer.quantizers.") and k.endswith(".codebook.embed") + ) + + quantizer = HiggsAudioRVQ( + num_quantizers=num_quantizers, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + hidden_size=hidden_size, + ).to(device) + for i in range(num_quantizers): + prefix = f"quantizer.quantizers.{i}" + embed_key = f"{prefix}.codebook.embed" + if embed_key in state_dict: + quantizer.quantizers[i].codebook.weight.data.copy_(state_dict[embed_key]) + proj_out_w = f"{prefix}.project_out.weight" + proj_out_b = f"{prefix}.project_out.bias" + if proj_out_w in state_dict: + quantizer.quantizers[i].project_out.weight.data.copy_(state_dict[proj_out_w]) + if proj_out_b in state_dict: + quantizer.quantizers[i].project_out.bias.data.copy_(state_dict[proj_out_b]) + + # Decide which acoustic decoder we need. A boson-layout state dict (after + # _remap_boson_model_pth_state_dict) carries ``acoustic_decoder.model.0.weight_v`` + # (the WNConv first conv) which the OmniVoice DacModel doesn't have; the + # OmniVoice layout carries ``acoustic_decoder.block.0.*`` instead. + is_boson_layout = "acoustic_decoder.model.0.weight_v" in state_dict + if is_boson_layout: + acoustic_decoder = build_boson_dac_decoder(device) + boson_sd = { + k[len("acoustic_decoder.") :]: v for k, v in state_dict.items() if k.startswith("acoustic_decoder.") + } + load_report = acoustic_decoder.load_state_dict(boson_sd, strict=False) + if load_report.missing_keys: + raise RuntimeError( + "Boson DAC decoder is missing keys after load: " + f"{load_report.missing_keys[:5]}... ({len(load_report.missing_keys)} total)" + ) + # ``unexpected_keys`` is informational only. + else: + acoustic_decoder = build_higgs_audio_acoustic_decoder(tokenizer_config, device) + for name, param in acoustic_decoder.named_parameters(): + higgs_name = f"acoustic_decoder.{name}" + if higgs_name in state_dict: + param.data.copy_(state_dict[higgs_name]) + acoustic_decoder.eval() + + if "fc2.weight" in state_dict and "fc2.bias" in state_dict: + fc2_w = state_dict["fc2.weight"] + fc2_b = state_dict["fc2.bias"] + fc2 = nn.Linear(fc2_w.shape[1], fc2_w.shape[0]).to(device) + fc2.weight.data.copy_(fc2_w) + fc2.bias.data.copy_(fc2_b) + else: + # Last-resort fallback: shape fc2 from the decoder's first conv input + # channels so the matmul produces compatible shapes. Weights stay + # uninitialised — PCM quality will be noise on this path. + first_conv = getattr(acoustic_decoder, "conv1", None) + if first_conv is None or not hasattr(first_conv, "in_channels"): + raise RuntimeError("acoustic_decoder has no conv1.in_channels; cannot size fc2 fallback") + fc2 = nn.Linear(hidden_size, int(first_conv.in_channels)).to(device) + + return quantizer, fc2, acoustic_decoder, tokenizer_config diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py new file mode 100644 index 000000000..6a14707ee --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage 1 (codec decoder) for higgs-audio v2. + +Two surfaces exposed: + +1. **Direct decode API** (kept for offline / unit tests): + ``decode_codes(audio_codes: [B, num_codebooks=8, T])`` -> ``[B, 1, T*hop]`` PCM. + ``forward_chunk(audio_codes, *, left_context_size, hop_length=960)`` + trims overlap from streamed PCM. ``forward(audio_codes)`` retains the same + single-tensor signature for backward compatibility with code that imports + the class directly. + +2. **vLLM stage runtime API** (used by the engine): + ``__init__(*, vllm_config, prefix)``, ``embed_input_ids``, + ``compute_logits=None``, plus a runtime ``forward(input_ids=..., positions=..., + runtime_additional_information=...)`` that takes flat codebook-major + ``input_ids`` (``[Q * num_frames]``) per request and returns an + :class:`vllm_omni.model_executor.models.output_templates.OmniOutput`. + +The kernel is the HiggsAudio codec helper in +``vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py``; both +surfaces share the same RVQ + DAC weights. +""" + +from __future__ import annotations + +import os +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import ( + HiggsAudioV2Config, +) +from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import ( + HiggsAudioRVQ, + load_higgs_audio_codec, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +__all__ = [ + "HiggsAudioV2Code2Wav", + "HiggsAudioV2Code2WavForConditionalGeneration", +] + +logger = init_logger(__name__) + + +class HiggsAudioV2Code2Wav(nn.Module): + """Stage-1 codec decoder for higgs-audio v2. + + Constructor accepts either the vLLM-runtime signature (``*, vllm_config, + prefix``) or the direct config-object form (``HiggsAudioV2Code2Wav(config)``) + so both engine-side and unit-test callers work. When the runtime form is + used we read the higgs_audio_v2 config out of ``vllm_config.model_config.hf_config`` + and the model_path off ``vllm_config.model_config.model`` so the codec can + be loaded directly from disk. + """ + + input_modalities = "audio" + + def __init__( + self, + config: HiggsAudioV2Config | None = None, + *, + vllm_config: VllmConfig | None = None, + prefix: str = "", + ): + super().__init__() + + # Resolve config from either positional (unit-test) or kwarg (engine) form. + if vllm_config is not None: + hf_config = vllm_config.model_config.hf_config + if isinstance(hf_config, HiggsAudioV2Config): + self.config = hf_config + else: + self.config = HiggsAudioV2Config(**hf_config.to_dict()) + self._model_path: str | None = vllm_config.model_config.model + self.vllm_config: VllmConfig | None = vllm_config + else: + if config is None: + raise TypeError( + "HiggsAudioV2Code2Wav: provide either positional `config` " + "(HiggsAudioV2Config) or keyword `vllm_config` (VllmConfig)." + ) + self.config = config + self._model_path = None + self.vllm_config = None + + self.sample_rate: int = int(self.config.sample_rate) + self.num_codebooks: int = int(self.config.num_codebooks) + self.num_real_codes: int = int(self.config.num_real_codes) + # Each codec frame upsamples to 960 24 kHz samples (= 25 fps * 960 = 24000). + self.hop_length: int = 960 + + # Engine-runner hooks (Stage 1 has no token sampling). + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + # Populated by load_weights(). + self.quantizer: HiggsAudioRVQ | None = None + self.fc2: nn.Linear | None = None + self.acoustic_decoder: nn.Module | None = None + self._loaded: bool = False + + # When constructed via the engine path, eagerly load codec weights from + # the model directory. The engine's subsequent ``load_weights(weights)`` + # call (with the Stage-0 safetensors iterator) is then a cheap no-op. + if self._model_path is not None: + try: + self.load_codec_from_disk(self._model_path) + except FileNotFoundError as exc: + logger.warning("HiggsAudioV2Code2Wav: eager codec load deferred (%s)", exc) + + # ------------------------------------------------------------- engine hooks + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Stage 1 ignores embeddings; vLLM's runner still needs a stable shape.""" + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None: + return None + + # ------------------------------------------------------------------ load + def load_weights(self, weights_or_model_dir, device: torch.device | None = None): + """Dual-signature load: + + - ``load_weights(weights_iterator)`` -- the vLLM engine path. The iterator + yields ``(name, tensor)`` pairs from the Stage-0 safetensors which DO NOT + contain codec weights, so we ignore them and eagerly load the codec from + ``self._model_path``. Returns the set of all our params as "initialized" + so the engine's loader passes its sanity check. + - ``load_weights(model_dir: str)`` -- the legacy file-based call used by + offline unit tests and ``end2end.py --mode stage1_only``. Returns ``None``. + """ + if not isinstance(weights_or_model_dir, (str, bytes, os.PathLike)): + # Engine call: consume the iterator (we don't need it) and eagerly + # load the codec from disk. + try: + for _ in weights_or_model_dir: + pass + except TypeError: + pass + if self._model_path is None: + raise RuntimeError( + "HiggsAudioV2Code2Wav.load_weights called via the engine path but " + "no model directory is known. Construct with vllm_config so the " + "constructor caches vllm_config.model_config.model." + ) + self.load_codec_from_disk(self._model_path, device=device) + return {name for name, _ in self.named_parameters()} + + # Legacy file-based call. + self.load_codec_from_disk(weights_or_model_dir, device=device) + return None + + def load_codec_from_disk(self, model_dir: str, device: torch.device | None = None) -> None: + """Load codec weights from disk (RVQ + fc2 + DAC decoder). + + Resolution order for the codec directory: + 1. If :attr:`HiggsAudioV2Config.audio_tokenizer_id` is set, treat it as + either a local path (if it exists on disk) or an HF repo id that we + resolve via ``huggingface_hub.snapshot_download``. This is the + default for the boson-ai release (codec in a separate repo). + 2. Otherwise fall back to ``/``. + """ + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + audio_tokenizer_path = self._resolve_audio_tokenizer_path(model_dir) + quantizer, fc2, acoustic_decoder, _tokenizer_config = load_higgs_audio_codec(audio_tokenizer_path, device) + if len(quantizer.quantizers) != self.num_codebooks: + raise ValueError( + f"checkpoint has {len(quantizer.quantizers)} quantizers but config.num_codebooks={self.num_codebooks}" + ) + self.quantizer = quantizer + self.fc2 = fc2 + self.acoustic_decoder = acoustic_decoder + self._loaded = True + logger.info( + "Loaded HiggsAudioV2Code2Wav: %d quantizers, fc2(%d->%d), sample_rate=%d", + len(self.quantizer.quantizers), + self.fc2.in_features, + self.fc2.out_features, + self.sample_rate, + ) + + def _resolve_audio_tokenizer_path(self, model_dir: str) -> str: + """Return a local filesystem path containing ``config.json`` + codec + weights. See :meth:`load_codec_from_disk` for the resolution rules.""" + tokenizer_id = getattr(self.config, "audio_tokenizer_id", None) + if tokenizer_id: + if os.path.isdir(tokenizer_id): + return tokenizer_id + return self._resolve_hf_id_to_local(tokenizer_id) + subdir = self.config.audio_tokenizer_subdir or "" + if os.path.isdir(model_dir): + return os.path.join(model_dir, subdir) if subdir else model_dir + local = self._resolve_hf_id_to_local(model_dir) + return os.path.join(local, subdir) if subdir else local + + @staticmethod + def _resolve_hf_id_to_local(repo_id: str) -> str: + """Resolve a HF repo id to a local snapshot directory. + + Prefers a read-only cache lookup (no write locks) so a quota-constrained + Lustre cache can still serve already-downloaded weights. Falls back to + ``snapshot_download`` only when nothing is cached locally. + """ + from huggingface_hub import try_to_load_from_cache + + cached = try_to_load_from_cache(repo_id=repo_id, filename="config.json") + if isinstance(cached, str) and os.path.isfile(cached): + return os.path.dirname(cached) + + from huggingface_hub.constants import HF_HUB_CACHE + + safe = repo_id.replace("/", "--") + snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots") + if os.path.isdir(snapshots_dir): + for rev in os.listdir(snapshots_dir): + candidate = os.path.join(snapshots_dir, rev) + if os.path.isfile(os.path.join(candidate, "config.json")): + return candidate + + from huggingface_hub import snapshot_download + + return snapshot_download(repo_id=repo_id) + + # ------------------------------------------------------ direct decode API + @torch.inference_mode() + def decode_codes(self, audio_codes: torch.Tensor) -> torch.Tensor: + """Decode a ``[B, num_codebooks=8, T]`` code tensor to PCM ``[B, 1, T*hop]``.""" + if not self._loaded: + raise RuntimeError("HiggsAudioV2Code2Wav not loaded. Call load_weights() first.") + + codes = self._validate_codes(audio_codes) + device = codes.device + + rvq_codes = codes.transpose(0, 1).long() # [num_codebooks, B, T] + quantized = self.quantizer.decode(rvq_codes) # [B, hidden, T] + # Match fc2's parameter dtype — vLLM may auto-cast the codec to the + # stage's model dtype (e.g. bf16), while the quantizer's RVQ codebooks + # often stay in float32 because they're embedded lookups. A defensive + # cast keeps the matmul homogeneous. + quantized = quantized.to(dtype=self.fc2.weight.dtype) + quantized = self.fc2(quantized.transpose(1, 2)).transpose(1, 2) + # Acoustic decoder may have parameters in yet another dtype — match + # its first parameter so the linear/conv ops stay consistent. + first_param = next(self.acoustic_decoder.parameters(), None) + if first_param is not None and quantized.dtype != first_param.dtype: + quantized = quantized.to(dtype=first_param.dtype) + audio = self.acoustic_decoder(quantized) + if audio.dim() == 2: + audio = audio.unsqueeze(1) + return audio.to(device) + + @torch.inference_mode() + def forward_chunk( + self, + audio_codes: torch.Tensor, + *, + left_context_size: int = 0, + hop_length: int | None = None, + ) -> torch.Tensor: + """Chunked decode that trims ``left_context_size * hop_length`` samples + off the leading edge of the upsampled PCM. Mirrors the qwen3_tts + ``talker2code2wav_async_chunk`` overlap contract. + """ + if left_context_size < 0: + raise ValueError(f"left_context_size must be >= 0; got {left_context_size}") + hop = int(hop_length) if hop_length is not None else self.hop_length + if hop <= 0: + raise ValueError(f"hop_length must be > 0; got {hop}") + pcm = self.decode_codes(audio_codes) + if left_context_size == 0: + return pcm + trim = left_context_size * hop + if pcm.shape[-1] <= trim: + return pcm[..., :0] + return pcm[..., trim:] + + # ---------------------------------------------------- vLLM runtime forward + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> torch.Tensor | OmniOutput: + """Dual-mode forward. + + - When ``input_ids`` is a 3-D ``[B, num_codebooks, T]`` tensor we treat + this as a direct decode call (legacy API) and return raw PCM. + - When ``input_ids`` is a flat 1-D / 2-D tensor we treat it as the + engine runtime payload (codebook-major flat per request) and return + an :class:`OmniOutput` carrying multimodal audio. + """ + # Legacy direct-decode signature: caller passed a 3-D code tensor as input_ids. + if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 3: + return self.decode_codes(input_ids) + + sr_val = int(self.sample_rate) + sr_tensor = torch.tensor(sr_val, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) + + left_context_size = [0] * len(request_ids_list) + if runtime_additional_information is not None: + for i, info in enumerate(runtime_additional_information): + if i >= len(left_context_size): + break + meta = info.get("meta", {}) if isinstance(info, dict) else {} + if "left_context_size" in meta: + left_context_size[i] = int(meta["left_context_size"]) + + wavs: list[torch.Tensor] = [] + for i, req_ids in enumerate(request_ids_list): + n = int(req_ids.numel()) + if n == 0: + wavs.append(empty) + continue + if n % self.num_codebooks != 0: + logger.warning( + "HiggsAudioV2Code2Wav: flat code length %d not divisible by num_codebooks=%d; dropping", + n, + self.num_codebooks, + ) + wavs.append(empty) + continue + frames = n // self.num_codebooks + codes_qf = req_ids.reshape(self.num_codebooks, frames) + codes_bqf = codes_qf.unsqueeze(0) + try: + pcm = self.forward_chunk( + codes_bqf, + left_context_size=left_context_size[i], + hop_length=self.hop_length, + ) + except ValueError as exc: + logger.warning("HiggsAudioV2Code2Wav: decode skipped (%s)", exc) + wavs.append(empty) + continue + wavs.append(pcm.squeeze(0).squeeze(0).to(torch.float32).cpu()) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": wavs, + "sr": [sr_tensor] * len(wavs), + }, + ) + + # --------------------------------------------------------------- helpers + def _validate_codes(self, audio_codes: torch.Tensor) -> torch.Tensor: + """Ensure shape and value range; reject stream specials with ValueError.""" + if not isinstance(audio_codes, torch.Tensor): + raise TypeError(f"audio_codes must be a torch.Tensor, got {type(audio_codes)!r}") + if audio_codes.ndim != 3: + raise ValueError( + f"audio_codes must have shape [B, num_codebooks={self.num_codebooks}, T]; " + f"got shape {tuple(audio_codes.shape)}" + ) + if int(audio_codes.shape[1]) != self.num_codebooks: + raise ValueError( + f"audio_codes second dim must equal num_codebooks={self.num_codebooks}; got {int(audio_codes.shape[1])}" + ) + if audio_codes.numel() > 0: + max_val = int(audio_codes.max().item()) + min_val = int(audio_codes.min().item()) + if max_val >= self.num_real_codes or min_val < 0: + raise ValueError( + "audio_codes contains stream-special or out-of-range IDs: " + f"min={min_val}, max={max_val}; real code range is " + f"[0, {self.num_real_codes - 1}]. Filter audio_stream_bos_id=" + f"{self.config.audio_stream_bos_id} and audio_stream_eos_id=" + f"{self.config.audio_stream_eos_id} (and anything above) at " + "Stage 0 before sending codes to the codec decoder." + ) + return audio_codes + + @staticmethod + def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]: + """Split a concatenated flat-codes tensor into per-request segments.""" + if seq_token_counts is not None and len(seq_token_counts) > 1: + boundaries = [0] + for count in seq_token_counts: + boundaries.append(boundaries[-1] + int(count)) + n = int(ids.numel()) + return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))] + return [ids] + + +# Engine-side architecture identifier alias (mirrors Qwen3TTSCode2Wav usage). +HiggsAudioV2Code2WavForConditionalGeneration = HiggsAudioV2Code2Wav diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py new file mode 100644 index 000000000..5e3883fd2 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py @@ -0,0 +1,1684 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-0 talker for higgs-audio v2 (vLLM-native, DualFFN-aware). + +Architecture summary (verified against `transformers.models.higgs_audio_v2` +source and the boson-ai checkpoint config.json): + +- Backbone: Llama-3.2-3B (hidden=3072, 28 layers, GQA 24Q/8KV, head_dim=128, + vocab=128256, max_position_embeddings=2048, RoPE llama3 scaling factor=32). + We reuse vLLM's compiled :class:`vllm.model_executor.models.llama.LlamaModel` + for the attention path so PagedAttention scheduling stays intact. + +- DualFFN: every transformer block carries a parallel audio expert + (``audio_input_layernorm`` + ``audio_post_attention_layernorm`` + ``audio_mlp``) + next to the standard text path (``input_layernorm`` + ``post_attention_layernorm`` + + ``mlp``). A per-position ``audio_token_mask`` of shape ``[B, S]`` selects + between the two. The mask is the union of positions where the input token id + equals ``audio_token_id=128016`` or ``audio_delay_token_id=128014``. + +- Multi-codebook output head: at audio positions, Stage 0 emits one ID per + codebook (8 codebooks of vocabulary 1026 each) in parallel via a single + fused ``audio_lm_head: Linear(hidden, num_codebooks * codebook_size)`` + whose output is reshaped to ``[N_audio, num_codebooks, codebook_size]``. + +- Delay pattern: codebook k is staggered by k frames using + ``audio_delay_token_id=128014`` as a filler at the LM-token positions where + codebook k has not yet started emitting real codes. This matches + ``HiggsAudioV2DelayPatternLogitsProcessor`` upstream and is the canonical + MusicGen pattern ``[0, 1, 2, 3, 4, 5, 6, 7]``. + +This file delivers the structural pieces (module classes, weight mapping, +forward pass, AR sampler dispatch, and the upstream delay-pattern logits +processor) to reproduce the upstream behavior under vLLM's compiled forward +path. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaMLP + +from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import ( + HiggsAudioV2Config, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +__all__ = [ + "HiggsAudioV2DecoderLayer", + "HiggsAudioV2TalkerForConditionalGeneration", +] + +logger = init_logger(__name__) + + +class HiggsAudioRMSNorm(nn.Module): + """Minimal RMSNorm matching the upstream HiggsAudioV2RMSNorm semantics. + + Kept self-contained so the parameter names (``weight``) line up with the + upstream state dict for direct load. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states.to(input_dtype) + + +class HiggsAudioV2DecoderLayer(LlamaDecoderLayer): + """One transformer block: vLLM-native attention + DualFFN-routed MLPs. + + Wraps a stock :class:`vllm.model_executor.models.llama.LlamaDecoderLayer` + for the self-attention path (so PagedAttention KV caches keep working) and + overlays the upstream HiggsAudioV2 DualFFN routing for the layernorm + MLP + pairs. The implementation follows the upstream rule from + ``transformers.models.higgs_audio_v2.HiggsAudioV2DecoderLayer.forward``: + + 1. Pre-attention norm: per-position split into ``audio_input_layernorm`` + (audio mask True) vs ``input_layernorm`` (audio mask False). The mixed + output is fed to a single shared self-attention. + 2. Post-attention residual + dual MLP: text positions go through + ``mlp(post_attention_layernorm(.))`` and audio positions go through + ``audio_mlp(audio_post_attention_layernorm(.))``. Both deltas are added + to the residual. + + The classical (non-fused-residual) transformer pattern is used so the per- + position split is straightforward. vLLM's compiled forward path passes + ``audio_token_mask`` through as an extra layer kwarg via + :class:`HiggsAudioV2TalkerForConditionalGeneration.forward`. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + # Subclass LlamaDecoderLayer so self_attn registers at the + # CANONICAL ``layers.{i}.self_attn.attn`` prefix (not nested under + # ``.base``). This is the architecturally-correct registration path + # for vLLM's compilation_config.static_forward_context. + super().__init__(vllm_config=vllm_config, prefix=prefix) + config = vllm_config.model_config.hf_config + hidden_size = int(config.hidden_size) + intermediate_size = int(config.intermediate_size) + hidden_act = str(config.hidden_act) + mlp_bias = bool(getattr(config, "mlp_bias", False)) + rms_norm_eps = float(config.rms_norm_eps) + # Parallel audio-expert norms + MLP (mirrors upstream parameter names). + self.audio_input_layernorm = HiggsAudioRMSNorm(hidden_size, eps=rms_norm_eps) + self.audio_post_attention_layernorm = HiggsAudioRMSNorm(hidden_size, eps=rms_norm_eps) + self.audio_mlp = LlamaMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_act=hidden_act, + bias=mlp_bias, + prefix=f"{prefix}.audio_mlp", + ) + + # ------------------------------------------------------------------ helpers + def _routed_norm( + self, + hidden: torch.Tensor, + text_norm: nn.Module, + audio_norm: HiggsAudioRMSNorm, + audio_mask: torch.Tensor | None, + residual: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Per-position routed RMSNorm with residual handling. + + Matches the vLLM fused-residual pattern: takes (hidden, residual) and + returns (normed_hidden, post_add_residual). Routes per position: + audio positions through ``audio_norm`` (HiggsAudioRMSNorm, plain RMS), + text positions through ``text_norm`` (vLLM compiled RMSNorm). + + For consistency we explicitly compute the residual sum once, then call + the vLLM text_norm on text positions WITHOUT its residual argument + (else it'd add the residual again). + """ + # Fold residual into hidden once. + if residual is not None: + hidden = hidden + residual + new_residual = hidden + if audio_mask is None: + return audio_norm(hidden), new_residual + mask_flat = audio_mask.reshape(-1) + if hidden.ndim == 3: + hidden_flat = hidden.reshape(-1, hidden.shape[-1]) + else: + hidden_flat = hidden + out = hidden_flat.clone() + if (~mask_flat).any(): + text_out = text_norm(hidden_flat[~mask_flat]) + # vLLM compiled RMSNorm.forward(x) returns Tensor; some variants + # return (Tensor, residual). Normalize to Tensor. + if isinstance(text_out, tuple): + text_out = text_out[0] + out[~mask_flat] = text_out.to(out.dtype) + if mask_flat.any(): + audio_out = audio_norm(hidden_flat[mask_flat]) + out[mask_flat] = audio_out.to(out.dtype) + return out.reshape_as(hidden), new_residual + + def _routed_mlp( + self, + hidden: torch.Tensor, + audio_mask: torch.Tensor | None, + ) -> torch.Tensor: + """Apply ``audio_mlp`` to mask positions and ``mlp`` to non-mask positions.""" + if audio_mask is None: + return self.audio_mlp(hidden) + mask_flat = audio_mask.reshape(-1) + if hidden.ndim == 3: + hidden_flat = hidden.reshape(-1, hidden.shape[-1]) + else: + hidden_flat = hidden + out = torch.zeros_like(hidden_flat) + if (~mask_flat).any(): + out[~mask_flat] = self.mlp(hidden_flat[~mask_flat]).to(out.dtype) + if mask_flat.any(): + out[mask_flat] = self.audio_mlp(hidden_flat[mask_flat]).to(out.dtype) + return out.reshape_as(hidden) + + # ------------------------------------------------------------------ forward + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + *, + audio_token_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Fused-residual pattern (compatible with vLLM's LlamaModel forward + # loop which calls ``hidden, residual = self.norm(hidden, residual)`` + # at the end). + attn_input, residual = self._routed_norm( + hidden_states, + self.input_layernorm, + self.audio_input_layernorm, + audio_token_mask, + residual, + ) + attn_out = self.self_attn(positions=positions, hidden_states=attn_input) + + mlp_input, residual = self._routed_norm( + attn_out, + self.post_attention_layernorm, + self.audio_post_attention_layernorm, + audio_token_mask, + residual, + ) + mlp_out = self._routed_mlp(mlp_input, audio_token_mask) + return mlp_out, residual + + +class HiggsAudioV2TalkerForConditionalGeneration(nn.Module): + """Stage-0 talker class registered under + ``HiggsAudioV2ForConditionalGeneration`` (canonical HF arch identifier) + AND ``HiggsAudioV2TalkerForConditionalGeneration`` (explicit alias). + + The class wires together: + 1. vLLM-native ``LlamaModel`` backbone (PagedAttention scheduling). + 2. A parallel set of DualFFN modules (one per decoder layer) for the + audio-expert path. + 3. Multi-codebook output via a single fused ``audio_lm_head`` of shape + ``[num_codebooks * codebook_size, hidden]`` (all codebooks emitted + in parallel per audio position). + 4. ``load_weights`` HF -> vLLM mapping (GQA reshape, llama3 RoPE, + text/audio MLP split, audio code heads, audio token embeddings). + + The forward path interleaves DualFFN routing with vLLM's compiled + attention path so the talker integrates with the rest of the wiring + (pipeline_registry, serving_speech, deploy yaml, stage_input_processor). + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + hf_config = vllm_config.model_config.hf_config + if isinstance(hf_config, HiggsAudioV2Config): + self.config: HiggsAudioV2Config = hf_config + else: + # When loaded via AutoConfig the class is the upstream + # transformers config; wrap it in our typed shell preserving + # every attribute via PretrainedConfig. + self.config = HiggsAudioV2Config(**hf_config.to_dict()) + + # ------------------------------------------------------------------ embedding + final norm + # directly instantiate the text-side embedding and final norm + # so we don't pull in a full LlamaModel.layers list that the engine's + # weight loader would require initialized. Each HiggsAudioV2DecoderLayer + # below already owns its own LlamaDecoderLayer for the attention path; + # duplicating those via a parallel LlamaModel made the loader fail. + self.embed_tokens = VocabParallelEmbedding( + num_embeddings=self.config.vocab_size, + embedding_dim=self.config.hidden_size, + org_num_embeddings=self.config.vocab_size, + quant_config=None, + prefix=f"{prefix}.embed_tokens", + ) + self.norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + + # ------------------------------------------------------------------ DualFFN-aware layer stack + # use vLLM's ``make_layers`` helper to register attention layers + # at canonical ``layers.{i}.self_attn.attn`` paths in + # ``compilation_config.static_forward_context``. This is what gets + # picked up by ``bind_kv_cache`` so prefill writes K/V to the cache + # (root cause for empty cache + zero attention output). + from vllm.model_executor.models.utils import make_layers + + _, _, self.layers = make_layers( + self.config.num_hidden_layers, + lambda prefix: HiggsAudioV2DecoderLayer(vllm_config=vllm_config, prefix=prefix), + prefix=f"{prefix}.layers", + ) + + # ------------------------------------------------------------------ audio embedding (shared) + # Audio frames are embedded by summing per-codebook lookups; the LM + # input_ids stream uses ``audio_token_id`` / ``audio_delay_token_id`` + # placeholders at the corresponding positions, and the talker + # substitutes the audio embedding via the audio_codebook_embeddings + # table declared further below. + self.register_buffer( + "audio_tokens_offsets", + torch.arange(self.config.num_codebooks) * self.config.codebook_size, + persistent=False, + ) + + # ------------------------------------------------------------------ heads + # Standard text LM head (Llama 128256-wide vocab). + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + org_num_embeddings=self.config.vocab_size, + quant_config=None, + prefix=f"{prefix}.lm_head", + ) + + # Audio side: the upstream ``audio_decoder_proj.audio_lm_head`` is a + # single fused ``Linear(hidden, num_codebooks * (codebook_size + 2))``. + # In our v2 config ``codebook_size = 1026`` ALREADY includes the two + # stream specials, so the fused head is shape ``(8 * 1026, 3072) = + # (8208, 3072)``. At sample time we run it once and reshape + # ``[N, 8208] -> [N, 8, 1026]`` to read each codebook's logits. + self.audio_lm_head = nn.Linear( + self.config.hidden_size, + self.config.num_codebooks * self.config.codebook_size, + bias=False, + ) + # Multi-codebook input embedding used both for prompt-side audio frame + # substitution and for next-step audio-input-embedding feedback during + # decode. Mirrors upstream ``HiggsAudioModel.audio_codebook_embeddings``. + self.audio_codebook_embeddings = nn.Embedding( + self.config.num_codebooks * self.config.codebook_size, + self.config.hidden_size, + ) + + # vLLM's logits processor wires sampling metadata into the lm_head. + self.logits_processor = LogitsProcessor(self.config.vocab_size) + + # ----------------------------------------------------------------- masks + @torch.inference_mode() + def audio_token_mask(self, input_ids: torch.Tensor) -> torch.Tensor: + """Compute the per-position audio_token_mask used by DualFFN routing. + + Matches ``HiggsAudioV2Model.get_placeholder_mask``: + + mask = (input_ids == audio_token_id) | (input_ids == audio_delay_token_id) + """ + return (input_ids == self.config.audio_token_id) | (input_ids == self.config.audio_delay_token_id) + + # ------------------------------------------------------------ ref-audio + @torch.inference_mode() + def _maybe_apply_ref_audio_substitution( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + info_dicts: Any, + ) -> torch.Tensor: + """Substitute reference-audio embeddings at prompt-side audio placeholders. + + Voice-clone path: the serving layer threads per-request + ``audio_input_ids`` (shape ``[T_frames, num_codebooks]``) + + ``audio_input_ids_mask`` (``[T_frames]``) into + ``model_intermediate_buffer`` (one dict per request in batch order). + For each request that has those tensors, we look up its span in the + flat ``input_ids`` via ``query_start_loc`` from the forward context, + embed the codes through ``audio_codebook_embeddings`` (with the + per-codebook ``audio_tokens_offsets`` shift, mirroring + ``HiggsAudioV2Embeddings.forward`` upstream), and overwrite the + text-side embedding at every audio_token / audio_delay placeholder + position in that span. + + Decode steps (span length 1) are skipped here — the audio-feedback + path in ``_maybe_apply_audio_feedback`` takes over from the moment + the talker emits its first codebook frame via ``sample()``. + """ + if not info_dicts: + return hidden_states + + # Pull the per-request (codes, mask) tensors in batch order. Values + # may be on CPU after IPC serialization (msgspec ships tensors as raw + # bytes); the loop below moves them to the input_ids device on use. + per_req_codes: list[torch.Tensor | None] = [] + any_codes = False + for info in info_dicts: + if not isinstance(info, dict): + per_req_codes.append(None) + continue + ids = info.get("audio_input_ids") + mask = info.get("audio_input_ids_mask") + if isinstance(ids, list): + ids = ids[0] if ids else None + if isinstance(mask, list): + mask = mask[0] if mask else None + if not isinstance(ids, torch.Tensor) or not isinstance(mask, torch.Tensor): + per_req_codes.append(None) + continue + if ids.ndim == 3: + ids = ids[0] + if mask.ndim == 2: + mask = mask[0] + valid = ids[mask.to(dtype=torch.bool)] # [T_valid, num_codebooks] + if valid.numel() == 0: + per_req_codes.append(None) + continue + per_req_codes.append(valid) + any_codes = True + + if not any_codes: + return hidden_states + + # Per-request slices in the flat input_ids are recovered from the + # forward context's attention metadata (mimo_audio_llm.py:1022-1025 + # uses the same pattern). Without spans we cannot route codes to the + # right placeholder positions, so we no-op instead of risking a + # cross-request mix-up. + from vllm.forward_context import get_forward_context + + try: + attn_metadata = get_forward_context().attn_metadata + except Exception: + return hidden_states + if isinstance(attn_metadata, dict): + if not attn_metadata: + return hidden_states + attn = next(iter(attn_metadata.values())) + else: + attn = attn_metadata + q_start = getattr(attn, "query_start_loc", None) + if not isinstance(q_start, torch.Tensor): + return hidden_states + q_start_cpu = q_start.detach().to("cpu").tolist() + if len(q_start_cpu) != len(info_dicts) + 1: + return hidden_states + + audio_token_id = int(self.config.audio_token_id) + audio_delay_id = int(self.config.audio_delay_token_id) + offsets = self.audio_tokens_offsets.to(input_ids.device) + + new_hidden: torch.Tensor | None = None + for i, codes in enumerate(per_req_codes): + if codes is None: + continue + s = int(q_start_cpu[i]) + e = int(q_start_cpu[i + 1]) + if e - s <= 1: + # Decode step — let _maybe_apply_audio_feedback handle this row. + continue + slice_ids = input_ids[s:e] + mask = (slice_ids == audio_token_id) | (slice_ids == audio_delay_id) + placeholders = mask.nonzero(as_tuple=True)[0] + n_real = int(codes.shape[0]) + if int(placeholders.numel()) < n_real: + # Mismatch (likely a truncated prompt): skip this request rather + # than corrupt the batch. The user-visible failure surfaces as + # the model not honoring the ref-audio rather than a crash. + continue + target = placeholders[:n_real] + s + embeds = self.audio_codebook_embeddings(codes.to(device=input_ids.device, dtype=torch.long) + offsets).sum( + dim=-2 + ) + if new_hidden is None: + new_hidden = hidden_states.clone() + new_hidden[target] = embeds.to(new_hidden.dtype) + + return new_hidden if new_hidden is not None else hidden_states + + # ----------------------------------------------------------------- weights + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Map HF state dict names to the vLLM-native layout. + + The upstream tensor names live under ``model.layers....`` plus a few + top-level entries. The vLLM-native layout uses fused QKV/MLP tensors, + so we transform the relevant pairs at load time: + + * Attention: HF ``q_proj/k_proj/v_proj`` -> vLLM ``qkv_proj`` (packed + ``[hidden + 2 * kv_dim, hidden]``; for GQA with 24 Q heads, 8 KV + heads, head_dim=128 the slabs are + ``q_proj[3072, 3072]``, ``k_proj[1024, 3072]``, ``v_proj[1024, 3072]`` + -> ``qkv_proj[5120, 3072]``). + * RoPE parameters are read from ``self.config.rope_parameters`` and + consumed by vLLM's LlamaAttention without any weight transform + (``rope_type="llama3"``, factor=32, low_freq_factor=0.125, + high_freq_factor=0.5, original_max_position_embeddings=1024). + * Text MLP: HF ``mlp.gate_proj`` + ``mlp.up_proj`` -> + vLLM ``mlp.gate_up_proj`` (concatenated on dim 0). + * Audio MLP: HF ``audio_mlp.gate_proj`` + ``audio_mlp.up_proj`` -> + vLLM ``dual_ffns..audio_mlp.gate_up_proj`` (same shape transform). + * Layer norms: HF ``input_layernorm`` / ``post_attention_layernorm`` + -> ``dual_ffns..{input,post_attention}_layernorm``. Audio norms + -> ``dual_ffns..{audio_input,audio_post_attention}_layernorm``. + * Audio token embedding: HF ``model.embed_audio_tokens.weight`` (with + optional outer ``embed_audio_tokens.`` prefix from upstream's + HiggsAudioV2Embeddings) -> ``self.audio_codebook_embeddings.weight`` + (see ``_map_simple_name``). + * Audio codebook heads: HF ``audio_lm_head.weight`` -> the fused 8x1026 + head from which Stage 0 reads codebook 0 directly and codebook 1..7 + via the fast-AR code predictor's residual heads. (Boson-ai's actual + checkpoint exposes a single fused ``audio_lm_head`` of shape + ``[num_codebooks * codebook_size, hidden]``, not one head per + codebook; we split it into the per-codebook heads at load time.) + + this routine covers the simple, fused-QKV, and fused-MLP + transcriptions. The remaining gap (full Stage-0 forward integration with + vLLM's PagedAttention scheduler so the loaded weights drive a working + greedy decode) is tracked as an Open Issue in the goal tracker. + """ + loaded: set[str] = set() + own_params = dict(self.named_parameters()) + # Debug: log the first few HF keys we receive so we can verify the + # mapper hits them. Also log our own param-name shape so missing + # mappings are obvious. + logger.info( + "higgs_audio_v2 load_weights: %d own params (sample: %s)", + len(own_params), + sorted(list(own_params.keys()))[:5], + ) + debug_seen: list[str] = [] + + # First pass: collect tensors that need fusing (q/k/v, gate/up for each + # MLP slot). We accumulate into per-(layer, slot, kind) buckets and + # commit the fused tensor once all parts have arrived. + fuse_buckets: dict[tuple[int, str, str], dict[str, torch.Tensor]] = {} + + def _strip_model_prefix(name: str) -> str: + """vLLM's DefaultModelLoader strips the leading ``model.`` from HF + state-dict keys before invoking ``load_weights``, but we also want + to accept raw HF names (for unit tests that pass the state dict + directly). Normalize by stripping at most one ``model.`` prefix. + """ + if name.startswith("model."): + return name[len("model.") :] + return name + + def _try_simple(name: str, tensor: torch.Tensor) -> bool: + normalized = _strip_model_prefix(name) + # Direct match: many keys (embed_tokens.weight, norm.weight, + # layers..{audio_,}input_layernorm.weight, layers..{audio_,}post_attention_layernorm.weight, + # layers..self_attn.o_proj.weight, layers..{,audio_}mlp.down_proj.weight) already match + # our named_parameters() keys after the model-prefix strip. + if normalized in own_params: + target = own_params[normalized] + if tuple(target.shape) == tuple(tensor.shape): + with torch.no_grad(): + target.copy_(tensor) + loaded.add(normalized) + return True + logger.warning( + "higgs_audio_v2 load_weights: shape mismatch %s -> %s: %s vs %s", + name, + normalized, + tuple(tensor.shape), + tuple(target.shape), + ) + return False + + # After making HiggsAudioV2DecoderLayer inherit from + # LlamaDecoderLayer, text-side params live directly at + # ``layers..`` (no ``.base`` indirection). audio side + # remains at ``layers..audio_<...>``. + parts = normalized.split(".") + if len(parts) >= 3 and parts[0] == "layers": + tail = ".".join(parts[2:]) + # Audio_mlp.down_proj lives at our top-level audio_mlp slot. + if tail == "audio_mlp.down_proj.weight": + audio_name = f"layers.{parts[1]}.audio_mlp.down_proj.weight" + if audio_name in own_params: + target = own_params[audio_name] + if tuple(target.shape) == tuple(tensor.shape): + with torch.no_grad(): + target.copy_(tensor) + loaded.add(audio_name) + return True + + # Special simple cases (nested audio embedding, text_lm_head alias, etc.) + mapped = self._map_simple_name(name) + if mapped is None: + # Also try after stripping the model. prefix so the _map_simple_name + # branches written for un-stripped names continue to work. + mapped = self._map_simple_name("model." + normalized) + if mapped is None or mapped not in own_params: + return False + target = own_params[mapped] + if tuple(target.shape) != tuple(tensor.shape): + logger.warning( + "higgs_audio_v2 load_weights: shape mismatch %s -> %s: %s vs %s", + name, + mapped, + tuple(tensor.shape), + tuple(target.shape), + ) + return False + with torch.no_grad(): + target.copy_(tensor) + loaded.add(mapped) + return True + + def _stash_fusion(name: str, tensor: torch.Tensor) -> bool: + normalized = _strip_model_prefix(name) + parts = normalized.split(".") + if len(parts) < 4 or parts[0] != "layers": + return False + try: + layer_idx = int(parts[1]) + except ValueError: + return False + tail = ".".join(parts[2:]) + for slot in ("mlp", "audio_mlp"): + for kind in ("gate_proj", "up_proj"): + if tail == f"{slot}.{kind}.weight": + fuse_buckets.setdefault((layer_idx, slot, "gate_up_proj"), {})[kind] = tensor + return True + for kind in ("q_proj", "k_proj", "v_proj"): + if tail == f"self_attn.{kind}.weight": + fuse_buckets.setdefault((layer_idx, "self_attn", "qkv_proj"), {})[kind] = tensor + return True + return False + + unhandled: list[str] = [] + for name, tensor in weights: + if len(debug_seen) < 30: + debug_seen.append(name) + if _try_simple(name, tensor): + continue + if _stash_fusion(name, tensor): + continue + # Audio code head: upstream has a fused [num_codebooks*codebook_size, hidden] + # head; split into codebook-0 head + residual heads. Accept both + # ``audio_lm_head.weight`` and ``model.audio_lm_head.weight``. + if name in ( + "audio_lm_head.weight", + "model.audio_lm_head.weight", + "audio_decoder_proj.audio_lm_head.weight", + "model.audio_decoder_proj.audio_lm_head.weight", + ): + self._consume_fused_audio_head(tensor, loaded) + continue + unhandled.append(name) + + # Second pass: emit fused tensors. Targets are the per-layer slots on + # the HiggsAudioV2DecoderLayer stack: text-side params at + # ``layers..`` directly (HiggsAudioV2DecoderLayer inherits + # from LlamaDecoderLayer); ``layers..audio_mlp...`` for the audio + # expert. + for (layer_idx, slot, fused_kind), parts in fuse_buckets.items(): + if fused_kind == "gate_up_proj": + gate = parts.get("gate_proj") + up = parts.get("up_proj") + if gate is None or up is None: + continue + fused = torch.cat([gate, up], dim=0) + if slot == "mlp": + target_name = f"layers.{layer_idx}.mlp.gate_up_proj.weight" + else: # audio_mlp + target_name = f"layers.{layer_idx}.audio_mlp.gate_up_proj.weight" + elif fused_kind == "qkv_proj": + q = parts.get("q_proj") + k = parts.get("k_proj") + v = parts.get("v_proj") + if q is None or k is None or v is None: + continue + fused = torch.cat([q, k, v], dim=0) + target_name = f"layers.{layer_idx}.self_attn.qkv_proj.weight" + else: + continue + if target_name in own_params and tuple(own_params[target_name].shape) == tuple(fused.shape): + with torch.no_grad(): + own_params[target_name].copy_(fused) + loaded.add(target_name) + else: + logger.warning( + "higgs_audio_v2 load_weights: fused %s not found or shape mismatch (%s vs target %s)", + target_name, + tuple(fused.shape), + tuple(own_params[target_name].shape) if target_name in own_params else "", + ) + logger.info( + "higgs_audio_v2 load_weights: %d/%d params initialized " + "(sample seen: %s; unhandled count=%d, sample unhandled=%s)", + len(loaded), + len(own_params), + debug_seen[:10], + len(unhandled), + unhandled[:20], + ) + return loaded + + def _consume_fused_audio_head(self, fused: torch.Tensor, loaded: set[str]) -> None: + """Load the fused ``audio_lm_head[num_codebooks*codebook_size, hidden]`` + directly into ``self.audio_lm_head.weight``. + + Upstream uses a single big head; there is no separate codebook-0 + head or residual predictor in the actual checkpoint. + """ + own_params = dict(self.named_parameters()) + num_codebooks = int(self.config.num_codebooks) + codebook_size = int(self.config.codebook_size) + hidden = int(self.config.hidden_size) + target = own_params.get("audio_lm_head.weight") + if target is None: + logger.warning("higgs_audio_v2: no audio_lm_head.weight slot to load into") + return + expected = (num_codebooks * codebook_size, hidden) + if tuple(fused.shape) != expected: + logger.warning( + "higgs_audio_v2: unexpected audio_lm_head shape %s; expected %s", + tuple(fused.shape), + expected, + ) + return + with torch.no_grad(): + target.copy_(fused) + loaded.add("audio_lm_head.weight") + + # --------------------------------------------------------------- helpers + def _map_simple_name(self, hf_name: str) -> str | None: + """Translate the simple (non-fused) HF parameter names to vLLM names.""" + # Upstream nests the audio embedding as + # ``model.embed_audio_tokens.embed_audio_tokens.weight``; the boson-ai + # checkpoint also uses this exact key. Accept both nested and flat. + # the upstream multi-codebook embedding is ``audio_codebook_embeddings`` + # (a flat ``Embedding(8 * codebook_size, hidden)``). also had a separate + # ``embed_audio_tokens`` (used in prefill-time DualFFN routing helpers); + # we keep that slot so the existing prompt-side embedding logic stays + # functional. Both slots are loaded from the same upstream tensor. + if hf_name in ( + "embed_audio_tokens.weight", + "embed_audio_tokens.embed_audio_tokens.weight", # stripped form + "model.embed_audio_tokens.embed_audio_tokens.weight", # raw HF form + "model.embed_audio_tokens.weight", + # Actual key used in the boson-ai checkpoint: + "audio_codebook_embeddings.weight", + "model.audio_codebook_embeddings.weight", + ): + return "audio_codebook_embeddings.weight" + if hf_name in ( + "lm_head.weight", + "text_lm_head.weight", + "model.text_lm_head.weight", + # Actual key used in the boson-ai checkpoint: + "audio_decoder_proj.text_lm_head.weight", + "model.audio_decoder_proj.text_lm_head.weight", + ): + return "lm_head.weight" + # Dropped LlamaModel wrapper; text embedding + final norm now + # live directly on the talker (self.embed_tokens / self.norm). + if hf_name == "model.embed_tokens.weight": + return "embed_tokens.weight" + if hf_name == "model.norm.weight": + return "norm.weight" + # Per-layer audio norms -> our HiggsAudioV2DecoderLayer side. + if hf_name.startswith("model.layers.") and hf_name.endswith(".audio_input_layernorm.weight"): + layer_idx = hf_name.split(".")[2] + return f"layers.{layer_idx}.audio_input_layernorm.weight" + if hf_name.startswith("model.layers.") and hf_name.endswith(".audio_post_attention_layernorm.weight"): + layer_idx = hf_name.split(".")[2] + return f"layers.{layer_idx}.audio_post_attention_layernorm.weight" + # Per-layer text-side params now live directly at + # ``layers.{i}.`` (no ``.base`` indirection) since + # HiggsAudioV2DecoderLayer inherits from LlamaDecoderLayer. + if hf_name.startswith("model.layers.") and hf_name.endswith(".input_layernorm.weight"): + layer_idx = hf_name.split(".")[2] + return f"layers.{layer_idx}.input_layernorm.weight" + if hf_name.startswith("model.layers.") and hf_name.endswith(".post_attention_layernorm.weight"): + layer_idx = hf_name.split(".")[2] + return f"layers.{layer_idx}.post_attention_layernorm.weight" + if hf_name.startswith("model.layers.") and hf_name.endswith(".self_attn.o_proj.weight"): + layer_idx = hf_name.split(".")[2] + return f"layers.{layer_idx}.self_attn.o_proj.weight" + if hf_name.startswith("model.layers.") and hf_name.endswith(".mlp.down_proj.weight"): + layer_idx = hf_name.split(".")[2] + return f"layers.{layer_idx}.mlp.down_proj.weight" + if hf_name.startswith("model.layers.") and hf_name.endswith(".audio_mlp.down_proj.weight"): + layer_idx = hf_name.split(".")[2] + return f"layers.{layer_idx}.audio_mlp.down_proj.weight" + return None + + # ----------------------------------------------------------------- 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: + """DualFFN-routed Stage-0 forward. + + Drives the custom :class:`HiggsAudioV2DecoderLayer` stack with a + per-position ``audio_token_mask`` derived from the input ``input_ids`` + (matching :func:`HiggsAudioV2Model.get_placeholder_mask`). Each layer + applies the upstream DualFFN routing internally; this method only + composes the embedding and final-norm steps around the layer loop. + """ + if inputs_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = inputs_embeds + + # Stash the step's input_ids for :meth:`_apply_audio_mode_bias` to read + # in :meth:`sample`. ``sampling_metadata.prompt_token_ids`` is only + # populated when penalties or token-id-aware logits processors are + # active, so for our use case we have to keep our own copy. + if input_ids is not None: + self._last_step_input_ids = input_ids + + # Stash query_start_loc here (inside the forward context) so + # _apply_audio_mode_bias (which runs LATER in sample(), after the + # forward context exits) can recover each request's last token + # position from the flat input_ids — see the per-row prev-token + # lookup below. Without this, mixed prefill+decode batches would + # fall back to ``input_ids[-num_rows:]``, which silently misroutes + # the audio_out_bos detection in batch>1 voice clone (was the root + # cause of the slot-shifted gibberish at batch=4). + 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 + + # Voice-clone ref-audio substitution: at the prompt-side audio + # placeholder span we overwrite the text-embed lookups with the + # multi-codebook embedding of the encoded reference clip carried in + # ``model_intermediate_buffer``. No-op for plain-text requests and + # for decode steps (which fall through to the audio-feedback path). + if input_ids is not None and inputs_embeds is None: + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") + hidden_states = self._maybe_apply_ref_audio_substitution( + hidden_states, + input_ids, + info_dicts, + ) + + # Audio-feedback substitution: vLLM's AR runner calls forward() + # with ``inputs_embeds=None`` and lets the model embed its own + # input_ids inline. ``embed_input_ids`` is only invoked via the + # supports_mm_inputs/prompt_embeds branches, neither of which fires + # for the talker. So we apply the substitution directly here: at + # decode positions where input_ids == audio_token_id, swap the text + # embedding with embed_audio_codes(last audio frame stashed in + # ``self._audio_state[req_id]`` by :meth:`sample`). + if input_ids is not None and inputs_embeds is None: + hidden_states = self._maybe_apply_audio_feedback(hidden_states, input_ids) + + audio_mask = self.audio_token_mask(input_ids) if input_ids is not None else None + + # Fused-residual loop: layers return (hidden, residual) + # where ``residual`` carries everything up to the post-MLP add. + residual: torch.Tensor | None = None + for layer in self.layers: + hidden_states, residual = layer( + positions, + hidden_states, + residual, + audio_token_mask=audio_mask, + ) + + # Final fused norm: combine final mlp_out with carried residual. + norm_out = self.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: + """Sampler-compatible logits tensor. + + Returns a single ``[N, vocab_size]`` tensor that the stock + vLLM AR runner pipes through ``.contiguous()`` and a generic + sampler. The contract is intentionally tensor-only: + + - Text-position logits come from ``lm_head`` (128256-wide Llama vocab). + - Audio-position logits at audio_token_id placeholders are also + read from ``lm_head``; the per-position codebook routing path + lives in :meth:`audio_codebook_logits` and is consumed by a + separate codebook sampling adapter. + """ + # Stash for the audio-sampler dispatch in :meth:`sample`. Limited to a + # single step at a time; cleared in :meth:`sample` after consumption + # so a missed call doesn't leak stale tensors. + self._last_logits_hidden = hidden_states + return self.logits_processor(self.lm_head, hidden_states, sampling_metadata) + + def audio_codebook_logits( + self, + hidden_states: torch.Tensor, + audio_token_mask: torch.Tensor, + ) -> torch.Tensor: + """Per-codebook logits at audio positions, ``[N_audio, num_codebooks, codebook_size]``. + + Runs the single fused ``audio_lm_head: Linear(hidden, num_codebooks * + codebook_size)`` and reshapes the output. This mirrors upstream's + ``HiggsAudioDecoderProjector.audio_lm_head`` + the + ``view(-1, num_codebooks, codebook_size)`` reshape in + ``HiggsAudioModel.forward``. + + Returns an empty tensor with the correct trailing shape when the mask + selects no positions. + """ + num_codebooks = int(self.config.num_codebooks) + codebook_size = int(self.config.codebook_size) + mask = audio_token_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, num_codebooks, codebook_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + flat = self.audio_lm_head(hidden_flat[mask]) + return flat.view(-1, num_codebooks, codebook_size) + + def embed_audio_codes(self, audio_ids: torch.Tensor) -> torch.Tensor: + """Embed a ``[num_codebooks, T]`` audio-ids slice into a ``[T, hidden]`` tensor. + + Mirrors upstream ``HiggsAudioModel._embed_audio_ids``: + + shift = arange(num_codebooks) * codebook_size # 1026 per slot + audio_embed = sum_over_codebooks(audio_codebook_embeddings(audio_ids + shift)) + + The 8 codebooks share a single ``Embedding(8 * codebook_size, hidden)`` + table; each codebook indexes into its own per-codebook slot via the + shift. The output is the per-frame mixture of 8 codebook embeddings + and is what the model re-feeds at the next AR step's + ``<|AUDIO_OUT|>`` placeholder position. + """ + if audio_ids.ndim != 2: + raise ValueError(f"audio_ids must be [num_codebooks, T]; got shape {tuple(audio_ids.shape)}") + codebook_size = int(self.config.codebook_size) + num_codebooks = int(self.config.num_codebooks) + shift = torch.arange(num_codebooks, device=audio_ids.device) * codebook_size + shifted = audio_ids + shift.unsqueeze(-1) # [num_codebooks, T] + embed = self.audio_codebook_embeddings(shifted) # [num_codebooks, T, hidden] + return embed.sum(dim=0) # [T, hidden] + + def _maybe_apply_audio_feedback(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + """Return ``hidden_states`` with audio-token positions replaced by the + multi-codebook embedding of the LAST audio frame stashed by + :meth:`sample`. + + Called from :meth:`forward` when ``inputs_embeds`` is ``None`` (vLLM's + AR runner path) — the equivalent of upstream's + ``_embed_audio_ids(audio_out_ids)`` injection at audio positions. + We use ``index_copy`` (out-of-place) instead of mutating the original + ``hidden_states`` because the embed-tokens output may be a view of a + compiled buffer that doesn't tolerate in-place writes. + """ + audio_token_id = int(self.config.audio_token_id) + flat_ids = input_ids.reshape(-1) + audio_positions = (flat_ids == audio_token_id).nonzero(as_tuple=False).reshape(-1) + if audio_positions.numel() == 0: + return hidden_states + state = getattr(self, "_audio_state", None) + if not state: + return hidden_states + # Build a list of (pos, audio_emb) overrides and apply via index_copy + # on a cloned hidden_states. + rep_positions: list[int] = [] + rep_embeds: list[torch.Tensor] = [] + for pos in audio_positions.tolist(): + req_state = state.get(int(pos)) + if req_state is None: + continue + audio_out_ids = req_state.get("audio_out_ids") + if not isinstance(audio_out_ids, torch.Tensor) or audio_out_ids.numel() == 0: + continue + last_frame = audio_out_ids[:, -1:].to(hidden_states.device) + audio_emb = self.embed_audio_codes(last_frame) # [1, hidden] + rep_positions.append(int(pos)) + rep_embeds.append(audio_emb[0].to(dtype=hidden_states.dtype)) + if not rep_positions: + return hidden_states + new_hidden = hidden_states.clone() + flat_hidden = new_hidden.reshape(-1, new_hidden.shape[-1]) + idx = torch.tensor(rep_positions, dtype=torch.long, device=new_hidden.device) + rep = torch.stack(rep_embeds, dim=0) + flat_hidden.index_copy_(0, idx, rep) + return new_hidden + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Embed the LM token stream, swapping in audio embeddings at placeholders. + + Text positions use the standard ``embed_tokens`` lookup. At audio + positions (input id equals ``audio_token_id``) we substitute the + text embedding with the multi-codebook embedding of the LAST audio + frame stashed by :meth:`sample` in ``self._audio_state[req_id]``. + This mirrors upstream's ``HiggsAudioModel.forward`` flow where + ``_embed_audio_ids(audio_out_ids)`` is plugged into the input at + ``<|AUDIO_OUT|>`` placeholder positions. + + Without this feedback the model receives the SAME embedding (just + the text ``<|AUDIO_OUT|>`` embedding) at every audio step and has no + way to attend to its own audio history — the resulting codes + collapse to a fixed loop and the ASR cannot recover the original + prompt. + """ + text_embed = self.embed_tokens(input_ids) + audio_token_id = int(self.config.audio_token_id) + flat_ids = input_ids.reshape(-1) + audio_positions = (flat_ids == audio_token_id).nonzero(as_tuple=False).reshape(-1) + if audio_positions.numel() == 0: + return text_embed + state = getattr(self, "_audio_state", None) + if not state: + return text_embed + # We expect one audio position per active request per decode step. + # Match rows by walking through the per-request audio_out_ids buffers. + # The runner concatenates requests in batch order; ``audio_positions`` + # lists their indices into the flat ``input_ids``. For each position we + # look up the corresponding request's last audio frame. + # NOTE: this O(N_audio) loop is small (N_audio == active requests in a + # decode step) and runs once per step. + flat_embed = text_embed.reshape(-1, text_embed.shape[-1]) + for pos in audio_positions.tolist(): + # ``state`` keys are the absolute batch row indices the sampler + # saw, which match the positions in flat ``input_ids`` during a + # decode step (one token per active request). For a prefill + # step's prompt-side audio_token_id positions, no state entry + # exists yet — that's fine, we leave the text embed in place. + req_state = state.get(int(pos)) + if req_state is None: + continue + audio_out_ids = req_state.get("audio_out_ids") + if not isinstance(audio_out_ids, torch.Tensor) or audio_out_ids.numel() == 0: + continue + last_frame = audio_out_ids[:, -1:].to(text_embed.device) # [Q, 1] + audio_emb = self.embed_audio_codes(last_frame) # [1, hidden] + flat_embed[pos] = audio_emb[0].to(dtype=text_embed.dtype) + return text_embed + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + """Wrap raw decoder outputs into the :class:`OmniOutput` contract. + + Mirrors the canonical Qwen3-TTS / Fish-Speech recovery pattern: + the runner threads per-request ``codes.audio`` into + ``model_intermediate_buffer`` (a list of dicts in batch order); we + concatenate those and trim ``text_hidden_states`` to the emitted + audio span. Falls back to the deprecated ``runtime_additional_information`` + kwarg for older runners, and finally to ``audio_codes`` / + ``model_kwargs[audio_codes]`` for direct callers. + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + hidden = model_outputs + + # Primary contract: model_intermediate_buffer (Qwen3-TTS / Fish-Speech). + 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 = [] + + # Build a per-request audio-codes list whose length exactly equals + # ``len(info_dicts)`` (== batch size). Requests with no emitted codes + # this step get an empty tensor placeholder — that way the runner's + # ``to_payload_element`` (which indexes ``element[idx]`` per request) + # can never silently fall back to ``element[0]`` for higher slots, + # which was the root of the batch>1 voice-clone slot-shift bug + # (slot k ended up speaking slot k-1's text). + 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}}, + ) + + # Fallbacks: explicit kwarg or wrapped model_kwargs dicts (preserved for + # direct-API callers that don't go through the runner buffer). + audio_codes = kwargs.get("audio_codes") + if audio_codes is None: + for source_name in ("model_kwargs", "model_kwargs_extra"): + source = kwargs.get(source_name) + if isinstance(source, dict) and "audio_codes" in source: + audio_codes = source["audio_codes"] + break + if audio_codes is None: + audio_codes = torch.empty(0, dtype=torch.long) + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={"codes": {"audio": audio_codes}}, + ) + + # ----------------------------------------------------- model-owned sampler + # Opt into the AR runner's model-sampler hook (see + # vllm_omni/worker/gpu_ar_model_runner.py:_sample). The runner calls + # ``self.model.sample(logits, sampling_metadata)`` when this flag is True; + # we delegate to the stock vLLM sampler for now and document the + # per-position audio-codebook-0 dispatch as the immediate follow-up. The + # hook itself is what unblocks an external integrator from plugging in the + # final audio-side sampler without touching the runner. + prefer_model_sampler: bool = True + + # Tell the runner's :func:`extract_multimodal_outputs` to forward the OmniOutput + # produced by :meth:`make_omni_output`. Without this flag the multimodal + # payload (audio codes) is dropped and the downstream stage receives only + # text hidden states. + have_multimodal_outputs: bool = True + + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> Any: + """Model-owned sampler with delay-pattern audio dispatch. + + Each generation step the model emits one LM token. At audio positions + (where the previous LM token was ``audio_bos_token_id`` or + ``audio_token_id``), the same hidden state also drives the fused + ``audio_lm_head`` to produce ``[num_codebooks, codebook_size]`` + per-codebook logits. We argmax those per-codebook, then APPLY THE + UPSTREAM DELAY PATTERN: + + - For the first ``num_codebooks - 1`` audio steps, force codebooks + ``[num_delay + 1 :]`` to ``audio_stream_bos_id``; only codebooks + up to ``num_delay`` are real. + - When ``audio_stream_eos_id`` first appears in a codebook k, force + all earlier codebooks to ``eos`` too and start the trailing + ``num_codebooks - k - 1`` "ramp-down" steps where successive + codebooks also get ``eos``. + - When all 8 codebooks have ramped down, switch the LM next-token + from ``audio_token_id`` to ``audio_eos_token_id`` so the engine + sees a stop signal at the right place. + + Per-request state (``num_delay``, ``num_remaining_delays``, + cumulative ``audio_out_ids: [num_codebooks, T_so_far]``) lives in + ``self._audio_state[req_id]``; ``embed_input_ids`` reads it on the + next step to substitute the audio-token embedding via + :meth:`embed_audio_codes`. + """ + sampler = getattr(self, "_stock_sampler", None) + if sampler is None: + from vllm.v1.sample.sampler import Sampler + + sampler = Sampler() + self._stock_sampler = sampler + + # Bias the text-vocab logits so the LM stays on the audio path. This + # forces ``audio_token_id`` at audio positions; the delay-pattern + # logic below later overrides the LM token to ``audio_eos_token_id`` + # when the ramp-down completes. + self._apply_audio_mode_bias(logits, sampling_metadata) + sampler_output = sampler(logits=logits, sampling_metadata=sampling_metadata) + + hidden = getattr(self, "_last_logits_hidden", None) + self._last_logits_hidden = None + if hidden 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 + + audio_token_id = int(self.config.audio_token_id) + audio_delay_id = int(self.config.audio_delay_token_id) + is_audio = (sampled_flat == audio_token_id) | (sampled_flat == audio_delay_id) + + if not hasattr(self, "_audio_state"): + self._audio_state: dict[int, dict[str, Any]] = {} + if not hasattr(self, "_slot_output_len"): + self._slot_output_len: dict[int, int] = {} + + # Skip the audio-code sampling for rows that just transitioned + # OUT of audio_out_bos. The hidden state at this prefill-end step is + # at the audio_out_bos position, where upstream's audio_out_mask is + # False and no audio_logits are computed. The LM bias above already + # forced ``audio_token_id`` at this step, and the next forward will + # produce a hidden at the new audio_token_id position — that is what + # upstream samples its first audio frame from. Also reset any stale + # ``_audio_state`` and ``_slot_output_len`` for the slot, so a fresh + # request reusing a finished slot starts clean. + first_after_bos = getattr(self, "_last_first_audio_after_bos", None) + self._last_first_audio_after_bos = None + num_codebooks_local = int(self.config.num_codebooks) + bos_local = int(self.config.audio_stream_bos_id) + seeded_rows: list[int] = [] + seeded_frame: torch.Tensor | None = None + # HF seeds frame 0 with all-BOS at the audio_bos position and ignores + # the model's prediction there; the first real audio frame is sampled + # one step later at the audio_token_id position, so we skip the + # prediction at the bos row and emit the all-BOS frame instead. + if isinstance(first_after_bos, torch.Tensor) and first_after_bos.numel() == is_audio.shape[0]: + first_after_bos = first_after_bos.to(is_audio.device) + skip_rows = first_after_bos & is_audio + if bool(skip_rows.any()): + # Upstream's AUDIO_INIT step seeds ``audio_out_ids`` with + # a single all-``audio_stream_bos_id`` frame and the NEXT + # forward feeds ``embed_audio_codes(all_bos)`` at the new + # audio_out_token_idx position. Without this, our audio MLP + # receives a text embedding at the audio_token_id position + # and produces NaN logits. Seed per-slot state with the + # all-BOS frame here, AND emit it as the first output frame + # so the downstream delay-pattern revert sees the leading + # BOS-only frame upstream emits. + seeded_frame = torch.full((num_codebooks_local,), bos_local, dtype=torch.long, device=hidden.device) + for bi in torch.nonzero(skip_rows, as_tuple=False).reshape(-1).tolist(): + bi = int(bi) + # Reset stale state from a finished prior request. + self._slot_output_len[bi] = 0 + self._audio_state[bi] = { + "num_delay": 0, + "num_remaining_delays": None, + "audio_out_ids": seeded_frame.unsqueeze(-1).clone(), # [Q, 1] + } + seeded_rows.append(bi) + is_audio = is_audio & ~first_after_bos + + if not bool(is_audio.any()) and not seeded_rows: + self._last_audio_codes = None + return sampler_output + + # Map audio-position rows to their absolute batch indices. The + # delay-pattern + state-update loop below uses this for per-request + # state lookup; the RAS check between the head sample and the + # delay-pattern apply also reads it. + audio_row_indices = torch.nonzero(is_audio, as_tuple=False).reshape(-1).tolist() + + # Per-codebook logits at audio positions only. + cb_logits = self.audio_codebook_logits(hidden, is_audio) # [N_audio, Q, V] + + # Emulate upstream's ``HiggsAudioV2DelayPatternLogitsProcessor``. + # The raw audio_lm_head has very large weights at index 1025 + # (audio_stream_eos): cb1 |w[1025]|=9.25 vs mean content 1.53 (6x). + # Without masking, the model naturally fires 1025 at codebook 1..6 + # with huge confidence (logit 32 vs runner-up 6) and triggers + # premature ramp-down at frame 2. Upstream's processor MASKS + # forbidden vocab BEFORE sampling so each codebook k is bound to: + # - emit ``audio_stream_bos_id`` (1024) for k leading frames, + # - emit content tokens [0..1023] thereafter, + # - emit ``audio_stream_eos_id`` (1025) only during ramp-down. + bos_pre = int(self.config.audio_stream_bos_id) + eos_pre = int(self.config.audio_stream_eos_id) + for local_i, batch_i in enumerate(audio_row_indices): + state_pre = self._audio_state.get(int(batch_i)) + num_delay_pre = int(state_pre.get("num_delay", 0)) if state_pre else 0 + num_rem_pre = state_pre.get("num_remaining_delays") if state_pre else None + if num_rem_pre is not None: + # Ramp-down: mask all but EOS for the EOS-locked codebooks; the + # remaining ones still emit content. Locked range: codebooks + # [0 : num_codebooks - num_rem_pre]. + lock_until = int(self.config.num_codebooks) - int(num_rem_pre) + for q in range(int(self.config.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: + # Disallow stream specials during ramp-down content frames + cb_logits[local_i, q, bos_pre] = float("-inf") + cb_logits[local_i, q, eos_pre] = float("-inf") + else: + # Leading delay phase. Codebook k starts emitting real content + # at step k (delay = k). Steps with k > num_delay_pre are still + # in BOS phase — force 1024. Steps with k <= num_delay_pre are + # content — disallow 1024 (BOS) and 1025 (EOS). + for q in range(int(self.config.num_codebooks)): + row = cb_logits[local_i, q] + if q > num_delay_pre: + 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") + # Allow EOS at codebook 0 only (it gates ramp-down). + # Disallow at all other content codebooks. + if q != 0: + cb_logits[local_i, q, eos_pre] = float("-inf") + + # Sample per-codebook, following the upstream + # ``HiggsAudioModel._sample_audio_tokens`` pipeline byte-for-byte. + cb_logits_2d = cb_logits.reshape(-1, cb_logits.shape[-1]) + codes_2d = self._sample_audio_codes_upstream( + cb_logits_2d, + temperature=1.0, + top_k=50, + top_p=0.95, + ) + codes_flat = codes_2d.view(cb_logits.shape[0], cb_logits.shape[1]).to(torch.long) + + # (PyTorch align): Repetition-Aware Sampling (RAS). Upstream + # ``HiggsAudioModel._sample_audio_tokens`` re-samples any codebook + # that just emitted the same token >= ras_win_max_num_repeat times + # in the last ras_win_len steps. This breaks codebook collapse loops + # — a 3B model with greedy/low-temp sampling often locks into one + # value otherwise. + ras_win_len = 7 + ras_win_max_num_repeat = 2 + # Codes_flat is [N_audio, Q]; we re-sample PER (audio_position, codebook) + # row that hits the repeat threshold over the last ras_win_len audio + # frames of that codebook's history (from state.audio_out_ids). + if codes_flat.shape[0] > 0: + audio_idx_to_batch = audio_row_indices # already a list[int] + for local_i in range(codes_flat.shape[0]): + bi = int(audio_idx_to_batch[local_i]) + req_state = self._audio_state.get(bi) + if req_state is None: + continue + aoi = req_state.get("audio_out_ids") + if not isinstance(aoi, torch.Tensor) or aoi.shape[-1] < ras_win_len: + continue + window = aoi[:, -ras_win_len:].to(codes_flat.device) # [Q, W] + this_codes = codes_flat[local_i] # [Q] + rep_count = (window == this_codes.unsqueeze(-1)).sum(dim=-1) # [Q] + row_indices = torch.nonzero(rep_count >= ras_win_max_num_repeat, as_tuple=False).reshape(-1) + if row_indices.numel() == 0: + continue + # Re-sample those codebooks from cb_logits directly (no top-k/top-p + # in RAS resample per upstream — pure softmax + multinomial). + resample_logits = cb_logits[local_i, row_indices].float() + resample_probs = resample_logits.softmax(dim=-1) + resampled = torch.multinomial(resample_probs, num_samples=1).squeeze(-1) + codes_flat[local_i, row_indices] = resampled.to(codes_flat.dtype) + + # Apply the upstream delay pattern + EOS ramp-down per request. + num_codebooks = int(self.config.num_codebooks) + bos = int(self.config.audio_stream_bos_id) + eos_stream = int(self.config.audio_stream_eos_id) + audio_eos_vocab = int(getattr(self.config, "audio_eos_token_id", -1)) + + # ``audio_row_indices`` and ``self._audio_state`` were initialised + # earlier (right after ``is_audio`` was computed) so the RAS check + # can reference them. + + # Per-request reset: detect "new request occupied this slot" by + # tracking the per-slot output_token_ids length across sample() + # calls. Within a request the length monotonically grows; when it + # resets to 0 (or goes BELOW the previously-seen high-water mark), + # the previous occupant has finished and a new request now owns + # the slot. Drop the stale ``_audio_state`` for that slot. + if not hasattr(self, "_slot_output_len"): + self._slot_output_len: dict[int, int] = {} + output_ids = getattr(sampling_metadata, "output_token_ids", None) + for batch_i in audio_row_indices: + bi = int(batch_i) + current_len = len(output_ids[bi]) if output_ids is not None and bi < len(output_ids) else 0 + prior_len = self._slot_output_len.get(bi, -1) + # A fresh request has either never been seen (prior_len == -1 + # and we're at step 0) or has length less than the prior + # high-water mark. + is_new_request = (prior_len > current_len) or (prior_len == -1 and current_len == 0) + if is_new_request and bi in self._audio_state: + self._audio_state.pop(bi, None) + self._slot_output_len[bi] = current_len + + # Simplification: we DON'T override ``sampler_output.sampled_token_ids`` + # back to ``audio_eos_token_id`` at ramp-down completion. The vLLM + # ``SamplerOutput.sampled_token_ids`` is plumbed through several + # downstream structures that may have already cached the GPU pointer; + # an in-place mutation here was crashing the worker with exit code + # None (silent CUDA error). Instead, ``max_tokens`` (set in the deploy + # yaml) is the stop bound while the audio span ramps down naturally. + 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}, + ) + num_delay: int = state["num_delay"] + num_remaining_delays: int | None = state["num_remaining_delays"] + this_codes = codes_flat[local_i].clone() # [Q] + + # 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: + # Audio ramp-down finished. Set a per-slot termination flag + # so the NEXT step's bias forces audio_eos_token_id (upstream's + # ``next_tokens[...] = audio_eos_token_id`` override at the + # ramp-completion point — see modeling_higgs_audio.py:1564). + _ = audio_eos_vocab + num_delay = 0 + num_remaining_delays = None + state["num_delay"] = num_delay + state["num_remaining_delays"] = num_remaining_delays + 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 + # Append to the per-request audio_out_ids buffer. + 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) + + # Reassemble the per-position codes tensor with the post-delay-pattern + # values; `-1` rows mark "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 + # Emit the seeded all-BOS frame at the prefill-end (AUDIO_INIT) + # rows so the talker output has the leading delay-pattern frame the + # downstream revert helper expects. + if seeded_rows and seeded_frame is not None: + seeded_rows_idx = torch.tensor(seeded_rows, dtype=torch.long, device=hidden.device) + codes_full.index_copy_( + 0, + seeded_rows_idx, + seeded_frame.unsqueeze(0).expand(len(seeded_rows), -1).to(hidden.device), + ) + + self._last_audio_codes = codes_full + self._postprocess_cursor = 0 + return sampler_output + + @staticmethod + def _sample_audio_codes_upstream( + logits: torch.Tensor, + temperature: float = 1.0, + top_k: int | None = None, + top_p: float | None = None, + ) -> torch.Tensor: + """Replicate the upstream HF ``LogitsProcessorList`` pipeline that + ``HiggsAudioModel._sample_audio_tokens`` calls via + ``logits_processor(None, next_audio_token_logits)``. + + Order (matching transformers' standard ``TemperatureLogitsWarper``, + ``TopKLogitsWarper``, ``TopPLogitsWarper``): + + 1. Divide by temperature. + 2. Top-k: keep the K largest, set others to -inf. + 3. Top-p (nucleus): keep the smallest prefix whose cumulative + softmax probability is <= top_p; mask the rest. Always keep at + least the top-1. + 4. Softmax → multinomial sample. + + ``logits`` shape: ``[N, V]``. Returns ``[N]`` of int sampled indices. + """ + x = logits.float() + if temperature is not None and temperature > 0.0 and temperature != 1.0: + x = x / temperature + if top_k is not None and 0 < top_k < x.shape[-1]: + kth = x.topk(top_k, dim=-1).values[..., -1:] + x = x.masked_fill(x < kth, float("-inf")) + if top_p is not None and 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 + # Shift right by one so we always keep at least the top-1 token. + sorted_mask[..., 1:] = sorted_mask[..., :-1].clone() + sorted_mask[..., 0] = False + # Scatter back to original index space and apply. + mask = torch.zeros_like(x, dtype=torch.bool) + mask.scatter_(-1, sorted_idx, sorted_mask) + x = x.masked_fill(mask, float("-inf")) + if temperature is None or temperature <= 0.0: + return torch.argmax(x, dim=-1) + probs = x.softmax(dim=-1) + # Guard against any all-masked rows (shouldn't happen with the shift + # above) — fall back to argmax of original logits. + valid = probs.sum(dim=-1) > 0 + sampled = torch.empty(x.shape[0], dtype=torch.long, device=x.device) + if valid.any(): + sampled[valid] = torch.multinomial(probs[valid], num_samples=1).squeeze(-1) + if (~valid).any(): + sampled[~valid] = torch.argmax(logits[~valid].float(), dim=-1) + return sampled + + def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> int: + """Mask non-audio tokens at audio-mode positions, in-place on ``logits``. + + Heuristic: per-request, find the last token seen so far (last of + ``output_token_ids[i]`` if non-empty, else last of + ``prompt_token_ids[i]``). If that token is ``audio_bos_token_id`` or + ``audio_token_id``, force the next emit to be one of + ``{audio_token_id, audio_eos_token_id}``. This unblocks live audio + generation since the un-biased argmax over the 128k vocab favours + text-token-id ranges (in particular the global ``eos_token_id``). + + We deliberately do NOT include the standard ``eos_token_id`` in the + allowed set — letting it through would end the whole sequence at the + first step. Stopping inside the audio span is the ``audio_eos`` token's + job; once it fires we drop the bias and the next call falls back to + the stock sampler. + """ + if logits is None or logits.ndim != 2: + return 0 + audio_bos = int(self.config.audio_bos_token_id) + audio_id = int(self.config.audio_token_id) + # (PyTorch align, corrected): upstream NEVER lets the model + # naturally emit ``audio_eos_token_id`` in the LM stream. The LM + # next-token is force-set to ``audio_out_token_idx`` every step; + # the audio span only ends when the delay-pattern ramp-down completes + # (driven by ``audio_stream_eos_id`` appearing in a codebook), at + # which point upstream OVERRIDES the LM token to ``audio_eos``. + # Our diag showed the model's ``audio_eos`` logit is very high right + # after the first audio step; allowing it via sampling stops the + # audio span with 1 frame and the rest of the LM stream goes off-rails. + # Force ONLY ``audio_token_id``. + allowed_extra: list[int] = [] + # Walk per-request to decide which rows to mask. + prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None) + output_ids = getattr(sampling_metadata, "output_token_ids", None) + num_rows = int(logits.shape[0]) + biased = 0 + # Fallback "previous token" source: read each request's LAST input + # position from ``self._last_step_input_ids`` using ``query_start_loc`` + # from the forward context. The previous version of this code took + # ``stash_ids[-num_rows:]`` which only works for pure-decode batches + # (where each row contributes exactly one token to the flat input_ids). + # In MIXED prefill+decode batches the last N tokens of flat input_ids + # all belong to the LAST prefilling request, so the "prev token" for + # other rows came back as random tokens from inside that request's + # prompt — which caused the audio_out_bos detection to fail for some + # rows, the audio_token_id bias to skip them, and the per-row codes to + # collapse to -1 silently (see the batch>1 voice-clone gibberish bug). + stash_ids = getattr(self, "_last_step_input_ids", None) + stash_tail: list[int] | None = None + if isinstance(stash_ids, torch.Tensor) and stash_ids.numel() > 0: + # forward() stashed query_start_loc for us inside the forward + # context — get_forward_context() would fail here because we're + # past the with-block by the time sample() runs. + q_start = getattr(self, "_last_step_query_start_loc", None) + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) == num_rows + 1: + # Per-request last position: q_start_loc[i+1] - 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: + # Pure-decode fallback when query_start_loc is unavailable: + # input_ids has exactly num_rows tokens (one per request) and + # the slice is order-preserving. + stash_tail = stash_ids[-num_rows:].detach().to("cpu").tolist() + # Also stash a per-row "is this the first audio-mode step right + # after audio_out_bos" flag. Upstream NEVER samples audio codes at + # the audio_out_bos position (its audio_out_mask is False there); + # the first audio frame is sampled one step later, at the new + # audio_out_token_idx position. ``sample()`` reads this flag to + # suppress the audio-code sampling at the prefill-end step so the + # first audio frame is sampled from the correct hidden state. + first_after_bos = 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: + # Prompt_token_ids may be a tensor or a list-of-lists + 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 + if prev not in (audio_bos, audio_id): + continue + if prev == audio_bos: + first_after_bos[i] = True + # If this slot's audio ramp finished last step, force LM token + # to audio_eos so the engine stops naturally. Matches upstream's + # ``next_tokens[...] = audio_eos_token_id`` override. + audio_state = getattr(self, "_audio_state", {}) or {} + slot_state = audio_state.get(i) + should_terminate = bool(isinstance(slot_state, dict) and slot_state.get("should_terminate")) + audio_eos_vocab = int(getattr(self.config, "audio_eos_token_id", -1)) + if should_terminate and 0 <= audio_eos_vocab < int(logits.shape[-1]): + # Force audio_eos at this position; upstream override at + # ramp-down completion (modeling_higgs_audio.py:1564). + row = logits[i] + mask = torch.full_like(row, float("-inf")) + mask[audio_eos_vocab] = row[audio_eos_vocab] + logits[i].copy_(mask) + slot_state["should_terminate"] = False + biased += 1 + continue + # Always force ``audio_token_id`` — upstream mirror. + allowed: set[int] = {audio_id, *allowed_extra} + row = logits[i] + mask = torch.full_like(row, float("-inf")) + for tok in allowed: + if 0 <= tok < row.shape[-1]: + mask[tok] = row[tok] + logits[i].copy_(mask) + biased += 1 + self._last_first_audio_after_bos = first_after_bos + return biased + + has_postprocess: bool = True + + 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. + + The runner calls this once per request in batch order. We index into + ``self._last_audio_codes`` (a per-batch ``[num_requests, num_codebooks]`` + tensor produced by :meth:`sample`, one row per request slot regardless + of prompt span length) by a running cursor — exactly one row per + request. Using ``hidden_states_slice.shape[0]`` as the stride is + WRONG at prefill: prompt span length is hundreds of tokens, but each + request still contributes exactly one sampled row to ``_last_audio_codes``. + That bug used to drop the first audio frame at prefill-end (the + all-BOS seed) and, in mixed prefill+decode batches, scrambled the + per-request → codes mapping so batch=4 output came back shifted / + cross-talked across slots. + """ + _ = multimodal_outputs # consumed by the runner directly + codes_full = getattr(self, "_last_audio_codes", None) + if codes_full is None: + return {} + + # The runner walks requests in batch order; pop from the cursor as we go. + # ``_last_audio_codes`` is one row per request slot (== one logits_indices + # row), not one row per prompt token, so advance by 1 per call. + cursor = int(getattr(self, "_postprocess_cursor", 0)) + if cursor >= int(codes_full.shape[0]): + # Defensive: shape drift between forward and postprocess. + self._postprocess_cursor = 0 + return {} + slice_codes = codes_full[cursor : cursor + 1] + self._postprocess_cursor = cursor + 1 + # Drop placeholder rows (-1) — only emit codes for actual audio positions. + audio_rows = slice_codes[:, 0] >= 0 + if not bool(audio_rows.any()): + return {} + new_codes = slice_codes[audio_rows].to(torch.int32) + + # Return ONLY this step's frames. The engine's + # MultimodalOutputProcessor.add_multimodal_tensor accumulates per-step + # OmniOutput.multimodal_outputs across all steps and consolidates at + # request finish via torch.cat(dim=0). Returning cumulative codes here + # caused N(N+1)/2 duplication (prior bug: engine internal accumulation). + return {"codes": {"audio": new_codes}} diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py new file mode 100644 index 000000000..a51d692b9 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TTS prompt builder + scope validators for higgs-audio v2. + +vllm-omni's higgs path supports two request shapes: + +* **Plain text -> 24 kHz speech.** :func:`build_plain_text_prompt` runs the + upstream HF processor with a bare ``"Generate audio following instruction."`` + system prompt; the emitted ``input_ids`` are byte-identical to upstream. +* **Voice clone (shallow).** :func:`build_voice_clone_prompt` runs the + upstream HF processor with both a target text turn and a reference + ``(ref_audio, ref_text)`` ICL turn; the processor returns ``input_ids`` + pre-expanded with audio placeholders plus an ``audio_input_ids`` tensor + carrying the encoded reference codes (HF tokenizer encodes audio on the + spot — no encoder vendored in vllm-omni). + +Still out of scope and rejected with explicit 4xx: multi-speaker +``[SPEAKERn]`` dialogue, ``profile:`` text-only speaker descriptions, the +``ref_audio_in_system_message`` system-block variant, and chunked long-form +generation. +""" + +from __future__ import annotations + +import os +import re +from typing import Any + +import numpy as np +import torch + +__all__ = [ + "UnsupportedInputError", + "MULTI_SPEAKER_TAG_PATTERN", + "validate_plain_text_input", + "build_plain_text_conversation", + "build_plain_text_prompt", + "build_voice_clone_conversation", + "build_voice_clone_prompt", + "input_ids_to_python_list", +] + + +class UnsupportedInputError(ValueError): + """Raised when a request asks for an out-of-scope higgs_audio_v2 feature.""" + + +# Matches the upstream multi-speaker SPEAKERn tag, e.g. [SPEAKER0], [SPEAKER12]. +MULTI_SPEAKER_TAG_PATTERN = re.compile(r"\[SPEAKER\d+\]", re.IGNORECASE) + + +def validate_plain_text_input(text: str) -> None: + """Reject multi-speaker tags inside the user text body. + + Phase-1 explicitly does NOT support multi-speaker dialogue. Catching the + pattern here means the rejection happens at the tokenizer boundary and is + visible to both offline (`pipeline.py`) and online (`serving_speech.py`) + code paths. + """ + if not isinstance(text, str): + raise UnsupportedInputError(f"higgs_audio_v2 expects plain text input; got {type(text).__name__}") + if MULTI_SPEAKER_TAG_PATTERN.search(text): + raise UnsupportedInputError( + "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; received text contains a speaker tag" + ) + + +def build_plain_text_conversation(text: str) -> list[dict[str, Any]]: + """Build the canonical single-speaker plain-text conversation. + + Uses the bare system prompt ``"Generate audio following instruction."`` + that matches the upstream HF reference's input formatting; this exact + wording is required for input-token parity with the upstream processor. + """ + validate_plain_text_input(text) + system_prompt = "Generate audio following instruction." + return [ + { + "role": "system", + "content": [{"type": "text", "text": system_prompt}], + }, + { + "role": "user", + "content": [{"type": "text", "text": text}], + }, + ] + + +def build_plain_text_prompt( + processor: Any, + text: str, + *, + sampling_rate: int = 24000, + return_tensors: str | None = "pt", +) -> dict[str, Any]: + """Run the upstream processor's chat template on a plain-text input. + + Returns the processor output dict (``input_ids`` plus any auxiliary tensors + such as ``attention_mask``). The serving layer passes ``input_ids`` to + Stage 0 as ``prompt_token_ids`` after a ``.tolist()``. + + Using the upstream processor verbatim (no system-prompt rewriting) is + what preserves input-token parity with the HF reference. + """ + conversation = build_plain_text_conversation(text) + inputs = processor.apply_chat_template( + conversation, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + sampling_rate=sampling_rate, + return_tensors=return_tensors, + ) + if "input_ids" not in inputs: + raise RuntimeError(f"HiggsAudioV2 processor returned no input_ids; got keys {list(inputs.keys())!r}") + return inputs + + +def input_ids_to_python_list(inputs: dict[str, Any]) -> list[int]: + """Convenience: pull a flat ``list[int]`` of token IDs from a processor output.""" + ids = inputs["input_ids"] + if isinstance(ids, torch.Tensor): + if ids.ndim == 2 and int(ids.shape[0]) != 1: + raise ValueError(f"expected batch=1 prompt; got input_ids shape {tuple(ids.shape)}") + return ids.reshape(-1).tolist() + return list(ids) + + +_AUDIO_OUT_TOKEN = "<|AUDIO_OUT|>" +_AUDIO_OUT_BOS_TOKEN = "<|audio_out_bos|>" +_AUDIO_EOS_TOKEN = "<|audio_eos|>" +_AUDIO_DELAY_TOKEN = "<|reserved_special_token_6|>" +_AUDIO_STREAM_BOS_ID = 1024 +_AUDIO_STREAM_EOS_ID = 1025 + + +def _build_delay_pattern(codes: torch.Tensor) -> torch.Tensor: + """Apply the upstream delay-pattern wrap to a ``[num_codebooks, T]`` code tensor. + + Mirrors ``HiggsAudioV2Processor.build_delay_pattern``: prepend BOS, append + EOS, then arrange in a triangular pattern that stretches the sequence by + ``num_codebooks - 1`` frames so codebook k starts emitting real codes at + frame k. + """ + num_codebooks, seq_len = codes.shape + bos = codes.new_full((num_codebooks, 1), _AUDIO_STREAM_BOS_ID) + eos = codes.new_full((num_codebooks, 1), _AUDIO_STREAM_EOS_ID) + wrapped = torch.cat([bos, codes, eos], dim=1) + wrapped_len = wrapped.shape[1] + new_seq_len = wrapped_len + num_codebooks - 1 + + output = torch.ones((1, num_codebooks, new_seq_len), dtype=codes.dtype, device=codes.device) + bos_mask = torch.tril(output, -1) > 0 + eos_mask = torch.triu(output, wrapped_len) > 0 + data_mask = ~(bos_mask | eos_mask) + output[bos_mask] = _AUDIO_STREAM_BOS_ID + output[data_mask] = wrapped.reshape(-1) + output[eos_mask] = _AUDIO_STREAM_EOS_ID + return output[0] + + +_ENCODER_CACHE: Any | None = None +_K2_OMNIVOICE_REPO = "k2-fsa/OmniVoice" +_K2_OMNIVOICE_SUBDIR = "audio_tokenizer" + + +def _load_audio_tokenizer(): + """Load the higgs-audio v2 codec via HF ``HiggsAudioV2TokenizerModel``. + + The boson-ai standalone tokenizer repo (``bosonai/higgs-audio-v2-tokenizer``) + ships a ``model.safetensors`` that is actually a copy of the 3B talker LM, + not the codec — pointing HF's ``from_pretrained`` at it yields an + encoder/decoder with randomly initialized structural weights, producing + noise codes at encode time. The k2-fsa OmniVoice repo bundles the same + codec weights repackaged under + ``audio_tokenizer/{config.json,model.safetensors,preprocessor_config.json}`` + with key naming aligned to ``HiggsAudioV2TokenizerModel``'s class layout + (``acoustic_encoder.*`` / ``acoustic_decoder.*`` / ``quantizer.quantizers.*`` + / ``fc/fc2`` / ``semantic_model.*``), so HF can load it directly. + + Only the ``audio_tokenizer/`` subdirectory (~806 MB) is downloaded. + """ + from huggingface_hub import snapshot_download + from transformers import HiggsAudioV2TokenizerModel + + repo_path = snapshot_download( + _K2_OMNIVOICE_REPO, + allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], + ) + audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR) + device = "cuda" if torch.cuda.is_available() else "cpu" + model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device) + return model.eval() + + +def _encode_ref_audio_codes( + wav: np.ndarray, + sr: int, +) -> torch.Tensor: + """Encode a single ref clip to codec codes via HF ``HiggsAudioV2TokenizerModel``. + + Returns shape ``[num_codebooks, T_raw]`` (before BOS/EOS + delay-pattern wrap). + """ + import torchaudio + + global _ENCODER_CACHE + if _ENCODER_CACHE is None: + try: + _ENCODER_CACHE = _load_audio_tokenizer() + except ImportError as exc: + raise RuntimeError( + "higgs_audio_v2 voice clone needs `transformers>=5.3.0` " + "(which ships the HiggsAudioV2TokenizerModel class). " + "Install via `pip install -U 'transformers>=5.3.0'`." + ) from exc + + target_sr = int(getattr(_ENCODER_CACHE.config, "sample_rate", 24000)) + wav_t = torch.as_tensor(wav, dtype=torch.float32).reshape(-1) + if sr and sr != target_sr: + wav_t = torchaudio.functional.resample(wav_t, sr, target_sr) + # HF encode() expects (batch=1, channels=1, num_samples). + input_values = wav_t.unsqueeze(0).unsqueeze(0).to(_ENCODER_CACHE.device) + with torch.inference_mode(): + codes = _ENCODER_CACHE.encode(input_values, return_dict=False) + # codes shape: (batch=1, num_quantizers, codes_length) -> (num_quantizers, T). + if codes.dim() == 3: + codes = codes.squeeze(0) + return codes.detach().to("cpu").long() + + +def build_voice_clone_conversation( + text: str, + ref_text: str, +) -> list[dict[str, Any]]: + """ChatML conversation for voice clone (HF jinja-template-compatible). + + The assistant turn is a list with a single ``{"type": "audio"}`` content + block — no ``"audio"``/``"url"``/``"path"`` key, so ``apply_chat_template`` + won't try to extract audio data (that path would collide with our explicit + audio encoder). The template renders the assistant block as the literal + ``<|audio_out_bos|><|AUDIO_OUT|><|audio_eos|>``; we expand the single + ``<|AUDIO_OUT|>`` token to ``N × audio_token + (num_codebooks-1) × delay_token`` + in :func:`build_voice_clone_prompt` to match the reference clip's frame count. + """ + validate_plain_text_input(text) + validate_plain_text_input(ref_text) + return [ + {"role": "system", "content": "Generate audio following instruction."}, + {"role": "user", "content": ref_text}, + {"role": "assistant", "content": [{"type": "audio"}]}, + {"role": "user", "content": text}, + ] + + +def build_voice_clone_prompt( + processor: Any, + text: str, + ref_audio_wav: np.ndarray | torch.Tensor, + ref_audio_sr: int, + ref_text: str, + *, + return_tensors: str | None = "pt", +) -> dict[str, Any]: + """Build a voice-clone prompt using the upstream boson audio tokenizer. + + Returns a dict carrying: + - ``prompt_token_ids``: ``list[int]`` — input_ids with the ref-audio + ``<|AUDIO_OUT|>`` + ``<|reserved_special_token_6|>`` placeholders + already expanded to match the encoded reference clip's frame count. + - ``audio_input_ids``: ``Tensor[T_frames, num_codebooks]`` — encoded + reference codes with BOS/EOS + delay pattern, transposed to ``[T, Q]`` + for parity with the HF processor output shape. + - ``audio_input_ids_mask``: ``Tensor[T_frames]`` — all-True bool mask. + """ + if isinstance(ref_audio_wav, torch.Tensor): + wav = ref_audio_wav.detach().to("cpu").float().reshape(-1).numpy() + else: + wav = np.asarray(ref_audio_wav, dtype=np.float32).reshape(-1) + + # 1. Encode via the upstream boson tokenizer (real weights from model.pth). + codes_qt = _encode_ref_audio_codes(wav, ref_audio_sr) # [num_codebooks, T_raw] + if codes_qt.ndim == 3: + codes_qt = codes_qt[0] + num_codebooks = int(codes_qt.shape[0]) + + # 2. Apply BOS/EOS + delay-pattern wrap. + audio_input_ids = _build_delay_pattern(codes_qt) # [num_codebooks, T_full] + T_full = int(audio_input_ids.shape[1]) + + # 3. Render the chat template, then expand the single <|AUDIO_OUT|> marker + # in the assistant turn into the full placeholder block. After expansion + # the text contains exactly T_full audio-mask positions, matching the + # audio_input_ids time axis. + conversation = build_voice_clone_conversation(text, ref_text) + rendered = processor.tokenizer.apply_chat_template( + conversation, + add_generation_prompt=True, + tokenize=False, + ) + n_delay = num_codebooks - 1 + n_audio = T_full - n_delay + if n_audio < 0: + raise RuntimeError(f"ref clip too short ({T_full} frames) for delay pattern with num_codebooks={num_codebooks}") + placeholders = _AUDIO_OUT_TOKEN * n_audio + _AUDIO_DELAY_TOKEN * n_delay + expanded_assistant = f"{_AUDIO_OUT_BOS_TOKEN}{placeholders}{_AUDIO_EOS_TOKEN}" + if _AUDIO_OUT_TOKEN not in rendered: + raise RuntimeError( + f"Voice-clone chat-template render is missing the assistant " + f"audio placeholder marker. conversation={conversation!r}" + ) + rendered = rendered.replace(_AUDIO_OUT_TOKEN, expanded_assistant, 1) + + # 4. Tokenize the rendered prompt. ``apply_chat_template`` already emits + # ``<|begin_of_text|>``, so disable add_special_tokens. + encoded = processor.tokenizer( + rendered, + add_special_tokens=False, + return_tensors=return_tensors, + ) + prompt_token_ids = encoded["input_ids"] + if isinstance(prompt_token_ids, torch.Tensor): + prompt_token_ids = prompt_token_ids.reshape(-1).tolist() + + # 5. Transpose codes to [T_full, num_codebooks] for HF parity. + audio_input_ids_t = audio_input_ids.transpose(0, 1).contiguous().to(torch.long) + audio_input_ids_mask = torch.ones(T_full, dtype=torch.bool) + + return { + "prompt_token_ids": prompt_token_ids, + "audio_input_ids": audio_input_ids_t, + "audio_input_ids_mask": audio_input_ids_mask, + } diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py b/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py new file mode 100644 index 000000000..fb6832569 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""higgs-audio v2 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM).""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v2" + +HIGGS_AUDIO_V2_PIPELINE = PipelineConfig( + model_type="higgs_audio_v2", + model_arch="HiggsAudioV2ForConditionalGeneration", + hf_architectures=("HiggsAudioV2ForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="higgs_audio_v2", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"), + sampling_constraints={ + "detokenize": False, + # Two stop signals for Stage-0: + # * 128009 = standard LM eos_token_id (sequence-level stop) + # * 128012 = audio_eos_token_id (forced by the talker's + # _apply_audio_mode_bias at audio-ramp completion; matches + # upstream HiggsAudioModel._sample_audio_tokens override). + "stop_token_ids": [128009, 128012], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="code2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="HiggsAudioV2Code2WavForConditionalGeneration", + sync_process_input_func=f"{_PROC}.talker2code2wav", + sampling_constraints={"detokenize": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 4c49afc36..ff61d6e6c 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -102,6 +102,22 @@ _OMNI_MODELS = { "qwen3_tts_code2wav", "Qwen3TTSCode2Wav", ), + ## higgs-audio v2 + "HiggsAudioV2ForConditionalGeneration": ( + "higgs_audio_v2", + "higgs_audio_v2_talker", + "HiggsAudioV2TalkerForConditionalGeneration", + ), + "HiggsAudioV2TalkerForConditionalGeneration": ( + "higgs_audio_v2", + "higgs_audio_v2_talker", + "HiggsAudioV2TalkerForConditionalGeneration", + ), + "HiggsAudioV2Code2WavForConditionalGeneration": ( + "higgs_audio_v2", + "higgs_audio_v2_code2wav", + "HiggsAudioV2Code2WavForConditionalGeneration", + ), ## mimo_audio "MiMoAudioModel": ( "mimo_audio", diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py new file mode 100644 index 000000000..0d6781586 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-input processor for higgs-audio v2: Talker -> Code2Wav. + +Mirrors the qwen3_tts contract: + +- ``talker2code2wav`` is the sync (non-async-chunk) adapter that runs once, + after Stage 0 has finished, collecting every audio frame into a single + ``OmniTokensPrompt`` for Stage 1. +- ``talker2code2wav_async_chunk`` is the streaming adapter invoked once per + emitted codec frame; it accumulates frames in a per-request buffer and + flushes an ``OmniPayloadStruct`` once a chunk is ready (or at end-of-stream). + +Stage 1 strictly rejects any code id >= ``audio_stream_bos_id`` (1024). This +processor filters frames that still carry stream specials so the codec sees +only real codes in [0, 1023], regardless of where Stage 0 emitted them. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import ( + CodesStruct, + MetaStruct, + OmniPayload, + OmniPayloadStruct, +) + +__all__ = ["talker2code2wav", "talker2code2wav_async_chunk"] + +logger = init_logger(__name__) + +# Default codec chunk-shape knobs for async_chunk mode. These match the +# canonical Qwen3-TTS defaults (25 frames = 1 s at 25 fps). Real values +# come from the connector config; these are only fallbacks. +_DEFAULT_CODEC_CHUNK_FRAMES = 25 +_DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 25 + +# Boson-ai constants. Kept local (rather than imported from the model package) +# so this module remains usable without circular imports during +# stage-input-processor discovery. +_NUM_CODEBOOKS = 8 +_AUDIO_STREAM_BOS_ID = 1024 +_AUDIO_STREAM_EOS_ID = 1025 +_NUM_REAL_CODES = _AUDIO_STREAM_BOS_ID # codes in [0, 1023] are real + + +def _filter_real_code_frames(audio_codes: torch.Tensor) -> torch.Tensor: + """Keep only frames whose codes are entirely in [0, 1023]. + + ``audio_codes`` has shape ``[num_frames, num_codebooks]`` after the + standard Stage-0 transpose. We drop frames containing stream specials + or anything padded outside the real-code range. + """ + if audio_codes.numel() == 0: + return audio_codes + if audio_codes.ndim != 2: + raise ValueError(f"expected [num_frames, num_codebooks] audio_codes; got shape {tuple(audio_codes.shape)}") + valid = (audio_codes >= 0).all(dim=1) & (audio_codes < _NUM_REAL_CODES).all(dim=1) + return audio_codes[valid] + + +def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: + """Reverse the delay-pattern shift, mirroring upstream's + ``boson_multimodal.model.higgs_audio.utils.revert_delay_pattern``. + + Input shape: ``[num_codebooks, seq_len + num_codebooks - 1]``. + Output shape: ``[num_codebooks, seq_len]``. + + For each codebook ``i``, slice rows ``[i:i+1]`` from columns + ``[i : seq_len + i]`` to remove ``i`` leading BOS pads and ``Q-1-i`` + trailing PAD entries. + """ + if audio_codes_qt.ndim != 2: + raise ValueError(f"_revert_delay_pattern expects [Q, T] input; got {tuple(audio_codes_qt.shape)}") + q, t = audio_codes_qt.shape + if t < q: + # Not enough frames to revert delay pattern; return as-is. + return audio_codes_qt + seq_len = t - q + 1 + out_l = [] + for i in range(q): + out_l.append(audio_codes_qt[i : i + 1, i : seq_len + i]) + return torch.cat(out_l, dim=0) + + +def talker2code2wav( + source_outputs: list[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list[Any]: + """Sync: collect all talker codes, then pass to code2wav at once.""" + from vllm_omni.inputs.data import OmniTokensPrompt + + code2wav_inputs: list[OmniTokensPrompt] = [] + for talker_output in source_outputs: + if not talker_output.finished: + continue + output = talker_output.outputs[0] + mm = output.multimodal_output + mm_codes = mm.get("codes", {}) + + audio_codes = mm_codes.get("audio") + if audio_codes is None or not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + # Nothing to decode for this request; emit an empty payload so the + # downstream stage can still close the response. + code2wav_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[], + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=None, + ) + ) + continue + + audio_codes = audio_codes.to(torch.long) + if audio_codes.ndim == 1: + # Stage 0 emitted a flat sequence; reshape to [num_frames, num_codebooks]. + if audio_codes.numel() % _NUM_CODEBOOKS != 0: + raise ValueError( + f"flat audio_codes length {audio_codes.numel()} not divisible by num_codebooks={_NUM_CODEBOOKS}" + ) + audio_codes = audio_codes.reshape(-1, _NUM_CODEBOOKS) + + if audio_codes.ndim != 2: + raise ValueError(f"audio_codes must be 1D or 2D; got shape {tuple(audio_codes.shape)}") + + # Audio_codes from Stage-0 are emitted in DELAY-PATTERN layout. + # Revert the shift to recover the canonical ``[Q, seq_len]`` form, + # then clip to the real-code range and trim the first/last frame + # (the codec doesn't consume the BOS/EOS edge frames). This mirrors + # the upstream ``HiggsAudioServeEngine.generate`` decode line: + # vq_code = revert_delay_pattern(out).clip(0, real-1)[:, 1:-1] + codes_qt = audio_codes.transpose(0, 1).contiguous().cpu() + codes_qt = _revert_delay_pattern(codes_qt) + codes_qt = codes_qt.clamp_(min=0, max=_NUM_REAL_CODES - 1) + if codes_qt.shape[-1] >= 3: + codes_qt = codes_qt[:, 1:-1] + # Code2Wav expects codebook-major flat: [Q * num_frames]. + codec_codes = codes_qt.reshape(-1).tolist() + + code2wav_inputs.append( + OmniTokensPrompt( + prompt_token_ids=codec_codes, + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=None, + ) + ) + return code2wav_inputs + + +def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None: + audio_codes = pooling_output.get("codes", {}).get("audio") + if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + return None + if audio_codes.ndim == 2: + frame = audio_codes[-1] + elif audio_codes.ndim == 1: + frame = audio_codes + else: + raise ValueError(f"unexpected audio_codes shape for higgs_audio_v2 async_chunk: {tuple(audio_codes.shape)}") + if frame.numel() == 0: + return None + frame = frame.to(torch.long).reshape(-1) + # Filter frames that still carry stream specials. + if int(frame.min().item()) < 0 or int(frame.max().item()) >= _NUM_REAL_CODES: + return None + return frame + + +def talker2code2wav_async_chunk( + transfer_manager: Any, + pooling_output: OmniPayload | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Streaming adapter: buffer frames per request and flush a chunk at a time. + + Emits codes in codebook-major flat layout (``[Q * num_frames]``) plus a + ``MetaStruct(left_context_size, finished)`` so Stage 1 can stitch chunks + together with the conventional left-context overlap. + """ + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + if isinstance(pooling_output, dict): + frame = _extract_last_frame(pooling_output) + if frame is not None: + transfer_manager.code_prompt_token_ids[request_id].append(frame.cpu().tolist()) + elif not finished: + return None + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + chunk_size = int(cfg.get("codec_chunk_frames", _DEFAULT_CODEC_CHUNK_FRAMES)) + left_context_size_config = int(cfg.get("codec_left_context_frames", _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES)) + + if chunk_size <= 0 or left_context_size_config < 0: + raise ValueError( + f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, " + f"codec_left_context_frames={left_context_size_config}" + ) + + length = len(transfer_manager.code_prompt_token_ids[request_id]) + if length <= 0: + if finished: + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + return None + + chunk_remainder = length % chunk_size + if not finished and chunk_remainder != 0: + # Wait until we have a clean chunk boundary. + return None + + context_length = chunk_remainder if (finished and chunk_remainder > 0) else chunk_size + end_index = min(length, left_context_size_config + context_length) + left_context_size = max(0, end_index - context_length) + window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + + if not window_frames: + # Defensive: should not happen because length > 0, but keeps the contract. + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(finished, dtype=torch.bool)), + ) + + num_codebooks = len(window_frames[0]) + if num_codebooks != _NUM_CODEBOOKS: + raise ValueError(f"expected {_NUM_CODEBOOKS} codebooks per frame; got {num_codebooks}") + num_frames = len(window_frames) + flat = torch.tensor( + [window_frames[f][q] for q in range(num_codebooks) for f in range(num_frames)], + dtype=torch.long, + ) + + return OmniPayloadStruct( + codes=CodesStruct(audio=flat), + meta=MetaStruct( + left_context_size=left_context_size, + finished=torch.tensor(finished, dtype=torch.bool), + ), + )