update

Author: sayakpaul

src/diffusers/hooks/group_offloading.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import glob
 import hashlib
 import os
 from contextlib import contextmanager, nullcontext
@@ -37,8 +36,7 @@
 _GROUP_OFFLOADING = "group_offloading"
 _LAYER_EXECUTION_TRACKER = "layer_execution_tracker"
 _LAZY_PREFETCH_GROUP_OFFLOADING = "lazy_prefetch_group_offloading"
-_GROUP_ID_LAZY_LEAF = "lazy_leafs"
-_GROUP_ID_UNMATCHED_GROUP = "top_level_unmatched_modules"
+GROUP_ID_LAZY_LEAF = "lazy_leafs"
 _SUPPORTED_PYTORCH_LAYERS = (
     torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d,
     torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d,
@@ -609,7 +607,7 @@ 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]
-            group_id = f"{name}_{i}_{i+len(current_modules)-1}"
+            group_id = f"{name}_{i}_{i + len(current_modules) - 1}"
             group = ModuleGroup(
                 modules=current_modules,
                 offload_device=offload_device,
@@ -622,7 +620,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,
+                group_id=group_id,
             )
             matched_module_groups.append(group)
             for j in range(i, i + len(current_modules)):
@@ -657,7 +655,7 @@ def _apply_group_offloading_block_level(
         stream=None,
         record_stream=False,
         onload_self=True,
-        _group_id=_GROUP_ID_UNMATCHED_GROUP,
+        group_id=f"{module.__class__.__name__}_unmatched_group",
     )
     if stream is None:
         _apply_group_offloading_hook(module, unmatched_group, None)
@@ -724,7 +722,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,
+            group_id=name,
         )
         _apply_group_offloading_hook(submodule, group, None)
         modules_with_group_offloading.add(name)
@@ -772,7 +770,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,
+            group_id=name,
         )
         _apply_group_offloading_hook(parent_module, group, None)
 
@@ -794,7 +792,7 @@ def _apply_group_offloading_leaf_level(
             record_stream=False,
             low_cpu_mem_usage=low_cpu_mem_usage,
             onload_self=True,
-            _group_id=_GROUP_ID_LAZY_LEAF,
+            group_id=GROUP_ID_LAZY_LEAF,
         )
         _apply_lazy_group_offloading_hook(module, unmatched_group, None)
 
@@ -908,90 +906,3 @@ def _compute_group_hash(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]
-
-
-def _get_expected_safetensors_files(
-    module: torch.nn.Module,
-    offload_to_disk_path: str,
-    offload_type: str,
-    num_blocks_per_group: Optional[int] = None,
-) -> Set[str]:
-    expected_files = set()
-
-    def get_hashed_filename(group_id: str) -> str:
-        short_hash = _compute_group_hash(group_id)
-        return os.path.join(offload_to_disk_path, f"group_{short_hash}.safetensors")
-
-    if offload_type == "block_level":
-        if num_blocks_per_group is None:
-            raise ValueError("num_blocks_per_group must be provided for 'block_level' offloading.")
-
-        # Handle groups of ModuleList and Sequential blocks
-        for name, submodule in module.named_children():
-            if not isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
-                continue
-
-            for i in range(0, len(submodule), num_blocks_per_group):
-                current_modules = submodule[i : i + num_blocks_per_group]
-                if not current_modules:
-                    continue
-                start_idx = i
-                end_idx = i + len(current_modules) - 1
-                group_id = f"{name}.{start_idx}_to_{end_idx}"
-                expected_files.add(get_hashed_filename(group_id))
-
-        # Handle the group for unmatched top-level modules and parameters
-        group_id = _GROUP_ID_UNMATCHED_GROUP
-        expected_files.add(get_hashed_filename(group_id))
-
-    elif offload_type == "leaf_level":
-        # Handle leaf-level module groups
-        for name, submodule in module.named_modules():
-            if isinstance(submodule, _SUPPORTED_PYTORCH_LAYERS):
-                # These groups will always have parameters, so a file is expected
-                expected_files.add(get_hashed_filename(name))
-
-        # Handle groups for non-leaf parameters/buffers
-        modules_with_group_offloading = {
-            name for name, sm in module.named_modules() if isinstance(sm, _SUPPORTED_PYTORCH_LAYERS)
-        }
-        parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
-        buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)
-
-        all_orphans = parameters + buffers
-        if all_orphans:
-            parent_to_tensors = {}
-            module_dict = dict(module.named_modules())
-            for tensor_name, _ in all_orphans:
-                parent_name = _find_parent_module_in_module_dict(tensor_name, module_dict)
-                if parent_name not in parent_to_tensors:
-                    parent_to_tensors[parent_name] = []
-                parent_to_tensors[parent_name].append(tensor_name)
-
-            for parent_name in parent_to_tensors:
-                # A file is expected for each parent that gathers orphaned tensors
-                expected_files.add(get_hashed_filename(parent_name))
-        expected_files.add(get_hashed_filename(_GROUP_ID_LAZY_LEAF))
-
-    else:
-        raise ValueError(f"Unsupported offload_type: {offload_type}")
-
-    return expected_files
-
-
-def _check_safetensors_serialization(
-    module: torch.nn.Module,
-    offload_to_disk_path: str,
-    offload_type: str,
-    num_blocks_per_group: Optional[int] = None,
-) -> bool:
-    if not os.path.isdir(offload_to_disk_path):
-        return False, None, None
-
-    expected_files = _get_expected_safetensors_files(module, offload_to_disk_path, offload_type, num_blocks_per_group)
-    actual_files = set(glob.glob(os.path.join(offload_to_disk_path, "*.safetensors")))
-    missing_files = expected_files - actual_files
-    extra_files = actual_files - expected_files
-
-    is_correct = not missing_files and not extra_files
-    return is_correct, extra_files, missing_files

