diff --git a/src/diffusers/hooks/hooks.py b/src/diffusers/hooks/hooks.py index 1d047251cede..6ba3d7dba695 100644 --- a/src/diffusers/hooks/hooks.py +++ b/src/diffusers/hooks/hooks.py @@ -13,6 +13,7 @@ # limitations under the License. import functools +import inspect from typing import Any import torch @@ -192,6 +193,7 @@ def new_forward(module, *args, **kwargs): return new_forward forward = self._module_ref.forward + forward_signature = inspect.signature(forward) fn_ref = HookFunctionReference() fn_ref.pre_forward = hook.pre_forward @@ -208,6 +210,7 @@ def new_forward(module, *args, **kwargs): self._module_ref.forward = functools.update_wrapper( functools.partial(rewritten_forward, self._module_ref), rewritten_forward ) + self._module_ref.forward.__signature__ = forward_signature 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..fe3009ea10a5 100644 --- a/tests/hooks/test_hooks.py +++ b/tests/hooks/test_hooks.py @@ -13,6 +13,7 @@ # limitations under the License. import gc +import inspect import pytest import torch @@ -62,6 +63,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +class ExportModel(torch.nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + ) -> torch.Tensor: + return hidden_states + timestep.float().reshape(1, 1, 1) + + class AddHook(ModelHook): def __init__(self, value: int): super().__init__() @@ -178,6 +189,54 @@ def test_hook_registry(self): assert len(registry.hooks) == 1 assert registry._hook_order == ["multiply_hook"] + def test_hook_registry_preserves_forward_signature(self): + model = ExportModel().to(torch_device) + expected_signature = inspect.signature(model.forward) + + registry = HookRegistry.check_if_exists_or_initialize(model) + registry.register_hook(ModelHook(), "noop_hook") + + assert inspect.signature(model.forward) == expected_signature + + def test_hook_registry_preserves_forward_signature_with_multiple_hooks(self): + model = ExportModel().to(torch_device) + original_signature = inspect.signature(model.forward) + + registry = HookRegistry.check_if_exists_or_initialize(model) + registry.register_hook(ModelHook(), "noop_hook") + registry.register_hook(ModelHook(), "noop_hook_2") + + assert inspect.signature(model.forward) == original_signature + + def test_hook_registry_keeps_torch_export_keyword_binding(self): + model = ExportModel().to(torch_device) + registry = HookRegistry.check_if_exists_or_initialize(model) + registry.register_hook(ModelHook(), "noop_hook") + + hidden_states = torch.randn(1, 4, 8, device=torch_device) + timestep = torch.tensor([1], device=torch_device) + + torch.export.export( + model, + args=(), + kwargs={"hidden_states": hidden_states, "timestep": timestep}, + ) + + def test_hook_registry_multiple_hooks_keep_torch_export_keyword_binding(self): + model = ExportModel().to(torch_device) + registry = HookRegistry.check_if_exists_or_initialize(model) + registry.register_hook(ModelHook(), "noop_hook") + registry.register_hook(ModelHook(), "noop_hook_2") + + hidden_states = torch.randn(1, 4, 8, device=torch_device) + timestep = torch.tensor([1], device=torch_device) + + torch.export.export( + model, + args=(), + kwargs={"hidden_states": hidden_states, "timestep": timestep}, + ) + def test_stateful_hook(self): registry = HookRegistry.check_if_exists_or_initialize(self.model) registry.register_hook(StatefulAddHook(1), "stateful_add_hook")