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

Code before:
        self._is_offloaded_to_disk = True

        # We do this to free up the RAM which is still holding the up tensor data.
        for tensor_obj in self.tensor_to_key.keys():
            tensor_obj.data = torch.empty_like(tensor_obj.data, device=self.offload_device)

    def _offload_to_memory(self):
        if self.stream is not None:
            if not self.record_stream:
                self._torch_accelerator_module.current_stream().synchronize()

            current_stream = self._torch_accelerator_module.current_stream()

            for group_module in self.modules:
                for param in group_module.parameters():
                    if self.record_stream and param.device.type != 'cpu':
                        param.data.record_stream(current_stream)
                    param.data = self.cpu_param_dict[param]
            for param in self.parameters:
                if self.record_stream and param.device.type != 'cpu':
                    param.data.record_stream(current_stream)
                param.data = self.cpu_param_dict[param]
            for buffer in self.buffers:
                if self.record_stream and buffer.device.type != 'cpu':
                    buffer.data.record_stream(current_stream)
                buffer.data = self.cpu_param_dict[buffer]
        else:
            for group_module in self.modules:
                group_module.to(self.offload_device, non_blocking=False)
            for param in self.parameters:
                param.data = param.data.to(self.offload_device, non_blocking=False)
            for buffer in self.buffers:
                buffer.data = buffer.data.to(self.offload_device, non_blocking=False)

    @torch.compiler.disable()
    def onload_(self):
        r"""Onloads the group of parameters to the onload_device."""
        if self.offload_to_disk_path is not None:
            self._onload_from_disk()
        else:
            self._onload_from_memory()

    @torch.compiler.disable()
    def offload_(self):
        r"""Offloads the group of parameters to the offload_device."""
        if self.offload_to_disk_path:
            self._offload_to_disk()
        else:
            self._offload_to_memory()



Review comment (suggestion):
Small suggestion to minimize duplication

```python
        current_stream = self._torch_accelerator_module.current_stream()

        def offload_tensor(tensor):
            if self.record_stream and tensor.device.type != 'cpu':
                tensor.data.record_stream(current_stream)
            tensor.data = self.cpu_param_dict[tensor]

        for group_module in self.modules:
            for param in group_module.parameters():
                offload_tensor(param)
        for param in self.parameters:
            offload_tensor(param)
        for buffer in self.buffers:
            offload_tensor(buffer)
```

Code after:
            tensors_to_save = {key: tensor.data.to(self.offload_device) for tensor, key in self.tensor_to_key.items()}
            safetensors.torch.save_file(tensors_to_save, self.safetensors_file_path)

        # The group is now considered offloaded to disk for the rest of the session.
        self._is_offloaded_to_disk = True

        # We do this to free up the RAM which is still holding the up tensor data.
        for tensor_obj in self.tensor_to_key.keys():
            tensor_obj.data = torch.empty_like(tensor_obj.data, device=self.offload_device)

    def _offload_to_memory(self):
        if self.stream is not None:
            if not self.record_stream:
                self._torch_accelerator_module.current_stream().synchronize()

            for group_module in self.modules:
                for param in group_module.parameters():
                    param.data = self.cpu_param_dict[param]
            for param in self.parameters:
                param.data = self.cpu_param_dict[param]
            for buffer in self.buffers:
                buffer.data = self.cpu_param_dict[buffer]
        else:
            for group_module in self.modules:
                group_module.to(self.offload_device, non_blocking=False)
            for param in self.parameters:
                param.data = param.data.to(self.offload_device, non_blocking=False)
            for buffer in self.buffers:
                buffer.data = buffer.data.to(self.offload_device, non_blocking=False)

    @torch.compiler.disable()
    def onload_(self):
        r"""Onloads the group of parameters to the onload_device."""
        if self.offload_to_disk_path is not None:
            self._onload_from_disk()
        else:
            self._onload_from_memory()

    @torch.compiler.disable()
    def offload_(self):
        r"""Offloads the group of parameters to the offload_device."""
        if self.offload_to_disk_path:
            self._offload_to_disk()
        else:
            self._offload_to_memory()


class GroupOffloadingHook(ModelHook):
    r"""
    A hook that offloads groups of torch.nn.Module to the CPU for storage and onloads to accelerator device for
    computation. Each group has one "onload leader" module that is responsible for onloading, and an "offload leader"