src/diffusers/utils/testing_utils.py

@@ -1,4 +1,5 @@
 import functools
+import glob
 import importlib
 import importlib.metadata
 import inspect
@@ -18,7 +19,7 @@
 from contextlib import contextmanager
 from io import BytesIO, StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import PIL.Image
@@ -1377,6 +1378,103 @@ def get_device_properties() -> DeviceProperties:
 else:
     DevicePropertiesUserDict = UserDict
 
+if is_torch_available():
+    from diffusers.hooks.group_offloading import (
+        _GROUP_ID_LAZY_LEAF,
+        _SUPPORTED_PYTORCH_LAYERS,
+        _compute_group_hash,
+        _find_parent_module_in_module_dict,
+        _gather_buffers_with_no_group_offloading_parent,
+        _gather_parameters_with_no_group_offloading_parent,
+    )
+
+    def _get_expected_safetensors_files(
+        module: torch.nn.Module,
+        offload_to_disk_path: str,
+        offload_type: str,
+        num_blocks_per_group: Optional[int] = None,
+    ) -> Set[str]:
+        expected_files = set()
+
+        def get_hashed_filename(group_id: str) -> str:
+            short_hash = _compute_group_hash(group_id)
+            return os.path.join(offload_to_disk_path, f"group_{short_hash}.safetensors")
+
+        if offload_type == "block_level":
+            if num_blocks_per_group is None:
+                raise ValueError("num_blocks_per_group must be provided for 'block_level' offloading.")
+
+            # Handle groups of ModuleList and Sequential blocks
+            unmatched_modules = []
+            for name, submodule in module.named_children():
+                if not isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
+                    unmatched_modules.append(module)
+                    continue
+
+                for i in range(0, len(submodule), num_blocks_per_group):
+                    current_modules = submodule[i : i + num_blocks_per_group]
+                    if not current_modules:
+                        continue
+                    group_id = f"{name}_{i}_{i + len(current_modules) - 1}"
+                    expected_files.add(get_hashed_filename(group_id))
+
+            # Handle the group for unmatched top-level modules and parameters
+            for module in unmatched_modules:
+                expected_files.add(get_hashed_filename(f"{module.__class__.__name__}_unmatched_group"))
+
+        elif offload_type == "leaf_level":
+            # Handle leaf-level module groups
+            for name, submodule in module.named_modules():
+                if isinstance(submodule, _SUPPORTED_PYTORCH_LAYERS):
+                    # These groups will always have parameters, so a file is expected
+                    expected_files.add(get_hashed_filename(name))
+
+            # Handle groups for non-leaf parameters/buffers
+            modules_with_group_offloading = {
+                name for name, sm in module.named_modules() if isinstance(sm, _SUPPORTED_PYTORCH_LAYERS)
+            }
+            parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
+            buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)
+
+            all_orphans = parameters + buffers
+            if all_orphans:
+                parent_to_tensors = {}
+                module_dict = dict(module.named_modules())
+                for tensor_name, _ in all_orphans:
+                    parent_name = _find_parent_module_in_module_dict(tensor_name, module_dict)
+                    if parent_name not in parent_to_tensors:
+                        parent_to_tensors[parent_name] = []
+                    parent_to_tensors[parent_name].append(tensor_name)
+
+                for parent_name in parent_to_tensors:
+                    # A file is expected for each parent that gathers orphaned tensors
+                    expected_files.add(get_hashed_filename(parent_name))
+            expected_files.add(get_hashed_filename(_GROUP_ID_LAZY_LEAF))
+
+        else:
+            raise ValueError(f"Unsupported offload_type: {offload_type}")
+
+        return expected_files
+
+    def _check_safetensors_serialization(
+        module: torch.nn.Module,
+        offload_to_disk_path: str,
+        offload_type: str,
+        num_blocks_per_group: Optional[int] = None,
+    ) -> bool:
+        if not os.path.isdir(offload_to_disk_path):
+            return False, None, None
+
+        expected_files = _get_expected_safetensors_files(
+            module, offload_to_disk_path, offload_type, num_blocks_per_group
+        )
+        actual_files = set(glob.glob(os.path.join(offload_to_disk_path, "*.safetensors")))
+        missing_files = expected_files - actual_files
+        extra_files = actual_files - expected_files
+
+        is_correct = not missing_files and not extra_files
+        return is_correct, extra_files, missing_files
+
 
 class Expectations(DevicePropertiesUserDict):
     def get_expectation(self) -> Any:

tests/models/test_modeling_common.py

@@ -41,7 +41,6 @@
 from pytest import skip
 from requests.exceptions import HTTPError
 
-from diffusers.hooks.group_offloading import _check_safetensors_serialization
 from diffusers.models import SD3Transformer2DModel, UNet2DConditionModel
 from diffusers.models.attention_processor import (
     AttnProcessor,
@@ -62,6 +61,7 @@
 from diffusers.utils.hub_utils import _add_variant
 from diffusers.utils.testing_utils import (
     CaptureLogger,
+    _check_safetensors_serialization,
     backend_empty_cache,
     backend_max_memory_allocated,
     backend_reset_peak_memory_stats,