vllm-omni/tests/engine/test_output_modality.py

"""Unit tests for Phase 1 foundation types (RFC #1601).

Note: Uses importlib to load modules directly, bypassing the vllm_omni
package __init__ which requires the vllm base package.
"""

import importlib.util
import sys
from pathlib import Path

import pytest
import torch

# ── Load modules without triggering vllm_omni.__init__ ─────────────
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

_OUTPUTS_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "outputs"


def _load_module(name: str, filepath: Path):
    spec = importlib.util.spec_from_file_location(name, filepath)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod


_om_mod = _load_module(
    "vllm_omni.outputs.output_modality",
    _OUTPUTS_DIR / "output_modality.py",
)
_load_module(
    "vllm_omni.outputs.utils",
    _OUTPUTS_DIR / "utils.py",
)
_mm_mod = _load_module(
    "vllm_omni.outputs.mm_outputs",
    _OUTPUTS_DIR / "mm_outputs.py",
)

OutputModality = _om_mod.OutputModality
TensorAccumulationStrategy = _om_mod.TensorAccumulationStrategy
get_accumulation_strategy = _om_mod.get_accumulation_strategy
register_key_accumulation_strategy = _om_mod.register_key_accumulation_strategy
MultimodalPayload = _mm_mod.MultimodalPayload
MultimodalCompletionOutput = _mm_mod.MultimodalCompletionOutput


def test_output_modality_parsing_and_flags():
    """Test OutputModality enum: from_string, aliases, compounds, properties, and accumulation strategy."""
    # Defaults
    assert OutputModality.from_string(None) == OutputModality.TEXT
    assert OutputModality.from_string("") == OutputModality.TEXT

    # Direct names and case insensitivity
    assert OutputModality.from_string("image") == OutputModality.IMAGE
    assert OutputModality.from_string("Audio") == OutputModality.AUDIO

    # Aliases
    assert OutputModality.from_string("speech") == OutputModality.AUDIO
    assert OutputModality.from_string("latents") == OutputModality.LATENT
    assert OutputModality.from_string("pixel_values") == OutputModality.IMAGE

    # Compound
    compound = OutputModality.from_string("text+image")
    assert compound.has_text and compound.has_multimodal

    # Flag properties
    assert OutputModality.TEXT.has_text and not OutputModality.TEXT.has_multimodal
    assert OutputModality.IMAGE.has_multimodal and not OutputModality.IMAGE.has_text

    # Accumulation strategy
    assert get_accumulation_strategy(OutputModality.AUDIO) == TensorAccumulationStrategy.CONCAT_LAST
    assert get_accumulation_strategy(OutputModality.IMAGE) == TensorAccumulationStrategy.CONCAT_DIM0

    # Unknown raises
    with pytest.raises(ValueError, match="Unknown modality"):
        OutputModality.from_string("video")


def test_multimodal_payload_and_completion_output():
    """Test MultimodalPayload and MultimodalCompletionOutput wrapper."""
    # Payload from_dict separates tensors and metadata
    data = {"waveform": torch.ones(1, 16000), "sample_rate": 16000}
    p = MultimodalPayload.from_dict(data)
    assert p is not None
    assert "waveform" in p.tensors and torch.equal(p.primary_tensor, data["waveform"])
    assert p.metadata["sample_rate"] == 16000
    assert not p.is_empty and len(p) == 2  # 1 tensor + 1 metadata

    # None/empty returns None
    assert MultimodalPayload.from_dict(None) is None
    assert MultimodalPayload.from_dict({}) is None

    wrapper = MultimodalCompletionOutput(
        multimodal_output=p,
        index=0,
        text="hello",
        token_ids=[],
        cumulative_logprob=None,
        logprobs=None,
    )
    assert wrapper.text == "hello"
    assert wrapper.multimodal_output is p


def test_output_modality_printed_examples(capsys):
    """Printed examples for output modality types."""
    print("\n=== OutputModality Parsing ===")
    for s in [None, "", "image", "Audio", "speech", "latents", "pixel_values", "text+image"]:
        print(f"  from_string({s!r:20s}) -> {OutputModality.from_string(s)}")

    print("\n=== Flag Properties ===")
    for m in [
        OutputModality.TEXT,
        OutputModality.IMAGE,
        OutputModality.AUDIO,
        OutputModality.TEXT | OutputModality.IMAGE,
    ]:
        print(f"  {str(m):40s} has_text={m.has_text}  has_multimodal={m.has_multimodal}")

    print("\n=== Accumulation Strategies ===")
    for m in [OutputModality.AUDIO, OutputModality.IMAGE, OutputModality.LATENT]:
        print(f"  {str(m):30s} -> {get_accumulation_strategy(m)}")

    print("\n=== MultimodalPayload ===")
    data = {"waveform": torch.ones(1, 16000), "sample_rate": 16000}
    p = MultimodalPayload.from_dict(data)
    print("  from_dict({waveform: tensor, sample_rate: 16000})")
    print(f"    tensors keys : {list(p.tensors.keys())}")
    print(f"    primary_tensor: shape={p.primary_tensor.shape}, dtype={p.primary_tensor.dtype}")
    print(f"    metadata      : {p.metadata}")
    print(f"    is_empty={p.is_empty}, len={len(p)}")  # len = tensors + metadata
    print(f"  from_dict(None) -> {MultimodalPayload.from_dict(None)}")
    print(f"  from_dict({{}})   -> {MultimodalPayload.from_dict({})}")

    print("\n=== MultimodalCompletionOutput ===")
    wrapper = MultimodalCompletionOutput(
        multimodal_output=p,
        index=0,
        text="hello",
        token_ids=[],
        cumulative_logprob=None,
        logprobs=None,
    )
    print(f"  text             : {wrapper.text}")
    print(f"  index            : {wrapper.index}")
    print(f"  multimodal_output: {wrapper.multimodal_output}")
    print(f"  repr             : {wrapper!r}")

    print("\n=== Unknown Modality ===")
    try:
        OutputModality.from_string("video")
    except ValueError as e:
        print(f'  from_string("video") raised ValueError: {e}')

    captured = capsys.readouterr()
    assert "OutputModality Parsing" in captured.out
    assert "MultimodalPayload" in captured.out


