Function _apply_group_offloading

src/diffusers/hooks/group_offloading.py:706–712

(module: torch.nn.Module, config: GroupOffloadingConfig)

def _apply_group_offloading(module: torch.nn.Module, config: GroupOffloadingConfig) -> None:
    if config.offload_type == GroupOffloadingType.BLOCK_LEVEL:
        _apply_group_offloading_block_level(module, config)
    elif config.offload_type == GroupOffloadingType.LEAF_LEVEL:
        _apply_group_offloading_leaf_level(module, config)
    else:
        assert False


def _apply_group_offloading_block_level(module: torch.nn.Module, config: GroupOffloadingConfig) -> None: