[Frontend] Add midway prompt update for streaming video generation (#4652)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -14,7 +14,7 @@ import copy
 import time
 from collections.abc import Callable
 from contextlib import AbstractContextManager, nullcontext
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from torch.profiler import record_function
@@ -32,7 +32,7 @@ from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import supports_step_execution
+from vllm_omni.diffusion.models.interface import SupportsPromptUpdate, supports_prompt_update, supports_step_execution
 from vllm_omni.diffusion.offloader import get_offload_backend
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -53,6 +53,9 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
+if TYPE_CHECKING:
+    from vllm_omni.inputs.data import OmniInteractionPrompt
+
 logger = init_logger(__name__)
 
 
@@ -803,3 +806,48 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 self._update_states_after(states, input_batch, pipeline_interrupted)
 
                 return BatchRunnerOutput.from_list(runner_output_list)
+
+    def submit_interaction(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+    ) -> None:
+        """Route a midway interaction to the matching active stepwise request feature."""
+        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")
+
+        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
+
+        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):
+            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}")
+
+        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)