Skip to content

feat: add Dropout to Flax UNet #3894

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Jul 7, 2023
Merged

feat: add Dropout to Flax UNet #3894

merged 12 commits into from
Jul 7, 2023

Conversation

SauravMaheshkar
Copy link
Contributor

@SauravMaheshkar SauravMaheshkar commented Jun 28, 2023

What does this PR do?

Fixes #522

Before submitting

Who can review?

Request for Review: @patrickvonplaten

@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Jun 28, 2023

The documentation is not available anymore as the PR was closed or merged.

@patrickvonplaten
Copy link
Contributor

cc @pcuenca

Copy link
Member

@pcuenca pcuenca left a comment

Choose a reason for hiding this comment

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

Thanks for this! I just have a question about the use of nn.compact and then we can merge.

@@ -169,6 +169,7 @@ def reshape_batch_dim_to_heads(self, tensor):
tensor = tensor.reshape(batch_size // head_size, seq_len, dim * head_size)
return tensor

@nn.compact
Copy link
Member

Choose a reason for hiding this comment

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

There's already a setup method in this class, it's a bit weird (to me) that we also use nn.compact. Would it be feasible to add nn.Dropout as a module in setup instead?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Addressed in 364432b

Copy link
Member

@pcuenca pcuenca left a comment

Choose a reason for hiding this comment

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

Awesome, many thanks!

@pcuenca pcuenca merged commit 03d829d into huggingface:main Jul 7, 2023
@SauravMaheshkar SauravMaheshkar deleted the saurav/add-dropout-flax-unet branch July 7, 2023 09:39
yoonseokjin pushed a commit to yoonseokjin/diffusers that referenced this pull request Dec 25, 2023
* feat: add Dropout to Flax UNet

* feat: add @compact decorator

* fix: drop nn.compact
AmericanPresidentJimmyCarter pushed a commit to AmericanPresidentJimmyCarter/diffusers that referenced this pull request Apr 26, 2024
* feat: add Dropout to Flax UNet

* feat: add @compact decorator

* fix: drop nn.compact
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.

[Community] Add dropout to Flax UNet
4 participants