[utils] preserve method metadata in apply_forward_hook#13803
Open
javierdejesusda wants to merge 1 commit into
Open
[utils] preserve method metadata in apply_forward_hook#13803javierdejesusda wants to merge 1 commit into
javierdejesusda wants to merge 1 commit into
Conversation
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
9d4632e to
fbdae29
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #6575
@apply_forward_hookwraps methods likeAutoencoderKL.encodeandAutoencoderKL.decode, but the innerwrapperdid not usefunctools.wraps. As a result the wrapped methods lost their__name__,__doc__,__module__and signature, sohuggingface/doc-buildercould not render them through the[[autodoc]]directive and they show up aswrapperwith 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'sinspect.unwrap+inspect.signatureflow picks up the real method. Call-time behaviour is unchanged.Before
After
Regression tests in
tests/others/test_utils.py::ApplyForwardHookTesterassert 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
AutoencoderKLin online docs #6575Who can review?
@stevhliu @sayakpaul