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

Code before:
                            param.data.record_stream(current_stream)

                    for buffer in self.buffers:
                        buffer.data = pinned_memory[buffer].to(self.onload_device, non_blocking=self.non_blocking)
                        if self.record_stream:
                            buffer.data.record_stream(current_stream)

            else:
                for group_module in self.modules:
                    for param in group_module.parameters():
                        param.data = param.data.to(self.onload_device, non_blocking=self.non_blocking)
                    for buffer in group_module.buffers():
                        buffer.data = buffer.data.to(self.onload_device, non_blocking=self.non_blocking)

                for param in self.parameters:
                    param.data = param.data.to(self.onload_device, non_blocking=self.non_blocking)

                for buffer in self.buffers:
                    buffer.data = buffer.data.to(self.onload_device, non_blocking=self.non_blocking)
                    if self.record_stream:
                        buffer.data.record_stream(current_stream)

    @torch.compiler.disable()
    def offload_(self):
        r"""Offloads the group of modules to the offload_device."""
        if self.offload_to_disk:
            if not self._is_offloaded_to_disk:
                os.makedirs(os.path.dirname(self.safetensors_file_path), exist_ok=True)
                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)
                self._is_offloaded_to_disk = True

            for tensor_obj in self.tensor_to_key.keys():
                tensor_obj.data = torch.empty_like(tensor_obj.data, device=self.offload_device)
            return

        torch_accelerator_module = (
            getattr(torch, torch.accelerator.current_accelerator().type)
            if hasattr(torch, "accelerator")
            else torch.cuda
        )
        if self.stream is not None:
            if not self.record_stream:
                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]

Review comment (refactor):
(nit): we probably need to refactor this a bit and break into smaller methods so we don't have to branch and do early-returns every time a new feature is added (we can do refactor once we have everything working, so not urgent)

Code after:
                        if self.record_stream:
                            buffer.data.record_stream(current_stream)

            else:
                for group_module in self.modules:
                    for param in group_module.parameters():
                        param.data = param.data.to(self.onload_device, non_blocking=self.non_blocking)
                    for buffer in group_module.buffers():
                        buffer.data = buffer.data.to(self.onload_device, non_blocking=self.non_blocking)

                for param in self.parameters:
                    param.data = param.data.to(self.onload_device, non_blocking=self.non_blocking)

                for buffer in self.buffers:
                    buffer.data = buffer.data.to(self.onload_device, non_blocking=self.non_blocking)
                    if self.record_stream:
                        buffer.data.record_stream(current_stream)

    @torch.compiler.disable()
    def offload_(self):
        r"""Offloads the group of modules to the offload_device."""
        if self.offload_to_disk_path:
            # TODO: we can potentially optimize this code path by checking if the _all_ the desired
            # safetensor files exist on the disk and if so, skip this step entirely, reducing IO
            # overhead. Currently, we just check if the given `safetensors_file_path` exists and if not
            # we perform a write.
            # Check if the file has been saved in this session or if it already exists on disk.
            if not self._is_offloaded_to_disk and not os.path.exists(self.safetensors_file_path):
                os.makedirs(os.path.dirname(self.safetensors_file_path), exist_ok=True)
                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)
            return

        torch_accelerator_module = (
            getattr(torch, torch.accelerator.current_accelerator().type)
            if hasattr(torch, "accelerator")
            else torch.cuda
        )
        if self.stream is not None:
            if not self.record_stream:
                torch_accelerator_module.current_stream().synchronize()
            for group_module in self.modules: