Skip to content

vllm_omni.model_executor.models.personaplex.personaplex_code2wav

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 duplex subpackage of this model folder for the same pattern).

logger module-attribute 

logger = init_logger(__name__)

PersonaPlexCode2Wav

Bases: 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).

config instance-attribute 

config = vllm_config.model_config.hf_config

enable_update_additional_information instance-attribute 

enable_update_additional_information = True

has_postprocess instance-attribute 

has_postprocess = False

has_preprocess instance-attribute 

has_preprocess = False

have_multimodal_outputs instance-attribute 

have_multimodal_outputs = True

input_modalities class-attribute instance-attribute 

input_modalities = 'audio'

mimi instance-attribute 

mimi: Module | None = None

model_path instance-attribute 

model_path = vllm_config.model_config.model

requires_raw_input_tokens instance-attribute 

requires_raw_input_tokens = True

requires_request_ids class-attribute instance-attribute 

requires_request_ids = True

vllm_config instance-attribute 

vllm_config = vllm_config

compute_logits 

compute_logits(
    hidden_states: Tensor | OmniOutput,
    sampling_metadata: Any = None,
) -> None

embed_input_ids 

embed_input_ids(input_ids: Tensor, **_: Any) -> Tensor

forward 

forward(
    input_ids: Tensor | None = None,
    positions: Tensor | None = None,
    intermediate_tensors: Any = None,
    inputs_embeds: 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.

get_dummy_runtime_additional_information 

get_dummy_runtime_additional_information(
    num_reqs: int,
) -> list[dict[str, object]]

load_weights 

load_weights(
    weights: Iterable[tuple[str, 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.

make_omni_output 

make_omni_output(
    model_outputs: Tensor | OmniOutput | tuple,
    **kwargs: Any,
) -> OmniOutput

on_requests_finished 

on_requests_finished(
    finished_req_ids: set[str] | list[str],
) -> None