Method register_hook

src/diffusers/hooks/hooks.py:177–215

(self, hook: ModelHook, name: str)

        self._fn_refs = []

    def register_hook(self, hook: ModelHook, name: str) -> None:
        if name in self.hooks.keys():
            raise ValueError(
                f"Hook with name {name} already exists in the registry. Please use a different name or "
                f"first remove the existing hook and then add a new one."
            )

        self._module_ref = hook.initialize_hook(self._module_ref)

        def create_new_forward(function_reference: HookFunctionReference):
            def new_forward(module, *args, **kwargs):
                args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
                output = function_reference.forward(*args, **kwargs)
                return function_reference.post_forward(module, output)

            return new_forward

        forward = self._module_ref.forward

        fn_ref = HookFunctionReference()
        fn_ref.pre_forward = hook.pre_forward
        fn_ref.post_forward = hook.post_forward
        fn_ref.forward = forward

        if hasattr(hook, "new_forward"):
            fn_ref.original_forward = forward
            fn_ref.forward = functools.update_wrapper(
                functools.partial(hook.new_forward, self._module_ref), hook.new_forward
            )

        rewritten_forward = create_new_forward(fn_ref)
        self._module_ref.forward = functools.update_wrapper(
            functools.partial(rewritten_forward, self._module_ref), rewritten_forward
        )

        hook.fn_ref = fn_ref
        self.hooks[name] = hook
        self._hook_order.append(name)
        self._fn_refs.append(fn_ref)

    def get_hook(self, name: str) -> ModelHook | None:
        return self.hooks.get(name, None)