Skip to content

vllm_omni.diffusion.offloader.base

logger module-attribute 

logger = init_logger(__name__)

OffloadBackend

Bases: ABC

Base class for CPU offload backends

config instance-attribute 

config = config

device instance-attribute 

device = device

enabled instance-attribute 

enabled = False

disable abstractmethod 

disable() -> None

Disable offloading and cleanup resources.

Removes all registered hooks. Does NOT move modules back to original devices (caller responsible for that).

enable abstractmethod 

enable(pipeline: Module) -> None

Enable offloading on the pipeline.

Discovers modules, moves them to appropriate devices, and registers forward hooks for swapping/prefetching.

Parameters:

Name Type Description Default
pipeline Module

Diffusion pipeline model (e.g., Wan22Pipeline)

 required

is_enabled 

is_enabled() -> bool

OffloadConfig dataclass

pin_cpu_memory class-attribute instance-attribute 

pin_cpu_memory: bool = True

strategy instance-attribute 

strategy: OffloadStrategy

use_hsdp class-attribute instance-attribute 

use_hsdp: bool = False

from_od_config classmethod 

from_od_config(
    od_config: OmniDiffusionConfig,
) -> OffloadConfig

Extract and validate offload settings from OmniDiffusionConfig.

For now, enforces mutual exclusion between model-level and layer-wise offloading. Layer-wise takes priority if both are enabled.

Parameters:

Name Type Description Default
od_config OmniDiffusionConfig

OmniDiffusionConfig with offload settings

 required

Returns:

Type Description
OffloadConfig

OffloadConfig with validated settings

OffloadStrategy

Bases: Enum

LAYER_WISE class-attribute instance-attribute 

LAYER_WISE = 'layer_wise'

MODEL_LEVEL class-attribute instance-attribute 

MODEL_LEVEL = 'model_level'

NONE class-attribute instance-attribute 

NONE = 'none'