From e14295f8b3bfbef01a97f022ad5cc89f5fca4710 Mon Sep 17 00:00:00 2001 From: sergioperezcheco Date: Sun, 19 Jul 2026 11:47:22 +0800 Subject: [PATCH] fix(hooks): preserve original forward signature in register_hook HookRegistry.register_hook wraps the module forward in a (module, *args, **kwargs) closure and bound it via functools.update_wrapper pointing __wrapped__ at that closure. As a result inspect.signature(module.forward) resolved to (module, *args, **kwargs) instead of the real parameters, and torch.export/torch.compile failed with "missing a required argument: 'module'". Point __wrapped__ (via update_wrapper's wrapped argument) at the original forward captured before it is overwritten, so inspect.signature follows the chain back to the real signature. The chain works for multiple stacked hooks because each wrapper's __wrapped__ points at the previous (also-correct) wrapper. Call-time behaviour is unchanged. Regression test in tests/hooks/test_hooks.py asserts the signature and __name__ are preserved after registering one and two hooks. Fixes #14135 --- src/diffusers/hooks/hooks.py | 11 ++++++++--- tests/hooks/test_hooks.py | 21 +++++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/src/diffusers/hooks/hooks.py b/src/diffusers/hooks/hooks.py index 1d047251cede..752e69d01acb 100644 --- a/src/diffusers/hooks/hooks.py +++ b/src/diffusers/hooks/hooks.py @@ -205,9 +205,14 @@ def new_forward(module, *args, **kwargs): ) rewritten_forward = create_new_forward(fn_ref) - self._module_ref.forward = functools.update_wrapper( - functools.partial(rewritten_forward, self._module_ref), rewritten_forward - ) + # Preserve the original forward's signature/name/docstring so that + # inspect.signature(module.forward) and torch.export see the real + # parameters instead of the hook wrapper's (module, *args, **kwargs). + # We point __wrapped__ at the original forward (captured above before it + # was overwritten), which is what inspect.signature follows. + wrapped = functools.partial(rewritten_forward, self._module_ref) + functools.update_wrapper(wrapped, forward) + self._module_ref.forward = wrapped hook.fn_ref = fn_ref self.hooks[name] = hook diff --git a/tests/hooks/test_hooks.py b/tests/hooks/test_hooks.py index b40d9478355d..ec62ccb9ecfb 100644 --- a/tests/hooks/test_hooks.py +++ b/tests/hooks/test_hooks.py @@ -178,6 +178,27 @@ def test_hook_registry(self): assert len(registry.hooks) == 1 assert registry._hook_order == ["multiply_hook"] + def test_register_hook_preserves_forward_signature(self): + # register_hook rewires module.forward through a (module, *args, **kwargs) + # wrapper. Ensure the wrapped forward keeps the original signature/name so + # that inspect.signature and torch.export see the real parameters. + import inspect + + model = DummyModel(self.in_features, self.hidden_features, self.out_features, self.num_layers) + model.to(torch_device) + expected_params = list(inspect.signature(model.forward).parameters) + + registry = HookRegistry.check_if_exists_or_initialize(model) + registry.register_hook(AddHook(1), "add_hook") + registry.register_hook(MultiplyHook(2), "multiply_hook") + + assert list(inspect.signature(model.forward).parameters) == expected_params + assert model.forward.__name__ == "forward" + + # runtime behaviour is unchanged + input = torch.randn(1, 4, device=torch_device, generator=self.get_generator()) + assert model(input).shape == (1, self.out_features) + def test_stateful_hook(self): registry = HookRegistry.check_if_exists_or_initialize(self.model) registry.register_hook(StatefulAddHook(1), "stateful_add_hook")