Skip to content

vllm_omni.diffusion.diffusion_engine

logger module-attribute 

logger = init_logger(__name__)

DiffusionEngine

The diffusion engine for vLLM-Omni diffusion models.

abort_queue instance-attribute 

abort_queue: Queue[str] = queue.Queue()

action_post_process_func instance-attribute 

action_post_process_func = (
    get_diffusion_action_post_process_func(od_config)
)

default_diffusion_model_runner_cls class-attribute instance-attribute 

default_diffusion_model_runner_cls: str | None = None

execute_fn instance-attribute 

execute_fn = self.executor.execute_step

executor instance-attribute 

executor = executor_class(od_config)

main_loop instance-attribute 

main_loop: AbstractEventLoop | None = None

od_config instance-attribute 

od_config = od_config

post_process_func instance-attribute 

post_process_func = get_diffusion_post_process_func(
    od_config
)

pre_process_func instance-attribute 

pre_process_func = get_diffusion_pre_process_func(od_config)

scheduler instance-attribute 

scheduler: SchedulerInterface = scheduler or (
    StepScheduler()
    if self.step_execution
    else RequestScheduler()
)

step_execution instance-attribute 

step_execution = bool(
    getattr(od_config, "step_execution", False)
)

stop_event instance-attribute 

stop_event: Event | None = None

supports_request_batch instance-attribute 

supports_request_batch = (
    False
    if self.step_execution
    else supports_request_batch(od_config)
)

worker_thread instance-attribute 

worker_thread: Thread | None = None

abort 

abort(request_id: str | Iterable[str]) -> None

add_req_and_wait_for_response 

add_req_and_wait_for_response(
    request: OmniDiffusionRequest,
) -> DiffusionOutput

add_request 

add_request(request: OmniDiffusionRequest) -> str

async_add_req_and_stream_response 

async_add_req_and_stream_response(
    request: OmniDiffusionRequest,
) -> AsyncGenerator[DiffusionOutput, None]

async_add_req_and_wait_for_response async 

async_add_req_and_wait_for_response(
    request: OmniDiffusionRequest,
) -> DiffusionOutput

async_collective_rpc async 

async_collective_rpc(
    method: str,
    timeout: float | None = None,
    args: tuple = (),
    kwargs: dict | None = None,
    unique_reply_rank: int | None = None,
) -> Any

Async variant of :meth:collective_rpc for event-loop callers.

Mirrors :meth:async_add_req_and_wait_for_response: enqueue a task keyed by a future and await the result without blocking the loop.

close 

close() -> None

collective_rpc 

collective_rpc(
    method: str,
    timeout: float | None = None,
    args: tuple = (),
    kwargs: dict | None = None,
    unique_reply_rank: int | None = None,
) -> Any

Call a method on worker processes and get results immediately.

The call is enqueued and executed by the engine's busy loop between scheduler steps, so it is naturally serialized against per-request execute_fn() invocations without any explicit mutual-exclusion lock.

Parameters:

Name Type Description Default
method str

The method name (str) to execute on workers

 required
timeout float | None

Optional timeout in seconds

 None
args tuple

Positional arguments for the method

 ()
kwargs dict | None

Keyword arguments for the method

 None
unique_reply_rank int | None

If set, only get reply from this rank

 None

Returns:

Type Description
Any

Single result if unique_reply_rank is provided, otherwise list of results

get_result async 

get_result(request_id: str) -> DiffusionOutput

get_streaming_result async 

get_streaming_result(
    request_id: str,
) -> AsyncGenerator[DiffusionOutput, None]

Mirrors get_result() in non-streaming mode.

make_engine staticmethod 

make_engine(
    config: OmniDiffusionConfig,
    scheduler: SchedulerInterface | None = None,
) -> DiffusionEngine

Factory method to create the engine selected by config.engine_backend.

Parameters:

Name Type Description Default
config OmniDiffusionConfig

The configuration for the diffusion engine.

 required

Returns:

Type Description
DiffusionEngine

An instance of the resolved DiffusionEngine (sub)class.

postprocess_output 

postprocess_output(
    request: OmniDiffusionRequest,
    output: DiffusionOutput,
    diffusion_engine_start_time: float,
    preprocess_time: float,
    exec_total_time: float,
) -> list[OmniRequestOutput]

Convert a DiffusionOutput to a list of OmniRequestOutput, attaching profiling metrics.

profile 

profile(
    is_start: bool = True, profile_prefix: str | None = None
) -> None

Start or stop profiling on all diffusion workers.

Parameters:

Name Type Description Default
is_start bool

True to start profiling, False to stop.

 True
profile_prefix str | None

Optional prefix for trace filename.

 None

resolve_engine_class staticmethod 

resolve_engine_class(
    config: OmniDiffusionConfig,
) -> type[DiffusionEngine]

Resolve the engine class selected by config.engine_backend.

Mirrors DiffusionExecutor.get_class: accepts "default", a DiffusionEngine subclass, or an import-path string (e.g. a deploy config's engine_backend). Kept separate from :meth:make_engine so the selection is testable without constructing an engine (which runs a dummy forward).

Parameters:

Name Type Description Default
config OmniDiffusionConfig

The configuration for the diffusion engine.

 required

Returns:

Type Description
type[DiffusionEngine]

The DiffusionEngine (sub)class to instantiate.

step async 

step(
    request: OmniDiffusionRequest,
) -> list[OmniRequestOutput]

step_streaming async 

step_streaming(
    request: OmniDiffusionRequest,
) -> AsyncGenerator[list[OmniRequestOutput], None]

get_dummy_run_num_frames 

get_dummy_run_num_frames(
    model_class_name: str, supports_audio_input: bool
) -> int

Get num_frames for the dummy warmup run. Returns 0 to skip warmup.

image_color_format 

image_color_format(model_class_name: str) -> str

supports_audio_output 

supports_audio_output(model_class_name: str) -> bool

supports_multimodal_input 

supports_multimodal_input(
    od_config: OmniDiffusionConfig,
) -> tuple[bool, bool]

supports_request_batch 

supports_request_batch(
    od_config: OmniDiffusionConfig,
) -> bool