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")