fix memory address problem

Author: sayakpaul

src/diffusers/hooks/group_offloading.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import hashlib
 import os
 from contextlib import contextmanager, nullcontext
 from typing import Dict, List, Optional, Set, Tuple, Union
@@ -62,6 +63,7 @@ def __init__(
         low_cpu_mem_usage: bool = False,
         onload_self: bool = True,
         offload_to_disk_path: Optional[str] = None,
+        _group_id: Optional[int] = None,
     ) -> None:
         self.modules = modules
         self.offload_device = offload_device
@@ -80,7 +82,9 @@ def __init__(
         self._is_offloaded_to_disk = False
 
         if self.offload_to_disk_path:
-            self.safetensors_file_path = os.path.join(self.offload_to_disk_path, f"group_{id(self)}.safetensors")
+            self._group_id = _group_id
+            short_hash = self._compute_group_hash(self._group_id)
+            self.safetensors_file_path = os.path.join(self.offload_to_disk_path, f"group_{short_hash}.safetensors")
 
             all_tensors = []
             for module in self.modules:
@@ -260,6 +264,11 @@ def offload_(self):
             for buffer in self.buffers:
                 buffer.data = buffer.data.to(self.offload_device, non_blocking=self.non_blocking)
 
+    def _compute_group_hash(self, group_id):
+        hashed_id = hashlib.sha256(group_id.encode("utf-8")).hexdigest()
+        # first 16 characters for a reasonably short but unique name
+        return hashed_id[:16]
+
 
 class GroupOffloadingHook(ModelHook):
     r"""
@@ -603,6 +612,9 @@ def _apply_group_offloading_block_level(
 
         for i in range(0, len(submodule), num_blocks_per_group):
             current_modules = submodule[i : i + num_blocks_per_group]
+            start_idx = i
+            end_idx = i + len(current_modules) - 1
+            group_id = f"{name}.{start_idx}_to_{end_idx}"
             group = ModuleGroup(
                 modules=current_modules,
                 offload_device=offload_device,
@@ -615,6 +627,7 @@ def _apply_group_offloading_block_level(
                 record_stream=record_stream,
                 low_cpu_mem_usage=low_cpu_mem_usage,
                 onload_self=True,
+                _group_id=group_id,
             )
             matched_module_groups.append(group)
             for j in range(i, i + len(current_modules)):
@@ -649,6 +662,7 @@ def _apply_group_offloading_block_level(
         stream=None,
         record_stream=False,
         onload_self=True,
+        _group_id="top_level_unmatched_modules",
     )
     if stream is None:
         _apply_group_offloading_hook(module, unmatched_group, None)
@@ -715,6 +729,7 @@ def _apply_group_offloading_leaf_level(
             record_stream=record_stream,
             low_cpu_mem_usage=low_cpu_mem_usage,
             onload_self=True,
+            _group_id=name,
         )
         _apply_group_offloading_hook(submodule, group, None)
         modules_with_group_offloading.add(name)
@@ -762,6 +777,7 @@ def _apply_group_offloading_leaf_level(
             record_stream=record_stream,
             low_cpu_mem_usage=low_cpu_mem_usage,
             onload_self=True,
+            _group_id=name,
         )
         _apply_group_offloading_hook(parent_module, group, None)
 
@@ -783,6 +799,7 @@ def _apply_group_offloading_leaf_level(
             record_stream=False,
             low_cpu_mem_usage=low_cpu_mem_usage,
             onload_self=True,
+            name="lazy_leafs",
         )
         _apply_lazy_group_offloading_hook(module, unmatched_group, None)