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

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

vllm_omni/model_executor/models/higgs_audio_v2/__init__.py

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

vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py

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

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py

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

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py

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

vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py

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

vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py

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

vllm_omni/model_executor/models/registry.py

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