vllm-omni/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py

from __future__ import annotations

import base64
from copy import deepcopy
from typing import Any

import numpy as np
from vllm.multimodal.media import MediaConnector

from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45DuplexPolicy
from vllm_omni.experimental.fullduplex.openai.protocol import DuplexSessionConfig
from vllm_omni.experimental.fullduplex.openai.runtime_adapter import (
    ServingRuntimeConfigError,
)


class MiniCPMO45ClientRuntimeConfigError(ServingRuntimeConfigError):
    pass


class MiniCPMO45NativeDuplexServingAdapter:
    """Serving-side MiniCPM-o 4.5 native duplex session preparation.

    The generic duplex WebSocket handler should not let client-supplied local
    paths reach workers.  This adapter follows the existing media connector
    boundary: serving resolves client media URIs, then workers receive only
    normalized PCM payloads.  Server-owned model assets remain local paths here.
    """

    PRIVATE_RUNTIME_CONFIG_KEYS = frozenset(
        {
            "duplex_stage_sampling_params",
            "duplex_stage_max_tokens",
            "duplex_stage0_max_tokens",
            "duplex_scheduler_token_id",
            "duplex_first_append_context_tokens",
            "ref_audio_data",
            "ref_audio_format",
            "ref_audio_sample_rate_hz",
            "initial_user_text",
        }
    )

    @classmethod
    def is_enabled(cls, config: DuplexSessionConfig) -> bool:
        return config.extra_body.get("minicpmo45_native_duplex") is True

    @classmethod
    def validate_client_config(cls, config: DuplexSessionConfig) -> None:
        cls.validate_client_extra_body(config.extra_body)

    @classmethod
    def validate_client_extra_body(cls, extra_body: object) -> None:
        if not isinstance(extra_body, dict):
            return
        private_keys = sorted(cls.PRIVATE_RUNTIME_CONFIG_KEYS.intersection(extra_body))
        if private_keys:
            raise MiniCPMO45ClientRuntimeConfigError(
                "native duplex runtime configuration is server-owned: " + ", ".join(private_keys)
            )

    @classmethod
    def runtime_config_for_update(
        cls,
        config: DuplexSessionConfig,
        current: object,
    ) -> dict[str, object]:
        runtime_config = deepcopy(dict(current)) if isinstance(current, dict) else {}
        runtime_config["instructions"] = config.instructions
        stage_max_tokens = runtime_config.get("duplex_stage_max_tokens")
        stage_max_tokens = deepcopy(stage_max_tokens) if isinstance(stage_max_tokens, dict) else {}
        stage_max_tokens["0"] = (
            config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20
        )
        stage_max_tokens.setdefault("1", 8192)
        runtime_config["duplex_stage_max_tokens"] = stage_max_tokens

        stage_sampling = runtime_config.get("duplex_stage_sampling_params")
        stage_sampling = deepcopy(stage_sampling) if isinstance(stage_sampling, dict) else {}
        stage0 = stage_sampling.get("0")
        stage0 = deepcopy(stage0) if isinstance(stage0, dict) else {}
        stage0["temperature"] = config.temperature if config.temperature is not None else 0.7
        stage_sampling["0"] = stage0
        runtime_config["duplex_stage_sampling_params"] = stage_sampling
        return runtime_config

    @classmethod
    async def prepare_runtime_config(cls, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]:
        extra_body = dict(config.extra_body)
        if any(key in extra_body for key in ("ref_audio_path", "tts_ref_audio_path")):
            raise ValueError("ref_audio_path is not accepted by native duplex; use ref_audio URI instead")
        cls.validate_client_config(config)
        runtime_config: dict[str, object] = {"instructions": config.instructions}
        initial_user_text = extra_body.pop("duplex_initial_user_text", None)
        if isinstance(initial_user_text, str) and initial_user_text:
            runtime_config["initial_user_text"] = initial_user_text
        cls._apply_default_scheduler_policy(runtime_config, config=config, model_config=model_config)

        ref_audio = config.ref_audio
        if ref_audio is None and isinstance(extra_body.get("ref_audio"), str):
            ref_audio = extra_body.pop("ref_audio")
        if ref_audio is None and isinstance(extra_body.get("tts_ref_audio"), str):
            ref_audio = extra_body.pop("tts_ref_audio")

        if ref_audio is None:
            if any(str(modality).lower() == "audio" for modality in config.modalities):
                raise MiniCPMO45ClientRuntimeConfigError(
                    "MiniCPM-o native duplex audio output requires ref_audio",
                    code="ref_audio_required",
                )
            cls._apply_first_append_context_tokens(
                runtime_config,
                model_config=model_config,
                instructions=config.instructions,
                initial_user_text=initial_user_text,
                ref_sample_count=None,
            )
            config.extra_body = extra_body
            return runtime_config
        else:
            wav_np, sr = await cls.resolve_ref_audio(ref_audio, model_config=model_config)

        wav_np = cls.normalize_ref_audio(wav_np, int(sr), target_sr=16000)
        # Trim to a whole number of pooled audio embeddings (100 ms frames) so
        # the first-append scheduler reserve can count them exactly.
        usable = (len(wav_np) // MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN) * (
            MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN
        )
        wav_np = wav_np[:usable]
        ref_audio_bytes = np.ascontiguousarray(wav_np, dtype=np.float32).tobytes()
        runtime_config["ref_audio_data"] = base64.b64encode(ref_audio_bytes).decode("ascii")
        runtime_config["ref_audio_format"] = "pcm_f32le"
        runtime_config["ref_audio_sample_rate_hz"] = 16000
        cls._apply_first_append_context_tokens(
            runtime_config,
            model_config=model_config,
            instructions=config.instructions,
            initial_user_text=initial_user_text,
            ref_sample_count=len(wav_np),
        )
        config.extra_body = extra_body
        config.ref_audio = None
        return runtime_config

    @classmethod
    def _apply_default_scheduler_policy(
        cls,
        runtime_config: dict[str, object],
        *,
        config: DuplexSessionConfig,
        model_config: Any,
    ) -> None:
        stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20
        runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192}
        stage0_params: dict[str, object] = {
            "temperature": config.temperature if config.temperature is not None else 0.7,
            "top_p": 0.8,
            "top_k": 20,
            "repetition_penalty": 1.05,
        }
        stop_token_ids = cls._native_stage0_stop_token_ids(model_config)
        if stop_token_ids:
            stage0_params["stop_token_ids"] = stop_token_ids
        runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params}
        scheduler_token_id = cls._native_scheduler_token_id(model_config)
        if scheduler_token_id is not None:
            runtime_config["duplex_scheduler_token_id"] = scheduler_token_id

    @classmethod
    def _apply_first_append_context_tokens(
        cls,
        runtime_config: dict[str, object],
        *,
        model_config: Any,
        instructions: object,
        initial_user_text: object,
        ref_sample_count: int | None,
    ) -> None:
        """Precompute the exact session-context token count for the engine.

        The first data-plane append carries the system template and optional
        reference-audio embeddings ahead of the first unit. The engine
        reserves scheduler slots from this count; an inexact count turns into
        pad embeddings inside the model KV (surplus) or truncated context
        (deficit), so it is computed with the same template and pooling math
        the worker uses.
        """
        if "duplex_first_append_context_tokens" in runtime_config:
            return
        tokenizer = cls._load_native_tokenizer(model_config)
        if tokenizer is None:
            return
        prefix, suffix = MiniCPMO45DuplexPolicy.session_context_texts(
            instructions,
            ref_sample_count is not None,
            initial_user_text,
        )
        try:
            prefix_ids = tokenizer.encode(prefix, add_special_tokens=False)
            suffix_ids = tokenizer.encode(suffix, add_special_tokens=False)
        except Exception:
            return
        ref_tokens = MiniCPMO45DuplexPolicy.audio_token_count(ref_sample_count or 0)
        runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids)

    @staticmethod
    def _native_stage0_stop_token_ids(model_config: Any) -> list[int]:
        tokenizer = MiniCPMO45NativeDuplexServingAdapter._load_native_tokenizer(model_config)
        if tokenizer is None:
            return []
        out: list[int] = []
        stop_token_fields = (
            "chunk_eos_token_id",
            "chunk_tts_eos_token_id",
            "listen_token_id",
            "turn_eos_token_id",
        )
        for field in stop_token_fields:
            token = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS[field]
            token_id = MiniCPMO45NativeDuplexServingAdapter._convert_token_to_id(tokenizer, token)
            if token_id is not None and token_id not in out:
                out.append(token_id)
        return out

    @staticmethod
    def _native_scheduler_token_id(model_config: Any) -> int | None:
        tokenizer = MiniCPMO45NativeDuplexServingAdapter._load_native_tokenizer(model_config)
        if tokenizer is None:
            return None
        scheduler_tokens = (
            MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS["unit_token_id"],
            MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS["audio_placeholder_token_id"],
        )
        for token in scheduler_tokens:
            token_id = MiniCPMO45NativeDuplexServingAdapter._convert_token_to_id(tokenizer, token)
            if token_id is not None:
                return token_id
        eos_id = getattr(tokenizer, "eos_token_id", None)
        try:
            return int(eos_id)
        except (TypeError, ValueError):
            return None

    @staticmethod
    def _load_native_tokenizer(model_config: Any) -> Any | None:
        model_path = getattr(model_config, "model", None)
        if not isinstance(model_path, str) or not model_path:
            return None
        try:
            from transformers import AutoTokenizer

            return AutoTokenizer.from_pretrained(
                model_path,
                trust_remote_code=True,
                local_files_only=True,
            )
        except Exception:
            return None

    @staticmethod
    def _convert_token_to_id(tokenizer: Any, token: str) -> int | None:
        convert = getattr(tokenizer, "convert_tokens_to_ids", None)
        value = None
        if callable(convert):
            value = convert(token)
            if isinstance(value, list):
                value = value[0] if len(value) == 1 else None
        try:
            token_id = int(value)
        except (TypeError, ValueError):
            token_id = -1
        unk_token_id = getattr(tokenizer, "unk_token_id", None)
        if token_id >= 0 and token_id != unk_token_id:
            return token_id
        encode = getattr(tokenizer, "encode", None)
        if callable(encode):
            try:
                ids = list(encode(token, add_special_tokens=False))
            except TypeError:
                ids = list(encode(token))
            if len(ids) == 1:
                try:
                    token_id = int(ids[0])
                except (TypeError, ValueError):
                    token_id = -1
                if token_id >= 0 and token_id != unk_token_id:
                    return token_id
        return None

    @staticmethod
    async def resolve_ref_audio(ref_audio: str, *, model_config: Any) -> tuple[np.ndarray, int]:
        connector = MediaConnector(
            allowed_local_media_path=getattr(model_config, "allowed_local_media_path", None),
            allowed_media_domains=getattr(model_config, "allowed_media_domains", None),
        )
        wav_np, sr = await connector.fetch_audio_async(ref_audio)
        return np.asarray(wav_np, dtype=np.float32), int(sr)

    @staticmethod
    def normalize_ref_audio(wav_np: np.ndarray, sample_rate: int, *, target_sr: int) -> np.ndarray:
        wav_np = np.asarray(wav_np, dtype=np.float32)
        if wav_np.ndim > 1:
            wav_np = wav_np.mean(axis=-1)
        wav_np = wav_np.reshape(-1)
        if sample_rate <= 0 or sample_rate == target_sr or wav_np.size == 0:
            return wav_np.astype(np.float32, copy=False)
        import torch
        import torchaudio

        audio = torch.from_numpy(wav_np).to(dtype=torch.float32).unsqueeze(0)
        resampled = torchaudio.functional.resample(audio, int(sample_rate), int(target_sr))
        return resampled.squeeze(0).cpu().numpy().astype(np.float32, copy=False)