[theontho]

What does this PR do?

HookRegistry.remove_hook rewires the hook registered after the removed one by assigning to its
fn_ref.forward. For a hook that defines new_forward, though, fn_ref.forward holds that
new_forward and the rest of the chain lives in fn_ref.original_forward (see register_hook).
Overwriting forward silently drops the next hook's new_forward, so that hook stops running even
though it is still registered.

This PR rewires original_forward when the next hook has one, and keeps the previous behaviour
(rewiring forward) for hooks without new_forward.

How it shows up in practice

Lazy group offloading registers a temporary tracing hook and removes it after the first forward. When
FirstBlockCache is applied on top, removing that hook bypasses the cache's tail-block new_forward
while the head block's cache hook stays active. The next step that decides to reuse the cache then
fails with:

diffusers/hooks/first_block_cache.py, in new_forward
    hidden_states = shared_state.tail_block_residuals[0] + output
TypeError: 'NoneType' object is not subscriptable

We hit this running Wan 2.2 with group offloading plus FirstBlockCacheConfig over a long sequence of
jobs. It reproduces on 0.42.0 and on current main.

Minimal reproduction

import torch
from diffusers.hooks import HookRegistry, ModelHook

class AddHook(ModelHook):
    def pre_forward(self, module, *args, **kwargs):
        return tuple(x + 1 for x in args), kwargs

class SkipHook(ModelHook):  # same shape as SkipLayerHook(skip_layer=True) in tests/hooks/test_hooks.py
    def new_forward(self, module, *args, **kwargs):
        return args[0]

model = torch.nn.Linear(4, 4)
registry = HookRegistry.check_if_exists_or_initialize(model)
registry.register_hook(AddHook(), "inner")
registry.register_hook(SkipHook(), "outer")
x = torch.zeros(1, 4)
print(model(x).abs().sum())      # 0: the outer hook skips the layer
registry.remove_hook("inner")
print(model(x).abs().sum())      # non-zero before this fix: the outer hook's new_forward is gone

Tests

Added test_remove_hook_keeps_next_hook_new_forward to tests/hooks/test_hooks.py, which fails
before the change and passes after it. Existing hook tests are unaffected; a chain of hooks without
new_forward still rewires forward exactly as before.

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests?

🤖 Generated with Claude Code

[github-actions]

Hi @theontho, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

[ZihaoW123]

Could this also cover a new_forward that calls original_forward? SkipLayerHook(skip_layer=True) never follows that link, so this test cannot detect a stale link that still executes the removed inner hook.