[utils] preserve method metadata in apply_forward_hook#13803
[utils] preserve method metadata in apply_forward_hook#13803javierdejesusda 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
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
stevhliu
left a comment
There was a problem hiding this comment.
nice, thanks for this! imo seems a bit much to add a test for this simple fix, but i'll defer to @sayakpaul on whether we need to include this
|
Happy to defer on the test. I only added it because the #6575 breakage was silent: call-time behaviour was always fine, it only showed up in |
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