125 lines
4.6 KiB
Python
125 lines
4.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Interaction coordinator to be co-owned by a pipeline and a runner."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm_omni.diffusion.interaction.registry import STRUCTURED_HANDLER_REGISTRY
|
|
from vllm_omni.diffusion.interaction.types import (
|
|
InteractionChunkMetadata,
|
|
InteractionPayload,
|
|
merge_interaction_metadata,
|
|
)
|
|
from vllm_omni.diffusion.models.interface import supports_interaction_apply
|
|
from vllm_omni.diffusion.worker.utils import StepRequestState
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler
|
|
|
|
|
|
class InteractionCoordinator:
|
|
"""Handles mid-way interaction inputs for one loaded pipeline.
|
|
|
|
Request-local session state lives on ``StepRequestState.interaction_sessions``;
|
|
this object is request-agnostic.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
handlers: dict[str, InteractionHandler],
|
|
*,
|
|
model_class_name: str | None = None,
|
|
) -> None:
|
|
self._handlers = handlers
|
|
self.model_class_name = model_class_name
|
|
|
|
@classmethod
|
|
def build(cls, pipeline: object, od_config: OmniDiffusionConfig) -> InteractionCoordinator:
|
|
model_class_name = od_config.model_class_name
|
|
handlers: dict[str, InteractionHandler] = {}
|
|
|
|
if supports_interaction_apply(pipeline):
|
|
for modality, handler_cls in STRUCTURED_HANDLER_REGISTRY.get(model_class_name or "", {}).items():
|
|
handlers[modality] = handler_cls.from_pipeline(pipeline)
|
|
|
|
return cls(handlers, model_class_name=model_class_name)
|
|
|
|
def has_modality(self, modality: str) -> bool:
|
|
return modality in self._handlers
|
|
|
|
@property
|
|
def needs_chunk_media(self) -> bool:
|
|
"""Whether any registered handler needs chunk ``num_frames``/``fps``."""
|
|
return any(handler.needs_chunk_media for handler in self._handlers.values())
|
|
|
|
def get_handler(self, modality: str) -> InteractionHandler:
|
|
handler = self._handlers.get(modality)
|
|
if handler is None:
|
|
raise ValueError(
|
|
f"interaction modality {modality!r} is not supported by pipeline {self.model_class_name!r}"
|
|
)
|
|
return handler
|
|
|
|
def enqueue(
|
|
self,
|
|
state: StepRequestState,
|
|
*,
|
|
modality: str,
|
|
event_id: str,
|
|
received_at: float,
|
|
payload: InteractionPayload,
|
|
transition_chunks: int | None,
|
|
) -> None:
|
|
"""Fan out enqueue to handlers in stable order (prompt first)."""
|
|
handler = self.get_handler(modality)
|
|
handler.enqueue(
|
|
state,
|
|
event_id=event_id,
|
|
received_at=received_at,
|
|
payload=payload,
|
|
transition_chunks=transition_chunks,
|
|
)
|
|
|
|
def apply_at_chunk_boundary(
|
|
self,
|
|
state: StepRequestState,
|
|
*,
|
|
boundary_at: float,
|
|
chunk_index: int | None = None,
|
|
num_frames: int | None = None,
|
|
fps: float | None = None,
|
|
) -> InteractionChunkMetadata:
|
|
"""Fan out chunk-boundary apply to handlers in stable order (prompt first)."""
|
|
if chunk_index is None:
|
|
chunk_index = state.chunk_index
|
|
metas: list[InteractionChunkMetadata] = []
|
|
for handler in self._handlers_in_apply_order():
|
|
if handler.modality not in state.interaction_sessions:
|
|
# This function is always called at chunk boundary without knowing if an interaction is really enqueued.
|
|
# `state.interaction_sessions[this modality]` is lazy-populated after the first interaction is enqueued.
|
|
# So, check here and skip a modality handler if this modality's session is not yet created.
|
|
continue
|
|
meta = handler.apply_at_chunk_boundary(
|
|
state,
|
|
boundary_at=boundary_at,
|
|
chunk_index=chunk_index,
|
|
num_frames=num_frames,
|
|
fps=fps,
|
|
)
|
|
if meta is not None:
|
|
metas.append(meta)
|
|
return merge_interaction_metadata(metas)
|
|
|
|
def _handlers_in_apply_order(self) -> list[InteractionHandler]:
|
|
ordered: list[InteractionHandler] = []
|
|
if "prompt" in self._handlers:
|
|
ordered.append(self._handlers["prompt"])
|
|
for modality, handler in self._handlers.items():
|
|
if modality == "prompt":
|
|
continue
|
|
ordered.append(handler)
|
|
return ordered
|