from dataclasses import dataclass, field from typing import Any import torch from PIL import Image from vllm.outputs import RequestOutput from vllm.v1.outputs import ModelRunnerOutput from vllm_omni.inputs.data import OmniPromptType @dataclass class OmniConnectorOutput: """Communication results from Model Runner to Scheduler. Carries transfer readiness signals so the Scheduler can make scheduling decisions without ever calling connector.put()/get() directly. Attributes: chunk_ready_req_ids: Request IDs with newly arrived chunks this cycle. chunk_finished_req_ids: Request IDs whose final chunk has arrived. request_metadata: Lightweight scheduling metadata keyed by request ID (e.g. next_stage_prompt_len, code_predictor_codes, left_context_size). Full payloads are owned by the Model Runner's local cache. kv_sent_req_ids: Request IDs whose KV cache was successfully sent. stage_recv_req_ids: Request IDs that received batch stage inputs. has_pending_kv_work: True if the mixin has pending, active, or completed KV transfers that the scheduler should account for. """ chunk_ready_req_ids: set[str] = field(default_factory=set) chunk_finished_req_ids: set[str] = field(default_factory=set) request_metadata: dict[str, dict[str, Any]] = field(default_factory=dict) kv_sent_req_ids: list[str] = field(default_factory=list) stage_recv_req_ids: set[str] = field(default_factory=set) has_pending_kv_work: bool = False class OmniModelRunnerOutput(ModelRunnerOutput): """Model runner output for omni models. Extends the base ModelRunnerOutput with support for multimodal outputs that may be produced by non-autoregressive stages. Attributes: multimodal_outputs: Optional dictionary mapping modality names to output tensors (e.g., {"image": tensor, "audio": tensor}) """ multimodal_outputs: dict[str, torch.Tensor] | None = None # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU. # The Scheduler can safely free the block tables for these requests. kv_extracted_req_ids: list[str] | None = None omni_connector_output: OmniConnectorOutput | None = None @dataclass class OmniRequestOutput: """Unified request output for both pipeline stages and diffusion models. This class handles outputs from: 1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output) 2. Diffusion models (with images, prompt, metrics) Attributes: request_id: Unique identifier for this request finished: Whether generation is complete stage_id: Identifier of the stage that produced this output (pipeline mode) final_output_type: Type of output ("text", "image", "audio", "latents") request_output: The underlying RequestOutput from the stage (pipeline mode) images: List of generated PIL images (diffusion mode) prompt: The prompt used for generation (diffusion mode) latents: Optional tensor of latent representations (diffusion mode) metrics: Optional dictionary of generation metrics """ request_id: str = "" finished: bool = True # Pipeline stage fields stage_id: int | None = None final_output_type: str = "text" request_output: RequestOutput | None = None # Diffusion model fields images: list[Image.Image] = field(default_factory=list) prompt: OmniPromptType | None = None latents: torch.Tensor | None = None trajectory_latents: torch.Tensor | None = None trajectory_timesteps: torch.Tensor | None = None trajectory_log_probs: torch.Tensor | None = None trajectory_decoded: list | 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) # profiling data stage_durations: dict[str, float] = field(default_factory=dict) # memory usage info peak_memory_mb: float = 0.0 # error handling error: str | None = None @classmethod def from_error( cls, request_id: str, error_message: str, ) -> "OmniRequestOutput": """Create a terminal error output. Args: request_id: Request identifier error_message: Human-readable error description Returns: OmniRequestOutput with ``finished=True`` and the ``error`` field set. """ return cls( request_id=request_id, finished=True, error=error_message, ) @classmethod def from_pipeline( cls, stage_id: int, final_output_type: str, request_output: RequestOutput, ) -> "OmniRequestOutput": """Create output from pipeline stage. Args: stage_id: Stage identifier final_output_type: Type of output request_output: The stage's output Returns: OmniRequestOutput configured for pipeline mode """ return cls( request_id=getattr(request_output, "request_id", ""), stage_id=stage_id, final_output_type=final_output_type, request_output=request_output, finished=True, ) @classmethod def from_diffusion( cls, request_id: str, images: list[Image.Image], prompt: OmniPromptType | None = None, metrics: dict[str, Any] | None = None, latents: torch.Tensor | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_timesteps: torch.Tensor | None = None, trajectory_log_probs: torch.Tensor | None = None, trajectory_decoded: list | None = None, multimodal_output: dict[str, Any] | None = None, custom_output: dict[str, Any] | None = None, final_output_type: str = "image", stage_durations: dict[str, float] | None = None, peak_memory_mb: float = 0.0, ) -> "OmniRequestOutput": """Create output from diffusion model. Args: request_id: Request identifier images: Generated images prompt: The prompt used metrics: Generation metrics latents: Optional latent tensors trajectory_latents: Optional stacked trajectory latent tensors trajectory_timesteps: Optional stacked trajectory timestep tensors trajectory_log_probs: Optional stacked trajectory log-probability tensors trajectory_decoded: Optional list of decoded trajectory images multimodal_output: Optional multimodal output dict custom_output: Optional custom output dict (e.g. prompt embeds) stage_durations: Optional stage durations (execution time of each stage) dict peak_memory_mb: Peak memory usage in MB Returns: OmniRequestOutput configured for diffusion mode """ return cls( request_id=request_id, final_output_type=final_output_type, images=images, prompt=prompt, latents=latents, trajectory_latents=trajectory_latents, trajectory_timesteps=trajectory_timesteps, trajectory_log_probs=trajectory_log_probs, trajectory_decoded=trajectory_decoded, metrics=metrics or {}, _multimodal_output=multimodal_output or {}, _custom_output=custom_output or {}, stage_durations=stage_durations or {}, peak_memory_mb=peak_memory_mb, finished=True, ) @property def multimodal_output(self) -> dict[str, Any]: """Return multimodal output from the underlying request output or local field. For pipeline outputs, this checks completion outputs first, then request_output. For diffusion outputs, this returns the local _multimodal_output field. """ if self.request_output is None: return self._multimodal_output # Check completion outputs first (where multimodal_output is attached) for output in getattr(self.request_output, "outputs", []): if mm := getattr(output, "multimodal_output", None): return mm if mm := getattr(self.request_output, "multimodal_output", None): 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.""" return len(self.images) # Pass-through properties keep vLLM serving codepaths compatible with # OmniRequestOutput for pipeline outputs (Issue #345). @property def prompt_token_ids(self) -> list[int] | None: """Return prompt token IDs from the underlying request output. This property is required for compatibility with vLLM's streaming chat completion generator which checks res.prompt_token_ids. """ if self.request_output is not None: return getattr(self.request_output, "prompt_token_ids", None) return None @property def outputs(self) -> list[Any]: """Return outputs from the underlying request output. This property is required for compatibility with vLLM's streaming and non-streaming chat completion generators. """ if self.request_output is not None: return getattr(self.request_output, "outputs", []) return [] @property def encoder_prompt_token_ids(self) -> list[int] | None: """Return encoder prompt token IDs from the underlying request output.""" if self.request_output is not None: return getattr(self.request_output, "encoder_prompt_token_ids", None) return None @property def prompt_logprobs(self) -> Any: """Return prompt logprobs from the underlying request output.""" if self.request_output is not None: return getattr(self.request_output, "prompt_logprobs", None) return None @property def num_cached_tokens(self) -> int | None: """Return number of cached tokens from the underlying request output.""" if self.request_output is not None: return getattr(self.request_output, "num_cached_tokens", None) return None @property def kv_transfer_params(self) -> Any: """Return KV transfer params from the underlying request output.""" if self.request_output is not None: return getattr(self.request_output, "kv_transfer_params", None) return None @property def is_diffusion_output(self) -> bool: """Check if this is a diffusion model output.""" return len(self.images) > 0 or self.final_output_type == "image" @property def is_pipeline_output(self) -> bool: """Check if this is a pipeline stage output.""" return self.stage_id is not None and self.request_output is not None def to_dict(self) -> dict[str, Any]: """Convert to dictionary for JSON serialization.""" result = { "request_id": self.request_id, "finished": self.finished, "final_output_type": self.final_output_type, } if self.is_diffusion_output: result.update( { "num_images": self.num_images, "prompt": self.prompt, "metrics": self.metrics, } ) if self.is_pipeline_output: result.update( { "stage_id": self.stage_id, } ) return result def __repr__(self) -> str: """Custom repr to properly show image count instead of image objects.""" # For images, show count instead of full list images_repr = f"[{len(self.images)} PIL Images]" if self.images else "[]" # Build repr string parts = [ f"request_id={self.request_id!r}", f"finished={self.finished}", f"stage_id={self.stage_id}", f"final_output_type={self.final_output_type!r}", f"request_output={self.request_output}", f"images={images_repr}", f"prompt={self.prompt!r}", f"latents={self.latents}", f"metrics={self.metrics}", f"multimodal_output={self._multimodal_output}", f"custom_output={self._custom_output}", f"stage_durations={self.stage_durations}", f"peak_memory_mb={self.peak_memory_mb}", ] return f"OmniRequestOutput({', '.join(parts)})"