Skip to content

vllm_omni.diffusion.worker.utils

Per-request mutable state for step-wise diffusion execution.

BaseRunnerOutput

Bases: ABC

get_request_output abstractmethod 

get_request_output(request_id: str) -> RunnerOutput | None

BatchRunnerOutput dataclass

Bases: BaseRunnerOutput

request_ids property 

request_ids: list[str]

runner_outputs instance-attribute 

runner_outputs: list[RunnerOutput]

from_list classmethod 

from_list(
    runner_output_list: list[RunnerOutput],
) -> BatchRunnerOutput

get_request_output 

get_request_output(request_id: str) -> RunnerOutput | None

DiffusionRequestState dataclass

Per-request mutable state across all pipeline stages.

Owned by Runner and passed through all step-execution stages: prepare_encode() initializes/updates fields, denoise_step() and step_scheduler() mutate per-step fields, and post_decode() consumes final latents. This state object is also the cache unit for future continuous batching.

This dataclass keeps only the minimal cross-model state required by the step-execution contract. Pipeline-specific state should be stored in extra and promoted here only when it becomes shared across models.

Examples:

  • Wan-style pipelines may keep condition, first_frame_mask, or image_embeds in extra.
  • Bagel-style pipelines may keep gen_context, cfg_text_context, cfg_img_context, or image_shape in extra.

chunk_denoise_completed property 

chunk_denoise_completed: bool

chunk_index class-attribute instance-attribute 

chunk_index: int = 0

chunk_num_steps class-attribute instance-attribute 

chunk_num_steps: int | None = None

current_timestep property 

current_timestep: Tensor | None

denoise_completed property 

denoise_completed: bool

do_true_cfg class-attribute instance-attribute 

do_true_cfg: bool = False

extra class-attribute instance-attribute 

extra: dict[str, Any] = field(default_factory=dict)

guidance class-attribute instance-attribute 

guidance: Tensor | None = None

img_shapes class-attribute instance-attribute 

img_shapes: list | None = None

kv_sender_info class-attribute instance-attribute 

kv_sender_info: dict | None = None

latents class-attribute instance-attribute 

latents: Tensor | None = None

negative_prompt_embeds class-attribute instance-attribute 

negative_prompt_embeds: Tensor | None = None

negative_prompt_embeds_mask class-attribute instance-attribute 

negative_prompt_embeds_mask: Tensor | None = None

negative_txt_seq_lens class-attribute instance-attribute 

negative_txt_seq_lens: list[int] | None = None

new_request property 

new_request: bool

peak_memory_mb class-attribute instance-attribute 

peak_memory_mb: float = 0.0

prompt class-attribute instance-attribute 

prompt: OmniPromptType | None = None

prompt_embeds class-attribute instance-attribute 

prompt_embeds: Tensor | None = None

prompt_embeds_mask class-attribute instance-attribute 

prompt_embeds_mask: Tensor | None = None

request_denoise_completed property 

request_denoise_completed: bool

request_id instance-attribute 

request_id: str

sampling instance-attribute 

sampling: OmniDiffusionSamplingParams

scheduler class-attribute instance-attribute 

scheduler: Any | None = None

stage_durations class-attribute instance-attribute 

stage_durations: dict[str, float] = field(
    default_factory=dict
)

step_in_chunk class-attribute instance-attribute 

step_in_chunk: int = 0

step_index class-attribute instance-attribute 

step_index: int = 0

timesteps class-attribute instance-attribute 

timesteps: Tensor | list[Tensor] | None = None

total_chunks class-attribute instance-attribute 

total_chunks: int = 1

total_steps property 

total_steps: int

txt_seq_lens class-attribute instance-attribute 

txt_seq_lens: list[int] | None = None

RunnerOutput dataclass

Bases: BaseRunnerOutput

Output of a single denoising step for a request.

NOTE: latents may be None when returned through IPC to avoid serialization overhead. The actual latents are kept in Worker's _request_state_cache.

finished class-attribute instance-attribute 

finished: bool = False

request_id instance-attribute 

request_id: str

result class-attribute instance-attribute 

result: DiffusionOutput | None = None

step_index class-attribute instance-attribute 

step_index: int | None = None

get_request_output 

get_request_output(request_id: str) -> RunnerOutput | None

attach_stage_durations 

attach_stage_durations(
    state: DiffusionRequestState, output: DiffusionOutput
) -> None

clear_pipeline_stage_durations 

clear_pipeline_stage_durations(pipeline: Any) -> None

consume_pipeline_stage_durations 

consume_pipeline_stage_durations(
    pipeline: Any,
) -> dict[str, float]

merge_stage_durations 

merge_stage_durations(
    state: DiffusionRequestState,
    stage_durations: dict[str, float],
) -> None