Skip to content

[utils] preserve method metadata in apply_forward_hook#13803

Open
javierdejesusda wants to merge 1 commit into
huggingface:mainfrom
javierdejesusda:fix/apply-forward-hook-preserves-metadata
Open

[utils] preserve method metadata in apply_forward_hook#13803
javierdejesusda wants to merge 1 commit into
huggingface:mainfrom
javierdejesusda:fix/apply-forward-hook-preserves-metadata

Conversation

@javierdejesusda
Copy link
Copy Markdown

What does this PR do?

Fixes #6575

@apply_forward_hook wraps methods like AutoencoderKL.encode and AutoencoderKL.decode, but the inner wrapper did not use functools.wraps. As a result the wrapped methods lost their __name__, __doc__, __module__ and signature, so huggingface/doc-builder could not render them through the [[autodoc]] directive and they show up as wrapper with a generic (*args, **kwargs) signature on the live API docs:

The same affects every autoencoder that uses the decorator (~29 files under src/diffusers/models/autoencoders/). Adding @functools.wraps(method) to the inner wrapper restores the original metadata so doc-builder's inspect.unwrap + inspect.signature flow picks up the real method. Call-time behaviour is unchanged.

Before

AutoencoderKL.decode.__name__   == "wrapper"
AutoencoderKL.decode.__doc__    is None
inspect.signature(AutoencoderKL.decode) == (self, *args, **kwargs)

After

AutoencoderKL.decode.__name__   == "decode"
AutoencoderKL.decode.__doc__    == "Decode a batch of images. ..."
inspect.signature(AutoencoderKL.decode) == (self, z: torch.FloatTensor, return_dict: bool = True, generator=None) -> ...

Regression tests in tests/others/test_utils.py::ApplyForwardHookTester assert that the wrapped method exposes the right metadata (name, docstring, signature including default values, __wrapped__) and that the underlying offload hook still fires through the wrapper.

The HF doc-builder CI preview will show the corrected rendering on the autoencoder pages once this PR is open.

Before submitting

Who can review?

@stevhliu @sayakpaul

Add @functools.wraps to the inner wrapper so it inherits __name__,
__doc__, __module__ and the signature of the wrapped method. Without
this, methods decorated with @apply_forward_hook (e.g. AutoencoderKL's
encode and decode) lose their docstring and signature, so doc-builder
cannot render them through the [[autodoc]] directive and they show up
as "wrapper" with a generic (*args, **kwargs) signature on the live
API docs. The same affects VQModel, ConsistencyDecoderVAE and the
other autoencoders that use the decorator.

Add regression tests asserting that the wrapped method exposes the
right metadata (name, docstring, signature including default values,
__wrapped__) and that the underlying offload hook still fires through
the wrapper.

Fixes huggingface#6575
@javierdejesusda javierdejesusda force-pushed the fix/apply-forward-hook-preserves-metadata branch from 9d4632e to fbdae29 Compare May 24, 2026 23:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Missing methods of AutoencoderKL in online docs

1 participant