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

Code before:
        low_cpu_mem_usage (`bool`, defaults to `False`):
            If True, the CPU memory usage is minimized by pinning tensors on-the-fly instead of pre-pinning them. This
            option only matters when using streamed CPU offloading (i.e. `use_stream=True`). This can be useful when
            the CPU memory is a bottleneck but may counteract the benefits of using streams.
    """
    if stream is not None and num_blocks_per_group != 1:
        logger.warning(
            f"Using streams is only supported for num_blocks_per_group=1. Got {num_blocks_per_group=}. Setting it to 1."
        )
        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), num_blocks_per_group):
            current_modules = submodule[i : i + num_blocks_per_group]
            start_idx = i
            end_idx = i + len(current_modules) - 1
            group_id = f"{name}.{start_idx}_to_{end_idx}"
            group = ModuleGroup(
                modules=current_modules,
                offload_device=offload_device,
                onload_device=onload_device,
                offload_to_disk_path=offload_to_disk_path,
                offload_leader=current_modules[-1],
                onload_leader=current_modules[0],
                non_blocking=non_blocking,
                stream=stream,
                record_stream=record_stream,
                low_cpu_mem_usage=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}")

    # Apply group offloading hooks to the module groups
    for i, group in enumerate(matched_module_groups):
        for group_module in group.modules:
            _apply_group_offloading_hook(group_module, group, None)

    # Parameters and Buffers of the top-level module need to be offloaded/onloaded separately
    # when the forward pass of this module is called. This is because the top-level module is not

Review comment (suggestion):
One liner is good here
```suggestion
            group_id = f"{name}_{i}_{i+len(current_modules)-1}"
```

Code after:
    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_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}")

    # Apply group offloading hooks to the module groups
    for i, group in enumerate(matched_module_groups):
        for group_module in group.modules:
            _apply_group_offloading_hook(group_module, group, None, config=config)