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)

vaes instance-attribute 

vaes: list[Module]