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