HookRegistry.register_hook breaks module forward signature

[jloftin-nv]

Describe the bug

Similar to #13803

HookRegistry.register_hook breaks a modules forward signature, and downstream prevents torch.compile / torch.export - specifically affecting context parallel hooks

Reproduction

import inspect
import torch
from diffusers.hooks.hooks import HookRegistry, ModelHook

class DummyModel(torch.nn.Module):
    def forward(self, hidden_states, timestep, encoder_hidden_states=None):
        return hidden_states + timestep.sum()

def main():
    model = DummyModel()

    sig_before = inspect.signature(model.forward)
    print("forward sig BEFORE hook:", sig_before)

    HookRegistry.check_if_exists_or_initialize(model).register_hook(ModelHook(), "noop")

    sig_after = inspect.signature(model.forward)
    print("forward sig AFTER  hook:", sig_after)

    hs = torch.randn(1, 4, 8)
    ts = torch.tensor([1.0])

    print("runtime call output shape:", tuple(model(hs, ts).shape))

    try:
        torch.export.export(model, args=(), kwargs={"hidden_states": hs, "timestep": ts})
        print("EXPORT: OK")
    except Exception as e: 
        print(f"EXPORT: FAIL -> {type(e).__name__}: {e}")


if __name__ == "__main__":
    main()

Logs

forward sig BEFORE hook: (hidden_states, timestep, encoder_hidden_states=None)
forward sig AFTER  hook: (module, *args, **kwargs)
runtime call output shape: (1, 4, 8)
EXPORT: FAIL -> TypeError: missing a required argument: 'module'

System Info

• 🤗 Diffusers version: 0.39.0
• Platform: Linux-6.11.0-17-generic-x86_64-with-glibc2.28
• Running on Google Colab?: No
• Python version: 3.12.13
• PyTorch version (GPU?): 2.12.1+cu130 (False)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 1.22.0
• Transformers version: not installed
• Accelerate version: not installed
• PEFT version: not installed
• Safetensors version: 0.8.0
• xFormers version: not installed
• Accelerator: NA
• Using GPU in script?: No
• Using distributed or parallel set-up in script?: No

Who can help?

No response