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

Code before:
            yield pinned_dict

        finally:
            pinned_dict = None

    @torch.compiler.disable()
    def onload_(self):
        r"""Onloads the group of modules to the onload_device."""
        torch_accelerator_module = (
            getattr(torch, torch.accelerator.current_accelerator().type)
            if hasattr(torch, "accelerator")
            else torch.cuda
        )
        context = nullcontext() if self.stream is None else torch_accelerator_module.stream(self.stream)
        current_stream = torch_accelerator_module.current_stream() if self.record_stream else None

        if self.offload_to_disk:
            if self.stream is not None:
                # Wait for previous Host->Device transfer to complete
                self.stream.synchronize()

            with context:
                if self.stream is not None:
                    # Load to CPU, pin, and async copy to device for overlapping transfer and compute
                    loaded_cpu_tensors = safetensors.torch.load_file(self.safetensors_file_path, device="cpu")
                    for key, tensor_obj in self.key_to_tensor.items():
                        pinned_tensor = loaded_cpu_tensors[key].pin_memory()
                        tensor_obj.data = pinned_tensor.to(self.onload_device, non_blocking=self.non_blocking)
                        if self.record_stream:
                            tensor_obj.data.record_stream(current_stream)
                else:
                    # Load directly to the target device (synchronous)
                    onload_device = (
                        self.onload_device.type if isinstance(self.onload_device, torch.device) else self.onload_device
                    )
                    loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device=onload_device)
                    for key, tensor_obj in self.key_to_tensor.items():
                        tensor_obj.data = loaded_tensors[key]
            return

        if self.stream is not None:
            # Wait for previous Host->Device transfer to complete
            self.stream.synchronize()

        with context:
            if self.stream is not None:
                with self._pinned_memory_tensors() as pinned_memory:
                    for group_module in self.modules:
                        for param in group_module.parameters():
                            param.data = pinned_memory[param].to(self.onload_device, non_blocking=self.non_blocking)
                            if self.record_stream:

Review comment (suggestion):
I think cleaner approach would be to provide a callable to `map_location` (assuming we were using torch.load instead of safetensors), which for each tensor can pin and move to device. Do we know if there is a equivalent to passing a callable with safetensors? If not, this is okay too

Code after:

    @torch.compiler.disable()
    def onload_(self):
        r"""Onloads the group of modules to the onload_device."""
        torch_accelerator_module = (
            getattr(torch, torch.accelerator.current_accelerator().type)
            if hasattr(torch, "accelerator")
            else torch.cuda
        )
        context = nullcontext() if self.stream is None else torch_accelerator_module.stream(self.stream)
        current_stream = torch_accelerator_module.current_stream() if self.record_stream else None

        if self.offload_to_disk_path:
            if self.stream is not None:
                # Wait for previous Host->Device transfer to complete
                self.stream.synchronize()

            with context:
                if self.stream is not None:
                    # Load to CPU, pin, and async copy to device for overlapping transfer and compute
                    loaded_cpu_tensors = safetensors.torch.load_file(self.safetensors_file_path, device="cpu")
                    for key, tensor_obj in self.key_to_tensor.items():
                        pinned_tensor = loaded_cpu_tensors[key].pin_memory()
                        tensor_obj.data = pinned_tensor.to(self.onload_device, non_blocking=self.non_blocking)
                        if self.record_stream:
                            tensor_obj.data.record_stream(current_stream)
                else:
                    # Load directly to the target device (synchronous)
                    onload_device = (
                        self.onload_device.type if isinstance(self.onload_device, torch.device) else self.onload_device
                    )
                    loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device=onload_device)
                    for key, tensor_obj in self.key_to_tensor.items():
                        tensor_obj.data = loaded_tensors[key]
            return

        if self.stream is not None:
            # Wait for previous Host->Device transfer to complete
            self.stream.synchronize()

        with context:
            if self.stream is not None:
                with self._pinned_memory_tensors() as pinned_memory:
                    for group_module in self.modules:
                        for param in group_module.parameters():
                            param.data = pinned_memory[param].to(self.onload_device, non_blocking=self.non_blocking)
                            if self.record_stream:
                                param.data.record_stream(current_stream)
                        for buffer in group_module.buffers():
                            buffer.data = pinned_memory[buffer].to(self.onload_device, non_blocking=self.non_blocking)
                            if self.record_stream: