vllm-omni/vllm_omni/model_executor/models/audex/audex_xcodec.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Audex TTA stage 1: XCodec1 RVQ codes → waveform.

Unlike the streaming causal speech decoder used for TTS, XCodec1 is a CNN
codec that the official flow decodes over the full sequence, so this stage
runs sync full-payload only: each request delivers its complete
de-interleaved ``[frames, 4]`` codec-id payload once at finish and the
decode happens in a single call.

The checkpoint (``hf-audio/xcodec-hubert-general-balanced``) is external to
the Audex repo and loads through transformers remote code; the model
skeleton is built from config at startup for vLLM's memory profiler and the
weights arrive via the standard ``load_weights`` iteration over the
snapshot's safetensors.
"""

from collections.abc import Iterable
from typing import Any

import torch
from torch import nn
from vllm.config import VllmConfig
from vllm.logger import init_logger

from vllm_omni.model_executor.models.audex.tta import (
    XCODEC1_CODEBOOK_SIZE,
    XCODEC1_NUM_CODEBOOKS,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput

logger = init_logger(__name__)

# Official cap on decoded frames (~10 s of audio at the codec frame rate).
_DEFAULT_MAX_TTA_FRAMES = 500


def _meta_str(value: Any) -> str | None:
    if isinstance(value, list | tuple):
        return _meta_str(value[0]) if value else None
    if isinstance(value, str):
        return value
    return None


def _codes_from_runtime_info(runtime_info: dict[str, Any] | None) -> torch.Tensor | None:
    """Extract the ``[frames, 4]`` codec-id payload from the connector info."""
    if not isinstance(runtime_info, dict):
        return None
    codes = runtime_info.get("codes")
    if not isinstance(codes, dict):
        return None
    audio = codes.get("audio")
    if isinstance(audio, torch.Tensor):
        return audio
    if isinstance(audio, list | tuple) and audio:
        return torch.as_tensor(audio, dtype=torch.long)
    return None


class AudexXCodec1(nn.Module):
    """Stage-1 model for Audex TTA (GenerationModelRunner)."""

    input_modalities = "audio"

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        self.vllm_config = vllm_config
        self.model_path = vllm_config.model_config.model

        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

        from transformers import AutoConfig, AutoModel

        config = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True)
        self.codec = AutoModel.from_config(config, trust_remote_code=True)
        self.codec.eval()
        self._sample_rate = int(getattr(config, "sampling_rate", 16000))
        self._max_frames = int(getattr(config, "max_tta_frames", _DEFAULT_MAX_TTA_FRAMES))

    # -------------------- vLLM model interface --------------------

    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
        # This stage ignores token embeddings; keep a stable dummy for the runner.
        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: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
        return None

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        state_dict = dict(weights)
        self.codec.load_state_dict(state_dict, strict=True)
        device = self.vllm_config.device_config.device
        self.codec.to(device=device, dtype=torch.float32)
        return {f"codec.{name}" for name in state_dict}

    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
        if isinstance(model_outputs, OmniOutput):
            return model_outputs
        # The runner may unpack the NamedTuple in transit; rebuild it.
        if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields):
            return OmniOutput(*model_outputs)
        raise TypeError(f"AudexXCodec1 expected OmniOutput, got {type(model_outputs)}")

    # -------------------- decoding --------------------

    def _decode_codes(self, req_id: str | None, codes: torch.Tensor) -> torch.Tensor:
        """Decode a ``[frames, 4]`` codec-id payload to a mono waveform."""
        if codes.numel() == 0:
            # Zero codec tokens for the whole request. Do NOT raise: an
            # exception in forward() kills the engine core for every request.
            # Returning empty audio lets the serving layer fail this request.
            logger.error("Audex TTA: request %s finished with no codec tokens; no audio generated", req_id)
            return torch.empty(0, dtype=torch.float32)
        if codes.ndim != 2 or codes.shape[1] != XCODEC1_NUM_CODEBOOKS:
            logger.error(
                "Audex TTA: request %s payload has shape %s, expected [frames, %d]; no audio generated",
                req_id,
                tuple(codes.shape),
                XCODEC1_NUM_CODEBOOKS,
            )
            return torch.empty(0, dtype=torch.float32)

        device = self.vllm_config.device_config.device
        with torch.inference_mode():
            frames = codes[: self._max_frames].to(device=device, dtype=torch.long)
            # [frames, 4] codec ids -> [1, 4, frames] per-layer indices.
            layer_offsets = (
                torch.arange(XCODEC1_NUM_CODEBOOKS, device=frames.device, dtype=frames.dtype) * XCODEC1_CODEBOOK_SIZE
            )
            per_layer = (frames - layer_offsets.view(1, -1)).clamp(0, XCODEC1_CODEBOOK_SIZE - 1)
            stacked = per_layer.transpose(0, 1).unsqueeze(0)
            decoded = self.codec.decode(stacked)
            audio = decoded.audio_values if hasattr(decoded, "audio_values") else decoded[0]
            return audio.reshape(-1).to(dtype=torch.float32, device="cpu")

    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,
    ) -> OmniOutput:
        runtime_infos = runtime_additional_information or []
        seq_token_counts = kwargs.get("seq_token_counts")
        if seq_token_counts is not None:
            num_reqs = len(seq_token_counts)
        elif input_ids is not None and input_ids.numel() > 0 and not runtime_infos:
            num_reqs = 1
        else:
            num_reqs = len(runtime_infos)

        audios: list[torch.Tensor] = []
        sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32)
        for i in range(num_reqs):
            info = runtime_infos[i] if i < len(runtime_infos) else {}
            info = info if isinstance(info, dict) else {}
            meta = info.get("meta", {})
            meta = meta if isinstance(meta, dict) else {}
            req_id = _meta_str(meta.get("req_id"))
            codes = _codes_from_runtime_info(info)
            if codes is None:
                audios.append(torch.empty(0, dtype=torch.float32))
                continue
            audios.append(self._decode_codes(req_id, codes))

        return OmniOutput(
            text_hidden_states=None,
            multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_reqs},
        )