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

Code before:
    defined block modules. In comparison to the "leaf_level" offloading, which is more fine-grained, this offloading
    is done at the top-level blocks and modules specified in block_modules.

    When block_modules is provided, only those modules will be treated as blocks for offloading. For each specified
    module, we either offload the entire submodule or recursively apply block offloading to it.
    """
    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

    block_modules = set(config.block_modules) if config.block_modules is not None else set()

    # Create module groups for ModuleList and Sequential blocks, and explicitly defined block modules
    modules_with_group_offloading = set()
    unmatched_modules = []
    matched_module_groups = []

    for name, submodule in module.named_children():
        # Check if this is an explicitly defined block module
        if name in block_modules:
            # Apply block offloading to the specified submodule
            _apply_block_offloading_to_submodule(
                submodule, name, config, modules_with_group_offloading, matched_module_groups
            )
        elif isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
            # Handle ModuleList and Sequential blocks as before
            for i in range(0, len(submodule), config.num_blocks_per_group):
                current_modules = list(submodule[i : i + config.num_blocks_per_group])
                if len(current_modules) == 0:
                    continue

                group_id = f"{name}_{i}_{i + len(current_modules) - 1}"
                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,
                    group_id=group_id,
                )
                matched_module_groups.append(group)
                for j in range(i, i + len(current_modules)):
                    modules_with_group_offloading.add(f"{name}.{j}")

Review comment (suggestion):
```suggestion
            )
            modules_with_group_offloading.add(name)
```


Code after:
    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,
        block_modules=block_modules,
        exclude_kwargs=exclude_kwargs,
    )
    _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, and explicitly
    defined block modules. In comparison to the "leaf_level" offloading, which is more fine-grained, this offloading is
    done at the top-level blocks and modules specified in block_modules.

    When block_modules is provided, only those modules will be treated as blocks for offloading. For each specified
    module, recursively apply block offloading to it.
    """
    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

    block_modules = set(config.block_modules) if config.block_modules is not None else set()

    # Create module groups for ModuleList and Sequential blocks, and explicitly defined block modules
    modules_with_group_offloading = set()
    unmatched_modules = []
    matched_module_groups = []

    for name, submodule in module.named_children():
        # Check if this is an explicitly defined block module
        if name in block_modules:
            # Track submodule using a prefix to avoid filename collisions during disk offload.