Public diffusion CPU-offload configuration helpers.
module-attribute ¶
DEFAULT_OFFLOAD_COMPONENTS = frozenset({DIT_COMPONENT})
module-attribute ¶
OFFLOAD_COMPONENTS = frozenset(
{DIT_COMPONENT, TEXT_ENCODER_COMPONENT}
)
dataclass ¶
Layer-mode settings for one selected diffusion component.
Resolved internal backend strategy.
class-attribute instance-attribute ¶
dataclass ¶
Validated internal form of the raw public offload mapping.
dataclass ¶
Canonical offload policy derived once from public or legacy inputs.
Return whether an enabled layer backend requires weight collectives.
component_uses_allgather(
config: Any, component: str = DIT_COMPONENT
) -> bool
Return whether one selected component uses AllGather transport.
materialize_legacy_offload_flags(
config: Any,
) -> OffloadStrategy
Keep existing strategy readers working after resolving the compact API.
Return whether any CPU-offload policy is active.
Return whether an active policy streams weights block by block.
parse_diffusion_offload_config(
value: Any,
) -> ParsedDiffusionOffloadConfig | None
Validate the compact diffusion_offload_config public API.
parse_dlo_transfer(
value: Mapping[str, str | DLOTransfer],
) -> dict[str, DLOTransfer]
Resolve an internal per-component transfer mapping.
parse_offload_components(
value: Collection[str],
) -> frozenset[str]
Validate an internal component collection.
resolve_offload(config: Any) -> ResolvedOffload
Resolve and cache the complete offload policy on its owning config.
resolve_offload_strategy(config: Any) -> OffloadStrategy
Resolve the compact config or compatibility boolean entry points.
Resolve selected components while preserving legacy topology defaults.