update

Author: DN6

src/diffusers/hooks/hooks.py

@@ -32,7 +32,7 @@ class ModelHook:
     _is_stateful = False
 
     def __init__(self):
-        self.fn_ref: "FunctionReference" = None
+        self.fn_ref: "HookFunctionReference" = None
 
     def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
         r"""
@@ -101,12 +101,17 @@ def reset_state(self, module: torch.nn.Module):
         return module
 
 
-class FunctionReference:
+class HookFunctionReference:
     def __init__(self) -> None:
+        """
+        Holding class for forward functions references used in Diffusers hooks. This struct allows you to easily swap
+        out the forward function in the when a hook is removed from a modules hook registry.
+
+        """
         self.pre_forward = None
         self.post_forward = None
-        self.old_forward = None
-        self.overwritten_forward = None
+        self.forward = None
+        self.original_forward = None
 
 
 class HookRegistry:
@@ -125,24 +130,24 @@ def register_hook(self, hook: ModelHook, name: str) -> None:
 
         self._module_ref = hook.initialize_hook(self._module_ref)
 
-        def create_new_forward(function_reference: FunctionReference):
+        def create_new_forward(function_reference: HookFunctionReference):
             def new_forward(module, *args, **kwargs):
                 args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
-                output = function_reference.old_forward(*args, **kwargs)
+                output = function_reference.forward(*args, **kwargs)
                 return function_reference.post_forward(module, output)
 
             return new_forward
 
         forward = self._module_ref.forward
 
-        fn_ref = FunctionReference()
+        fn_ref = HookFunctionReference()
         fn_ref.pre_forward = hook.pre_forward
         fn_ref.post_forward = hook.post_forward
-        fn_ref.old_forward = forward
+        fn_ref.forward = forward
 
         if hasattr(hook, "new_forward"):
-            fn_ref.overwritten_forward = forward
-            fn_ref.old_forward = functools.update_wrapper(
+            fn_ref.original_forward = forward
+            fn_ref.forward = functools.update_wrapper(
                 functools.partial(hook.new_forward, self._module_ref), hook.new_forward
             )
 
@@ -160,25 +165,28 @@ def get_hook(self, name: str) -> Optional[ModelHook]:
         return self.hooks.get(name, None)
 
     def remove_hook(self, name: str, recurse: bool = True) -> None:
-        num_hooks = len(self._hook_order)
-        if name in self.hooks.keys():
-            hook = self.hooks[name]
-            index = self._hook_order.index(name)
-            fn_ref = self._fn_refs[index]
+        if name not in self.hooks.keys():
+            logger.warning(f"hook: {name} was not found in HookRegistry")
+            return
 
-            old_forward = fn_ref.old_forward
-            if fn_ref.overwritten_forward is not None:
-                old_forward = fn_ref.overwritten_forward
-
-            if index == num_hooks - 1:
-                self._module_ref.forward = old_forward
-            else:
-                self._fn_refs[index + 1].old_forward = old_forward
-
-            self._module_ref = hook.deinitalize_hook(self._module_ref)
-            del self.hooks[name]
-            self._hook_order.pop(index)
-            self._fn_refs.pop(index)
+        num_hooks = len(self._hook_order)
+        hook = self.hooks[name]
+        index = self._hook_order.index(name)
+        fn_ref = self._fn_refs[index]
+
+        old_forward = fn_ref.forward
+        if fn_ref.original_forward is not None:
+            old_forward = fn_ref.original_forward
+
+        if index == num_hooks - 1:
+            self._module_ref.forward = old_forward
+        else:
+            self._fn_refs[index + 1].forward = old_forward
+
+        self._module_ref = hook.deinitalize_hook(self._module_ref)
+        del self.hooks[name]
+        self._hook_order.pop(index)
+        self._fn_refs.pop(index)
 
         if recurse:
             for module_name, module in self._module_ref.named_modules():

src/diffusers/hooks/pyramid_attention_broadcast.py

@@ -162,7 +162,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
         )
 
         if should_compute_attention:
-            output = self.fn_ref.overwritten_forward(*args, **kwargs)
+            output = self.fn_ref.original_forward(*args, **kwargs)
         else:
             output = self.state.cache