def test_get_accumulation_strategy_key_override_precedence():
    """A registered per-key override wins over the modality-wide default;
    any other key under the same modality keeps using that default."""
    assert get_accumulation_strategy(OutputModality.AUDIO, "unregistered.key") == TensorAccumulationStrategy.CONCAT_LAST

    register_key_accumulation_strategy("__test_override__.frames", TensorAccumulationStrategy.CONCAT_DIM0)
    assert (
        get_accumulation_strategy(OutputModality.AUDIO, "__test_override__.frames")
        == TensorAccumulationStrategy.CONCAT_DIM0
    )
    # A plain waveform key under the same AUDIO modality is unaffected.
    assert get_accumulation_strategy(OutputModality.AUDIO, "waveform") == TensorAccumulationStrategy.CONCAT_LAST
    assert get_accumulation_strategy(OutputModality.AUDIO) == TensorAccumulationStrategy.CONCAT_LAST


def _talker_codec_frame_payload(audio_key: str, ref_key: str) -> MultimodalPayload:
    """Synthetic stand-in for a TTS talker's accumulated per-step output.

    ``audio_key`` grows one ``[1, num_codebooks]`` row per decode step on top
    of a ``[5, num_codebooks]`` prefill block (Qwen3-TTS's ``codes.audio``);
    ``ref_key`` re-emits the same constant ``[100, num_codebooks]``
    reference-context matrix unchanged at every step (Qwen3-TTS's
    ``codes.ref``). Neither is a waveform chunk.
    """
    payload = MultimodalPayload()
    payload.tensors[audio_key] = [
        torch.zeros(5, 16),
        torch.full((1, 16), 1.0),
        torch.full((1, 16), 2.0),
        torch.full((1, 16), 3.0),
    ]
    ref_block = torch.arange(100 * 16, dtype=torch.float32).reshape(100, 16)
    payload.tensors[ref_key] = [ref_block.clone(), ref_block.clone(), ref_block.clone()]
    return payload


def test_consolidate_tensors_default_audio_strategy_corrupts_codec_frame_keys():
    """Reproduces both failure modes of forcing codec-frame keys through the
    AUDIO modality's waveform-tuned default (CONCAT_LAST) with no per-key
    override registered: the audio-frames key's dim-0 growth makes
    CONCAT_LAST's dim=-1 concat raise (mismatched dim 0), which the
    flatten-chunks fallback turns into a 1-D blob instead of the intended
    ``[8, 16]``; the ref-frames key's matching last dim lets CONCAT_LAST
    "succeed" by silently concatenating three identical ``[100, 16]`` blocks
    into ``[100, 48]`` instead of keeping a single copy.
    """
    payload = _talker_codec_frame_payload("unregistered.codes.audio", "unregistered.codes.ref")

    payload.consolidate_tensors(OutputModality.AUDIO)

    assert tuple(payload.tensors["unregistered.codes.audio"].shape) != (8, 16)
    assert tuple(payload.tensors["unregistered.codes.ref"].shape) == (100, 48)


def test_consolidate_tensors_with_qwen3_tts_key_overrides_is_correct():
    """With the per-key overrides Qwen3-TTS's pipeline module registers
    (``codes.audio`` -> CONCAT_DIM0, ``codes.ref`` -> REPLACE), both keys
    consolidate correctly under the AUDIO modality despite not being
    waveform chunks."""
    register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0)
    register_key_accumulation_strategy("codes.ref", TensorAccumulationStrategy.REPLACE)

    payload = _talker_codec_frame_payload("codes.audio", "codes.ref")
    expected_ref = payload.tensors["codes.ref"][0]

    payload.consolidate_tensors(OutputModality.AUDIO)

    assert tuple(payload.tensors["codes.audio"].shape) == (8, 16)
    assert tuple(payload.tensors["codes.ref"].shape) == (100, 16)
    assert torch.equal(payload.tensors["codes.ref"], expected_ref)


def test_consolidate_tensors_raises_with_key_name_instead_of_silently_keeping_last():
    """A concat failure under a non-CONCAT_LAST strategy must be surfaced
    with the offending key name rather than resolved by silently keeping
    only the last chunk. Previously the except-branch only special-cased
    the literal key ``"audio"``, so any other mismatched key (e.g. a
    codec-frame key routed to CONCAT_DIM0) was corrupted instead of
    failing loudly.
    """
    payload = MultimodalPayload()
    payload.tensors["latent.frames"] = [torch.zeros(2, 8), torch.zeros(3, 9)]

    with pytest.raises(RuntimeError, match="latent.frames"):
        payload.consolidate_tensors(OutputModality.LATENT)