diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 678797652..e2b318062 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -779,6 +779,31 @@ turning them on. For Qwen-Image-style serving examples, document `--step-execution` as the feature gate and `--max-num-seqs N` as the companion batching knob. +### Streaming Output and Mid-Stream Interactions + +To add streaming generation with optional mid-stream prompt updates to a new diffusion pipeline (often a video generation pipeline): + +0. **Model should internally support** chunked video emission and bounded VAE decode of a chunk. +1. **Implement stepwise execution** so interaction RPCs can be consumed at chunk + boundaries (see [Step Execution](#step-execution) above). +2. **Mix in** ``InteractionMixin`` on the pipeline class. +3. **Implement** ``peek_chunk_media(StepRequestState) -> ChunkMediaSpec`` to offer information + about the next chunk to generate. + This is called whenever one supported interaction modality needs frame-level calculation, + e.g., per-frame camera trajectory. If only chunk-level calculation is involved, peeking is skipped. +4. **Optionally implement** ``prepare_next_chunk(StepRequestState)`` if the + pipeline must rebuild something (e.g., per-chunk state) after an interaction is applied. + This function is called at chunk boundary. +5. **Register** supported modalities and handlers in ``vllm_omni/diffusion/interaction/registry.py`` + (keyed by ``od_config.model_class_name`` → modality key as in ``req["multi_modal_data"]``). + +Serve with ``--diffusion-streaming-output`` (auto-enables step execution). See +[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md#streaming-output) +and the example script at `examples/online_serving/streaming_video_generation/`. + +**Interaction payload contract:** Although ``OmniInteractionPrompt`` type annocation requires ``event_id``, +the WebSocket API payload may omit ``event_id`` and let the API layer assigns one. + ### Cache Acceleration #### TeaCache diff --git a/docs/user_guide/diffusion/execution_modes.md b/docs/user_guide/diffusion/execution_modes.md index 3aa625295..5eee3c4d9 100644 --- a/docs/user_guide/diffusion/execution_modes.md +++ b/docs/user_guide/diffusion/execution_modes.md @@ -8,7 +8,7 @@ are exposed. ## Choose a Mode | Goal | CLI configuration | -|---|---| +| --- | --- | | Serial request execution | `--max-num-seqs 1` | | Fused request-level batching | `--max-num-seqs N` | | Single-request step execution | `--step-execution --max-num-seqs 1` | @@ -177,8 +177,10 @@ wait The scheduler may batch these requests when their sampling parameters are compatible. See the [Image Generation API](../../serving/image_generation_api.md) for response -formats and additional client examples. Streaming models can use -model-specific streaming endpoints documented by their serving guide. +formats and additional client examples. + +Streaming models can use model-specific streaming endpoints documented at +[Streaming Video Output API](../../serving/streaming_video_output_api.md). ## Python API @@ -241,7 +243,7 @@ For step execution, set `step_execution: true` and remove ## CLI Reference | Flag | Default | Effect | -|---|---:|---| +| --- | ---: | --- | | `--step-execution` | disabled | Select step-wise scheduling | | `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity | | `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode | diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py similarity index 100% rename from tests/diffusion/test_diffusion_streaming_output.py rename to tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py diff --git a/tests/diffusion/test_prompt_update.py b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py similarity index 59% rename from tests/diffusion/test_prompt_update.py rename to tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py index f08825d1a..2e5ceb8e2 100644 --- a/tests/diffusion/test_prompt_update.py +++ b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for midway prompt update across runner, batch, pipeline, and entrypoint layers.""" from __future__ import annotations @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio import queue import time +from contextlib import contextmanager from types import SimpleNamespace from typing import Any, cast from unittest.mock import MagicMock, patch @@ -21,6 +22,11 @@ from tests.engine.test_orchestrator import ( ) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient +from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator +from vllm_omni.diffusion.interaction.modality_handlers.prompt import ( + PromptInteractionHandler, + PromptSession, +) from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.input_batch import InputBatch @@ -54,30 +60,55 @@ def pipeline() -> HeliosPipeline: None, ) ) - pipeline._prepare_next_chunk = MagicMock() + pipeline.prepare_next_chunk = MagicMock() + od_config = SimpleNamespace(model_class_name="HeliosPipeline") + pipeline._interaction_coordinator = InteractionCoordinator.build(pipeline, od_config) return pipeline -def _make_diffusion_request_state(*, request_id: str = "req-1") -> StepRequestState: +def _make_diffusion_request_state(*, request_id: str = "req-1", fps: int | None = None) -> StepRequestState: state = StepRequestState( request_id=request_id, - sampling=SimpleNamespace(num_outputs_per_prompt=1, max_sequence_length=226), # pyright: ignore[reportArgumentType] + sampling=SimpleNamespace( # pyright: ignore[reportArgumentType] + num_outputs_per_prompt=1, + max_sequence_length=226, + fps=fps, + ), prompt="hello", ) state.prompt_embeds = torch.zeros(1, 4, 2) - state.extra = {} + state.extra = {"window_num_frames": 8} return state -def _make_diffusion_model_runner(*, pipeline, streaming_output: bool = True) -> DiffusionModelRunner: +def _make_diffusion_model_runner( + *, + pipeline, + streaming_output: bool = True, + model_class_name: str = "HeliosPipeline", +) -> DiffusionModelRunner: runner = object.__new__(DiffusionModelRunner) runner.pipeline = pipeline runner.state_cache = {} runner.od_config = SimpleNamespace( # pyright: ignore[reportAttributeAccessIssue] - model_class_name="HeliosPipeline", + model_class_name=model_class_name, streaming_output=streaming_output, + cache_backend=None, + parallel_config=SimpleNamespace(use_hsdp=False), ) + runner.vllm_config = object() + runner.device = torch.device("cpu") + runner.cache_backend = None + runner.offload_backend = None + runner.prompt_embed_cache = None + runner.kv_transfer_manager = SimpleNamespace( + receive_multi_kv_cache_distributed=lambda *a, **k: None, + ) + # HeliosPipeline skeletons from object.__new__ are not full SupportsStepExecution + # instances; unit tests of submit_interaction stub this. Stepwise path tests + # rebind to the real method. runner._supports_step_mode = lambda: True + runner._interaction_coordinator = None return runner @@ -108,10 +139,13 @@ class TestPromptUpdateExecution: runner.submit_interaction("req-1", _prompt_interaction(transition_chunks=2)) pipeline.encode_prompt.assert_called_once() # pyright: ignore[reportAttributeAccessIssue] - pending = state.extra["pending_prompt_update"] - assert pending["event_id"] == "ui-update-1" - assert pending["transition_chunks"] == 2 - assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0)) + session = state.interaction_sessions["prompt"] + assert isinstance(session, PromptSession) + pending = session.pending_event + assert pending is not None + assert pending.event_id == "ui-update-1" + assert pending.transition_chunks == 2 + assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0)) def test_runner_prompt_update_rejects_unsupported_pipeline(self) -> None: """Runner rejects prompt updates when the pipeline lacks prompt-update support.""" @@ -135,6 +169,9 @@ class TestPromptUpdateExecution: {}, {"multi_modal_data": {"camera": {"type": "pose"}}}, {"event": {"prompt": "updated", "multi_modal_data": {"camera": {"type": "pose"}}}}, + {"event": {"prompt": "updated", "multi_modal_data": {}}}, + {"event": {"prompt": "updated", "multi_modal_data": None}}, + {"event": {"prompt": "updated", "multi_modal_data": "bad"}}, ], ) def test_runner_interaction_rejects_structural_payloads_until_implemented( @@ -149,6 +186,34 @@ class TestPromptUpdateExecution: with pytest.raises(NotImplementedError, match="Only text-only prompt update interactions"): runner.submit_interaction("req-1", cast(Any, interaction)) + @pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HeliosPyramidPipeline"]) + def test_coordinator_registers_prompt_for_helios_aliases( + self, + pipeline: HeliosPipeline, + model_class_name: str, + ) -> None: + """Both Helios architecture names resolve the prompt handler registry.""" + od_config = SimpleNamespace(model_class_name=model_class_name) + coordinator = InteractionCoordinator.build(pipeline, od_config) + assert coordinator.has_modality("prompt") + + runner = _make_diffusion_model_runner(pipeline=pipeline, model_class_name=model_class_name) + runner.state_cache["req-1"] = _make_diffusion_request_state() + runner.submit_interaction("req-1", _prompt_interaction()) + session = runner.state_cache["req-1"].interaction_sessions["prompt"] + assert isinstance(session, PromptSession) + assert session.pending_event is not None + + def test_helios_peek_chunk_media_requires_positive_fps(self, pipeline: HeliosPipeline) -> None: + """Media-timeline peek must raise on missing/invalid fps instead of coercing to 0.0.""" + state = _make_diffusion_request_state(fps=None) + with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"): + pipeline.peek_chunk_media(state) + + state.sampling.fps = 0 # pyright: ignore[reportAttributeAccessIssue] + with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"): + pipeline.peek_chunk_media(state) + def test_input_batch_refreshes_prompt_embeds_on_version_change(self) -> None: """InputBatch rebuilds prompt_embeds when prompt_update_version changes.""" state = StepRequestState( @@ -159,28 +224,37 @@ class TestPromptUpdateExecution: state.prompt_embeds = torch.zeros(1, 2, 3) state.latents = torch.zeros(1, 2) state.timesteps = torch.tensor([1.0]) - state.extra["prompt_update_version"] = 0 + state.interaction_sessions["prompt"] = PromptSession(version=0) batch = InputBatch.make_batch([state]) assert torch.equal(batch.prompt_embeds, torch.zeros(1, 2, 3)) # pyright: ignore[reportArgumentType] state.prompt_embeds = torch.ones(1, 2, 3) - state.extra["prompt_update_version"] = 1 + state.interaction_sessions["prompt"] = PromptSession(version=1) refreshed = InputBatch.make_batch([state], cached_batch=batch) assert torch.equal(refreshed.prompt_embeds, torch.ones(1, 2, 3)) # pyright: ignore[reportArgumentType] - def test_helios_prepare_prompt_update_queues_pending_target(self, pipeline: HeliosPipeline) -> None: - """HeliosPipeline queues target embeds without mutating current prompt_embeds.""" + def test_helios_prompt_handler_enqueue_queues_pending_target(self, pipeline: HeliosPipeline) -> None: + """Prompt handler queues target embeds without mutating current prompt_embeds.""" state = _make_diffusion_request_state() - pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2) + PromptInteractionHandler.from_pipeline(pipeline).enqueue( + state, + event_id="ui-update-1", + received_at=0.0, + payload={"prompt": "new scene"}, + transition_chunks=2, + ) - pending = state.extra["pending_prompt_update"] - assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0)) - assert pending["transition_chunks"] == 2 + session = state.interaction_sessions["prompt"] + assert isinstance(session, PromptSession) + pending = session.pending_event + assert pending is not None + assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0)) + assert pending.transition_chunks == 2 assert torch.equal(state.prompt_embeds, torch.zeros(1, 4, 2)) # pyright: ignore[reportArgumentType] - def test_helios_prepare_prompt_update_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None: + def test_helios_prompt_handler_enqueue_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None: """Reject prompt updates submitted before initial prompt embeds exist.""" state = _make_diffusion_request_state() state.prompt_embeds = None @@ -189,43 +263,209 @@ class TestPromptUpdateExecution: ValueError, match="prompt_update is not allowed before initial generation has started", ): - pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2) + PromptInteractionHandler.from_pipeline(pipeline).enqueue( + state, + event_id="ui-update-1", + received_at=0.0, + payload={"prompt": "new scene"}, + transition_chunks=2, + ) - assert "pending_prompt_update" not in state.extra + session = state.interaction_sessions.get("prompt") + assert session is None or (isinstance(session, PromptSession) and session.pending_event is None) - def test_helios_apply_prompt_update_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> None: + def test_helios_apply_interaction_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> None: """At chunk boundary, starts transition state and bumps prompt_update_version.""" - state = _make_diffusion_request_state() - pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2) + state = _make_diffusion_request_state(fps=None) + PromptInteractionHandler.from_pipeline(pipeline).enqueue( + state, + event_id="ui-update-1", + received_at=0.0, + payload={"prompt": "new scene"}, + transition_chunks=2, + ) - pipeline._apply_prompt_update_at_chunk_boundary(state) + pipeline.apply_interaction_at_chunk_boundary(state) - assert state.extra.get("prompt_update_state") is not None + session = state.interaction_sessions["prompt"] + assert isinstance(session, PromptSession) + assert session.active_event is not None assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2)) # pyright: ignore[reportArgumentType] - assert state.extra["prompt_update_version"] == 1 - assert state.extra["prompt_update_chunk_metadata"] == { + assert session.version == 1 + assert state.interaction_chunk_metadata is not None + assert state.interaction_chunk_metadata.as_dict() == { "started_event_ids": ["ui-update-1"], "active_event_ids": ["ui-update-1"], "completed_event_ids": [], } - def test_helios_apply_prompt_update_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None: + def test_helios_apply_interaction_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None: """At chunk boundary, interpolates embeds until the target prompt is reached.""" - state = _make_diffusion_request_state() - pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=3) - pipeline._apply_prompt_update_at_chunk_boundary(state) + state = _make_diffusion_request_state(fps=None) + PromptInteractionHandler.from_pipeline(pipeline).enqueue( + state, + event_id="ui-update-1", + received_at=0.0, + payload={"prompt": "new scene"}, + transition_chunks=3, + ) + pipeline.apply_interaction_at_chunk_boundary(state) assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0 / 3.0)) # pyright: ignore[reportArgumentType] for _ in range(2): - pipeline._apply_prompt_update_at_chunk_boundary(state) + pipeline.apply_interaction_at_chunk_boundary(state) assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0)) # pyright: ignore[reportArgumentType] - assert state.extra["prompt_update_version"] == 3 + session = state.interaction_sessions["prompt"] + assert isinstance(session, PromptSession) + assert session.version == 3 # Further chunk boundaries after completion must not bump the version. - pipeline._apply_prompt_update_at_chunk_boundary(state) - assert state.extra["prompt_update_version"] == 3 + pipeline.apply_interaction_at_chunk_boundary(state) + assert session.version == 3 + + def test_stepwise_runner_applies_prompt_update_and_acks( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Drive the production runner loop: submit → apply(no media args) → ACK on next chunk. + + Also regresses ``sampling.fps is None`` on the prompt-only path (no peek). + ACK ids land on the *next* chunk's ``DiffusionOutput`` via + ``_attach_stepwise_metadata`` consuming the prior boundary metadata. + """ + import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module + from vllm_omni.diffusion.interaction.mixin import InteractionMixin + from vllm_omni.diffusion.interaction.types import ChunkMediaSpec + + chunks = [ + DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), # pyright: ignore[reportArgumentType] + DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), # pyright: ignore[reportArgumentType] + ] + + class _InteractiveChunkPipeline(InteractionMixin): + device = torch.device("cpu") + supports_step_execution = True + + def __init__(self) -> None: + self._outputs = chunks + self.decode_calls = 0 + self.transformer = SimpleNamespace(dtype=torch.float32) + self.encode_prompt = MagicMock(return_value=(torch.full((1, 4, 2), 2.0), None)) + self.prepare_next_chunk = MagicMock() + od_config = SimpleNamespace(model_class_name="HeliosPipeline") + self._interaction_coordinator = InteractionCoordinator.build(self, od_config) # pyright: ignore[reportArgumentType] + + def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec: + # Same contract as Helios: never coerce missing fps to 0.0. + fps = state.sampling.fps + if fps is None or float(fps) <= 0: + raise ValueError( + "sampling.fps is required and must be > 0 for interaction modalities " + f"that use the chunk media timeline, got {fps!r} " + f"(request_id={state.request_id!r})" + ) + return ChunkMediaSpec(num_frames=int(state.extra["window_num_frames"]), fps=float(fps)) + + def prepare_encode(self, state: StepRequestState) -> StepRequestState: + state.prompt_embeds = torch.zeros(1, 4, 2) + state.latents = torch.zeros(1, 1) + state.timesteps = torch.tensor([1.0, 0.0, 1.0, 0.0]) + state.step_index = 0 + state.step_in_chunk = 0 + state.chunk_num_steps = 2 + state.total_chunks = 2 + state.extra = {"window_num_frames": 8} + return state + + def denoise_step(self, input_batch: Any, states: Any) -> torch.Tensor: + del states + return torch.ones_like(input_batch.latents) + + def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor | None) -> None: + state.latents = noise_pred + state.step_index += 1 + state.step_in_chunk += 1 + + def post_decode(self, state: StepRequestState) -> DiffusionOutput: + output = self._outputs[self.decode_calls] + self.decode_calls += 1 + state.chunk_index += 1 + state.step_index = 0 + state.step_in_chunk = 0 + if not state.request_denoise_completed: + state.latents = torch.zeros(1, 1) + return output + + pipeline = _InteractiveChunkPipeline() + runner = _make_diffusion_model_runner(pipeline=pipeline) + runner._supports_step_mode = lambda: DiffusionModelRunner._supports_step_mode(runner) + runner.od_config.streaming_output = True + runner.od_config.step_execution = True + runner.diffusion_kv_backend = SimpleNamespace(remove_diffusion_kv_requests=lambda *_: 0) # pyright: ignore[reportAttributeAccessIssue] + + sampling = SimpleNamespace( + generator=None, + seed=None, + generator_device=None, + num_inference_steps=4, + num_outputs_per_prompt=1, + max_sequence_length=226, + fps=None, + ) + req = SimpleNamespace( + request_id="req", + prompt="a prompt", + sampling_params=sampling, + kv_sender_info=None, + ) + + @contextmanager + def _noop_forward_context(*args: Any, **kwargs: Any): + del args, kwargs + yield + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "is_available", lambda: False) + + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req, diffusion_kv_metadata=None)], + scheduled_cached_reqs=SimpleNamespace(request_ids=[]), + ) + # Admit request and run first denoise step (no chunk boundary yet). + DiffusionModelRunner.execute_stepwise(runner, scheduler_output) # pyright: ignore[reportArgumentType] + assert "req" in runner.state_cache + assert runner.state_cache["req"].sampling.fps is None + + runner.submit_interaction("req", _prompt_interaction("new scene", transition_chunks=2)) + + cached = SimpleNamespace( + finished_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]), + ) + # Second step completes chunk 0: apply runs (no peek), metadata staged. + chunk0 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req") # pyright: ignore[reportArgumentType] + assert chunk0 is not None and chunk0.result is not None + assert chunk0.result.started_event_ids == [] + state = runner.state_cache["req"] + assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2)) # pyright: ignore[reportArgumentType] + assert state.interaction_chunk_metadata is not None + assert state.interaction_chunk_metadata.started_event_ids == ["ui-update-1"] + pipeline.prepare_next_chunk.assert_called() # pyright: ignore[reportAttributeAccessIssue] + + # Next chunk decode attaches prior-boundary ACK ids onto DiffusionOutput. + DiffusionModelRunner.execute_stepwise(runner, cached) # pyright: ignore[reportArgumentType] + chunk1 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req") # pyright: ignore[reportArgumentType] + assert chunk1 is not None and chunk1.result is not None + assert chunk1.result.started_event_ids == ["ui-update-1"] + assert chunk1.result.active_event_ids == ["ui-update-1"] + assert chunk1.finished is True class TestPromptUpdateIntegration: @@ -266,8 +506,8 @@ class TestPromptUpdateIntegration: @pytest.mark.asyncio async def test_runner_prompt_update_failure_surfaces_non_fatal_error(self, pipeline: HeliosPipeline) -> None: - """Runner-side prepare_prompt_update rejection is reported through the orchestrator.""" - pipeline.prepare_prompt_update = MagicMock( # pyright: ignore[reportAttributeAccessIssue, reportMethodAssignment] + """Runner-side prompt encode rejection is reported through the orchestrator.""" + pipeline.encode_prompt = MagicMock( # pyright: ignore[reportAttributeAccessIssue] side_effect=ValueError("prompt embeds are not ready") ) runner = _make_diffusion_model_runner(pipeline=pipeline) @@ -293,12 +533,7 @@ class TestPromptUpdateIntegration: finally: inline_client.shutdown() - pipeline.prepare_prompt_update.assert_called_once_with( # pyright: ignore[reportAttributeAccessIssue] - runner.state_cache["req-1"], - "new scene", - "ui-update-1", - 2, - ) + pipeline.encode_prompt.assert_called_once() # pyright: ignore[reportAttributeAccessIssue] msg = output_queue.get_nowait() assert msg == ErrorMessage( error="Failed interaction for request req-1: prompt embeds are not ready", @@ -337,9 +572,10 @@ class TestPromptUpdateIntegration: outputs = await generate_task - pending = runner.state_cache[internal_request_id].extra["pending_prompt_update"] - assert pending["transition_chunks"] == 2 - assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0)) + pending = runner.state_cache[internal_request_id].interaction_sessions["prompt"].pending_event + assert pending is not None + assert pending.transition_chunks == 2 + assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0)) pipeline.encode_prompt.assert_called_once() # pyright: ignore[reportAttributeAccessIssue] assert [output.custom_output["chunk"] for output in outputs] == [0, 1] finally: @@ -544,7 +780,10 @@ class TestPromptUpdateIntegration: ] deadline = time.monotonic() + 5.0 - while "pending_prompt_update" not in state.extra: + while ( + not isinstance(state.interaction_sessions.get("prompt"), PromptSession) + or state.interaction_sessions["prompt"].pending_event is None + ): if time.monotonic() >= deadline: raise TimeoutError("timed out waiting for prompt update during streaming") await asyncio.sleep(0.01) diff --git a/vllm_omni/diffusion/interaction/__init__.py b/vllm_omni/diffusion/interaction/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/interaction/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/interaction/coordinator.py b/vllm_omni/diffusion/interaction/coordinator.py new file mode 100644 index 000000000..5b4986d88 --- /dev/null +++ b/vllm_omni/diffusion/interaction/coordinator.py @@ -0,0 +1,124 @@ +# 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 diff --git a/vllm_omni/diffusion/interaction/mixin.py b/vllm_omni/diffusion/interaction/mixin.py new file mode 100644 index 000000000..366e90588 --- /dev/null +++ b/vllm_omni/diffusion/interaction/mixin.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Interaction-handling mixin for diffusion pipelines. Process interactions at chunk boundaries.""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, cast + +from vllm_omni.diffusion.models.interface import SupportsInteractionApply +from vllm_omni.diffusion.worker.utils import StepRequestState + +if TYPE_CHECKING: + from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator + + +class InteractionMixin: + """Unified chunk-boundary interaction hook. + + Pipelines inherit this mixin and implement ``peek_chunk_media`` to offer playback information about the next chunk. + Diffusion runner calls ``apply_interaction_at_chunk_boundary`` without knowing which modality is being applied. + + Per-modality session state lives on ``StepRequestState.interaction_sessions``; + handler strategy objects live on the pipeline/runner ``InteractionCoordinator``. + """ + + _interaction_coordinator: InteractionCoordinator | None = None + + def apply_interaction_at_chunk_boundary(self, state: StepRequestState) -> None: + """Advance all active interaction tracks before the next chunk.""" + if self._interaction_coordinator is None: + raise RuntimeError( + "interaction coordinator is not initialized; " + "DiffusionModelRunner.load_model must wire InteractionCoordinator " + "onto the pipeline before chunked generation" + ) + + num_frames: int | None = None + fps: float | None = None + if self._interaction_coordinator.needs_chunk_media: + media = cast(SupportsInteractionApply, self).peek_chunk_media(state) + num_frames = media.num_frames + fps = media.fps + + merged = self._interaction_coordinator.apply_at_chunk_boundary( + state, + boundary_at=time.monotonic(), + num_frames=num_frames, + fps=fps, + ) + state.interaction_chunk_metadata = merged + + def prepare_next_chunk(self, state: StepRequestState) -> None: + """Set up pipeline state for the next chunk after interaction apply. + + Default no-op. Pipelines may override this if needed. + """ + pass diff --git a/vllm_omni/diffusion/interaction/modality_handlers/base.py b/vllm_omni/diffusion/interaction/modality_handlers/base.py new file mode 100644 index 000000000..61f548fcd --- /dev/null +++ b/vllm_omni/diffusion/interaction/modality_handlers/base.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Base interaction handler interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import ClassVar, Self + +from vllm_omni.diffusion.interaction.types import ( + InteractionChunkMetadata, + InteractionPayload, +) +from vllm_omni.diffusion.worker.utils import StepRequestState + + +class InteractionHandler(ABC): + """Strategy object for one interaction modality. + + Handler instances are pipeline/runner-owned and request-agnostic. Per-request + session state lives on ``StepRequestState.interaction_sessions``. + Each concrete modality handler decides how to handle timing from ``received_at``. + """ + + modality: ClassVar[str] + # When True, chunk-boundary apply needs ``ChunkMediaSpec`` (num_frames/fps) + # Those information are useful when interaction handler needs interpolation/integration on a frame-by-frame basis + needs_chunk_media: ClassVar[bool] = False + + @classmethod + def from_pipeline(cls, pipeline: object) -> Self: + """Bind a request-agnostic handler to ``pipeline`` when construction needs it. + + Default construction ignores the pipeline. Prompt handlers override this to + capture ``encode_prompt`` / device / dtype. + """ + del pipeline + return cls() + + @abstractmethod + def enqueue( + self, + state: StepRequestState, + *, + event_id: str, + received_at: float, + payload: InteractionPayload, + transition_chunks: int | None, + ) -> None: + """Validate and queue this track on request-local state.""" + + @abstractmethod + def apply_at_chunk_boundary( + self, + state: StepRequestState, + *, + boundary_at: float, + chunk_index: int | None = None, # defaults to state.chunk_index when omitted + num_frames: int | None = None, # only set when at least one interaction handler needs it + fps: float | None = None, # only set when at least one interaction handler needs it + ) -> InteractionChunkMetadata | None: + """Advance request-local state and materialize this chunk's effects.""" diff --git a/vllm_omni/diffusion/interaction/modality_handlers/prompt.py b/vllm_omni/diffusion/interaction/modality_handlers/prompt.py new file mode 100644 index 000000000..e87d9029b --- /dev/null +++ b/vllm_omni/diffusion/interaction/modality_handlers/prompt.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Prompt-track interaction handler (midway prompt updates).""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Any, ClassVar + +import torch +from typing_extensions import override + +from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler +from vllm_omni.diffusion.interaction.types import ( + InteractionChunkMetadata, + InteractionEvent, + InteractionMode, + InteractionPayload, + InteractionSession, +) +from vllm_omni.diffusion.worker.utils import StepRequestState + +DEFAULT_TRANSITION_CHUNKS = 3 + +logger = logging.getLogger(__name__) + + +@dataclass(kw_only=True) +class QueuedPromptEvent(InteractionEvent): + """Last-write-wins prompt update waiting for / active at a chunk boundary. + + ``source_prompt_embeds`` is the lerp start endpoint, staged when this event + is activated (not at enqueue). Live blended embeds live on + ``StepRequestState.prompt_embeds``. + """ + + mode: InteractionMode = "target" + prompt: str + target_prompt_embeds: torch.Tensor + source_prompt_embeds: torch.Tensor | None = None + + def blended_prompt_embeds(self) -> torch.Tensor: + assert self.source_prompt_embeds is not None + alpha = self._current_alpha() + if alpha >= 1.0: + return self.target_prompt_embeds + return (1.0 - alpha) * self.source_prompt_embeds + alpha * self.target_prompt_embeds + + def advance_transition(self) -> None: + assert self.source_prompt_embeds is not None + if self.transition_chunks <= 0: + self.source_prompt_embeds = self.target_prompt_embeds + return + self.elapsed_transition_chunks += 1 + if self.elapsed_transition_chunks >= self.transition_chunks: + self.source_prompt_embeds = self.target_prompt_embeds + self.elapsed_transition_chunks = float(self.transition_chunks) + + def _current_alpha(self) -> float: + if self.transition_chunks <= 0: + return 1.0 + return min(1.0, self.elapsed_transition_chunks / self.transition_chunks) + + +@dataclass +class PromptSession(InteractionSession): + """Request-local prompt-update session under ``state.interaction_sessions['prompt']``.""" + + pending_event: QueuedPromptEvent | None = None + active_event: QueuedPromptEvent | None = None + version: int = 0 + + +def prompt_update_versions(states: Sequence[StepRequestState]) -> tuple[int, ...]: + """Return per-request prompt-update versions for batch cache comparison.""" + versions: list[int] = [] + for state in states: + session = state.interaction_sessions.get("prompt") + versions.append(session.version if isinstance(session, PromptSession) else 0) + return tuple(versions) + + +class PromptInteractionHandler(InteractionHandler): + """Encode-and-queue prompt updates; lerp embeds at chunk boundaries. + + Mirrors the former ``PromptUpdateMixin`` behavior. Encoding happens at + enqueue time via an injected ``encode_prompt`` callable. + """ + + modality: ClassVar[str] = "prompt" + + def __init__( + self, + *, + encode_prompt: Callable[..., tuple[torch.Tensor, ...]], + device: torch.device | str, + dtype: torch.dtype, + ) -> None: + self._encode_prompt = encode_prompt + self._device = device + self._dtype = dtype + + @classmethod + @override + def from_pipeline(cls, pipeline: Any) -> PromptInteractionHandler: + """Build a handler from a diffusion pipeline that exposes ``encode_prompt``.""" + return cls( + encode_prompt=pipeline.encode_prompt, + device=pipeline.device, + dtype=pipeline.transformer.dtype, + ) + + @override + def enqueue( + self, + state: StepRequestState, + *, + event_id: str, + received_at: float, + payload: InteractionPayload, + transition_chunks: int | None, + ) -> None: + """Prompt updates are last-write-win and unbuffered at chunk boundary.""" + prompt = payload.get("prompt") + if not isinstance(prompt, str) or not prompt: + raise ValueError("prompt must be non-empty") + if not event_id: + raise ValueError("event_id must be non-empty") + if state.prompt_embeds is None: + raise ValueError( + f"prompt_update is not allowed before initial generation has started (request_id={state.request_id!r})" + ) + duration = DEFAULT_TRANSITION_CHUNKS if transition_chunks is None else transition_chunks + if duration < 0: + raise ValueError("transition_chunks must be >= 0") + + target_prompt_embeds, _ = self._encode_prompt( + prompt=prompt, + negative_prompt=None, + do_classifier_free_guidance=False, + num_videos_per_prompt=state.sampling.num_outputs_per_prompt, + max_sequence_length=state.sampling.max_sequence_length, + device=self._device, + dtype=self._dtype, + ) + session = state.interaction_sessions.setdefault("prompt", PromptSession()) + assert isinstance(session, PromptSession) + with session.lock: + # Chunk-level LWW: replace any prior pending event. + session.pending_event = QueuedPromptEvent( + event_id=event_id, + received_at=received_at, + transition_chunks=duration, + prompt=prompt, + target_prompt_embeds=target_prompt_embeds, + ) + + @override + 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 | None: + """Advance or start prompt interpolation before the next chunk.""" + # Prompt lerp is chunk-LWW; media timeline and caller chunk_index unused. + del chunk_index, num_frames, fps + session = state.interaction_sessions.get("prompt") + assert isinstance(session, PromptSession) + + embeds_changed = False + next_chunk_index = state.chunk_index + started_event_ids: list[str] = [] + active_event_ids: list[str] = [] + completed_event_ids: list[str] = [] + + with session.lock: + # Prompt ignores frame scheduling but still records the boundary clock. + session.last_boundary_at = boundary_at + pending_event = session.pending_event + session.pending_event = None + active_event = session.active_event + + # If current transition is not complete, advance it. + # After completion, leave active_event in place (so a later pending + # update can abort/overwrite), but do not keep bumping the version. + if active_event is not None: + in_transition = ( + active_event.transition_chunks > 0 + and active_event.elapsed_transition_chunks < active_event.transition_chunks + ) + if in_transition: + active_event.advance_transition() + state.prompt_embeds = active_event.blended_prompt_embeds() + embeds_changed = True + active_event_ids.append(active_event.event_id) + if active_event.elapsed_transition_chunks >= active_event.transition_chunks: + state.prompt_embeds = active_event.target_prompt_embeds + active_event.source_prompt_embeds = active_event.target_prompt_embeds + completed_event_ids.append(active_event.event_id) + logger.debug( + "prompt_update transition complete request_id=%s next_chunk_index=%d prompt=%.20s...", + state.request_id, + next_chunk_index, + active_event.prompt, + ) + + # If a new prompt update is pending, start a new transition. + if pending_event is not None: + if state.prompt_embeds is None: + raise RuntimeError( + "internal error: trying to apply a pending prompt update but " + f"current prompt_embeds is None (request_id={state.request_id!r})" + ) + # Stage lerp-start embeds at activation (after any same-boundary advance). + pending_event.source_prompt_embeds = state.prompt_embeds.detach().clone() + pending_event.elapsed_transition_chunks = 0.0 + session.active_event = pending_event + active_event = pending_event + event_id = pending_event.event_id + duration = pending_event.transition_chunks + prompt = pending_event.prompt + target = pending_event.target_prompt_embeds + started_event_ids.append(event_id) + active_event_ids.append(event_id) + if duration <= 0: + # A hard/immediate transition + state.prompt_embeds = target + pending_event.source_prompt_embeds = target + completed_event_ids.append(event_id) + logger.debug( + "prompt_update sharp transition request_id=%s next_chunk_index=%d prompt=%.20s...", + state.request_id, + next_chunk_index, + prompt, + ) + else: + # A transition that really takes some time + active_event.advance_transition() + state.prompt_embeds = active_event.blended_prompt_embeds() + if active_event.elapsed_transition_chunks >= active_event.transition_chunks: + state.prompt_embeds = active_event.target_prompt_embeds + active_event.source_prompt_embeds = active_event.target_prompt_embeds + completed_event_ids.append(event_id) + logger.debug( + "prompt_update transition start request_id=%s next_chunk_index=%d prompt=%.20s...", + state.request_id, + next_chunk_index, + prompt, + ) + embeds_changed = True + + # Indicate that prompt embeddings have changed---clear input batch cache. + if embeds_changed: + session.version += 1 + + return InteractionChunkMetadata( + started_event_ids=started_event_ids, + active_event_ids=active_event_ids, + completed_event_ids=completed_event_ids, + ) diff --git a/vllm_omni/diffusion/interaction/registry.py b/vllm_omni/diffusion/interaction/registry.py new file mode 100644 index 000000000..f86dab902 --- /dev/null +++ b/vllm_omni/diffusion/interaction/registry.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Registry of mid-way interaction handlers during a diffusion generation. +A pipeline only declares the modalities that it supports. +And a pipeline can choose which handler for each modality (there can be more than one handler for a modality). + +Outer key: pipeline architecture name (``od_config.model_class_name``). +Inner key: modality name (e.g. ``\"prompt\"``). +""" + +from __future__ import annotations + +from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler +from vllm_omni.diffusion.interaction.modality_handlers.prompt import PromptInteractionHandler + +STRUCTURED_HANDLER_REGISTRY: dict[str, dict[str, type[InteractionHandler]]] = { + # Pipeline class name -> modality -> handler class. + "HeliosPipeline": { + "prompt": PromptInteractionHandler, + }, + "HeliosPyramidPipeline": { + "prompt": PromptInteractionHandler, + }, +} diff --git a/vllm_omni/diffusion/interaction/types.py b/vllm_omni/diffusion/interaction/types.py new file mode 100644 index 000000000..904552129 --- /dev/null +++ b/vllm_omni/diffusion/interaction/types.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared types for midway interaction handlers.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from threading import Lock +from typing import Literal + +# Wire payload for one modality track. Handlers validate keys they need. +InteractionPayload = Mapping[str, object] + +# Shared command semantics across modalities that support them. +# Prompt specialization uses ``"target"`` only. +InteractionMode = Literal["target"] + + +@dataclass(frozen=True) +class ChunkMediaSpec: + """Media extent of one generation chunk, for interaction timeline mapping.""" + + num_frames: int + fps: float + + def __post_init__(self) -> None: + if self.num_frames <= 0: + raise ValueError(f"ChunkMediaSpec.num_frames must be > 0, got {self.num_frames}") + if self.fps <= 0: + raise ValueError(f"ChunkMediaSpec.fps must be > 0, got {self.fps}") + + @property + def duration_s(self) -> float: + return float(self.num_frames) / float(self.fps) + + +@dataclass(kw_only=True) +class InteractionEvent: + """Shared envelope for one queued or in-flight interaction command. **Intended to be subclassed**. + + Example: + - ``QueuedPromptEvent`` in ``vllm_omni.diffusion.interaction.modality_handlers.prompt`` + """ + + event_id: str + received_at: float + mode: InteractionMode + transition_chunks: int + elapsed_transition_chunks: float = 0.0 + + +@dataclass +class InteractionSession: + """Request-local session base for one interaction modality. **Intended to be subclassed**. + + Concrete modalities subclass this and store their own pending/active state. + ``last_boundary_at`` is available for frame-scheduled modalities; chunk-LWW + modalities may ignore it. + + Example: + - ``PromptSession`` in ``vllm_omni.diffusion.interaction.modality_handlers.prompt`` + """ + + lock: Lock = field(default_factory=Lock, repr=False) + last_boundary_at: float | None = None + + +@dataclass +class InteractionChunkMetadata: + """Per-modality (or merged) event-id bookkeeping for one chunk boundary.""" + + started_event_ids: list[str] = field(default_factory=list) + active_event_ids: list[str] = field(default_factory=list) + completed_event_ids: list[str] = field(default_factory=list) + + def as_dict(self) -> dict[str, list[str]]: + return { + "started_event_ids": list(self.started_event_ids), + "active_event_ids": list(self.active_event_ids), + "completed_event_ids": list(self.completed_event_ids), + } + + +def merge_interaction_metadata( + metas: list[InteractionChunkMetadata], +) -> InteractionChunkMetadata: + """Merge per-handler metadata while preserving order and de-duplicating ids.""" + started: list[str] = [] + active: list[str] = [] + completed: list[str] = [] + seen_started: set[str] = set() + seen_active: set[str] = set() + seen_completed: set[str] = set() + for meta in metas: + for event_id in meta.started_event_ids: + if event_id not in seen_started: + seen_started.add(event_id) + started.append(event_id) + for event_id in meta.active_event_ids: + if event_id not in seen_active: + seen_active.add(event_id) + active.append(event_id) + for event_id in meta.completed_event_ids: + if event_id not in seen_completed: + seen_completed.add(event_id) + completed.append(event_id) + return InteractionChunkMetadata( + started_event_ids=started, + active_event_ids=active, + completed_event_ids=completed, + ) diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index cb11eb847..45c6d0693 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from Helios (https://github.com/BestWishYsh/Helios) from __future__ import annotations @@ -18,18 +18,20 @@ from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel +from typing_extensions import override from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.interaction.mixin import InteractionMixin +from vllm_omni.diffusion.interaction.types import ChunkMediaSpec from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.prompt_update import PromptUpdateMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.platforms import current_omni_platform @@ -160,7 +162,7 @@ class HeliosPipeline( CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, - PromptUpdateMixin, + InteractionMixin, SupportsComponentDiscovery, ): """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni. @@ -533,10 +535,22 @@ class HeliosPipeline( "zero_steps": int(extra.get("zero_steps", 1)), } ) - self._prepare_next_chunk(state) + self.prepare_next_chunk(state) return state - def _prepare_next_chunk(self, state: StepRequestState) -> None: + def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec: + """Expose this chunk's decoded media extent for interaction timelines.""" + num_frames = int(state.extra["window_num_frames"]) + fps = state.sampling.fps + if fps is None or float(fps) <= 0: + raise ValueError( + "sampling.fps is required and must be > 0 for interaction modalities that use the " + f"chunk media timeline, got {fps!r} (request_id={state.request_id!r})" + ) + return ChunkMediaSpec(num_frames=num_frames, fps=float(fps)) + + @override + def prepare_next_chunk(self, state: StepRequestState) -> None: extra = state.extra k = state.chunk_index is_first_chunk = k == 0 @@ -900,14 +914,9 @@ class HeliosPipeline( output = current_latents if extra["output_type"] == "latent" else current_video completed_chunk_index = state.chunk_index - prompt_update_metadata = state.extra.pop("prompt_update_chunk_metadata", {}) state.chunk_index += 1 finished = state.request_denoise_completed - if not finished: - # Apply queued/advancing prompt updates before preparing the next chunk. - self._apply_prompt_update_at_chunk_boundary(state) - self._prepare_next_chunk(state) - else: + if finished: self._current_timestep = None if current_omni_platform.is_available(): current_omni_platform.empty_cache() @@ -917,9 +926,6 @@ class HeliosPipeline( stage_durations=self.stage_durations if hasattr(self, "stage_durations") else {}, chunk_index=completed_chunk_index, total_chunks=state.total_chunks, - started_event_ids=prompt_update_metadata.get("started_event_ids", []), - active_event_ids=prompt_update_metadata.get("active_event_ids", []), - completed_event_ids=prompt_update_metadata.get("completed_event_ids", []), finished=finished, ) diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index 7a30fa77c..75cc5a0d2 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations from dataclasses import dataclass @@ -19,6 +19,7 @@ if TYPE_CHECKING: from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend from vllm_omni.diffusion.data import DiffusionOutput + from vllm_omni.diffusion.interaction.types import ChunkMediaSpec from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import StepRequestState @@ -108,30 +109,34 @@ def supports_step_execution(pipeline: object) -> bool: @runtime_checkable -class SupportsPromptUpdate(Protocol): - """Optional protocol for pipelines that support midway prompt updates. +class SupportsInteractionApply(Protocol): + """Optional protocol for pipelines with unified mid-generation, chunk-boundary hooks.""" - Pipelines typically implement this via - :class:`~vllm_omni.diffusion.prompt_update.PromptUpdateMixin`. - """ + def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec: + """Return the media timeline represented by the upcoming/current chunk. - supports_prompt_update: ClassVar[bool] = True + Useful when interaction handler needs interpolation/integration on a frame-by-frame basis, + or for backpressure/pacing. + """ + ... - def prepare_prompt_update( - self, - state: StepRequestState, - prompt: str, - event_id: str, - transition_chunks: int | None = None, - ) -> None: - """Encode and queue a prompt update on request-local state.""" + def apply_interaction_at_chunk_boundary(self, state: StepRequestState) -> None: + """Advance queued interactions before the next generation chunk.""" + ... + + def prepare_next_chunk(self, state: StepRequestState) -> None: + """Set up pipeline state for the next chunk after interaction apply. + + Default implementations is a no-op; model-specific pipelines override + when chunk transitions require latent/history bookkeeping. + """ ... -def supports_prompt_update(pipeline: object) -> bool: - """Return whether ``pipeline`` implements :class:`SupportsPromptUpdate`.""" +def supports_interaction_apply(pipeline: object) -> bool: + """Return whether ``pipeline`` implements :class:`SupportsInteractionApply`.""" - return isinstance(pipeline, SupportsPromptUpdate) + return isinstance(pipeline, SupportsInteractionApply) @runtime_checkable diff --git a/vllm_omni/diffusion/prompt_update.py b/vllm_omni/diffusion/prompt_update.py deleted file mode 100644 index 400838035..000000000 --- a/vllm_omni/diffusion/prompt_update.py +++ /dev/null @@ -1,245 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Midway prompt updates for chunked diffusion step execution. - -``PromptUpdateMixin`` exposes -- a pipeline interface for the external to queue a prompt update, -- a pipeline internal method to apply prompt transitions at chunk boundaries. - -``prompt_update_versions``: a helper for ``InputBatch`` cache invalidation. -""" - -from __future__ import annotations - -import logging -from collections.abc import Sequence -from dataclasses import dataclass -from threading import Lock -from typing import ClassVar, TypedDict, cast - -import torch - -from vllm_omni.diffusion.worker.utils import StepRequestState - -DEFAULT_TRANSITION_CHUNKS = 3 - -logger = logging.getLogger(__name__) - - -class PromptUpdateExtra(TypedDict, total=False): - """A "protocol" for ``StepRequestState.extra`` keys used by the prompt-update feature.""" - - pending_prompt_update: _PendingPromptUpdate - prompt_update_state: _PromptUpdateState - prompt_update_chunk_metadata: _PromptUpdateChunkMetadata - prompt_update_version: int # Bumped when prompt_embeds change; invalidates InputBatch cache. - - # The lock for storing/updating pending prompt updates. - # Saved in state.extra to be naturally cleaned up with the request state. - # As PromptUpdateMixin (the pipeline) currently lacks an obvious request-finish hook - prompt_update_lock: Lock - - -def prompt_update_versions(states: Sequence[StepRequestState]) -> tuple[int, ...]: - """Return per-request prompt-update versions for batch cache comparison. - - To be used by ``InputBatch`` or other external to determine if the cached batch is still valid. - """ - return tuple(int(cast(PromptUpdateExtra, state.extra).get("prompt_update_version", 0)) for state in states) - - -class PromptUpdateMixin: - """Mixin for chunked streaming pipelines that support midway prompt updates. - - Implements the behavior expected by :class:`~vllm_omni.diffusion.models.interface.SupportsPromptUpdate`. - All per-request transition state lives on ``StepRequestState.extra``; this mixin itself is stateless. - """ - - supports_prompt_update: ClassVar[bool] = True - - @staticmethod - def _prompt_update_lock(extra: PromptUpdateExtra) -> Lock: - return extra.setdefault("prompt_update_lock", Lock()) - - def prepare_prompt_update( - self, - state: StepRequestState, - prompt: str, - event_id: str, - transition_chunks: int | None = None, - ) -> None: - """Encode and queue a prompt update to apply at the next chunk boundary. - It does not actually apply the prompt transition. - - This is an extended interface for the pipeline, to be called by the runner. - """ - if not prompt: - raise ValueError("prompt must be non-empty") - if not event_id: - raise ValueError("event_id must be non-empty") - if state.prompt_embeds is None: - raise ValueError( - f"prompt_update is not allowed before initial generation has started (request_id={state.request_id!r})" - ) - duration = DEFAULT_TRANSITION_CHUNKS if transition_chunks is None else transition_chunks - if duration < 0: - raise ValueError("transition_chunks must be >= 0") - - target_prompt_embeds, _ = self.encode_prompt( - prompt=prompt, - negative_prompt=None, - do_classifier_free_guidance=False, - num_videos_per_prompt=state.sampling.num_outputs_per_prompt, - max_sequence_length=state.sampling.max_sequence_length, - device=self.device, - dtype=self.transformer.dtype, - ) - extra = cast(PromptUpdateExtra, state.extra) - with self._prompt_update_lock(extra): - extra["pending_prompt_update"] = { - "event_id": event_id, - "prompt": prompt, - "target_prompt_embeds": target_prompt_embeds, - "transition_chunks": duration, - } - - def _apply_prompt_update_at_chunk_boundary(self, state: StepRequestState) -> None: - """Advance or start prompt interpolation before the next chunk.""" - extra = cast(PromptUpdateExtra, state.extra) - update_state: _PromptUpdateState | None = extra.get("prompt_update_state") - with self._prompt_update_lock(extra): - pending = extra.pop("pending_prompt_update", None) - embeds_changed = False - next_chunk_index = state.chunk_index - started_event_ids: list[str] = [] - active_event_ids: list[str] = [] - completed_event_ids: list[str] = [] - - # If current transition is not complete, advance it. - # After completion, leave prompt_update_state in place (so a later pending - # update can abort/overwrite), but do not keep bumping the version. - if update_state is not None: - in_transition = ( - update_state.transition_chunks > 0 - and update_state.elapsed_transition_chunks < update_state.transition_chunks - ) - if in_transition: - update_state.advance_transition() - state.prompt_embeds = update_state.blended_prompt_embeds() - embeds_changed = True - active_event_ids.append(update_state.event_id) - if update_state.elapsed_transition_chunks >= update_state.transition_chunks: - state.prompt_embeds = update_state.target_prompt_embeds - update_state.source_prompt_embeds = update_state.target_prompt_embeds - completed_event_ids.append(update_state.event_id) - if update_state.target_prompt is not None: - logger.debug( - "prompt_update transition complete request_id=%s next_chunk_index=%d prompt=%.20s...", - state.request_id, - next_chunk_index, - update_state.target_prompt, - ) - - # If a new prompt update is pending, start a new transition. - if pending is not None: - if state.prompt_embeds is None: - raise RuntimeError( - "internal error: trying to apply a pending prompt update but " - f"current prompt_embeds is None (request_id={state.request_id!r})" - ) - source = state.prompt_embeds.detach().clone() - target = pending["target_prompt_embeds"] - duration = int(pending["transition_chunks"]) - prompt = str(pending.get("prompt", "")) - event_id = pending.get("event_id") - update_state = _PromptUpdateState( - source_prompt_embeds=source, - target_prompt_embeds=target, - transition_chunks=duration, - event_id=event_id, - target_prompt=prompt, - ) - extra["prompt_update_state"] = update_state - started_event_ids.append(event_id) - active_event_ids.append(event_id) - if duration <= 0: - # A hard/immediate transition - state.prompt_embeds = target - completed_event_ids.append(event_id) - logger.debug( - "prompt_update sharp transition request_id=%s next_chunk_index=%d prompt=%.20s...", - state.request_id, - next_chunk_index, - prompt, - ) - else: - # A transition that really takes some time - update_state.advance_transition() - state.prompt_embeds = update_state.blended_prompt_embeds() - if update_state.elapsed_transition_chunks >= update_state.transition_chunks: - state.prompt_embeds = update_state.target_prompt_embeds - update_state.source_prompt_embeds = update_state.target_prompt_embeds - completed_event_ids.append(event_id) - logger.debug( - "prompt_update transition start request_id=%s next_chunk_index=%d prompt=%.20s...", - state.request_id, - next_chunk_index, - prompt, - ) - embeds_changed = True - - # Indicate that prompt embeddings have changed---clear input batch cache. - if embeds_changed: - new_version = int(extra.get("prompt_update_version", 0)) + 1 - extra["prompt_update_version"] = new_version - - extra["prompt_update_chunk_metadata"] = { - "started_event_ids": started_event_ids, - "active_event_ids": active_event_ids, - "completed_event_ids": completed_event_ids, - } - - -@dataclass -class _PromptUpdateState: - """Per-request prompt interpolation stored in ``state.extra["prompt_update_state"]``.""" - - source_prompt_embeds: torch.Tensor - target_prompt_embeds: torch.Tensor - transition_chunks: int - event_id: str - elapsed_transition_chunks: int = 0 - target_prompt: str | None = None - - def blended_prompt_embeds(self) -> torch.Tensor: - alpha = self._current_alpha() - if alpha >= 1.0: - return self.target_prompt_embeds - return (1.0 - alpha) * self.source_prompt_embeds + alpha * self.target_prompt_embeds - - def advance_transition(self) -> None: - if self.transition_chunks <= 0: - self.source_prompt_embeds = self.target_prompt_embeds - return - self.elapsed_transition_chunks += 1 - if self.elapsed_transition_chunks >= self.transition_chunks: - self.source_prompt_embeds = self.target_prompt_embeds - self.elapsed_transition_chunks = self.transition_chunks - - def _current_alpha(self) -> float: - if self.transition_chunks <= 0: - return 1.0 - return min(1.0, self.elapsed_transition_chunks / self.transition_chunks) - - -class _PendingPromptUpdate(TypedDict): - event_id: str - prompt: str - target_prompt_embeds: torch.Tensor - transition_chunks: int - - -class _PromptUpdateChunkMetadata(TypedDict): - started_event_ids: list[str] - active_event_ids: list[str] - completed_event_ids: list[str] diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 8a352fb7d..8804785fa 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -42,12 +42,13 @@ from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( ) from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_rank from vllm_omni.diffusion.forward_context import set_forward_context +from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import ( - SupportsPromptUpdate, + SupportsInteractionApply, adopt_request_scoped_cache_dit, is_request_scoped_cache_dit_enabled, - supports_prompt_update, + supports_interaction_apply, supports_step_execution, ) from vllm_omni.diffusion.offloader import enable_offload_backend @@ -176,6 +177,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.offload_backend: Any | None = None self.prompt_embed_cache: Any | None = None self.input_batch: InputBatch | None = None + self._interaction_coordinator: InteractionCoordinator | None = None self.model_memory_usage = 0 self.diffusion_kv_backend = DiffusionKVModelRunnerBackend( vllm_config=vllm_config, @@ -427,6 +429,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): model_tag=self.od_config.model_class_name, ) + self._interaction_coordinator = InteractionCoordinator.build(self.pipeline, self.od_config) + if hasattr(self.pipeline, "_interaction_coordinator"): + self.pipeline._interaction_coordinator = self._interaction_coordinator + logger.info("Model runner: Initialization complete.") def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: @@ -789,7 +795,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): return self._runner_output_from_outputs(reqs, outputs) - def _attach_stepwise_metrics( + def _attach_stepwise_metadata( self, state: StepRequestState, output: DiffusionOutput, @@ -800,6 +806,14 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) attach_stage_durations(state, output) + # In streaming output mode with interaction, acknowledge which events are handled in this chunk. + meta = state.interaction_chunk_metadata + state.interaction_chunk_metadata = None + if meta is not None: + output.started_event_ids = list(meta.started_event_ids) + output.active_event_ids = list(meta.active_event_ids) + output.completed_event_ids = list(meta.completed_event_ids) + def execute_model( self, req: OmniDiffusionRequest, @@ -1275,10 +1289,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): clear_pipeline_stage_durations(self.pipeline) result = self.pipeline.post_decode(req) if result is not None: - self._attach_stepwise_metrics( + self._attach_stepwise_metadata( req, result, ) + # After consuming this chunk's interaction metadata, apply pending interactions and + # prepare the next chunk (prepare_next_chunk may be a no-op---depending on pipeline) + if supports_interaction_apply(self.pipeline) and not req.request_denoise_completed: + pipe = cast(SupportsInteractionApply, self.pipeline) + pipe.apply_interaction_at_chunk_boundary(req) + pipe.prepare_next_chunk(req) else: result = None # finished should be computed after post_decode() advanced chunk_index @@ -1354,42 +1374,48 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): request_id: str, interaction: OmniInteractionPrompt, ) -> None: - """Route a midway interaction to the matching active stepwise request feature.""" + """Route a midway interaction through the pipeline interaction coordinator.""" assert self.pipeline is not None, "Model not loaded. Call load_model() first." if not self.od_config.streaming_output: raise ValueError("submit_interaction requires streaming_output=True") if not self._supports_step_mode(): raise ValueError("submit_interaction requires step execution support") + coordinator = self._interaction_coordinator + if coordinator is None: + coordinator = InteractionCoordinator.build(self.pipeline, self.od_config) + self._interaction_coordinator = coordinator + if hasattr(self.pipeline, "_interaction_coordinator"): + self.pipeline._interaction_coordinator = coordinator + event = interaction.get("event") - if isinstance(event, dict) and "prompt" in event and "multi_modal_data" not in event: - # Is a prompt update interaction. - self._submit_prompt_update_interaction(request_id, interaction) - return + has_mm = isinstance(event, dict) and "multi_modal_data" in event + has_prompt = isinstance(event, dict) and "prompt" in event and event.get("prompt") is not None - raise NotImplementedError( - "Only text-only prompt update interactions with 'event.prompt' and optional " - "'transition_chunks' are supported in this release" - ) - - def _submit_prompt_update_interaction( - self, - request_id: str, - interaction: OmniInteractionPrompt, - ) -> None: - """Queue a prompt-update interaction for an active stepwise request.""" - if not supports_prompt_update(self.pipeline): + # Prompt-only interactions in this release; multi_modal_data lands with camera support. + if not isinstance(event, dict) or has_mm or not has_prompt: + raise NotImplementedError( + "Only text-only prompt update interactions with 'event.prompt' and optional " + "'transition_chunks' are supported in this release" + ) + if not coordinator.has_modality("prompt"): raise ValueError(f"prompt_update is not supported by pipeline {self.od_config.model_class_name!r}") state = self.state_cache.get(request_id) if state is None: - raise ValueError(f"No active request state for prompt_update: {request_id!r}") + raise ValueError(f"No active request state for interaction: {request_id!r}") - event = cast(dict[str, Any], interaction.get("event")) - prompt = event["prompt"] - transition_chunks = interaction.get("transition_chunks") event_id = interaction.get("event_id") if not isinstance(event_id, str) or not event_id: raise ValueError("event_id must be non-empty") - pipeline = cast(SupportsPromptUpdate, self.pipeline) - pipeline.prepare_prompt_update(state, prompt, event_id, transition_chunks) + prompt = event["prompt"] + if not isinstance(prompt, str) or not prompt: + raise ValueError("prompt must be non-empty") + coordinator.enqueue( + state, + modality="prompt", + event_id=event_id, + received_at=time.monotonic(), + payload={"prompt": prompt}, + transition_chunks=interaction.get("transition_chunks"), + ) diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py index d21e73287..0e9a2bde1 100644 --- a/vllm_omni/diffusion/worker/input_batch.py +++ b/vllm_omni/diffusion/worker/input_batch.py @@ -17,7 +17,7 @@ from dataclasses import dataclass, field import numpy as np import torch -from vllm_omni.diffusion.prompt_update import prompt_update_versions +from vllm_omni.diffusion.interaction.modality_handlers.prompt import prompt_update_versions from vllm_omni.diffusion.worker.utils import StepRequestState diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index 62def6a0c..5158dbfbf 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -12,6 +12,10 @@ import torch if TYPE_CHECKING: from vllm_omni.diffusion.data import DiffusionOutput + from vllm_omni.diffusion.interaction.types import ( + InteractionChunkMetadata, + InteractionSession, + ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType @@ -100,6 +104,10 @@ class StepRequestState: total_chunks: int = 1 chunk_num_steps: int | None = None + # ── Optional interaction information in streaming output mode ── + interaction_sessions: dict[str, InteractionSession] = field(default_factory=dict) # Modality -> transition progress + interaction_chunk_metadata: InteractionChunkMetadata | None = None # Acknowledging completion of each interaction + # ── Per-request scheduler instance (set once by prepare_encode) ── scheduler: Any | None = None diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 3ef934bf5..f9f0a9a9a 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -46,11 +46,14 @@ class OmniInteractionEvent(OmniTextPrompt, total=False): pass -class OmniInteractionPrompt(TypedDict, total=False): - """Mid-stream interaction payload.""" +class OmniInteractionPrompt(TypedDict): + """Normalized mid-stream interaction payload passed to ``submit_interaction``. - event_id: NotRequired[str] - event: NotRequired[OmniInteractionEvent] + WebSocket clients payload may omit ``event_id``; the API layer assigns a fallback ID in this case. + """ + + event_id: str + event: OmniInteractionEvent transition_chunks: NotRequired[int]