anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
4882a36c41c481f70a3211c72f4e36e7bc800983
vllm-omni/vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py
T

434 lines
19 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Stage-1 Code2Wav model for PersonaPlex (Moshi finetune).
This is the ``LLM_GENERATION`` codec stage of the 2-stage PersonaPlex pipeline.
It consumes the per-frame audio codebooks produced by the talker (stage 0) and
the depformer (built by the lead), and turns them into 24 kHz PCM by calling the
external Mimi neural codec from the ``moshi`` package.
Layout contract (mirrors ``Qwen3TTSCode2Wav``):
* Per request, ``input_ids`` holds a flat codebook-major codec sequence
``[k * F]`` where ``k`` is the number of *active* codebooks Mimi decodes
(``num_codebooks``, i.e. the ``cb 0..7`` slice) and ``F`` is the number of
codec frames. The talker emits a 17-row token stack per frame; the input
processor (built by the lead) keeps only rows ``1:9`` (the 8 PCM-bearing audio
codebooks) and flattens them codebook-major before this stage.
* ``forward(...)`` returns an :class:`OmniOutput` whose ``multimodal_outputs``
carries ``{"model_outputs": [wav_per_request], "sr": [sr_per_request]}`` — the
exact shape ``Qwen3TTSCode2Wav`` returns, so the downstream audio packer is
unchanged.
The Mimi decoder is the transformers ``MimiModel`` (kyutai/mimi weights), not
in vLLM's safetensors loader, so they are loaded eagerly in ``load_weights``
part of the vLLM weights iterator (the codec owns its own checkpoint; see the
experimental fullduplex package for the same pattern).
"""
from __future__ import annotations
from collections.abc import Iterable, Mapping
from pathlib import Path
from typing import Any
import torch
import torch.nn as nn
from vllm.config import VllmConfig
from vllm.forward_context import get_forward_context, is_forward_context_available
from vllm.logger import init_logger
from vllm_omni.model_executor.models.output_templates import OmniOutput
logger = init_logger(__name__)
_MIMI_DECODE_BATCH_FRAMES = 5
def _codec_ids_from_payload_or_input(
input_ids: torch.Tensor,
runtime_info: Mapping[str, Any] | None,
) -> torch.Tensor:
"""Prefer connector-delivered codec ids over scheduler placeholders."""
if isinstance(runtime_info, Mapping):
codes = runtime_info.get("codes")
if isinstance(codes, Mapping):
audio = codes.get("audio")
if isinstance(audio, torch.Tensor) and audio.numel() > 0:
return audio.reshape(-1).to(device=input_ids.device, dtype=torch.long)
if isinstance(audio, (list, tuple)) and audio:
return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1)
return input_ids.reshape(-1).to(dtype=torch.long)
class PersonaPlexCode2Wav(nn.Module):
"""Stage-1 Code2Wav model for PersonaPlex (GenerationModelRunner).
Wraps the external Mimi decoder. The vLLM generation runner only calls a
handful of methods on this module: ``embed_input_ids`` (dummy embeddings),
``compute_logits`` (none -- this stage never samples), ``forward`` (the
actual codec->PCM decode), ``make_omni_output`` (output normalization), and
``load_weights`` (eager Mimi construction).
"""
input_modalities = "audio"
requires_request_ids = True
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
self.vllm_config = vllm_config
self.model_path = vllm_config.model_config.model
self.config = vllm_config.model_config.hf_config
# Runner-facing capability flags, matching Qwen3TTSCode2Wav so the
# GenerationModelRunner drives this stage the same way.
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
mimi_cfg = getattr(self.config, "mimi_config", None)
# Number of *active* codebooks Mimi decodes to PCM (cb 0..7).
self._num_codebooks = int(getattr(mimi_cfg, "num_codebooks", 8))
self._output_sample_rate = int(getattr(mimi_cfg, "sample_rate", 24000))
self._samples_per_frame = int(getattr(mimi_cfg, "samples_per_frame", 1920))
self._mimi_name = getattr(mimi_cfg, "mimi_name", None) or getattr(self.config, "mimi_name", None)
self._max_codec_sessions = int(getattr(vllm_config.model_config, "duplex_max_sessions", 1))
# The Mimi module is constructed in load_weights() (it owns its own
# weight format) and assigned here so vLLM's memory profiler can see it.
self.mimi: nn.Module | None = None
self._additional_mimi = nn.ModuleList()
self._mimi_device: torch.device | None = None
self._request_codes: dict[str, torch.Tensor] = {}
self._request_codec_slots: dict[str, int] = {}
# ------------------------------------------------------------------
# Runner-facing no-op / placeholder hooks (mirror Qwen3TTSCode2Wav).
# ------------------------------------------------------------------
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
# This stage ignores token embeddings; keep a stable dummy embedding so
# the vLLM runner has a valid hidden-state placeholder.
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:
# Code2Wav never samples; the runner short-circuits when logits are None.
return None
def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
return [{"meta": {"personaplex_dummy_profile": True}} for _ in range(num_reqs)]
def _split_request_ids(
self,
ids: torch.Tensor,
seq_token_counts: list[int] | None = None,
) -> list[torch.Tensor]:
"""Split concatenated ``input_ids`` into per-request segments.
Uses ``seq_token_counts`` (injected by the runner via model_kwargs) when
available, falling back to forward-context ``ubatch_slices`` when
micro-batching is active. Returns ``[ids]`` for single-request batches.
Mirrors ``Qwen3TTSCode2Wav._split_request_ids``.
"""
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] + count)
n = ids.numel()
return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))]
if is_forward_context_available():
slices = get_forward_context().ubatch_slices
if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):
boundaries = [0]
for s in slices:
boundaries.append(boundaries[-1] + s)
return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
return [ids]
# ------------------------------------------------------------------
# Decode.
# ------------------------------------------------------------------
@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,
) -> OmniOutput:
"""Decode flat codebook-major codec ids into PCM via Mimi.
``input_ids`` per request is ``[k * F]`` (codebook-major), where ``k``
is ``self._num_codebooks``. The connector may send either a new delta
chunk or the cumulative prefix of a resumable request. Request-local
code history identifies the new suffix, and the streaming Mimi decoder
consumes each new frame exactly once.
"""
sr_val = int(self._output_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"))
num_req = len(request_ids_list)
state_ids = self._resolve_request_ids(
num_req,
kwargs.get("request_ids"),
runtime_additional_information,
)
if self.mimi is None:
raise RuntimeError("PersonaPlexCode2Wav.forward called before Mimi was loaded in load_weights().")
k = int(self._num_codebooks)
device = self._mimi_device or ids.device
audios: list[torch.Tensor] = [empty] * num_req
srs = [sr_tensor] * num_req
for i, req_ids in enumerate(request_ids_list):
runtime_info = (
runtime_additional_information[i]
if runtime_additional_information is not None and i < len(runtime_additional_information)
else None
)
meta = runtime_info.get("meta") if isinstance(runtime_info, Mapping) else None
if isinstance(meta, Mapping) and meta.get("personaplex_dummy_profile") is True:
continue
flat = _codec_ids_from_payload_or_input(req_ids, runtime_info)
n = int(flat.numel())
if n == 0 or n % k != 0:
if n > 0:
logger.warning(
"PersonaPlex Code2Wav input_ids length %d not divisible by "
"num_codebooks %d; skipping malformed request.",
n,
k,
)
continue
frames = n // k
codes_kf = flat.reshape(k, frames)
state_id = state_ids[i]
delta_kf = self._new_code_suffix(state_id, codes_kf)
if delta_kf.shape[1] == 0:
continue
wav = self._decode_streaming_frames(state_id, delta_kf.to(device=device))
if wav.numel() > 0:
audios[i] = wav.to(dtype=torch.float32).reshape(-1)
return OmniOutput(
text_hidden_states=None,
multimodal_outputs={"model_outputs": audios, "sr": srs},
)
@staticmethod
def _resolve_request_ids(
count: int,
request_ids: object,
runtime_additional_information: list[dict[str, Any]] | None,
) -> list[str | None]:
if isinstance(request_ids, list | tuple):
if len(request_ids) != count:
raise ValueError(
"PersonaPlex Code2Wav request id count does not match inputs: "
f"request_ids={len(request_ids)}, inputs={count}"
)
return [str(request_id) for request_id in request_ids]
infos = runtime_additional_information or []
resolved: list[str | None] = []
for index in range(count):
info = infos[index] if index < len(infos) else None
request_id = None
if isinstance(info, Mapping):
request_id = info.get("request_id")
if request_id is None:
meta = info.get("meta")
if isinstance(meta, Mapping):
request_id = meta.get("request_id")
resolved.append(str(request_id) if request_id is not None else None)
return resolved
def _new_code_suffix(self, request_id: str | None, codes_kf: torch.Tensor) -> torch.Tensor:
if request_id is None:
return codes_kf
incoming = codes_kf.detach().to(device="cpu", dtype=torch.long)
previous = self._request_codes.get(request_id)
if (
previous is not None
and incoming.shape[1] >= previous.shape[1]
and torch.equal(incoming[:, : previous.shape[1]], previous)
):
delta = incoming[:, previous.shape[1] :]
self._request_codes[request_id] = incoming
return delta
self._request_codes[request_id] = incoming if previous is None else torch.cat([previous, incoming], dim=1)
return incoming
def _decode_streaming_frames(self, request_id: str | None, codes_kf: torch.Tensor) -> torch.Tensor:
codec, ephemeral = self._codec_for_request(request_id)
if not hasattr(codec, "decode_frame"):
raise RuntimeError("PersonaPlex Code2Wav requires a streaming Mimi decoder")
decode_frames = getattr(codec, "decode_frames", None)
chunks: list[torch.Tensor] = []
for start in range(0, codes_kf.shape[1], _MIMI_DECODE_BATCH_FRAMES):
frame_batch = codes_kf[:, start : start + _MIMI_DECODE_BATCH_FRAMES]
if frame_batch.shape[1] > 1 and callable(decode_frames):
chunks.append(self._flatten_wav(decode_frames(frame_batch.unsqueeze(0))))
else:
chunks.extend(
self._flatten_wav(codec.decode_frame(frame_batch[:, frame].unsqueeze(0)))
for frame in range(frame_batch.shape[1])
)
wav = torch.cat(chunks, dim=0) if chunks else codes_kf.new_empty(0, dtype=torch.float32)
if ephemeral:
codec.reset_streaming()
return wav
def _set_mimi_codecs(self, codecs: list[nn.Module]) -> None:
if not codecs:
raise ValueError("PersonaPlex Code2Wav requires at least one Mimi decoder")
self.mimi = codecs[0]
self._additional_mimi = nn.ModuleList(codecs[1:])
self._request_codec_slots.clear()
def _mimi_codecs(self) -> list[nn.Module]:
if self.mimi is None:
return []
return [self.mimi, *self._additional_mimi]
def _codec_for_request(self, request_id: str | None) -> tuple[nn.Module, bool]:
codecs = self._mimi_codecs()
if not codecs:
raise RuntimeError("PersonaPlexCode2Wav.forward called before Mimi was loaded in load_weights().")
occupied = set(self._request_codec_slots.values())
if request_id is not None:
slot = self._request_codec_slots.get(request_id)
if slot is None:
slot = next((index for index in range(len(codecs)) if index not in occupied), None)
if slot is None:
raise RuntimeError(f"PersonaPlex Code2Wav decoder capacity {len(codecs)} is exhausted")
codecs[slot].streaming_init(1)
self._request_codec_slots[request_id] = slot
return codecs[slot], False
slot = next((index for index in range(len(codecs)) if index not in occupied), None)
if slot is None:
raise RuntimeError(f"PersonaPlex Code2Wav decoder capacity {len(codecs)} is exhausted")
codecs[slot].streaming_init(1)
return codecs[slot], True
def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
for request_id in finished_req_ids:
state_id = str(request_id)
self._request_codes.pop(state_id, None)
slot = self._request_codec_slots.pop(state_id, None)
codecs = self._mimi_codecs()
if slot is not None and slot < len(codecs):
codecs[slot].reset_streaming()
@staticmethod
def _flatten_wav(wav: torch.Tensor) -> torch.Tensor:
"""Normalize Mimi's decode output to a 1D PCM tensor.
Mimi may return ``[B, samples]`` or ``[B, C, samples]`` depending on the
moshi version; PersonaPlex is mono (``C == 1``), so collapse to 1D.
"""
if wav.dim() == 3:
# [B, C, T] -> take batch 0, channel 0.
return wav[0, 0]
if wav.dim() == 2:
# [B, T] -> batch 0.
return wav[0]
return wav.reshape(-1)
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
return model_outputs
if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields):
return OmniOutput(*model_outputs)
if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2):
raise TypeError(
"PersonaPlexCode2Wav expected OmniOutput, OmniOutput tuple, "
f"or (audio_tensor, sr) outputs, got {type(model_outputs)}"
)
audio_tensor, sr = model_outputs
return OmniOutput(
text_hidden_states=None,
multimodal_outputs={
"model_outputs": audio_tensor,
"sr": sr,
},
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Construct and load the external Mimi decoder.
The primary vLLM weights iterator carries no Code2Wav parameters (Mimi
owns its own checkpoint format), so it is drained and the Mimi module is
built from the moshi package's loader.
"""
# Drain the primary iterator so callers don't hang on an unconsumed
# generator; none of these weights belong to this stage.
for _ in weights:
pass
device = self.vllm_config.device_config.device
from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
PersonaPlexMimiCodec,
)
checkpoint = Path(self.model_path) / (self._mimi_name or "tokenizer-e351c8d8-checkpoint125.safetensors")
codecs = [
PersonaPlexMimiCodec(
checkpoint=str(checkpoint) if checkpoint.is_file() else None,
device=str(device),
).eval()
for _ in range(self._max_codec_sessions)
]
self._set_mimi_codecs(codecs)
self._mimi_device = torch.device(str(device))
reported_sr = getattr(codecs[0].model.config, "sampling_rate", None)
if reported_sr is not None:
self._output_sample_rate = int(reported_sr)
bytes_per_mib = 1024**2
per_slot_weight_mib = (
sum(parameter.numel() * parameter.element_size() for parameter in codecs[0].parameters()) / bytes_per_mib
)
total_weight_mib = (
sum(parameter.numel() * parameter.element_size() for codec in codecs for parameter in codec.parameters())
/ bytes_per_mib
)
logger.info(
"PersonaPlex Code2Wav loaded %d Mimi stream(s): "
"per_slot_weight_mib=%.2f total_weight_mib=%.2f "
"num_codebooks=%d sample_rate=%d samples_per_frame=%d",
len(codecs),
per_slot_weight_mib,
total_weight_mib,
self._num_codebooks,
self._output_sample_rate,
self._samples_per_frame,
)
return {name for name, _ in self.named_parameters() if name.startswith(("mimi.", "_additional_mimi."))}
Reference in new issue View Git Blame Copy Permalink