stable-diffusion

stable-diffusion copied to clipboard

use torch's built-in Multihead attention class for self-attention

3

CompVis stable-diffusion [uses](https://github.com/CompVis/stable-diffusion/blob/69ae4b35e0a0f6ee1af8bb9a5d0016ccb27e36dc/ldm/modules/attention.py#L152) lucidrains' multi-head attention implementation [from `perceiver_pytorch`](https://github.com/lucidrains/perceiver-pytorch/blob/abbb5d5949d3509c57749bd134f5068f2761aac7/perceiver_pytorch/perceiver_pytorch.py#L82). there's actually a built-in [MultiheadAttention](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html) class in torch that we could be using. it even has a fast-path for self-attention, which...

Birch-san

Fixed a few things on your work on ToMe based on the author's paper. Seems to consistently generate about 20% faster when merging at 2048 tokens.

lalalune