Preserve HookRegistry forward signatures

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

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