anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Core][Diffusion] Camera interaction for diffusion streaming generation (LingBot World 2 as example) (#7198)

Browse Source 
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
This commit is contained in:
Zeyu Huang | 黃澤宇 authored and GitHub committed 2026-09-20 16:19:41 +00:00
1 parent 139a47a578
commit e36babd48f
26 files changed
+2071 -200
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/worker/diffusion_model_runner.py
+57 -52
View File
@@ -43,6 +43,7 @@ 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.interaction.types import InteractionPayload
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import (
SupportsInteractionApply,
@@ -1080,30 +1081,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
error_outputs: list[RunnerOutput] = []
for state in states:
if state.request_id in new_request_ids:
# Everything that runs before ``_dit_any_rank_failed`` must be
# inside the try: an exception in ``_initialize_generator`` or
# Everything that requires rank-synchronization must be called
# inside a try, record the exception and handle with `_dit_any_rank_failed`.
# Reason (example): An exception in ``_initialize_generator`` or
# ``clear_pipeline_stage_durations`` on one rank would skip the
# all-reduce here while every peer proceeds into it, and the
# peers then hang on the NCCL collective until timeout.
per_req_exc: BaseException | None = None
try:
self._initialize_generator(state.sampling)
clear_pipeline_stage_durations(pipeline)
# encode
pipeline.prepare_encode(state)
merge_stage_durations(
state,
consume_pipeline_stage_durations(pipeline),
)
except Exception as exc:
per_req_exc = exc
# Pipelines that do rank-0-only work (e.g. MiniMax H3
# reference-video prep) must broadcast per-request failures
# internally so downstream collectives stay in step; even so,
# cross-check that every DiT rank agrees so a rank-local error
# (or a future pipeline that omits the guard) does not leave
# the process group half-way through a new request.
if _dit_any_rank_failed(per_req_exc is not None):
def _abort_prep_failure(per_req_exc: BaseException | None) -> None:
self.state_cache.pop(state.request_id, None)
if per_req_exc is None:
per_req_exc = RuntimeError(
@@ -1123,6 +1107,41 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
result=DiffusionOutput.from_exception(per_req_exc),
)
)
per_req_exc: BaseException | None = None
try:
self._initialize_generator(state.sampling)
clear_pipeline_stage_durations(pipeline)
pipeline.prepare_encode(state)
except Exception as exc:
per_req_exc = exc
# Pipelines that do rank-0-only work (e.g. MiniMax H3
# reference-video prep) must broadcast per-request failures
# internally so downstream collectives stay in step; even so,
# cross-check that every DiT rank agrees so a rank-local error
# (or a future pipeline that omits the guard) does not leave
# the process group half-way through a new request.
if _dit_any_rank_failed(per_req_exc is not None):
_abort_prep_failure(per_req_exc)
continue
# If the pipeline supports interaction, the interaction session initialization also needs to call
# synchronized_monotonic_time(). Wrap in another try-block to not block on prepare_encode failures.
try:
if supports_interaction_apply(pipeline) and state.chunk_index == 0:
pipe = cast(SupportsInteractionApply, pipeline)
assert self._interaction_coordinator is not None, "Model not loaded. Call load_model() first."
state.interaction_chunk_metadata = self._interaction_coordinator.maybe_prepare_initial_session(
state, pipe
)
pipe.prepare_next_chunk(state)
merge_stage_durations(
state,
consume_pipeline_stage_durations(pipeline),
)
except Exception as exc:
per_req_exc = exc
if _dit_any_rank_failed(per_req_exc is not None):
_abort_prep_failure(per_req_exc)
continue
prepared_states.append(state)
@@ -1431,47 +1450,33 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
interaction: OmniInteractionPrompt,
) -> None:
"""Route a midway interaction through the pipeline interaction coordinator."""
assert self.pipeline is not None, "Model not loaded. Call load_model() first."
assert self.pipeline is not None and self._interaction_coordinator 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")
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
# 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 interaction: {request_id!r}")
event_id = interaction.get("event_id")
if not isinstance(event_id, str) or not event_id:
raise ValueError("event_id must be non-empty")
prompt = event["prompt"]
if not isinstance(prompt, str) or not prompt:
raise ValueError("prompt must be non-empty")
coordinator.enqueue(
event = interaction["event"]
parts: list[tuple[str, InteractionPayload]] = []
prompt = event.get("prompt")
if prompt is not None:
parts.append(("prompt", {"prompt": prompt}))
multi_modal_data = event.get("multi_modal_data")
if multi_modal_data:
parts.extend(
(str(modality), cast(InteractionPayload, payload)) for modality, payload in multi_modal_data.items()
)
self._interaction_coordinator.enqueue_parts(
state,
modality="prompt",
event_id=event_id,
parts=parts,
event_id=interaction["event_id"],
received_at=time.monotonic(),
payload={"prompt": prompt},
transition_chunks=interaction.get("transition_chunks"),
)