286 lines
11 KiB
Python
286 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Multimodal output data structures for vLLM-Omni.
|
|
|
|
This module defines structured types for multimodal outputs.
|
|
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterator, Mapping
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
from vllm.outputs import CompletionOutput
|
|
|
|
from vllm_omni.outputs.output_modality import (
|
|
OutputModality,
|
|
TensorAccumulationStrategy,
|
|
get_accumulation_strategy,
|
|
)
|
|
from vllm_omni.outputs.utils import _is_tensor_list, _to_cpu
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Keys whose values are metadata scalars (e.g. audio sample rate) but may
|
|
# arrive as 0-d torch.Tensors — from_dict routes all tensors into .tensors,
|
|
# so we relocate them to .metadata before consolidation to avoid a bogus
|
|
# torch.cat attempt and its warn-and-keep-last fallback.
|
|
_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"})
|
|
|
|
|
|
def _cat_tensors(
|
|
tensors: list[torch.Tensor],
|
|
strategy: TensorAccumulationStrategy,
|
|
) -> torch.Tensor:
|
|
"""Concatenate a list of tensors according to *strategy*."""
|
|
if strategy == TensorAccumulationStrategy.CONCAT_LAST:
|
|
return torch.cat(tensors, dim=-1)
|
|
if strategy == TensorAccumulationStrategy.REPLACE:
|
|
return tensors[-1]
|
|
# CONCAT_DIM0 / APPEND_LIST / default
|
|
return torch.cat(tensors, dim=0)
|
|
|
|
|
|
def _consolidate_tensor_list(
|
|
key: str,
|
|
tensor_list: list[torch.Tensor],
|
|
strategy: TensorAccumulationStrategy,
|
|
) -> torch.Tensor:
|
|
"""Concatenate a deferred tensor list, with a fallback for CONCAT_LAST only.
|
|
|
|
A shape mismatch under CONCAT_LAST is expected for genuine audio waveform
|
|
chunks (their non-last dims can vary between async-chunk steps), so that
|
|
case falls back to flattening each chunk before concatenating. Any other
|
|
strategy failing to concatenate means the tensors are not shaped the way
|
|
that strategy assumes (e.g. a key was assigned the wrong strategy), so it
|
|
is surfaced instead of silently keeping only the last chunk.
|
|
"""
|
|
try:
|
|
return _cat_tensors(tensor_list, strategy)
|
|
except RuntimeError as exc:
|
|
if strategy is not TensorAccumulationStrategy.CONCAT_LAST:
|
|
raise RuntimeError(
|
|
f"Error concatenating tensor list for key {key!r} under strategy {strategy}: {exc}"
|
|
) from exc
|
|
logger.warning("CONCAT_LAST failed for key %s; flattening each chunk before concatenating", key)
|
|
return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0)
|
|
|
|
|
|
def _append_entries(store: dict[str, Any], incoming: dict[str, Any]) -> None:
|
|
"""Append *incoming* entries into *store*.
|
|
|
|
Tensor values are collected into lists (deferred concatenation, avoiding
|
|
O(n²) repeated torch.cat); non-tensor values replace the existing entry.
|
|
"""
|
|
for key, new_value in incoming.items():
|
|
existing = store.get(key)
|
|
if isinstance(existing, list) and isinstance(new_value, torch.Tensor):
|
|
existing.append(new_value)
|
|
elif isinstance(existing, torch.Tensor) and isinstance(new_value, torch.Tensor):
|
|
store[key] = [existing, new_value]
|
|
else:
|
|
store[key] = new_value
|
|
|
|
|
|
@dataclass(eq=False)
|
|
class MultimodalPayload(Mapping):
|
|
"""Structured multimodal output payload.
|
|
|
|
Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)``
|
|
style checks in downstream code can be replaced with duck-typing, and
|
|
``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)``
|
|
all work seamlessly for both tensors and metadata.
|
|
|
|
Attributes:
|
|
tensors: Dictionary mapping modality/key names to their tensors.
|
|
metadata: Optional dictionary for non-tensor metadata
|
|
(e.g., sample rate for audio, image dimensions).
|
|
"""
|
|
|
|
tensors: dict[str, torch.Tensor] = field(default_factory=dict)
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
@property
|
|
def primary_tensor(self) -> torch.Tensor | None:
|
|
"""Return the first tensor in the payload, or None if empty."""
|
|
if self.tensors:
|
|
return next(iter(self.tensors.values()))
|
|
return None
|
|
|
|
@property
|
|
def is_empty(self) -> bool:
|
|
"""Return True if the payload has no tensors and no metadata."""
|
|
return len(self.tensors) == 0 and len(self.metadata) == 0
|
|
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
"""Get a value by key, searching tensors first then metadata."""
|
|
if key in self.tensors:
|
|
return self.tensors[key]
|
|
return self.metadata.get(key, default)
|
|
|
|
def __contains__(self, key: object) -> bool:
|
|
return key in self.tensors or key in self.metadata
|
|
|
|
def __getitem__(self, key: str) -> Any:
|
|
"""Dict-like indexing: search tensors first, then metadata."""
|
|
if key in self.tensors:
|
|
return self.tensors[key]
|
|
if key in self.metadata:
|
|
return self.metadata[key]
|
|
raise KeyError(key)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.tensors) + len(self.metadata)
|
|
|
|
def __iter__(self) -> Iterator[str]:
|
|
yield from self.tensors
|
|
yield from self.metadata
|
|
|
|
def __bool__(self) -> bool:
|
|
return bool(self.tensors) or bool(self.metadata)
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
"""Support equality with plain dicts and other Mappings."""
|
|
if isinstance(other, MultimodalPayload):
|
|
return self.tensors == other.tensors and self.metadata == other.metadata
|
|
if isinstance(other, Mapping):
|
|
return dict(self) == dict(other)
|
|
return NotImplemented
|
|
|
|
def merged_with(self, incoming: MultimodalPayload) -> MultimodalPayload:
|
|
"""Merge *incoming* onto this payload and return the result.
|
|
|
|
Content tensors accumulate into lists for deferred concatenation;
|
|
known sample-rate keys are snapshots replaced immediately, including
|
|
in DELTA streams that do not consolidate each emission. Missing keys
|
|
retain their previous value. Other values keep the existing merge
|
|
behavior. When this payload
|
|
is empty, *incoming* is returned as-is, so callers should use the
|
|
return value: ``accumulated = accumulated.merged_with(incoming)``.
|
|
"""
|
|
if self.is_empty:
|
|
return incoming
|
|
# Capture before merging: incoming may be this payload. Store these
|
|
# snapshots in only one partition even if their representation changes.
|
|
snapshots = {key: incoming[key] for key in _METADATA_TENSOR_KEYS if key in incoming}
|
|
_append_entries(self.tensors, incoming.tensors)
|
|
_append_entries(self.metadata, incoming.metadata)
|
|
for key, value in snapshots.items():
|
|
self.tensors.pop(key, None)
|
|
self.metadata[key] = value
|
|
return self
|
|
|
|
def consolidate_tensors(self, modality: OutputModality) -> None:
|
|
"""Concatenate deferred tensor lists into single tensors.
|
|
|
|
Tensors are generated content accumulated as chunks, so each key's
|
|
list is concatenated according to the strategy
|
|
``get_accumulation_strategy(modality, key)`` resolves for it (e.g.
|
|
audio waveform chunks along the time dimension, latent frames along
|
|
the batch dimension). Most keys share *modality*'s default, but a
|
|
key can be registered for a different strategy -- e.g. a codec-frame
|
|
matrix that grows along dim 0 rather than the waveform-tuned default
|
|
for its modality -- via ``register_key_accumulation_strategy``.
|
|
"""
|
|
# Relocate scalar metadata tensors (e.g. sample rate) that from_dict
|
|
# routed into .tensors, so they take the REPLACE path via .metadata
|
|
# instead of failing torch.cat as 0-d tensors.
|
|
for key in _METADATA_TENSOR_KEYS:
|
|
value = self.tensors.pop(key, None)
|
|
if value is None:
|
|
continue
|
|
self.metadata[key] = value[-1] if isinstance(value, list) else value
|
|
|
|
for key, value in list(self.tensors.items()):
|
|
if _is_tensor_list(value):
|
|
strategy = get_accumulation_strategy(modality, key)
|
|
self.tensors[key] = _consolidate_tensor_list(key, value, strategy)
|
|
|
|
def consolidate_metadata(self) -> None:
|
|
"""Resolve deferred tensor lists in metadata by keeping the latest value.
|
|
|
|
Metadata values are per-step snapshots (e.g. sample rate), not
|
|
content deltas, so the latest value supersedes earlier ones.
|
|
Nested dicts (unflattened payloads) are resolved one level down.
|
|
"""
|
|
for key, value in list(self.metadata.items()):
|
|
if _is_tensor_list(value):
|
|
self.metadata[key] = value[-1]
|
|
elif isinstance(value, dict):
|
|
for sub_key, sub_value in list(value.items()):
|
|
if _is_tensor_list(sub_value):
|
|
value[sub_key] = sub_value[-1]
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""Convert back to a plain dict (tensors + metadata merged)."""
|
|
result: dict[str, Any] = dict(self.tensors)
|
|
result.update(self.metadata)
|
|
return result
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None:
|
|
"""Create a MultimodalPayload from a raw dictionary.
|
|
|
|
Separates torch.Tensor values into tensors and everything
|
|
else into metadata.
|
|
"""
|
|
if not data:
|
|
return None
|
|
tensors: dict[str, torch.Tensor] = {}
|
|
metadata: dict[str, Any] = {}
|
|
for k, v in data.items():
|
|
if isinstance(v, torch.Tensor):
|
|
tensors[k] = v
|
|
else:
|
|
metadata[k] = v
|
|
if not tensors and not metadata:
|
|
return None
|
|
return cls(tensors=tensors, metadata=metadata)
|
|
|
|
@classmethod
|
|
def from_raw(cls, payload: Any, modality_key: str) -> MultimodalPayload | None:
|
|
"""Create a MultimodalPayload from a raw producer payload.
|
|
|
|
Accepts a MultimodalPayload (returned as-is), a dict, or a bare
|
|
tensor (stored under *modality_key*). Tensors are moved to CPU.
|
|
Producer-specific dict keys are remapped to the semantic modality
|
|
key (e.g. "audio", "latent"): AR runners produce {"hidden": ...}
|
|
and generation runners produce {"model_outputs": ...}.
|
|
"""
|
|
if isinstance(payload, MultimodalPayload):
|
|
return payload
|
|
|
|
if not isinstance(payload, dict):
|
|
return cls.from_dict({modality_key: _to_cpu(payload)})
|
|
|
|
remapped: dict[str, Any] = {}
|
|
for key, value in payload.items():
|
|
is_producer_key = key == "model_outputs" or (key == "hidden" and modality_key != "hidden")
|
|
remapped[modality_key if is_producer_key else key] = _to_cpu(value)
|
|
return cls.from_dict(remapped)
|
|
|
|
|
|
@dataclass
|
|
class MultimodalCompletionOutput(CompletionOutput):
|
|
"""CompletionOutput with multimodal support.
|
|
|
|
Inherits all CompletionOutput fields and adds multimodal_output.
|
|
As a CompletionOutput subclass, compatible with all existing vLLM consumers.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
multimodal_output: MultimodalPayload | None = None,
|
|
**kwargs: Any,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.multimodal_output = multimodal_output
|
|
|
|
def __repr__(self) -> str:
|
|
base = super().__repr__()
|
|
return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})"
|