vllm-omni/vllm_omni/experimental/fullduplex/personaplex/runtime.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Moshi-free PersonaPlex lockstep engine (native temporal + depformer + Mimi).

This is the real-time duplex backend built entirely from vllm-omni-native
components: :class:`PersonaPlexTemporalStreaming` (ring-KV Helium),
:class:`PersonaPlexInputEmbeddings`, :class:`PersonaPlexDepformer` and
:class:`PersonaPlexMimiCodec`. It reimplements the reference lockstep semantics
(staging ring, per-codebook delays, warmup gates, hybrid system-prompt prefill,
per-slot elastic recycle) as first-class policy code, so the external ``moshi``
package is no longer needed at runtime.

Lockstep contract per 80 ms tick (17 token rows: text, 8 agent codebooks, 8
user codebooks; delays ``[0, 1x8-ish from checkpoint]``):

1. Mimi-encode the user frame and write it (delayed) into the staging ring.
2. Read the model input at ``(o - 1) % CT``, embed, step the temporal.
3. Greedy text token; depformer expands the 16 audio rows (teacher-forced where
   the staging ring already holds provided values).
4. Write sampled rows at ``o % CT``; gather the de-delayed output frame; decode
   agent rows 1..8 through Mimi (rows still in warmup are sanitized to silence
   tokens: the LM's initial audio token id equals the codec cardinality and is
   out of range for the codebook tables).

