Fix HookRegistry.remove_hook dropping the next hook's new_forward

src/diffusers/hooks/hooks.py

@@ -261,7 +261,13 @@ def remove_hook(self, name: str, recurse: bool = True) -> None:
             if index == num_hooks - 1:
                 self._module_ref.forward = old_forward
             else:
-                self._fn_refs[index + 1].forward = old_forward
+                next_fn_ref = self._fn_refs[index + 1]
+                # A hook with `new_forward` keeps it in `forward` and reaches the rest of the chain
+                # through `original_forward`; rewire that, or its `new_forward` is silently dropped.
+                if next_fn_ref.original_forward is not None:
+                    next_fn_ref.original_forward = old_forward
+                else:
+                    next_fn_ref.forward = old_forward
 
             self._module_ref = hook.deinitalize_hook(self._module_ref)
             del self.hooks[name]

tests/hooks/test_hooks.py

@@ -259,6 +259,19 @@ def test_skip_layer_hook(self):
         output = self.model(input).mean().detach().cpu().item()
         assert output != 0.0
 
+    def test_remove_hook_keeps_next_hook_new_forward(self):
+        registry = HookRegistry.check_if_exists_or_initialize(self.model)
+        registry.register_hook(AddHook(1), "add_hook")
+        registry.register_hook(SkipLayerHook(skip_layer=True), "skip_layer_hook")
+
+        input = torch.zeros(1, 4, device=torch_device)
+        assert self.model(input).mean().detach().cpu().item() == 0.0
+
+        # Removing the inner hook must not bypass the outer hook's `new_forward`.
+        registry.remove_hook("add_hook")
+        assert registry._hook_order == ["skip_layer_hook"]
+        assert self.model(input).mean().detach().cpu().item() == 0.0
+
     def test_skip_layer_internal_block(self):
         registry = HookRegistry.check_if_exists_or_initialize(self.model.linear_1)
         input = torch.zeros(1, 4, device=torch_device)