new_forward support

Author: a-r-r-o-w

src/diffusers/hooks/hooks.py

@@ -146,7 +146,7 @@ def register_hook(self, hook: ModelHook, name: str) -> None:
         def create_new_forward(hook: ModelHook, function_reference: HookFunctionReference):
             def new_forward(module, *args, **kwargs):
                 if not hook._is_enabled:
-                    return function_reference.forward(*args, **kwargs)
+                    return function_reference.original_forward(*args, **kwargs)
                 args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
                 output = function_reference.forward(*args, **kwargs)
                 return function_reference.post_forward(module, output)
@@ -158,6 +158,7 @@ def new_forward(module, *args, **kwargs):
         fn_ref = HookFunctionReference()
         fn_ref.pre_forward = hook.pre_forward
         fn_ref.post_forward = hook.post_forward
+        fn_ref.original_forward = forward
         fn_ref.forward = forward
 
         if hasattr(hook, "new_forward"):

tests/hooks/test_hooks.py

@@ -450,6 +450,31 @@ def test_enable_disable_hook(self):
         self.assertNotEqual(output1, output2)
         self.assertEqual(output1, output3)
 
+    def test_enable_disable_hook_containing_new_forward(self):
+        registry = HookRegistry.check_if_exists_or_initialize(self.model)
+        registry.register_hook(AddHook(1), "add_hook")
+        for block in self.model.blocks:
+            block_registry = HookRegistry.check_if_exists_or_initialize(block)
+            block_registry.register_hook(SkipLayerHook(skip_layer=True), "skip_layer_hook")
+        registry.register_hook(MultiplyHook(2), "multiply_hook")
+
+        input = torch.randn(1, 4, device=torch_device, generator=self.get_generator())
+        output1 = self.model(input).mean().detach().cpu().item()
+
+        self.model._disable_hook("skip_layer_hook")
+        output2 = self.model(input).mean().detach().cpu().item()
+
+        self.model._disable_hook("add_hook")
+        output3 = self.model(input).mean().detach().cpu().item()
+
+        self.model._enable_hook("skip_layer_hook")
+        self.model._enable_hook("add_hook")
+        output4 = self.model(input).mean().detach().cpu().item()
+
+        self.assertNotEqual(output1, output2)
+        self.assertNotEqual(output2, output3)
+        self.assertEqual(output1, output4)
+
     def test_remove_all_hooks(self):
         registry = HookRegistry.check_if_exists_or_initialize(self.model)
         registry.register_hook(AddHook(1), "add_hook")