AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-29 07:22:12 +03:00
Code Activity

check if safetensors already exist.

Browse Source
This commit is contained in:
sayakpaul
2025-06-12 10:48:57 +05:30
parent d8179b10d3
commit 4e4842fb0b
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/diffusers/hooks/group_offloading.py
View File

@@ -216,13 +216,20 @@ class ModuleGroup:
def offload_(self):
r"""Offloads the group of modules to the offload_device."""
if self.offload_to_disk_path:
if not self._is_offloaded_to_disk:
# 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)
self._is_offloaded_to_disk = True
# The group is now considered offloaded to disk for the rest of the session.
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)