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

Code before:
                source = pinned_memory[param] if pinned_memory else param.data
                self._transfer_tensor_to_device(param, source)
            for buffer in group_module.buffers():
                source = pinned_memory[buffer] if pinned_memory else buffer.data
                self._transfer_tensor_to_device(buffer, source)

        for param in self.parameters:
            source = pinned_memory[param] if pinned_memory else param.data
            self._transfer_tensor_to_device(param, source)

        for buffer in self.buffers:
            source = pinned_memory[buffer] if pinned_memory else buffer.data
            self._transfer_tensor_to_device(buffer, source)

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

        context = nullcontext() if self.stream is None else self._torch_accelerator_module.stream(self.stream)
        current_stream = self._torch_accelerator_module.current_stream() if self.record_stream else None

        with context:
            # Load to CPU (if using streams) or directly to target device, pin, and async copy to device
            device = self.onload_device if self.stream is None else "cpu"
            loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device=device)

            if self.stream is not None:
                for key, tensor_obj in self.key_to_tensor.items():
                    pinned_tensor = loaded_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:
                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]

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

        context = nullcontext() if self.stream is None else self._torch_accelerator_module.stream(self.stream)
        with context:
            if self.stream is not None:
                with self._pinned_memory_tensors() as pinned_memory:
                    self._process_tensors_from_modules(pinned_memory)

Review comment (bug):
`device` when supplied as `torch.device("cuda")` to the `onload_device` argument, it would fail here from `safetensors` complaining `invalid device cuda`. Simply wrapping up it within `str()` solves the issue.

Code after:
                source = pinned_memory[param] if pinned_memory else param.data
                self._transfer_tensor_to_device(param, source)
            for buffer in group_module.buffers():
                source = pinned_memory[buffer] if pinned_memory else buffer.data
                self._transfer_tensor_to_device(buffer, source)

        for param in self.parameters:
            source = pinned_memory[param] if pinned_memory else param.data
            self._transfer_tensor_to_device(param, source)

        for buffer in self.buffers:
            source = pinned_memory[buffer] if pinned_memory else buffer.data
            self._transfer_tensor_to_device(buffer, source)

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

        context = nullcontext() if self.stream is None else self._torch_accelerator_module.stream(self.stream)
        current_stream = self._torch_accelerator_module.current_stream() if self.record_stream else None

        with context:
            # Load to CPU (if using streams) or directly to target device, pin, and async copy to device
            device = str(self.onload_device) if self.stream is None else "cpu"
            loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device=device)

            if self.stream is not None:
                for key, tensor_obj in self.key_to_tensor.items():
                    pinned_tensor = loaded_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:
                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]

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

        context = nullcontext() if self.stream is None else self._torch_accelerator_module.stream(self.stream)
        with context:
            if self.stream is not None:
                with self._pinned_memory_tensors() as pinned_memory:
                    self._process_tensors_from_modules(pinned_memory)