vllm-omni/vllm_omni/model_executor/models/voxcpm/voxcpm.py

from __future__ import annotations

import json
import os
import sys
import tempfile
import warnings
import wave
from collections.abc import Callable, Generator, Iterable
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from einops import rearrange
from tqdm import tqdm
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.sequence import IntermediateTensors

from vllm_omni.model_executor.models.output_templates import OmniOutput

from .voxcpm_loader import (
    _build_prompt_cache_with_soundfile,
    _device_to_string,
    _force_cuda_available_for_npu,
    _import_voxcpm_audio_vae_classes,
    _import_voxcpm_base_model_class,
    _is_torchcodec_load_error,
    _normalize_dtype_name,
    _prepare_runtime_model_dir,
    _resolve_runtime_device,
)
from .voxcpm_runtime_utils import resolve_voxcpm_model_dir
from .voxcpm_stage_wrappers import _DirectVoxCPMAudioVAE, _DirectVoxCPMLatentGenerator

logger = init_logger(__name__)
_VOXCPM_LATENT_MAGIC = 131071


def _make_voxcpm_model_for_omni(base: type[Any]) -> type[Any]:
    """Subclass upstream VoxCPMModel: local ``_inference`` + ``latents_only`` prompt-cache generation."""

    from voxcpm.model.utils import get_dtype

    class VoxCPMModelForOmni(base):
        @torch.inference_mode()
        def build_prompt_cache(self, *args: Any, **kwargs: Any):
            try:
                return super().build_prompt_cache(*args, **kwargs)
            except (ImportError, ModuleNotFoundError, RuntimeError) as exc:
                if not _is_torchcodec_load_error(exc):
                    raise
                return _build_prompt_cache_with_soundfile(self, *args, **kwargs)

        @torch.inference_mode()
        def _inference(
            self,
            text: torch.Tensor,
            text_mask: torch.Tensor,
            feat: torch.Tensor,
            feat_mask: torch.Tensor,
            min_len: int = 2,
            max_len: int = 2000,
            inference_timesteps: int = 10,
            cfg_value: float = 2.0,
            streaming: bool = False,
            streaming_prefix_len: int = 3,
        ) -> Generator[tuple[torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]:
            B, _, _, _ = feat.shape

            feat_embed = self.feat_encoder(feat)
            feat_embed = self.enc_to_lm_proj(feat_embed)

            scale_emb = self.config.lm_config.scale_emb if self.config.lm_config.use_mup else 1.0
            text_embed = self.base_lm.embed_tokens(text) * scale_emb
            combined_embed = text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed

            prefix_feat_cond = feat[:, -1, ...]
            pred_feat_seq: list[torch.Tensor] = []

            audio_patch_count = int(feat_mask.sum().item())
            if audio_patch_count > 0:
                context_len = min(streaming_prefix_len - 1, audio_patch_count)
                prompt_context_patches = list(feat[:, -context_len:, :, :].split(1, dim=1))
                pred_feat_seq = prompt_context_patches + pred_feat_seq

            enc_outputs, kv_cache_tuple = self.base_lm(
                inputs_embeds=combined_embed,
                is_causal=True,
            )
            self.base_lm.kv_cache.fill_caches(kv_cache_tuple)

            enc_outputs = self.fsq_layer(enc_outputs) * feat_mask.unsqueeze(-1) + enc_outputs * text_mask.unsqueeze(-1)
            lm_hidden = enc_outputs[:, -1, :]

            residual_enc_outputs, residual_kv_cache_tuple = self.residual_lm(
                inputs_embeds=enc_outputs + feat_mask.unsqueeze(-1) * feat_embed,
                is_causal=True,
            )
            self.residual_lm.kv_cache.fill_caches(residual_kv_cache_tuple)
            residual_hidden = residual_enc_outputs[:, -1, :]

            for step_idx in tqdm(range(max_len)):
                dit_hidden = self.lm_to_dit_proj(lm_hidden) + self.res_to_dit_proj(residual_hidden)
                pred_feat = self.feat_decoder(
                    mu=dit_hidden,
                    patch_size=self.patch_size,
                    cond=prefix_feat_cond.transpose(1, 2).contiguous(),
                    n_timesteps=inference_timesteps,
                    cfg_value=cfg_value,
                ).transpose(1, 2)

                curr_embed = self.enc_to_lm_proj(self.feat_encoder(pred_feat.unsqueeze(1)))
                pred_feat_seq.append(pred_feat.unsqueeze(1))
                prefix_feat_cond = pred_feat

                if streaming:
                    pred_feat_chunk = torch.cat(pred_feat_seq[-streaming_prefix_len:], dim=1)
                    feat_pred = rearrange(pred_feat_chunk, "b t p d -> b d (t p)", b=B, p=self.patch_size)
                    yield feat_pred, pred_feat_seq

                stop_flag = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden))).argmax(dim=-1)[0].cpu().item()
                if step_idx > min_len and stop_flag == 1:
                    break

                lm_hidden = self.base_lm.forward_step(
                    curr_embed[:, 0, :],
                    torch.tensor([self.base_lm.kv_cache.step()], device=curr_embed.device),
                ).clone()
                lm_hidden = self.fsq_layer(lm_hidden)
                residual_hidden = self.residual_lm.forward_step(
                    lm_hidden + curr_embed[:, 0, :],
                    torch.tensor([self.residual_lm.kv_cache.step()], device=curr_embed.device),
                ).clone()

            if not streaming:
                pred_feat_seq_cat = torch.cat(pred_feat_seq, dim=1)
                feat_pred = rearrange(pred_feat_seq_cat, "b t p d -> b d (t p)", b=B, p=self.patch_size)
                yield feat_pred, pred_feat_seq_cat.squeeze(0).cpu()

        @torch.inference_mode()
        def generate_latents_with_prompt_cache(
            self,
            target_text: str,
            prompt_cache: dict,
            min_len: int = 2,
            max_len: int = 2000,
            inference_timesteps: int = 10,
            cfg_value: float = 2.0,
            retry_badcase: bool = False,
            retry_badcase_max_times: int = 3,
            retry_badcase_ratio_threshold: float = 6.0,
            streaming_prefix_len: int = 3,
        ) -> tuple[None, torch.Tensor, torch.Tensor]:
            return next(
                self._generate_with_prompt_cache(
                    target_text=target_text,
                    prompt_cache=prompt_cache,
                    min_len=min_len,
                    max_len=max_len,
                    inference_timesteps=inference_timesteps,
                    cfg_value=cfg_value,
                    retry_badcase=retry_badcase,
                    retry_badcase_max_times=retry_badcase_max_times,
                    retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
                    streaming=False,
                    streaming_prefix_len=streaming_prefix_len,
                    latents_only=True,
                )
            )

        @torch.inference_mode()
        def generate_latents_with_prompt_cache_streaming(
            self,
            target_text: str,
            prompt_cache: dict,
            min_len: int = 2,
            max_len: int = 2000,
            inference_timesteps: int = 10,
            cfg_value: float = 2.0,
            retry_badcase: bool = False,
            retry_badcase_max_times: int = 3,
            retry_badcase_ratio_threshold: float = 6.0,
            streaming_prefix_len: int = 3,
        ) -> Generator[tuple[None, torch.Tensor, torch.Tensor], None, None]:
            return self._generate_with_prompt_cache(
                target_text=target_text,
                prompt_cache=prompt_cache,
                min_len=min_len,
                max_len=max_len,
                inference_timesteps=inference_timesteps,
                cfg_value=cfg_value,
                retry_badcase=retry_badcase,
                retry_badcase_max_times=retry_badcase_max_times,
                retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
                streaming=True,
                streaming_prefix_len=streaming_prefix_len,
                latents_only=True,
            )

        @torch.inference_mode()
        def _generate_with_prompt_cache(
            self,
            target_text: str,
            prompt_cache: dict,
            min_len: int = 2,
            max_len: int = 2000,
            inference_timesteps: int = 10,
            cfg_value: float = 2.0,
            retry_badcase: bool = False,
            retry_badcase_max_times: int = 3,
            retry_badcase_ratio_threshold: float = 6.0,
            streaming: bool = False,
            streaming_prefix_len: int = 3,
            latents_only: bool = False,
        ) -> Generator[tuple[torch.Tensor | None, torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]:
            if retry_badcase and streaming:
                warnings.warn("Retry on bad cases is not supported in streaming mode, setting retry_badcase=False.")
                retry_badcase = False
            if prompt_cache is None:
                prompt_audio_feat = torch.empty((0, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32)
                text = target_text
            else:
                prompt_audio_feat = prompt_cache["audio_feat"]
                prompt_text = prompt_cache["prompt_text"]
                text = prompt_text + target_text

            text_token = torch.LongTensor(self.text_tokenizer(text))
            text_token = torch.cat(
                [
                    text_token,
                    torch.tensor([self.audio_start_token], dtype=torch.int32, device=text_token.device),
                ],
                dim=-1,
            )
            target_text_token = torch.LongTensor(self.text_tokenizer(target_text))

            audio_length = prompt_audio_feat.size(0)
            text_length = text_token.shape[0]
            text_pad_token = torch.zeros(audio_length, dtype=torch.int32, device=text_token.device)
            audio_pad_feat = torch.zeros(
                (text_token.shape[0], self.patch_size, self.audio_vae.latent_dim),
                dtype=torch.float32,
                device=text_token.device,
            )
            text_token = torch.cat([text_token, text_pad_token])
            audio_feat = torch.cat([audio_pad_feat, prompt_audio_feat], dim=0)
            text_mask = (
                torch.cat([torch.ones(text_length), torch.zeros(audio_length)]).type(torch.int32).to(text_token.device)
            )
            audio_mask = (
                torch.cat([torch.zeros(text_length), torch.ones(audio_length)]).type(torch.int32).to(text_token.device)
            )

            text_token = text_token.unsqueeze(0).to(self.device)
            text_mask = text_mask.unsqueeze(0).to(self.device)
            audio_feat = audio_feat.unsqueeze(0).to(self.device).to(get_dtype(self.config.dtype))
            audio_mask = audio_mask.unsqueeze(0).to(self.device)

            target_text_length = len(self.text_tokenizer(target_text))
            retry_badcase_times = 0
            while retry_badcase_times < retry_badcase_max_times:
                inference_result = self._inference(
                    text_token,
                    text_mask,
                    audio_feat,
                    audio_mask,
                    min_len=min_len,
                    max_len=min(int(target_text_length * retry_badcase_ratio_threshold + 10), max_len),
                    inference_timesteps=inference_timesteps,
                    cfg_value=cfg_value,
                    streaming=streaming,
                    streaming_prefix_len=streaming_prefix_len,
                )
                if streaming:
                    patch_len = self.patch_size * self.chunk_size
                    for latent_pred, pred_audio_feat in inference_result:
                        if latents_only:
                            decode_audio = None
                            yield (decode_audio, target_text_token, latent_pred)
                        else:
                            decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
                            decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu()
                            yield (decode_audio, target_text_token, pred_audio_feat)
                    break

                latent_pred, pred_audio_feat = next(inference_result)
                if retry_badcase and pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold:
                    ratio = pred_audio_feat.shape[0] / target_text_length
                    print(f"  Badcase detected, audio_text_ratio={ratio}, retrying...", file=sys.stderr)
                    retry_badcase_times += 1
                    continue
                break

            if not streaming:
                if latents_only:
                    decode_audio = None
                else:
                    decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
                    patch_len = self.patch_size * self.chunk_size
                    if audio_mask.sum().item() > 0:
                        decode_audio = decode_audio[..., patch_len * (streaming_prefix_len - 1) :].squeeze(1).cpu()
                    else:
                        decode_audio = decode_audio[..., :].squeeze(1).cpu()
                yield (decode_audio, target_text_token, pred_audio_feat)

    VoxCPMModelForOmni.__name__ = "VoxCPMModelForOmni"
    VoxCPMModelForOmni.__qualname__ = "VoxCPMModelForOmni"
    return VoxCPMModelForOmni


def _import_voxcpm_model_class() -> type[Any]:
    base = _import_voxcpm_base_model_class()
    return _make_voxcpm_model_for_omni(base)


def _load_native_voxcpm_model(
    model_path: str,
    *,
    device: torch.device,
    dtype: str | None,
):
    VoxCPMModel = _import_voxcpm_model_class()
    model_dir = resolve_voxcpm_model_dir(model_path)
    runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype=dtype)

    if device.type == "npu" and hasattr(torch, "npu"):
        torch.npu.set_device(device)

    with _force_cuda_available_for_npu(device):
        return VoxCPMModel.from_local(
            runtime_model_path,
            optimize=device.type == "cuda",
        )


def _load_native_voxcpm_latent_generator(
    model_path: str,
    *,
    device: torch.device,
    dtype: str | None,
) -> _DirectVoxCPMLatentGenerator:
    return _DirectVoxCPMLatentGenerator(_load_native_voxcpm_model(model_path, device=device, dtype=dtype))


def _load_native_voxcpm_audio_vae(
    model_path: str,
    *,
    device: torch.device,
) -> _DirectVoxCPMAudioVAE:
    AudioVAE, AudioVAEConfig = _import_voxcpm_audio_vae_classes()
    model_dir = resolve_voxcpm_model_dir(model_path)
    runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype="float32")
    config_dict = json.loads((Path(runtime_model_path) / "config.json").read_text())
    audio_vae_config = config_dict.get("audio_vae_config")
    audio_vae = AudioVAE(config=AudioVAEConfig(**audio_vae_config)) if audio_vae_config is not None else AudioVAE()

    state_dict = torch.load(
        Path(runtime_model_path) / "audiovae.pth",
        map_location="cpu",
        weights_only=True,
    )["state_dict"]
    audio_vae.load_state_dict(state_dict, strict=True)
    audio_vae = audio_vae.to(device=device, dtype=torch.float32).eval()
    if device.type == "npu" and hasattr(torch, "npu"):
        torch.npu.set_device(device)
    patch_size = int(config_dict.get("patch_size", 2))
    return _DirectVoxCPMAudioVAE(audio_vae, patch_size=patch_size)


