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

Code before:
            offload_to_disk_path=offload_to_disk_path,
            non_blocking=non_blocking,
            stream=stream,
            record_stream=record_stream,
            low_cpu_mem_usage=low_cpu_mem_usage,
        )
    elif offload_type == "leaf_level":
        _apply_group_offloading_leaf_level(
            module=module,
            offload_device=offload_device,
            onload_device=onload_device,
            offload_to_disk_path=offload_to_disk_path,
            non_blocking=non_blocking,
            stream=stream,
            record_stream=record_stream,
            low_cpu_mem_usage=low_cpu_mem_usage,
        )
    else:
        raise ValueError(f"Unsupported offload_type: {offload_type}")


def _apply_group_offloading_block_level(
    module: torch.nn.Module,
    num_blocks_per_group: int,
    offload_device: torch.device,
    offload_to_disk_path: Optional[str],
    onload_device: torch.device,
    non_blocking: bool,
    stream: Union[torch.cuda.Stream, torch.Stream, None] = None,
    record_stream: Optional[bool] = False,
    low_cpu_mem_usage: bool = False,
) -> 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.

    Args:
        module (`torch.nn.Module`):
            The module to which group offloading is applied.
        offload_device (`torch.device`):
            The device to which the group of modules are offloaded. This should typically be the CPU.
        offload_to_disk_path: TODO
        onload_device (`torch.device`):
            The device to which the group of modules are onloaded.
        non_blocking (`bool`):
            If True, offloading and onloading is done asynchronously. This can be useful for overlapping computation
            and data transfer.
        stream (`torch.cuda.Stream`or `torch.Stream`, *optional*):
            If provided, offloading and onloading is done asynchronously using the provided stream. This can be useful
            for overlapping computation and data transfer.
        record_stream (`bool`, defaults to `False`): When enabled with `use_stream`, it marks the current tensor

Review comment (refactor):
(nit): let's move this to the bottom of argument list. some of my personal code will break with this change because I pass args instead of by name. Might be inconvenient for other users too

Code after:
            low_cpu_mem_usage=low_cpu_mem_usage,
        )
    elif offload_type == "leaf_level":
        _apply_group_offloading_leaf_level(
            module=module,
            offload_device=offload_device,
            onload_device=onload_device,
            offload_to_disk_path=offload_to_disk_path,
            non_blocking=non_blocking,
            stream=stream,
            record_stream=record_stream,
            low_cpu_mem_usage=low_cpu_mem_usage,
        )
    else:
        raise ValueError(f"Unsupported offload_type: {offload_type}")


def _apply_group_offloading_block_level(
    module: torch.nn.Module,
    num_blocks_per_group: int,
    offload_device: torch.device,
    onload_device: torch.device,
    non_blocking: bool,
    stream: Union[torch.cuda.Stream, torch.Stream, None] = None,
    record_stream: Optional[bool] = False,
    low_cpu_mem_usage: bool = False,
    offload_to_disk_path: Optional[str] = None,
) -> 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.

    Args:
        module (`torch.nn.Module`):
            The module to which group offloading is applied.
        offload_device (`torch.device`):
            The device to which the group of modules are offloaded. This should typically be the CPU.
        offload_to_disk_path (`str`, *optional*, defaults to `None`):
            The path to the directory where parameters will be offloaded. Setting this option can be useful in limited
            RAM environment settings where a reasonable speed-memory trade-off is desired.
        onload_device (`torch.device`):
            The device to which the group of modules are onloaded.
        non_blocking (`bool`):
            If True, offloading and onloading is done asynchronously. This can be useful for overlapping computation
            and data transfer.
        stream (`torch.cuda.Stream`or `torch.Stream`, *optional*):
            If provided, offloading and onloading is done asynchronously using the provided stream. This can be useful
            for overlapping computation and data transfer.
        record_stream (`bool`, defaults to `False`): When enabled with `use_stream`, it marks the current tensor
            as having been used by this stream. It is faster at the expense of slightly more memory usage. Refer to the
            [PyTorch official docs](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) more