Skip to content

vllm_omni.diffusion.offloader.module_collector

logger module-attribute 

logger = init_logger(__name__)

ModuleDiscovery

Discovers pipeline components.

If the pipeline implements :class:SupportsComponentDiscovery, its _dit_modules, _encoder_modules, and _vae_modules class variables are used directly. Otherwise, falls back to scanning well-known attribute names.

discover staticmethod 

discover(pipeline: Module) -> PipelineModules

Discover DiT, encoder, and VAE modules from pipeline.

Parameters:

Name Type Description Default
pipeline Module

Diffusion pipeline model

 required

Returns:

Type Description
PipelineModules

PipelineModules with lists of discovered modules and names

PipelineModules dataclass

dit_names instance-attribute 

dit_names: list[str]

dits instance-attribute 

dits: list[Module]

encoder_names instance-attribute 

encoder_names: list[str]

encoders instance-attribute 

encoders: list[Module]

resident_modules class-attribute instance-attribute 

resident_modules: list[Module] = field(default_factory=list)

resident_names class-attribute instance-attribute 

resident_names: list[str] = field(default_factory=list)

vae_names class-attribute instance-attribute 

vae_names: list[str] = field(default_factory=list)

vaes instance-attribute 

vaes: list[Module]

outermost_dits 

outermost_dits() -> tuple[list[str], list[Module]]

Discovered DiTs that are not nested inside another discovered DiT.

Some pipelines declare a DiT and one of its own submodules as separate DiTs (e.g. Cosmos3 lists both transformer and the nested transformer.language_model) so that offloading can treat them as independent rings. HSDP instead shards each DiT through that DiT's own _hsdp_shard_conditions; an inner DiT's blocks are already covered by its ancestor's conditions, so sharding it again would double-wrap them (and the inner stack may not declare conditions at all). Keeping only the outermost DiTs leaves the ancestor's conditions as the single source of truth for every nested block. Order follows dit_names.