Elastic recycle mirrors the verified design from the moshi-hosted path: the
temporal ring KV masks the recycled row (plus the one LM "sacrifice tick"
write), the staging row is wiped to initial tokens, and the system prompt
(voice embeddings, silence spacer, persona tokens, silence) is replayed per-row
while other slots keep conversing.
"""

from __future__ import annotations

import logging
import os
import tarfile
import threading
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np

from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
from vllm_omni.experimental.fullduplex.personaplex.policy import (
    AUDIO_INITIAL_TOKEN,
    AUDIO_SILENCE_FRAME_CNT,
    NON_TEXT_TOKEN_IDS,
    SILENCE_TOKENS,
    SINE_TOKENS,
    TEXT_INITIAL_TOKEN,
    TEXT_VOCAB_SIZE,
    ZERO_TEXT_TOKEN,
    PrefillStep,
    wrap_with_system_tags,
)

if TYPE_CHECKING:
    from numpy.typing import NDArray

logger = logging.getLogger(__name__)


class PersonaPlexEngine:
    """Frame-lockstep PersonaPlex on native components only (no moshi).

    One class serves both surfaces: the single-session ``FrameStepper`` protocol
    (``open_session`` / ``step``) and the batched slot API consumed by
    ``serving/batched.py`` (``open_batch`` / ``prefill_steps`` / ``recycle_slot``
    / ``voice_end`` / ``end_prefill`` / ``step_batch``).
    """

    def __init__(self, config: PersonaPlexConfig | None = None) -> None:
        self.config = config or PersonaPlexConfig()
        self.sample_rate = self.config.sample_rate
        self.frame_size = self.config.frame_size
        # Serializes the single-session FrameStepper surface (open_session/step)
        # across worker threads: a cancelled asyncio task cannot stop a running
        # engine call, so a racing open() must wait for the in-flight step.
        # The batched surface is serialized by BatchedSessionManager instead.
        self._fs_lock = threading.Lock()
        self._loaded = False
        self._opened = False
        self._tokenizer = None
        self._codec = None
        self._temporal = None
        self._emb = None
        self._dep = None
        self._voice_embeddings = None  # [T, 1, 1, H]
        self._voice_cache = None  # [1, K, CT]
        self._loaded_voice = None  # name/path of the currently loaded voice
        # lockstep state (built at open time)
        self._cache = None  # [B, K, CT] long
        self._provided = None  # [B, K, CT] bool
        self._initial = None  # [K] long
        self._delays = None  # [K] long (python list)
        self._phys = 0
        self._local = None  # [B] cpu long

    # -- loading --------------------------------------------------------------

    def load(self) -> PersonaPlexEngine:
        if self._loaded:
            return self
        import sentencepiece
        import torch
        from huggingface_hub import hf_hub_download
        from safetensors.torch import load_file

        from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
            PersonaPlexConfig as _PPModelConfig,
        )
        from vllm_omni.model_executor.models.personaplex.personaplex_depformer import (
            PersonaPlexDepformer,
        )
        from vllm_omni.model_executor.models.personaplex.personaplex_embeddings import (
            PersonaPlexInputEmbeddings,
        )
        from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
            PersonaPlexMimiCodec,
        )
        from vllm_omni.model_executor.models.personaplex.personaplex_temporal import (
            PersonaPlexTemporalStreaming,
        )

        cfg = self.config
        B = cfg.batch_size
        logger.info("PersonaPlex native: loading Mimi codec")
        self._codec = PersonaPlexMimiCodec(device=cfg.device)
        self._codec.streaming_init(B)

        tok_path = hf_hub_download(cfg.hf_repo, "tokenizer_spm_32k_3.model")
        self._tokenizer = sentencepiece.SentencePieceProcessor(tok_path)

        logger.info("PersonaPlex native: loading temporal + depformer + embeddings")
        sd = load_file(hf_hub_download(cfg.hf_repo, "model.safetensors"), device=cfg.device)
        dtype = torch.bfloat16
        self._temporal = PersonaPlexTemporalStreaming().to(cfg.device, dtype)
        self._temporal.load_weights(sd)
        self._temporal.eval()
        self._temporal.streaming_init(B)

        ppcfg = _PPModelConfig()
        self._emb = PersonaPlexInputEmbeddings(ppcfg).to(cfg.device, dtype)
        self._emb.load_weights(sd)
        self._emb.eval()
        self._dep = PersonaPlexDepformer(ppcfg.depformer_config, temporal_hidden_size=4096, text_card=32000).to(
            cfg.device, dtype
        )
        self._dep.load_weights(sd)
        self._dep.eval()
        del sd
        torch.accelerator.empty_cache()

        K = 17
        self._delays = [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1]  # text, agent cb0..7, user cb0..7
        self._initial = torch.tensor(
            [TEXT_INITIAL_TOKEN] + [AUDIO_INITIAL_TOKEN] * 16, dtype=torch.long, device=cfg.device
        )
        ct = max(self._delays) + 3
        self._cache = torch.full((B, K, ct), AUDIO_INITIAL_TOKEN, dtype=torch.long, device=cfg.device)
        self._provided = torch.zeros((B, K, ct), dtype=torch.bool, device=cfg.device)
        self._local = torch.zeros(B, dtype=torch.long)
        self._sine = torch.tensor(SINE_TOKENS, dtype=torch.long, device=cfg.device)
        self._silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long, device=cfg.device)

        self._load_voice(cfg.voice_prompt)
        self._loaded = True
        return self

    def _load_voice(self, voice: str) -> None:
        import torch

        if voice == self._loaded_voice:
            return
        path = self._resolve_voice_prompt(voice)
        if not path.endswith(".pt"):
            raise ValueError(f"native stack needs a .pt voice-embedding bundle; got {path!r}")
        state = torch.load(path, map_location=self.config.device, weights_only=True)
        self._voice_embeddings = state["embeddings"].to(self.config.device)
        self._voice_cache = state["cache"].to(self.config.device)
        self._loaded_voice = voice

    def _resolve_voice_prompt(self, voice: str) -> str:
        from huggingface_hub import hf_hub_download

        if os.path.exists(voice):
            return voice
        tgz = Path(hf_hub_download(self.config.hf_repo, "voices.tgz"))
        vdir = tgz.parent / "voices"
        if not vdir.exists():
            with tarfile.open(tgz, "r:gz") as tar:
                tar.extractall(path=tgz.parent)  # noqa: S202 - trusted first-party bundle
        cand = vdir / voice
        if not cand.exists():
            raise FileNotFoundError(f"voice prompt {voice!r} not found in voices.tgz")
        return str(cand)

    # -- lockstep core ----------------------------------------------------------

    def _reset_all_streaming(self) -> None:
        self._codec.reset_streaming()
        self._temporal.reset_streaming()
        self._cache.fill_(AUDIO_INITIAL_TOKEN)
        self._cache[:, 0, :] = TEXT_INITIAL_TOKEN
        self._provided.fill_(False)
        self._phys = 0
        self._local.zero_()

    def _tick(
        self,
        user_codes,  # [B, 8] long or None (prefill rows overridden separately)
        forced_moshi=None,  # [B, 8] long (values used on force_mask rows)
        forced_text=None,  # [B] long
        force_mask=None,  # [B] bool (device)
        embed_override=None,  # ([B, 1, H], [B] bool rows)
    ):
        """One lockstep step. Returns (ran, out[B,17] | None)."""
        import torch

        B, K, CT = self._cache.shape
        o = self._phys
        lo = self._local

        if user_codes is not None:
            for q in range(8):
                k = 9 + q
                wp = (o + self._delays[k]) % CT
                self._cache[:, k, wp] = user_codes[:, q]
                self._provided[:, k, wp] = True
        if forced_moshi is not None and force_mask is not None:
            for q in range(8):
                k = 1 + q
                wp = (o + self._delays[k]) % CT
                self._cache[force_mask, k, wp] = forced_moshi[force_mask, q]
                self._provided[force_mask, k, wp] = True
        if forced_text is not None and force_mask is not None:
            wp = (o + self._delays[0]) % CT
            self._cache[force_mask, 0, wp] = forced_text[force_mask]
            self._provided[force_mask, 0, wp] = True

        for k, delay in enumerate(self._delays):
            warm = lo <= delay  # [B] cpu bool
            if bool(warm.any()):
                wd = warm.to(self._cache.device)
                pos = o % CT
                self._cache[:, k, pos] = torch.where(wd, self._initial[k].expand(B), self._cache[:, k, pos])
                self._provided[:, k, pos] = self._provided[:, k, pos] | wd

        if o == 0:
            # Genuine cold start: no prior physical frame to read as model input.
            self._cache[:, :, 0] = self._initial.view(1, K).expand(B, K)
            self._phys += 1
            self._local += 1
            return False, None

        mip = (o - 1) % CT
        tp = o % CT
        input_ = self._cache[:, :, mip]  # [B, K]
        target_ = self._cache[:, :, tp]
        provided_ = self._provided[:, :, tp]

        emb = self._emb(input_.unsqueeze(-1))  # [B, 1, H]
        if embed_override is not None:
            ov, rows = embed_override
            emb = torch.where(rows.view(B, 1, 1), ov.to(emb.dtype), emb)
        hidden, text_logits = self._temporal.step(emb)

        sampled_text = text_logits.float().view(B, -1).argmax(-1)  # greedy
        next_text = torch.where(provided_[:, 0], target_[:, 0], sampled_text)
        audio = self._dep(next_text, hidden, audio_tokens=target_[:, 1:], audio_provided=provided_[:, 1:])  # [B, 16]

        self._provided[:, :, mip] = False
        self._cache[:, 0, tp] = torch.where(~provided_[:, 0], sampled_text, self._cache[:, 0, tp])
        self._cache[:, 1:, tp] = torch.where(~provided_[:, 1:], audio, self._cache[:, 1:, tp])

        max_delay = max(self._delays)
        valid = lo > max_delay
        if not bool(valid.any()):
            self._phys += 1
            self._local += 1
            return True, None

        idx = torch.tensor(
            [(self._phys - max_delay + d) % CT for d in self._delays],
            device=self._cache.device,
            dtype=torch.long,
        )
        out = self._cache.gather(2, idx.view(1, K, 1).expand(B, K, 1)).squeeze(-1)  # [B, K]
        self._last_valid = valid.clone()
        self._phys += 1
        self._local += 1
        return True, out

    def _boot_prefill(self, persona: str | None) -> None:
        """Batch-wide hybrid system prompt: voice -> cache restore -> silence ->
        persona -> silence (all rows identical)."""
        import torch

        B = self.config.batch_size
        dev = self.config.device
        dummy_user = self._initial[9:].view(1, 8).expand(B, 8)
        dummy_moshi = self._initial[1:9].view(1, 8).expand(B, 8)
        zero_text = torch.full((B,), ZERO_TEXT_TOKEN, dtype=torch.long, device=dev)
        all_rows = torch.ones(B, dtype=torch.bool, device=dev)

        for e in self._voice_embeddings:
            ov = e.view(1, 1, -1).expand(B, 1, -1)
            while True:
                ran, _ = self._tick(dummy_user, dummy_moshi, zero_text, all_rows, embed_override=(ov, all_rows))
                if ran:
                    break
        self._cache.copy_(self._voice_cache.expand_as(self._cache))

        sine = self._sine.view(1, 8).expand(B, 8)
        silence = self._silence.view(1, 8).expand(B, 8)
        text = persona if persona is not None else self.config.persona
        toks = self._tokenizer.encode(wrap_with_system_tags(text)) if text else []
        for phase_toks in (
            [ZERO_TEXT_TOKEN] * AUDIO_SILENCE_FRAME_CNT,
            toks,
            [ZERO_TEXT_TOKEN] * AUDIO_SILENCE_FRAME_CNT,
        ):
            for t in phase_toks:
                tt = torch.full((B,), int(t), dtype=torch.long, device=dev)
                self._tick(sine, silence, tt, all_rows)
        self._codec.reset_streaming()

    # -- batched surface (mirrors the previous engine's API) --------------------

    def open_batch(self) -> None:
        if not self._loaded:
            self.load()
        self._reset_all_streaming()
        self._boot_prefill(None)
        self._opened = True
        logger.info("PersonaPlex native: %d slots prefilled", self.config.batch_size)

    def prefill_steps(self, persona: str | None = None) -> list[PrefillStep]:
        zero = self._silence
        silence = [
            PrefillStep("tokens", moshi_tokens=zero, text_token=ZERO_TEXT_TOKEN, user_sine=True)
            for _ in range(AUDIO_SILENCE_FRAME_CNT)
        ]
        steps: list[PrefillStep] = [PrefillStep("sacrifice")]
        for e in self._voice_embeddings:
            steps.append(PrefillStep("voice", embedding=e.view(1, 1, -1)))
        steps.append(PrefillStep("voice_end"))
        steps.extend(silence)
        text = persona if persona is not None else self.config.persona
        if text:
            for tok in self._tokenizer.encode(wrap_with_system_tags(text)):
                steps.append(PrefillStep("tokens", moshi_tokens=zero, text_token=int(tok), user_sine=True))
        steps.extend(silence)
        return steps

    def recycle_slot(self, b: int) -> None:
        self._temporal.reset_slot(b)
        self._temporal.bump_slot_start(b)  # LM sacrifice-tick semantics
        self._local[b] = 0
        self._cache[b] = self._initial.view(-1, 1).expand(-1, self._cache.shape[2])
        self._provided[b] = False

    def voice_end(self, b: int) -> None:
        import torch

        ct = self._cache.shape[2]
        shift = int((self._phys - int(self._local[b])) % ct)
        self._cache[b] = torch.roll(self._voice_cache[0], shifts=shift, dims=-1)

    def end_prefill(self, b: int) -> None:
        self._codec.reset_slot(b)

    def step_batch(self, user_pcm: NDArray[np.float32], prefill: dict[int, PrefillStep] | None = None):
        import torch

        if not self._opened:
            raise RuntimeError("open_batch() must be called before step_batch()")
        cfg = self.config
        B = cfg.batch_size
        pcm = np.ascontiguousarray(user_pcm, dtype=np.float32).reshape(B, self.frame_size)
        codes = self._codec.encode_frame(torch.from_numpy(pcm).to(cfg.device)).long()
        return self._step_from_codes(codes, prefill)

    def step_codes(self, user_codes, prefill: dict[int, PrefillStep] | None = None):
        """Drive the lockstep loop from precomputed user codes (golden replays).

        Kept public deliberately: bit-exact golden-replay verification must
        bypass the Mimi encoder, whose streaming conv state makes encoded
        codes run-order dependent (valid but not reproducible token-for-token).
        """
        return self._step_from_codes(user_codes.to(self.config.device).long(), prefill)

    def _step_from_codes(self, codes, prefill):
        import torch

        from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput

        cfg = self.config
        B = cfg.batch_size
        dev = cfg.device
        forced_moshi = forced_text = force_mask = embed_override = None
        if prefill:
            force_mask = torch.zeros(B, dtype=torch.bool, device=dev)
            embed_rows = torch.zeros(B, dtype=torch.bool, device=dev)
            emb_full = None
            forced_moshi = self._initial[1:9].view(1, 8).expand(B, 8).clone()
            forced_text = torch.full((B,), ZERO_TEXT_TOKEN, dtype=torch.long, device=dev)
            for b, stp in prefill.items():
                if stp.kind == "voice_end":
                    raise ValueError("voice_end is not a tick; call engine.voice_end(b)")
                codes[b] = self._sine if stp.user_sine else self._initial[9:]
                if stp.kind == "voice":
                    embed_rows[b] = True
                    if emb_full is None:
                        emb_full = torch.zeros(B, 1, stp.embedding.shape[-1], dtype=stp.embedding.dtype, device=dev)
                    emb_full[b] = stp.embedding[0]
                    force_mask[b] = True  # dummy staging like step_embeddings
                else:
                    force_mask[b] = True
                    if stp.kind == "tokens":
                        forced_moshi[b] = stp.moshi_tokens
                        forced_text[b] = int(stp.text_token)
            if bool(embed_rows.any()):
                embed_override = (emb_full, embed_rows)

        ran, out = self._tick(codes, forced_moshi, forced_text, force_mask, embed_override)
        self.last_out_tokens = None if out is None else out.detach().cpu()
        if not ran or out is None:
            return [None] * B

        valid = self._last_valid  # [B] cpu bool
        emit = [bool(valid[b]) and (not prefill or b not in prefill) for b in range(B)]
        audio_tokens = out[:, 1:9].clone()
        if not all(emit):
            hide = torch.tensor([not e for e in emit], device=dev).view(B, 1)
            audio_tokens = torch.where(hide, self._silence.view(1, 8).expand(B, 8), audio_tokens)
        pcm_out = self._codec.decode_frame(audio_tokens)  # [B, frame]
        outs = []
        for b in range(B):
            if not emit[b]:
                outs.append(None)
                continue
            audio = pcm_out[b].float().detach().cpu().numpy().astype(np.float32)
            outs.append(FrameOutput(audio=audio, text=self._decode_text(int(out[b, 0].item()))))
        return outs

    # -- single-session FrameStepper surface -------------------------------------

    def open_session(self, voice_prompt: str | None = None, persona: str | None = None) -> None:
        with self._fs_lock:
            if not self._loaded:
                self.load()
            # Always resolve a voice: falling back to the configured default here
            # keeps a previous session's custom voice from leaking into this one.
            self._load_voice(voice_prompt or self.config.voice_prompt)
            self._reset_all_streaming()
            self._boot_prefill(persona)
            self._opened = True

    def step(self, user_pcm: NDArray[np.float32]):
        with self._fs_lock:
            pcm = np.ascontiguousarray(user_pcm, dtype=np.float32).reshape(-1)
            if pcm.shape[0] != self.frame_size:
                raise ValueError(f"expected {self.frame_size} samples per frame, got {pcm.shape[0]}")
            outs = self.step_batch(pcm.reshape(1, -1))
        from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput

        return outs[0] if outs[0] is not None else FrameOutput(audio=None, text=None)

    def _decode_text(self, token_id: int) -> str | None:
        if token_id in NON_TEXT_TOKEN_IDS or token_id >= TEXT_VOCAB_SIZE:
            return None
        return self._tokenizer.id_to_piece(token_id).replace("▁", " ")