diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index fd0be6c98..986bad9a8 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -96,6 +96,7 @@ class TestWorkerWrapperBaseInitialization: local_rank=0, rank=0, od_config=mock_od_config, + skip_load_model=False, ) diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index bf36b78ff..0627520c9 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -1,3 +1,4 @@ +import gc import sys from pathlib import Path @@ -20,8 +21,10 @@ models = ["riverclouds/qwen_image_random"] def inference(model_name: str, offload: bool = True): + gc.collect() current_omni_platform.empty_cache() device_index = torch.cuda.current_device() + torch.cuda.reset_peak_memory_stats(device=device_index) monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() m = Omni(model=model_name, enable_cpu_offload=offload) @@ -42,6 +45,10 @@ def inference(model_name: str, offload: bool = True): peak = monitor.peak_used_mb monitor.stop() + del m + gc.collect() + current_omni_platform.empty_cache() + return peak @@ -51,9 +58,9 @@ def inference(model_name: str, offload: bool = True): @pytest.mark.parametrize("model_name", models) def test_cpu_offload_diffusion_model(model_name: str): try: - no_offload_peak_memory = inference(model_name, offload=False) - cleanup_dist_env_and_memory() offload_peak_memory = inference(model_name, offload=True) + cleanup_dist_env_and_memory() + no_offload_peak_memory = inference(model_name, offload=False) except Exception: pytest.fail("Inference failed") print(f"Offload peak memory: {offload_peak_memory} MB") diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index da68adbdb..0a68c2e70 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -618,6 +618,10 @@ class DiffusionOutput: post_process_func: Callable[..., Any] | None = None + # Extra custom output data (e.g. latent trajectories, prompt embeds) + # passed through to OmniRequestOutput.custom_output + custom_output: dict[str, Any] = field(default_factory=dict) + # logged timings info, directly from Req.timings # timings: Optional["RequestTimings"] = None diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 066ba2bb8..103f0ca89 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -144,6 +144,7 @@ class DiffusionEngine: prompt=prompt, metrics=metrics, latents=output.trajectory_latents, + custom_output=output.custom_output or {}, multimodal_output=mm_output, ), ] @@ -197,6 +198,7 @@ class DiffusionEngine: prompt=prompt, metrics=metrics, latents=output.trajectory_latents, + custom_output=output.custom_output or {}, multimodal_output=mm_output, ), ) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index d4389f4d2..3fd29636e 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -63,6 +63,7 @@ class DiffusionWorker: local_rank: int, rank: int, od_config: OmniDiffusionConfig, + skip_load_model: bool = False, ): self.local_rank = local_rank self.rank = rank @@ -79,8 +80,9 @@ class DiffusionWorker: od_config=self.od_config, device=self.device, ) - self.load_model(load_format=self.od_config.diffusion_load_format) - self.init_lora_manager() + if not skip_load_model: + self.load_model(load_format=self.od_config.diffusion_load_format) + self.init_lora_manager() logger.info(f"Worker {self.rank}: Initialization complete.") def init_device(self) -> None: @@ -524,10 +526,15 @@ class WorkerWrapperBase: worker_class = self._prepare_worker_class() # Create the actual worker instance + # When custom_pipeline_args is provided, skip initial model loading + # since re_init_pipeline will handle it. This avoids allocating memory + # through CuMemAllocator twice, which causes assertion failures in + # sleep mode. self.worker = worker_class( local_rank=gpu_id, rank=gpu_id, od_config=od_config, + skip_load_model=(self.custom_pipeline_args is not None), ) # Re-initialize pipeline with custom pipeline if provided diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index e393f31be..f7e95aa32 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -4,8 +4,8 @@ import asyncio import copy import time import weakref -from collections.abc import AsyncGenerator, Iterable, Sequence -from typing import Any +from collections.abc import AsyncGenerator, Callable, Iterable, Sequence +from typing import Any, TypeVar from vllm.config import VllmConfig from vllm.inputs.preprocess import InputPreprocessor @@ -35,6 +35,8 @@ from vllm_omni.lora.request import LoRARequest from vllm_omni.metrics import OrchestratorAggregator from vllm_omni.outputs import OmniRequestOutput +_R = TypeVar("_R") + logger = init_logger(__name__) @@ -105,10 +107,17 @@ class AsyncOmni(OmniBase): self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False + # Sleep mode tracking + self._is_sleeping: bool = False + # Request state tracking self.request_states: dict[str, ClientRequestState] = {} self.output_handler: asyncio.Task | None = None + # RPC results storage: {stage_id: {rpc_id: result}} + # Used to avoid race condition between output_handler and collective_rpc + self._rpc_results: dict[int, dict[str, dict[str, Any]]] = {} + super().__init__(model, **kwargs) # Register weak reference cleanup (called on garbage collection) @@ -193,6 +202,8 @@ class AsyncOmni(OmniBase): "enforce_eager": kwargs.get("enforce_eager", False), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), + "worker_extension_cls": kwargs.get("worker_extension_cls", None), + "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), }, @@ -253,6 +264,26 @@ class AsyncOmni(OmniBase): self.io_processor = None self.model_config = None + def _setup_rpc_result_checkers(self) -> None: + """Override base class to use async-friendly RPC result checkers. + + In the async path the output_handler task drains the output queues + and stashes collective_rpc results into ``_rpc_results``. The + checker therefore only needs to look there (no queue draining). + """ + for stage in self.stage_list: + sid = stage.stage_id + + def make_rpc_checker(stage_id: int): + def rpc_checker(rpc_id: str) -> dict[str, Any] | None: + if stage_id in self._rpc_results and rpc_id in self._rpc_results[stage_id]: + return self._rpc_results[stage_id].pop(rpc_id) + return None + + return rpc_checker + + stage._rpc_result_checker = make_rpc_checker(sid) + def shutdown(self): """Shutdown, cleaning up the background proc and IPC. @@ -542,12 +573,17 @@ class AsyncOmni(OmniBase): images = engine_outputs.images if stage.final_output_type == "image": + # Propagate custom_output from inner diffusion output + custom_output = {} + if isinstance(engine_outputs, OmniRequestOutput): + custom_output = engine_outputs.custom_output or {} output_to_yield = OmniRequestOutput( stage_id=stage_id, final_output_type=stage.final_output_type, request_output=engine_outputs, images=images, finished=finished, + _custom_output=custom_output, ) else: output_to_yield = OmniRequestOutput( @@ -597,6 +633,15 @@ class AsyncOmni(OmniBase): # so we wait for a short time and try again await asyncio.sleep(0.05) continue + # Handle collective_rpc results separately to avoid + # race condition with the polling in collective_rpc + if result.get("type") == "collective_rpc_result": + rpc_id = result.get("rpc_id") + if rpc_id: + if stage_id not in self._rpc_results: + self._rpc_results[stage_id] = {} + self._rpc_results[stage_id][rpc_id] = result + continue req_id = result.get("request_id") req_state = request_states.get(req_id) if req_state is None: @@ -713,19 +758,73 @@ class AsyncOmni(OmniBase): async def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool: pass + async def collective_rpc( + self, + method: str | Callable[..., _R], + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + ) -> list[_R]: + """Execute an RPC call on all stages asynchronously. + + Args: + method: Name of the method to execute or callable + timeout: Optional timeout in seconds + args: Positional arguments for the method + kwargs: Keyword arguments for the method + + Returns: + List of results from each stage + """ + self._run_output_handler() + # Run synchronous collective_rpc in thread pool to avoid blocking event loop + loop = asyncio.get_event_loop() + + async def run_stage_rpc(stage: OmniStage) -> _R: + return await loop.run_in_executor( + None, + stage.collective_rpc, + method, + timeout, + args, + kwargs, + ) + + # Run all stages concurrently + results = await asyncio.gather(*[run_stage_rpc(stage) for stage in self.stage_list]) + return list(results) + async def sleep(self, level: int = 1) -> None: - pass + self._is_sleeping = True + await self.collective_rpc(method="sleep", args=(level,)) async def wake_up(self, tags: list[str] | None = None) -> None: - pass + self._is_sleeping = False + await self.collective_rpc(method="wake_up", args=(tags,)) async def is_sleeping(self) -> bool: """Check whether the engine is sleeping""" - return False + return getattr(self, "_is_sleeping", False) async def add_lora(self, lora_request: LoRARequest) -> bool: """Load a new LoRA adapter into the engine for future requests.""" - return False + result = await self.collective_rpc(method="add_lora", args=(lora_request,)) + return result[0][0] + + async def remove_lora(self, adapter_id: int) -> bool: + """Remove a LoRA adapter from the engine.""" + result = await self.collective_rpc(method="remove_lora", args=(adapter_id,)) + return result[0][0] + + async def list_loras(self) -> list[int]: + """List all loaded LoRA adapter IDs.""" + result = await self.collective_rpc(method="list_loras") + return result[0][0] + + async def pin_lora(self, adapter_id: int) -> bool: + """Pin a LoRA adapter so it is not evicted from the cache.""" + result = await self.collective_rpc(method="pin_lora", args=(adapter_id,)) + return result[0][0] async def encode( self, diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 9ecc5f762..598205971 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -9,7 +9,7 @@ import uuid import weakref from collections.abc import Callable, Generator, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Literal, overload +from typing import Any, Literal, TypeVar, overload import huggingface_hub import msgspec.msgpack @@ -46,12 +46,15 @@ from vllm_omni.entrypoints.utils import ( ) from vllm_omni.entrypoints.zmq_utils import ZmqQueue from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams +from vllm_omni.lora.request import LoRARequest from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) from vllm_omni.outputs import OmniRequestOutput +_R = TypeVar("_R") + logger = init_logger(__name__) @@ -177,6 +180,13 @@ class OmniBase: self._handshake_seen: set[int] = set() # Track which stage IDs have completed ZMQ handshake self._single_stage_id: int | None = None # Optional: deploy only a specific stage ID + # Sleep mode tracking + self._is_sleeping: bool = False + + # RPC results storage: {stage_id: {rpc_id: result}} + # Used by collective_rpc to retrieve results collected from the output queue + self._rpc_results: dict[int, dict[str, dict[str, Any]]] = {} + # Initialize stages - each stage will create appropriate instance based on stage_type # Stage workers will automatically create OmniLLM or OmniDiffusion instances # based on stage_type in YAML config (handled in omni_stage.py) @@ -362,6 +372,8 @@ class OmniBase: self._start_stages(model) # Wait for all stages to report readiness before seeding self._wait_for_stages_ready(timeout=init_timeout) + # Set up RPC result checkers so that collective_rpc works + self._setup_rpc_result_checkers() def _is_async_chunk_enable(self, stage_args: list) -> bool: """get async chunk flag""" @@ -517,6 +529,117 @@ class OmniBase: profiler = getattr(profiler_config, "profiler", None) return profiler is not None + def _setup_rpc_result_checkers(self) -> None: + """Set up RPC result checkers for all stages. + + Each checker reads from the shared ``_rpc_results`` dict so that + ``OmniStage.collective_rpc`` can retrieve results that were + collected by the orchestrator (or, in the sync path, by the + stage-level checker that drains the output queue). + + Uses a weak reference to ``self`` to avoid a circular reference + (OmniBase → stage_list → OmniStage → closure → OmniBase) that + would prevent the instance from being freed by reference counting. + """ + weak_self = weakref.ref(self) + for stage in self.stage_list: + sid = stage.stage_id + + def make_rpc_checker(stage_id: int): + def rpc_checker(rpc_id: str) -> dict[str, Any] | None: + _self = weak_self() + if _self is None: + return None + # First check the shared dict + if stage_id in _self._rpc_results and rpc_id in _self._rpc_results[stage_id]: + return _self._rpc_results[stage_id].pop(rpc_id) + # In the sync path there is no background output handler, + # so drain the output queue ourselves and stash any + # non-RPC results back. + out_q = _self._stage_out_queues[stage_id] if stage_id < len(_self._stage_out_queues) else None + if out_q is not None: + import queue as _queue + + try: + while True: + item = out_q.get_nowait() + if isinstance(item, dict) and item.get("type") == "collective_rpc_result": + item_rpc_id = item.get("rpc_id") + if item_rpc_id == rpc_id: + return item + # Stash for another caller + if stage_id not in _self._rpc_results: + _self._rpc_results[stage_id] = {} + _self._rpc_results[stage_id][item_rpc_id] = item + else: + # Non-RPC item — put it back + out_q.put(item) + break + except _queue.Empty: + pass + return None + + return rpc_checker + + stage._rpc_result_checker = make_rpc_checker(sid) + + def collective_rpc( + self, + method: str | Callable[..., _R], + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + ) -> list[_R]: + """Execute an RPC call on all stage workers. + + Args: + method: Name of the worker method to execute, or a callable. + timeout: Maximum time in seconds to wait for execution. + args: Positional arguments to pass to the worker method. + kwargs: Keyword arguments to pass to the worker method. + + Returns: + A list containing the results from each stage. + """ + results: list[_R] = [] + for stage in self.stage_list: + results.append( + stage.collective_rpc( + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + ) + ) + return results + + def sleep(self, level: int = 1) -> None: + """Put the engine to sleep to free up resources. + + Args: + level: Sleep level (1 = light sleep, higher = deeper sleep). + """ + self._is_sleeping = True + self.collective_rpc(method="sleep", args=(level,)) + + def wake_up(self, tags: list[str] | None = None) -> None: + """Wake the engine up from sleep. + + Args: + tags: Optional list of tags to selectively wake components. + """ + self._is_sleeping = False + self.collective_rpc(method="wake_up", args=(tags,)) + + def is_sleeping(self) -> bool: + """Check whether the engine is sleeping.""" + return getattr(self, "_is_sleeping", False) + + def add_lora(self, lora_request: LoRARequest) -> bool: + """Load a new LoRA adapter into the engine for future requests.""" + result = self.collective_rpc(method="add_lora", args=(lora_request,)) + return result[0][0] + def start_profile(self, stages: list[int] | None = None) -> None: """Start profiling for specified stages. diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 098cfa15d..1ce9dae48 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -15,10 +15,11 @@ import queue import sys import time import traceback -from collections.abc import Sequence +import uuid +from collections.abc import Callable, Sequence from contextlib import contextmanager from dataclasses import fields -from typing import Any, Literal, cast +from typing import Any, Literal, TypeVar, cast from vllm import PromptType, RequestOutput from vllm.inputs import TextPrompt @@ -60,6 +61,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, O from vllm_omni.metrics import count_tokens_from_outputs from vllm_omni.outputs import OmniRequestOutput +_R = TypeVar("_R") + logger = init_logger(__name__) @@ -306,6 +309,10 @@ class OmniStage: self._proc: mp.Process | None = None self._shm_threshold_bytes: int = 65536 self._stage_init_timeout: int = stage_init_timeout + # Callback used by the orchestrator's output handler to stash + # collective_rpc results so that ``collective_rpc`` can retrieve + # them without competing for the output queue. + self._rpc_result_checker: Callable[[str], dict | None] | None = None def set_engine(self, engine: LLMEngine) -> None: """Set the LLM engine for this stage. @@ -410,6 +417,71 @@ class OmniStage: logger.error(f"[Stage-{self.stage_id}] Timeout waiting for profiler results.") return {} + def collective_rpc( + self, + method: str | Callable[..., _R], + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + ) -> list[_R]: + """Execute an RPC call on all workers via the stage engine. + + Args: + method: Name of the worker method to execute, or a callable that + is serialized and sent to all workers to execute. + + If the method is a callable, it should accept an additional + ``self`` argument, in addition to the arguments passed in + ``args`` and ``kwargs``. The ``self`` argument will be the + worker object. + timeout: Maximum time in seconds to wait for execution. Raises a + :class:`TimeoutError` on timeout. ``None`` means wait + indefinitely. + args: Positional arguments to pass to the worker method. + kwargs: Keyword arguments to pass to the worker method. + + Returns: + A list containing the results from each worker. + + Note: + It is recommended to use this API to only pass control messages, + and set up data-plane communication to pass data. + """ + assert self._in_q is not None and self._out_q is not None, "Queues must be attached before collective_rpc" + + # Submit collective_rpc task to worker + rpc_id = str(uuid.uuid4()) + self._in_q.put( + { + "type": OmniStageTaskType.COLLECTIVE_RPC, + "rpc_id": rpc_id, + "method": method, + "timeout": timeout, + "args": args, + "kwargs": kwargs, + } + ) + + start_time = time.time() + while True: + if timeout is not None and (time.time() - start_time) > timeout: + raise TimeoutError(f"collective_rpc timed out after {timeout} seconds") + + # Check if result was already collected by the orchestrator's + # output handler (stored in a shared dict). + result = None + if self._rpc_result_checker is not None: + result = self._rpc_result_checker(rpc_id) + + if result is not None: + if result.get("type") == "collective_rpc_result": + if result.get("rpc_id") == rpc_id: + if "error" in result: + raise RuntimeError(f"collective_rpc failed: {result['error']}") + return result["result"] + + time.sleep(0.001) # Small sleep to avoid busy waiting + def init_stage_worker( self, model: str, @@ -542,14 +614,24 @@ class OmniStage: try: self._in_q.put_nowait(SHUTDOWN_TASK) except Exception as e: - logger.warning("Failed to send shutdown to in_q: %s", e) + # Queue may already be closed by the weak-ref cleanup + # (e.g. ZmqQueue socket closed) — this is expected. + logger.debug("Failed to send shutdown to in_q: %s", e) close_fn = getattr(self._in_q, "close", None) if callable(close_fn): - close_fn() + try: + close_fn() + except Exception: + pass + self._in_q = None if self._out_q is not None: close_fn = getattr(self._out_q, "close", None) if callable(close_fn): - close_fn() + try: + close_fn() + except Exception: + pass + self._out_q = None if hasattr(self, "_ray_actor") and self._ray_actor: kill_ray_actor(self._ray_actor) @@ -878,6 +960,40 @@ def _stage_worker( out_q.put({"type": "profiler_result", "data": profiler_data}) continue + # Handle collective_rpc commands + if task_type == OmniStageTaskType.COLLECTIVE_RPC: + rpc_id = task.get("rpc_id") + rpc_method = task.get("method") + rpc_timeout = task.get("timeout") + rpc_args = task.get("args", ()) + rpc_kwargs = task.get("kwargs") or {} + try: + rpc_result = stage_engine.collective_rpc( + method=rpc_method, + timeout=rpc_timeout, + args=rpc_args, + kwargs=rpc_kwargs, + ) + out_q.put( + { + "type": "collective_rpc_result", + "rpc_id": rpc_id, + "stage_id": stage_id, + "result": rpc_result, + } + ) + except Exception as e: + logger.exception("[Stage-%s] collective_rpc failed: %s", stage_id, e) + out_q.put( + { + "type": "collective_rpc_result", + "rpc_id": rpc_id, + "stage_id": stage_id, + "error": str(e), + } + ) + continue + batch_tasks: list[dict[str, Any]] = [task] tasks_failed_to_add_to_batch: list[dict[str, Any]] = [] start_time = _time.time() @@ -1406,6 +1522,54 @@ async def _stage_worker_async( # Send result back to orchestrator for STOP command if task_type == OmniStageTaskType.PROFILER_STOP: out_q.put({"type": "profiler_result", "data": profiler_data}) + elif task_type == OmniStageTaskType.COLLECTIVE_RPC: + rpc_id = task.get("rpc_id") + rpc_method = task.get("method") + rpc_timeout = task.get("timeout") + rpc_args = task.get("args", ()) + rpc_kwargs = task.get("kwargs") or {} + try: + if stage_type == "diffusion": + # DiffusionEngine.collective_rpc is synchronous + loop = asyncio.get_event_loop() + rpc_result = await loop.run_in_executor( + None, + lambda: stage_engine.engine.collective_rpc( + method=rpc_method, + timeout=rpc_timeout, + args=rpc_args, + kwargs=rpc_kwargs, + ), + ) + else: + rpc_result = await cast(AsyncLLM, stage_engine).collective_rpc( + method=rpc_method, + timeout=rpc_timeout, + args=rpc_args, + kwargs=rpc_kwargs, + ) + out_q.put( + { + "type": "collective_rpc_result", + "rpc_id": rpc_id, + "stage_id": stage_id, + "result": rpc_result, + } + ) + except Exception as e: + logger.exception( + "[Stage-%s] collective_rpc failed: %s", + stage_id, + e, + ) + out_q.put( + { + "type": "collective_rpc_result", + "rpc_id": rpc_id, + "stage_id": stage_id, + "error": str(e), + } + ) else: asyncio.create_task(generation_single_request(task)) @@ -1559,10 +1723,16 @@ def output_strip(r_output: RequestOutput | OmniRequestOutput, final_output: bool if mm_output is not None: r_output.multimodal_output = {} + custom_out = getattr(r_output, "_custom_output", None) + if custom_out is not None: + r_output._custom_output = {} + outputs = getattr(r_output, "outputs", None) if outputs is not None: for out in outputs: if getattr(out, "multimodal_output", None): out.multimodal_output = {} + if getattr(out, "_custom_output", None): + out._custom_output = {} return r_output diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 110bf48cd..de4a71719 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -36,6 +36,7 @@ class OmniStageTaskType(enum.Enum): SHUTDOWN = "shutdown" PROFILER_START = "profiler_start" PROFILER_STOP = "profiler_stop" + COLLECTIVE_RPC = "collective_rpc" SHUTDOWN_TASK = {"type": OmniStageTaskType.SHUTDOWN} diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 16d690bf3..c0f10af2b 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -1,7 +1,7 @@ import copy import pprint from dataclasses import asdict, dataclass, field -from typing import Any, TypeAlias +from typing import Any, TypeAlias, TypedDict from vllm.inputs import PromptType from vllm.sampling_params import SamplingParams @@ -103,13 +103,34 @@ class OmniEmbedsPrompt(EmbedsPrompt): additional_information: NotRequired[dict[str, Any]] +class OmniCustomPrompt(TypedDict, total=False): + """Custom prompt type for diffusion pipelines with pre-tokenized inputs. + + Allows users to pass pre-tokenized prompt IDs, attention masks, and extra + arguments directly, bypassing the tokenization stage in the pipeline. + + Attributes: + prompt_ids: Pre-tokenized prompt token IDs (single or batched) + negative_prompt_ids: Pre-tokenized negative prompt token IDs + prompt_mask: Attention mask tensor for the prompt + negative_prompt_mask: Attention mask tensor for the negative prompt + extra_args: Additional pipeline-specific arguments + """ + + prompt_ids: list[int] | list[list[int]] + negative_prompt_ids: list[int] | list[list[int]] + prompt_mask: torch.Tensor + negative_prompt_mask: torch.Tensor + extra_args: dict[str, Any] + + # Must ensure that all additional prompt types are inherited from vLLM prompt types # Because TypedDict doesn't support isinstance and are dict. Cannot distinguish them in runtime. # Inheritance ensure that there are only additional fields but not removing fields--safe to route to LLM.generate() OmniSingletonPrompt: TypeAlias = str | list[int] | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt """Omni singleton prompt type extending vLLM's SingletonPrompt.""" -OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt +OmniPromptType: TypeAlias = PromptType | OmniTextPrompt | OmniTokensPrompt | OmniEmbedsPrompt | OmniCustomPrompt def token_inputs_omni( diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 5bbd27a1d..6f035a81e 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -60,6 +60,7 @@ class OmniRequestOutput: latents: torch.Tensor | None = None metrics: dict[str, Any] = field(default_factory=dict) _multimodal_output: dict[str, Any] = field(default_factory=dict) + _custom_output: dict[str, Any] = field(default_factory=dict) @classmethod def from_pipeline( @@ -95,6 +96,7 @@ class OmniRequestOutput: metrics: dict[str, Any] | None = None, latents: torch.Tensor | None = None, multimodal_output: dict[str, Any] | None = None, + custom_output: dict[str, Any] | None = None, final_output_type: str = "image", ) -> "OmniRequestOutput": """Create output from diffusion model. @@ -105,6 +107,8 @@ class OmniRequestOutput: prompt: The prompt used metrics: Generation metrics latents: Optional latent tensors + multimodal_output: Optional multimodal output dict + custom_output: Optional custom output dict (e.g. latent trajectories, prompt embeds) Returns: OmniRequestOutput configured for diffusion mode @@ -117,6 +121,7 @@ class OmniRequestOutput: latents=latents, metrics=metrics or {}, _multimodal_output=multimodal_output or {}, + _custom_output=custom_output or {}, finished=True, ) @@ -140,6 +145,23 @@ class OmniRequestOutput: return mm return self._multimodal_output + @property + def custom_output(self) -> dict[str, Any]: + """Return custom output data from diffusion pipelines. + + For diffusion outputs, returns the local _custom_output field. + For pipeline outputs with an inner OmniRequestOutput, forwards + the custom_output from the inner request output. + """ + if self.request_output is not None: + if isinstance(self.request_output, OmniRequestOutput): + return self.request_output._custom_output + return self._custom_output + + @custom_output.setter + def custom_output(self, value: dict[str, Any]) -> None: + self._custom_output = value + @property def num_images(self) -> int: """Return the number of generated images.""" @@ -250,6 +272,7 @@ class OmniRequestOutput: f"latents={self.latents}", f"metrics={self.metrics}", f"multimodal_output={self._multimodal_output}", + f"custom_output={self._custom_output}", ] return f"OmniRequestOutput({', '.join(parts)})"