"""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)