vllm-omni/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Shared fixtures and argument builders for MiniCPM-o 4.5 duplex E2E tests."""

from __future__ import annotations

import hashlib
import uuid
import wave
from pathlib import Path
from types import SimpleNamespace

import pytest
from huggingface_hub import snapshot_download

from tests.helpers.runtime import OmniServerParams, get_model_prefix
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_duplex_max_sessions

MODEL = "openbmb/MiniCPM-o-4_5"
DEPLOY_CONFIG_REL = "minicpmo_4_5.yaml"
DEPLOY_CONFIG = get_deploy_config_path(DEPLOY_CONFIG_REL)
ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"
SOFT_INTERRUPT_WAV = ASSET_DIR / "soft_interrupt_16k.wav"
SOFT_INTERRUPT_SHA256 = "cadae6d0ddc510310f16f8775d6379f30a5195369ccb54ff10a8e3f9a1f4a2ea"
REF_AUDIO_RELATIVE_PATH = Path("assets") / "HT_ref_audio.wav"

SERVER_PARAMS = [
    pytest.param(
        OmniServerParams(
            model=MODEL,
            stage_config_path=DEPLOY_CONFIG,
            use_stage_cli=False,
            server_args=["--trust-remote-code"],
        ),
        id="three-stage-single-gpu",
    )
]


def deploy_max_sessions() -> int:
    """Concurrent duplex sessions the deploy config under test admits.

    The admission probe has to expect the capacity the server is actually
    started with. Hardcoding it silently drifts the moment the deploy config
    changes its capacity, turning that change into an unrelated probe timeout.
    """
    return get_deploy_duplex_max_sessions(DEPLOY_CONFIG_REL)


def validated_wav(path: Path, expected_sha256: str) -> Path:
    actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest()
    if actual_sha256 != expected_sha256:
        raise AssertionError(
            f"MiniCPM-o duplex fixture SHA256 mismatch: expected {expected_sha256}, got {actual_sha256}"
        )
    with wave.open(str(path), "rb") as wav_file:
        actual_format = (
            wav_file.getframerate(),
            wav_file.getnchannels(),
            wav_file.getsampwidth(),
            wav_file.getcomptype(),
        )
    expected_format = (16_000, 1, 2, "NONE")
    if actual_format != expected_format:
        raise AssertionError(f"MiniCPM-o duplex fixture must be 16 kHz mono PCM16, got {actual_format}")
    return path


def validated_input_wav() -> Path:
    return validated_wav(RESPONSE_REQUIRED_WAV, RESPONSE_REQUIRED_SHA256)


def validated_soft_interrupt_wav() -> Path:
    return validated_wav(SOFT_INTERRUPT_WAV, SOFT_INTERRUPT_SHA256)


def resolve_ref_audio() -> Path:
    # ``MODEL`` may be a local checkpoint directory instead of a Hub repo id
    # (``Path("/prefix") / "/abs/path"`` collapses to the absolute path), so only
    # fall back to the Hub cache when no such directory exists on disk.
    model_prefix = get_model_prefix()
    model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
    if not model_root.is_dir():
        model_root = Path(snapshot_download(MODEL, local_files_only=True))
    ref_audio = model_root / REF_AUDIO_RELATIVE_PATH
    if not ref_audio.is_file():
        raise FileNotFoundError(f"MiniCPM-o checkpoint ref audio is missing: {ref_audio}")
    return ref_audio


def duplex_camera_frames(*, seconds: int, cache_dir: Path) -> list[str]:
    """Base64 JPEG frames of a moving synthetic scene, one per second.

    Mirrors what the realtime web client captures (1 fps, no client-side
    resize: the server normalizes at ``scale_resolution=448``).
    """
    from tests.e2e.online_serving.helpers.minicpmo_realtime_duplex_scenarios import _video_frames_from_file
    from tests.helpers.media import generate_synthetic_video

    video = generate_synthetic_video(448, 448, 30 * seconds, cache_dir=cache_dir)
    return _video_frames_from_file(Path(video["file_path"]))


def realtime_url(omni_server) -> str:
    return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1"


def demo_args(
    *,
    omni_server,
    input_wav: Path,
    ref_audio: Path,
    output_dir: Path,
) -> SimpleNamespace:
    return SimpleNamespace(
        url=realtime_url(omni_server),
        model=omni_server.model,
        session_id=f"duplex-ci-single-{uuid.uuid4().hex}",
        input_wav=str(input_wav),
        ref_audio=str(ref_audio),
        turn_input_wav=[],
        output_dir=str(output_dir),
        output_audio_format="pcm16",
        chunk_ms=200,
        realtime_input=True,
        input_video=None,
        stack_frames=1,
        first_turn_ms=1400,
        turn_duration_ms=[],
        first_turn_transcript="duplex CI speech",
        omit_transcript_hints=False,
        validation_mode="response-required",
        temperature=0.0,
        scenario="sequential",
        require_audio=True,
        require_distinct_inputs=False,
        expect_empty_turn=[],
        short_ack_ms=350,
        turns=1,
        timeout_s=120.0,
        model_policy_settle_ms=2000,
    )


def multi_session_args(
    *,
    omni_server,
    input_wav: Path,
    ref_audio: Path,
    output_dir: Path,
    response_required: bool,
) -> SimpleNamespace:
    return SimpleNamespace(
        url=realtime_url(omni_server),
        model=omni_server.model,
        sessions=2,
        input_wav=str(input_wav),
        ref_audio=str(ref_audio),
        session_input_wav=[],
        session_expected_token=[],
        turn_input_wav=[],
        output_dir=str(output_dir),
        realtime_input=True,
        chunk_ms=200,
        turns=1,
        first_turn_ms=1400,
        turn_duration_ms=[],
        response_required=response_required,
        temperature=0.0 if response_required else None,
        disconnect_session_index=0,
        takeover_session_index=1,
        resume_after_ms=100,
        expire_session_index=None,
        expire_after_s=40.0,
        verify_admission_limit=None,
        model_policy_settle_ms=2000,
        timeout_s=180.0,
    )