Repository: huggingface/diffusers
File: src/diffusers/hooks/group_offloading.py (Python)

Code before:
    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:
    r"""
    This function applies offloading to groups of torch.nn.ModuleList or torch.nn.Sequential blocks. In comparison to
    the "leaf_level" offloading, which is more fine-grained, this offloading is done at the top-level blocks.
    """

    if config.stream is not None and config.num_blocks_per_group != 1:
        logger.warning(
            f"Using streams is only supported for num_blocks_per_group=1. Got {config.num_blocks_per_group=}. Setting it to 1."
        )
        config.num_blocks_per_group = 1

    # Create module groups for ModuleList and Sequential blocks
    modules_with_group_offloading = set()
    unmatched_modules = []
    matched_module_groups = []
    for name, submodule in module.named_children():
        if not isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
            unmatched_modules.append((name, submodule))
            modules_with_group_offloading.add(name)
            continue

        for i in range(0, len(submodule), config.num_blocks_per_group):
            current_modules = submodule[i : i + config.num_blocks_per_group]
            group = ModuleGroup(
                modules=current_modules,
                offload_device=config.offload_device,
                onload_device=config.onload_device,
                offload_to_disk_path=config.offload_to_disk_path,
                offload_leader=current_modules[-1],
                onload_leader=current_modules[0],
                non_blocking=config.non_blocking,
                stream=config.stream,
                record_stream=config.record_stream,
                low_cpu_mem_usage=config.low_cpu_mem_usage,
                onload_self=True,
            )
            matched_module_groups.append(group)
            for j in range(i, i + len(current_modules)):
                modules_with_group_offloading.add(f"{name}.{j}")

    # Apply group offloading hooks to the module groups
    for i, group in enumerate(matched_module_groups):

Review comment (suggestion):
Should this be turned into a sensible value error like previous? 👁️ 

Code after:
    if offload_type == GroupOffloadingType.BLOCK_LEVEL and num_blocks_per_group is None:
        raise ValueError("`num_blocks_per_group` must be provided when using `offload_type='block_level'.")

    _raise_error_if_accelerate_model_or_sequential_hook_present(module)

    config = GroupOffloadingConfig(
        onload_device=onload_device,
        offload_device=offload_device,
        offload_type=offload_type,
        num_blocks_per_group=num_blocks_per_group,
        non_blocking=non_blocking,
        stream=stream,
        record_stream=record_stream,
        low_cpu_mem_usage=low_cpu_mem_usage,
        offload_to_disk_path=offload_to_disk_path,
    )
    _apply_group_offloading(module, config)


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:
    r"""
    This function applies offloading to groups of torch.nn.ModuleList or torch.nn.Sequential blocks. In comparison to
    the "leaf_level" offloading, which is more fine-grained, this offloading is done at the top-level blocks.
    """

    if config.stream is not None and config.num_blocks_per_group != 1:
        logger.warning(
            f"Using streams is only supported for num_blocks_per_group=1. Got {config.num_blocks_per_group=}. Setting it to 1."
        )
        config.num_blocks_per_group = 1

    # Create module groups for ModuleList and Sequential blocks
    modules_with_group_offloading = set()
    unmatched_modules = []
    matched_module_groups = []
    for name, submodule in module.named_children():
        if not isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
            unmatched_modules.append((name, submodule))
            modules_with_group_offloading.add(name)
            continue

        for i in range(0, len(submodule), config.num_blocks_per_group):