Skip to content

Add missing MochiEncoder3D.gradient_checkpointing attribute#11146

Merged
a-r-r-o-w merged 4 commits into
huggingface:mainfrom
mjkvaak-amd:add-mochiencoder3d-grad-chkpt-attribute
Apr 5, 2025
Merged

Add missing MochiEncoder3D.gradient_checkpointing attribute#11146
a-r-r-o-w merged 4 commits into
huggingface:mainfrom
mjkvaak-amd:add-mochiencoder3d-grad-chkpt-attribute

Conversation

@mjkvaak-amd
Copy link
Copy Markdown
Contributor

What does this PR do?

This PR fixes MochiEncoder3D is missing the gradient_checkpointing attribute in current main:src/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py causing the forward pass to fail with AttributeError: 'MochiEncoder3D' object has no attribute 'gradient_checkpointing'. We simply added self.gradient_checkpointing = False in MochiEncoder3D.

Also a limited set of tests were added for AutoencoderKLMochi: there are a number of tests we skip, but we feel fixing those are not in the scope of this PR.

Before submitting

@DN6 DN6 requested a review from a-r-r-o-w March 31, 2025 08:04
@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

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.

@yiyixuxu
Copy link
Copy Markdown
Collaborator

@bot/ style

Comment on lines +83 to +89
@unittest.skip("Unsupported test.")
def test_effective_gradient_checkpointing(self):
""" Fails because of conv_cache:
tests/models/autoencoders/test_models_autoencoder_mochi.py::AutoencoderKLMochiTests::test_effective_gradient_checkpointing -
TypeError: ModelMixin.enable_gradient_checkpointing.<locals>._gradient_checkpointing_func() got an unexpected keyword argument 'conv_cache'
"""
pass
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this test should be made to pass. It fails because we don't accept kwargs in gradient checkpointing func:

def _gradient_checkpointing_func(module, *args):

Could you update the vae implementation to not pass conv_cache argument as a keyword arg and instead just use normal arg? LMK if you'd like me to take this up in a separate PR. Other than, changes LGTM, thanks!

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done ✅

@a-r-r-o-w
Copy link
Copy Markdown
Contributor

@bot /style

@github-actions
Copy link
Copy Markdown
Contributor

Style fixes have been applied. View the workflow run here.

@mjkvaak-amd mjkvaak-amd requested a review from a-r-r-o-w April 2, 2025 06:13
Copy link
Copy Markdown
Contributor

@a-r-r-o-w a-r-r-o-w left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks!

@a-r-r-o-w a-r-r-o-w merged commit 8ad68c1 into huggingface:main Apr 5, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants