Skip to content

Conversation

kig
Copy link
Contributor

@kig kig commented Dec 1, 2022

This PR adds xformers memory-efficient attention to the VAE to make the attention parts of the VAE require less memory.

This can help with issues like #1434 to some extent, but another part of the problem are the convolution layers. Running the VAE in channels_last memory format helps, but you're still dealing with RAM use in the tens of GB . At least it's not hundreds of GB. To run the VAE on low VRAM, switch the up_blocks / down_blocks to run on the CPU.

@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Dec 1, 2022

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

@blefaudeux
Copy link
Contributor

thoughts on #1493 @kig ? It's basically meant to help PRs like this one, checking just in case that it would also make sense to somebody else who spent time in these layers

@patil-suraj
Copy link
Contributor

Hey @kig , thanks a lot for the PR! The PR looks good, we just need to add the set_use_memory_efficient_attention_xformers only in AttentionBlock now thanks to #1493. So we can remove this method from other places. Should be good to merge after that :)

@kig
Copy link
Contributor Author

kig commented Dec 3, 2022

Hi @patil-suraj , thanks for the help! I applied the changes and it's a lot simpler now :)

@kig kig marked this pull request as ready for review December 3, 2022 05:01
Copy link
Contributor

@patil-suraj patil-suraj 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 lot for addressing the comments!

@patil-suraj
Copy link
Contributor

Tried it out, works really well, merging!

@patil-suraj patil-suraj merged commit daebee0 into huggingface:main Dec 3, 2022
tcapelle pushed a commit to tcapelle/diffusers that referenced this pull request Dec 12, 2022
* Add xformers attention to VAE

* Simplify VAE xformers code

* Update src/diffusers/models/attention.py

Co-authored-by: Ilmari Heikkinen <ilmari@fhtr.org>
Co-authored-by: Suraj Patil <surajp815@gmail.com>
sliard pushed a commit to sliard/diffusers that referenced this pull request Dec 21, 2022
* Add xformers attention to VAE

* Simplify VAE xformers code

* Update src/diffusers/models/attention.py

Co-authored-by: Ilmari Heikkinen <ilmari@fhtr.org>
Co-authored-by: Suraj Patil <surajp815@gmail.com>
yoonseokjin pushed a commit to yoonseokjin/diffusers that referenced this pull request Dec 25, 2023
* Add xformers attention to VAE

* Simplify VAE xformers code

* Update src/diffusers/models/attention.py

Co-authored-by: Ilmari Heikkinen <ilmari@fhtr.org>
Co-authored-by: Suraj Patil <surajp815@gmail.com>
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