vllm_omni.diffusion.offloader.component_utils ¶

Shared component-plan helpers for diffusion offload backends.

logger `module-attribute` ¶

logger = init_logger(__name__)

clear_encoder_layerwise_state ¶

clear_encoder_layerwise_state(module: Module) -> None

Clear encoder layerwise state after its backend hooks are removed.

encoder_component_type ¶

encoder_component_type(
    name: str, plan: OffloadPlan | None
) -> str | None

Map a discovered encoder path to its public offload component.

get_encoder_block_groups ¶

get_encoder_block_groups(
    module: Module,
    name: str,
    plan: OffloadPlan | None,
    *,
    strict: bool = False,
) -> list[ModuleList]

Resolve the streamable block lists declared for one encoder.

iter_streamable_dits ¶

iter_streamable_dits(
    resolved: ResolvedOffloadPlan, device: device
) -> Iterator[tuple[ResolvedComponent, BlockStack]]

Yield selected DiTs that resolved to a streamable block stack.

move_non_block_state_to_device ¶

move_non_block_state_to_device(
    module: Module,
    block_groups: Sequence[Sequence[Module]],
    device: device,
) -> None

Keep component state outside streamed block lists resident on device.

prepare_component ¶

prepare_component(
    module: Module,
    name: str,
    *,
    device: device,
    stage_on_demand: bool,
    blockwise: bool,
    staged_components: list[Module],
) -> None

Stage a selected component or keep its non-streamed form resident.

prepare_pipeline_components ¶

prepare_pipeline_components(
    resolved: ResolvedOffloadPlan,
    *,
    device: device,
    staged_components: list[Module],
    enable_encoder_blocks: Callable[
        [ResolvedComponent], bool
    ],
) -> None

Apply the shared encoder/VAE/resident placement policy.

set_encoder_layerwise_state ¶

set_encoder_layerwise_state(
    module: Module,
    hooks: list[Any],
    block_groups: Sequence[Sequence[Module]],
) -> None

Publish the backend-neutral state used by encoder stage lifecycles.

validate_on_demand_component ¶

validate_on_demand_component(
    module: Module, name: str
) -> None

Require the explicit lifecycle used by pipeline-managed components.