class VoxCPMForConditionalGeneration(nn.Module):
    input_modalities = "audio"
    _LATENT_STAGES = {"latent_generator", "latent", "ar_dit"}
    _VAE_STAGES = {"vae", "audio_vae"}

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        del prefix
        self.vllm_config = vllm_config
        self.model_path = vllm_config.model_config.model
        self.model_stage = getattr(vllm_config.model_config, "model_stage", "latent_generator")
        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
        self.inject_omni_request_id_into_runtime_info = True
        self._pipeline = None
        self._latent_stream_gens: dict[str, Any] = {}
        self._latent_stream_terminal_pending: dict[str, int] = {}
        self._latent_stream_completed: set[str] = set()
        self._next_local_stream_key = 0
        self._ar_emit_stop_token = True

    def _runner_hidden_device_dtype(self) -> tuple[torch.device, torch.dtype]:
        device = _resolve_runtime_device(self.vllm_config)
        model_config = getattr(self.vllm_config, "model_config", None)
        dtype = getattr(model_config, "dtype", torch.float32) if model_config is not None else torch.float32
        return device, dtype

    def _ensure_model_loaded(self):
        if self._pipeline is not None:
            return

        target_device = _resolve_runtime_device(self.vllm_config)
        model_dtype = getattr(self.vllm_config.model_config, "dtype", None)
        normalized_dtype = _normalize_dtype_name(model_dtype)
        if self.model_stage in self._LATENT_STAGES:
            self._pipeline = _load_native_voxcpm_latent_generator(
                self.model_path,
                device=target_device,
                dtype=normalized_dtype,
            )
        elif self.model_stage in self._VAE_STAGES:
            self._pipeline = _load_native_voxcpm_audio_vae(
                self.model_path,
                device=target_device,
            )
        else:
            raise ValueError(
                f"Unsupported VoxCPM model_stage: {self.model_stage}. "
                "pure_voxcpm only supports split-stage latent_generator/vae inference."
            )

        logger.info("Loaded VoxCPM stage '%s' on %s", self.model_stage, _device_to_string(target_device))

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        del weights
        self._ensure_model_loaded()
        return set()

    @staticmethod
    def _extract_val(info: dict[str, Any], key: str, default: Any) -> Any:
        value = info.get(key, default)
        if isinstance(value, list):
            return value[0] if value else default
        return value

    def _resolve_stream_request_key(self, info: dict[str, Any]) -> str:
        request_key = info.get("__voxcpm_stream_key")
        if request_key is not None:
            return str(request_key)

        request_key = info.get("_omni_req_id")
        if request_key is not None:
            request_key = str(request_key)
            info["__voxcpm_stream_key"] = request_key
            return request_key

        request_key = f"voxcpm-local-{self._next_local_stream_key}"
        self._next_local_stream_key += 1
        info["__voxcpm_stream_key"] = request_key
        return str(request_key)

    def _recover_latent_from_input_ids(self, input_ids: torch.Tensor | None) -> torch.Tensor | None:
        if input_ids is None or input_ids.numel() == 0:
            return None
        flat_ids = input_ids.detach().reshape(-1).to("cpu")
        if flat_ids.numel() < 4 or int(flat_ids[0].item()) != _VOXCPM_LATENT_MAGIC:
            return None
        latent_dim = int(flat_ids[1].item())
        time_dim = int(flat_ids[2].item())
        payload = flat_ids[3:]
        expected = latent_dim * time_dim
        if latent_dim <= 0 or time_dim <= 0:
            raise ValueError(f"Invalid VoxCPM latent header: latent_dim={latent_dim}, time_dim={time_dim}")
        if int(payload.numel()) != expected:
            raise ValueError(
                "Invalid VoxCPM latent payload size: "
                f"expected={expected}, actual={int(payload.numel())}, "
                f"latent_dim={latent_dim}, time_dim={time_dim}"
            )
        packed = payload.to(dtype=torch.int32).to(torch.uint16)
        return packed.view(torch.bfloat16).to(torch.float32).reshape(1, latent_dim, time_dim)

    def _maybe_recover_vae_infos(
        self,
        infos: list[dict[str, Any]],
        input_ids: torch.Tensor | None,
        *,
        async_chunk: bool,
    ) -> list[dict[str, Any]]:
        if not async_chunk:
            return infos
        if any(self._extract_val(info, "latent_audio_feat", None) is not None for info in infos):
            return infos
        recovered = self._recover_latent_from_input_ids(input_ids)
        if recovered is None:
            return infos
        return [{"latent_audio_feat": recovered}]

    @staticmethod
    def _normalize_audio_samples(samples: Any) -> np.ndarray:
        if isinstance(samples, torch.Tensor):
            return samples.detach().cpu().float().reshape(-1).numpy()
        return np.asarray(samples, dtype=np.float32).reshape(-1)

    @classmethod
    def _normalize_ref_audio(cls, ref_audio: Any) -> tuple[np.ndarray, int]:
        if isinstance(ref_audio, str):
            raise TypeError("String ref_audio should be handled as a path before waveform normalization.")

        if isinstance(ref_audio, dict):
            sample_rate = ref_audio.get("sample_rate") or ref_audio.get("sampling_rate") or ref_audio.get("sr")
            samples = None
            for key in ("audio", "wav", "samples", "array", "waveform"):
                if key in ref_audio and ref_audio[key] is not None:
                    samples = ref_audio[key]
                    break
            if sample_rate is None or samples is None:
                raise ValueError("ref_audio dict must contain waveform data and sample rate.")
            return cls._normalize_audio_samples(samples), int(sample_rate)

        if isinstance(ref_audio, (list, tuple)):
            if len(ref_audio) == 1:
                return cls._normalize_ref_audio(ref_audio[0])
            if len(ref_audio) == 2 and np.isscalar(ref_audio[1]):
                return cls._normalize_audio_samples(ref_audio[0]), int(ref_audio[1])

        raise TypeError(f"Unsupported ref_audio format: {type(ref_audio)!r}")

    @staticmethod
    def _write_temp_prompt_wav(waveform: np.ndarray, sample_rate: int) -> str:
        prompt_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        prompt_file.close()

        wav = np.asarray(waveform, dtype=np.float32).reshape(-1)
        wav = np.clip(wav, -1.0, 1.0)
        pcm16 = (wav * 32767.0).astype(np.int16)
        with wave.open(prompt_file.name, "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(int(sample_rate))
            wav_file.writeframes(pcm16.tobytes())

        return prompt_file.name

    @classmethod
    def _resolve_prompt_inputs(cls, info: dict[str, Any]) -> tuple[str | None, str | None, str | None]:
        prompt_text = cls._extract_val(info, "prompt_text", None)
        prompt_wav_path = cls._extract_val(info, "prompt_wav_path", None)
        if prompt_wav_path:
            if prompt_text is None:
                prompt_text = cls._extract_val(info, "ref_text", None)
            return prompt_wav_path, prompt_text, None

        ref_audio = cls._extract_val(info, "ref_audio", None)
        ref_text = cls._extract_val(info, "ref_text", None)
        if ref_audio is None or ref_text is None:
            return None, None, None
        if isinstance(ref_audio, str):
            return ref_audio, ref_text, None

        waveform, sample_rate = cls._normalize_ref_audio(ref_audio)
        temp_prompt_wav = cls._write_temp_prompt_wav(waveform, sample_rate)
        return temp_prompt_wav, ref_text, temp_prompt_wav

    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 _get_vocab_size(self) -> int:
        model_config = getattr(self.vllm_config, "model_config", None)
        if model_config is not None:
            getter = getattr(model_config, "get_vocab_size", None)
            if callable(getter):
                try:
                    return int(getter())
                except Exception:
                    pass
            hf_config = getattr(model_config, "hf_text_config", None)
            if hf_config is not None and hasattr(hf_config, "vocab_size"):
                return int(hf_config.vocab_size)
        return 32000

    def _make_empty_output(
        self,
        *,
        output_key: str,
        payload_factory: Callable[[], torch.Tensor],
        infos: list[dict[str, Any]],
        sample_rate: int,
        out_device: torch.device,
        out_dtype: torch.dtype,
        hidden_rows: int | None = None,
    ) -> OmniOutput:
        if hidden_rows is None:
            hidden_rows = len(infos)
        return OmniOutput(
            text_hidden_states=torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype),
            multimodal_outputs={
                output_key: [payload_factory() for _ in infos],
                "sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in infos],
            },
        )

    def _finalize_stage_output(
        self,
        *,
        output_key: str,
        outputs: list[torch.Tensor],
        sample_rates: list[torch.Tensor],
        out_device: torch.device,
        out_dtype: torch.dtype,
        hidden_rows: int | None = None,
    ) -> OmniOutput:
        multimodal_outputs: dict[str, Any] = {output_key: outputs, "sr": sample_rates}
        if hidden_rows is not None:
            text_hidden_states = torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype)
        elif outputs:
            outputs_tensor = torch.stack(outputs)
            text_hidden_states = (
                outputs_tensor.unsqueeze(-1)
                if outputs_tensor.ndim == 1
                else outputs_tensor.reshape(-1, outputs_tensor.shape[-1])
            )
        else:
            text_hidden_states = torch.zeros((0, 1), device=out_device, dtype=out_dtype)
        text_hidden_states = text_hidden_states.to(device=out_device, dtype=out_dtype)
        return OmniOutput(
            text_hidden_states=text_hidden_states,
            multimodal_outputs=multimodal_outputs,
        )

    def _forward_vae_stage(
        self,
        infos: list[dict[str, Any]],
        *,
        sample_rate: int,
        async_chunk: bool,
        out_device: torch.device,
        out_dtype: torch.dtype,
    ) -> OmniOutput:
        if all(self._extract_val(info, "latent_audio_feat", None) is None for info in infos):
            self._ar_emit_stop_token = True
            return self._make_empty_output(
                output_key="model_outputs",
                payload_factory=lambda: torch.zeros((0,), dtype=torch.float32),
                infos=infos,
                sample_rate=sample_rate,
                out_device=out_device,
                out_dtype=out_dtype,
            )

        outputs: list[torch.Tensor] = []
        sample_rates: list[torch.Tensor] = []
        for info in infos:
            latent_audio_feat = self._extract_val(info, "latent_audio_feat", None)
            audio_tensor = self._pipeline.decode(latent_audio_feat, trim_streaming_patch=async_chunk)
            outputs.append(audio_tensor.float().cpu())
            sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))

        self._ar_emit_stop_token = True
        return self._finalize_stage_output(
            output_key="model_outputs",
            outputs=outputs,
            sample_rates=sample_rates,
            out_device=out_device,
            out_dtype=out_dtype,
        )

    def _forward_latent_stage(
        self,
        infos: list[dict[str, Any]],
        *,
        sample_rate: int,
        async_chunk: bool,
        out_device: torch.device,
        out_dtype: torch.dtype,
        hidden_rows: int,
    ) -> OmniOutput:
        texts = [self._extract_val(info, "text", "") for info in infos]
        if all(not text for text in texts):
            self._ar_emit_stop_token = True
            return self._make_empty_output(
                output_key="latent_audio_feat",
                payload_factory=lambda: torch.zeros((0,), dtype=torch.float32),
                infos=infos,
                sample_rate=sample_rate,
                out_device=out_device,
                out_dtype=out_dtype,
                hidden_rows=hidden_rows,
            )

        outputs: list[torch.Tensor] = []
        sample_rates: list[torch.Tensor] = []
        last_chunk_flags: list[bool] | None = [] if async_chunk else None
        payload_finished_flags: list[bool] | None = [] if async_chunk else None
        for info in infos:
            text = self._extract_val(info, "text", "")
            cfg_value = float(self._extract_val(info, "cfg_value", 2.0))
            inference_timesteps = int(self._extract_val(info, "inference_timesteps", 10))
            min_len = int(self._extract_val(info, "min_len", 2))
            max_len = int(self._extract_val(info, "max_len", self._extract_val(info, "max_new_tokens", 4096)))
            retry_badcase = bool(self._extract_val(info, "retry_badcase", True))
            retry_badcase_max_times = int(self._extract_val(info, "retry_badcase_max_times", 3))
            retry_badcase_ratio_threshold = float(self._extract_val(info, "retry_badcase_ratio_threshold", 6.0))
            streaming_prefix_len = int(self._extract_val(info, "streaming_prefix_len", 3))

            request_key = self._resolve_stream_request_key(info)
            created_temp: str | None = None

            if async_chunk:
                terminal_pending = self._latent_stream_terminal_pending.get(request_key, 0)
                if terminal_pending > 0:
                    outputs.append(torch.zeros((0,), dtype=torch.float32))
                    assert last_chunk_flags is not None
                    last_chunk_flags.append(True)
                    assert payload_finished_flags is not None
                    payload_finished_flags.append(terminal_pending == 1)
                    if terminal_pending == 1:
                        self._latent_stream_terminal_pending.pop(request_key, None)
                    else:
                        self._latent_stream_terminal_pending[request_key] = terminal_pending - 1
                    sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
                    continue

                if request_key in self._latent_stream_completed:
                    outputs.append(torch.zeros((0,), dtype=torch.float32))
                    assert last_chunk_flags is not None
                    last_chunk_flags.append(True)
                    assert payload_finished_flags is not None
                    payload_finished_flags.append(False)
                    sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
                    continue

                if request_key not in self._latent_stream_gens:
                    prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info)
                    created_temp = temp_prompt_wav
                    self._latent_stream_gens[request_key] = self._pipeline.iter_latent_chunks_streaming(
                        text=text,
                        prompt_wav_path=prompt_wav_path,
                        prompt_text=prompt_text,
                        cfg_value=cfg_value,
                        inference_timesteps=inference_timesteps,
                        min_len=min_len,
                        max_len=max_len,
                        streaming_prefix_len=streaming_prefix_len,
                        retry_badcase=False,
                        retry_badcase_max_times=retry_badcase_max_times,
                        retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
                    )
                generator = self._latent_stream_gens[request_key]
                try:
                    chunk_latent, is_last = next(generator)
                except StopIteration:
                    self._latent_stream_gens.pop(request_key, None)
                    self._latent_stream_terminal_pending[request_key] = 1
                    self._latent_stream_completed.add(request_key)
                    outputs.append(torch.zeros((0,), dtype=torch.float32))
                    assert last_chunk_flags is not None
                    last_chunk_flags.append(True)
                    assert payload_finished_flags is not None
                    payload_finished_flags.append(True)
                else:
                    if is_last:
                        self._latent_stream_gens.pop(request_key, None)
                        self._latent_stream_terminal_pending[request_key] = 1
                        self._latent_stream_completed.add(request_key)
                    outputs.append(chunk_latent.detach().float().cpu())
                    assert last_chunk_flags is not None
                    last_chunk_flags.append(bool(is_last))
                    assert payload_finished_flags is not None
                    payload_finished_flags.append(False)
                finally:
                    if created_temp is not None and os.path.exists(created_temp):
                        os.unlink(created_temp)
                sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))
                continue

            prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info)
            try:
                latent_audio_feat = self._pipeline.generate_latents(
                    text=text,
                    prompt_wav_path=prompt_wav_path,
                    prompt_text=prompt_text,
                    cfg_value=cfg_value,
                    inference_timesteps=inference_timesteps,
                    min_len=min_len,
                    max_len=max_len,
                    retry_badcase=retry_badcase,
                    retry_badcase_max_times=retry_badcase_max_times,
                    retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
                )
                outputs.append(latent_audio_feat.float().cpu())
            finally:
                if temp_prompt_wav is not None and os.path.exists(temp_prompt_wav):
                    os.unlink(temp_prompt_wav)

            sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32))

        self._ar_emit_stop_token = all(last_chunk_flags) if async_chunk and last_chunk_flags else True
        output = self._finalize_stage_output(
            output_key="latent_audio_feat",
            outputs=outputs,
            sample_rates=sample_rates,
            out_device=out_device,
            out_dtype=out_dtype,
            hidden_rows=hidden_rows,
        )
        if async_chunk and payload_finished_flags is not None:
            output.multimodal_outputs["finished"] = [
                torch.tensor(flag, dtype=torch.bool) for flag in payload_finished_flags
            ]
        return output

    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor:
        del sampling_metadata
        if isinstance(hidden_states, OmniOutput):
            hidden_states = hidden_states.text_hidden_states
        if hidden_states is None:
            device, dtype = self._runner_hidden_device_dtype()
            hidden_states = torch.zeros((0, 1), device=device, dtype=dtype)
        if hidden_states.ndim == 1:
            hidden_states = hidden_states.unsqueeze(-1)
        elif hidden_states.ndim > 2:
            hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])

        vocab_size = self._get_vocab_size()
        num_rows = int(hidden_states.shape[0])
        logits = torch.zeros((num_rows, vocab_size), dtype=torch.float32, device=hidden_states.device)
        eos_id = 2 if vocab_size > 2 else 0
        safe_id = 1 if vocab_size > 1 and 1 != eos_id else 0
        emit_stop = getattr(self, "_ar_emit_stop_token", True)
        if num_rows > 0:
            if emit_stop:
                logits[:, eos_id] = 1.0e6
            else:
                logits[:, eos_id] = -1.0e9
                logits[:, safe_id] = 1.0e6
        return logits

    @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,
        model_intermediate_buffer: list[dict[str, Any]] | None = None,
        **kwargs: Any,
    ) -> OmniOutput:
        del positions, intermediate_tensors, inputs_embeds, kwargs
        self._ensure_model_loaded()
        out_device, out_dtype = self._runner_hidden_device_dtype()
        if input_ids is not None and input_ids.device.type == out_device.type:
            out_device = input_ids.device

        infos = model_intermediate_buffer or runtime_additional_information or [{}]
        hidden_rows = len(infos)
        if input_ids is not None and len(input_ids.shape) > 0:
            hidden_rows = max(hidden_rows, int(input_ids.shape[0]))
        sample_rate = int(getattr(self._pipeline, "sample_rate", 24000))
        async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False))
        if self.model_stage in self._VAE_STAGES:
            infos = self._maybe_recover_vae_infos(infos, input_ids, async_chunk=async_chunk)
            return self._forward_vae_stage(
                infos,
                sample_rate=sample_rate,
                async_chunk=async_chunk,
                out_device=out_device,
                out_dtype=out_dtype,
            )
        if self.model_stage in self._LATENT_STAGES:
            return self._forward_latent_stage(
                infos,
                sample_rate=sample_rate,
                async_chunk=async_chunk,
                out_device=out_device,
                out_dtype=out_dtype,
                hidden_rows=hidden_rows,
            )
        raise ValueError(f"Unsupported VoxCPM model_stage at runtime: {self.model_stage}")

    def make_empty_intermediate_tensors(
        self, batch_size: int, dtype: torch.dtype, device: torch.device
    ) -> IntermediateTensors:
        del batch_size, dtype, device
        return {}


__all__ = ["VoxCPMForConditionalGeneration"]