Skip to content

adapt attention.py to torch 2.0#2483

Closed
caiqi wants to merge 1 commit intohuggingface:mainfrom
caiqi:torch_2
Closed

adapt attention.py to torch 2.0#2483
caiqi wants to merge 1 commit intohuggingface:mainfrom
caiqi:torch_2

Conversation

@caiqi
Copy link

@caiqi caiqi commented Feb 24, 2023

The AttentionBlock is not adapted to the torch 2.0. When using StableDiffusionLatentUpscalePipeline with 768x768 images, it will raise OOM on 16GB GPU. This PR use F.scaled_dot_product_attention to decrease the memory usage. I tested on Colab this PR can fix the issue. https://colab.research.google.com/drive/1qMwzjweWSUHsYeG932OCECAeA-qkyUjb?usp=sharing .

@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Feb 24, 2023 •
edited
Loading

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

@patrickvonplaten
Copy link
Contributor

patrickvonplaten commented Mar 6, 2023

Hey @caiqi,

Thanks for the PR - I believe that we already merged PyTorch's 2.0 fast attention support in this PR: #2303

Think we can close this one, no? Very sorry for not replying earlier.

@caiqi
Copy link
Author

caiqi commented Mar 8, 2023 •
edited
Loading

Hey @caiqi,

Thanks for the PR - I believe that we already merged PyTorch's 2.0 fast attention support in this PR: #2303

Think we can close this one, no? Very sorry for not replying earlier.

@patrickvonplaten Thanks! It seems this PR #2303 updates the cross_attention.py but not attention.py?
image

I meet the memory issue due to the image decoder part, which replies on AttentionBlock.

@patrickvonplaten
Copy link
Contributor

patrickvonplaten commented Mar 8, 2023

Hey @caiqi,

yeah the naming is not great here, attention.py make use of cross_attention.py in all its attention computations.

@caiqi
Copy link
Author

caiqi commented Mar 9, 2023 •
edited
Loading

@patrickvonplaten I have tested the latest diffusers code and it seems that attention.py uses its own attention code. The following is the stack track:
image

I have tested in this colab notebook: https://colab.research.google.com/drive/1qMwzjweWSUHsYeG932OCECAeA-qkyUjb?usp=sharing

@patrickvonplaten
Copy link
Contributor

patrickvonplaten commented Mar 9, 2023

cc @williamberman we should clean this attention logic up to avoid confusion

@williamberman
Copy link
Contributor

williamberman commented Mar 16, 2023 •
edited
Loading

Appreciate it @caiqi :) we're in the process of deprecating AttentionBlock, see #1880

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@caiqi @HuggingFaceDocBuilderDev @patrickvonplaten @williamberman