update

Author: a-r-r-o-w

src/diffusers/hooks/group_offloading.py

@@ -448,18 +448,22 @@ def _apply_group_offloading_leaf_level(
     # of the module is called
     parameters = []
     buffers = []
+    module_dict = dict(module.named_modules())
 
-    def gather_non_module_parameters_and_buffers(m: torch.nn.Module):
-        if isinstance(m, _SUPPORTED_PYTORCH_LAYERS):
-            return
-        for parameter in m.parameters(recurse=False):
-            parameters.append(parameter)
-        for buffer in m.buffers(recurse=False):
-            buffers.append(buffer)
-        for submodule in m.children():
-            gather_non_module_parameters_and_buffers(submodule)
+    for name, parameter in module.named_parameters():
+        atoms = name.split(".")
+        parent_name = ".".join(atoms[:-1])
+        if parent_name in module_dict and isinstance(module_dict[parent_name], _SUPPORTED_PYTORCH_LAYERS):
+            continue
+        parameters.append(parameter)
+
+    for name, buffer in module.named_buffers():
+        atoms = name.split(".")
+        parent_name = ".".join(atoms[:-1])
+        if parent_name in module_dict and isinstance(module_dict[parent_name], _SUPPORTED_PYTORCH_LAYERS):
+            continue
+        buffers.append(buffer)
 
-    gather_non_module_parameters_and_buffers(module)
     unmatched_group = ModuleGroup(
         modules=[],
         offload_device=offload_device,

src/diffusers/hooks/hooks.py

@@ -31,6 +31,9 @@ class ModelHook:
 
     _is_stateful = False
 
+    def __init__(self) -> None:
+        self.fn_ref = None
+
     def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
         r"""
         Hook that is executed when a model is initialized.
@@ -103,6 +106,7 @@ def __init__(self) -> None:
         self.pre_forward = None
         self.post_forward = None
         self.old_forward = None
+        self.is_overwritten_forward = False
 
 
 class HookRegistry:
@@ -119,40 +123,36 @@ def register_hook(self, hook: ModelHook, name: str) -> None:
         if name in self.hooks.keys():
             logger.warning(f"Hook with name {name} already exists, replacing it.")
 
-        forward = self._module_ref.forward
-
-        fn_ref = FunctionReference()
-        fn_ref.pre_forward = hook.pre_forward
-        fn_ref.post_forward = hook.post_forward
-        fn_ref.old_forward = forward
-
         self._module_ref = hook.initialize_hook(self._module_ref)
 
-        def create_new_forward(function_reference: FunctionReference):
+        def create_new_forward(function_reference: FunctionReference, forward):
             def new_forward(module, *args, **kwargs):
                 args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
-                output = function_reference.old_forward(*args, **kwargs)
+                output = forward(*args, **kwargs)
                 return function_reference.post_forward(module, output)
 
             return new_forward
 
-        # if hasattr(hook, "new_forward"):
-        #     fn_ref.old_forward = hook.new_forward
+        forward = self._module_ref.forward
 
-        #     def new_forward(module, *args, **kwargs):
-        #         args, kwargs = hook.pre_forward(module, *args, **kwargs)
-        #         output = rewritten_forward(module, *args, **kwargs)
-        #         return hook.post_forward(module, output)
-        # else:
+        fn_ref = FunctionReference()
+        fn_ref.pre_forward = hook.pre_forward
+        fn_ref.post_forward = hook.post_forward
+        fn_ref.old_forward = forward
 
-        #     def new_forward(module, *args, **kwargs):
-        #         args, kwargs = hook.pre_forward(module, *args, **kwargs)
-        #         output = forward(*args, **kwargs)
-        #         return hook.post_forward(module, output)
+        if hasattr(hook, "new_forward"):
+            new_forward = hook.new_forward
+            fn_ref.is_overwritten_forward = True
+        else:
+            new_forward = forward
+            fn_ref.is_overwritten_forward = False
 
-        new_forward = create_new_forward(fn_ref)
-        self._module_ref.forward = functools.update_wrapper(functools.partial(new_forward, self._module_ref), forward)
+        rewritten_forward = create_new_forward(fn_ref, new_forward)
+        self._module_ref.forward = functools.update_wrapper(
+            functools.partial(rewritten_forward, self._module_ref), forward
+        )
 
+        hook.fn_ref = fn_ref
         self.hooks[name] = hook
         self._hook_order.append(name)
         self._fn_refs.append(fn_ref)
@@ -165,7 +165,6 @@ def remove_hook(self, name: str, recurse: bool = True) -> None:
         if name in self.hooks.keys():
             hook = self.hooks[name]
             index = self._hook_order.index(name)
-
             fn_ref = self._fn_refs[index]
 
             if index == num_hooks - 1: