Should we add xformers efficient memory attention mechanisms?

[Warvito]

In the recent stable diffusion 2.0 (https://github.com/Stability-AI/stablediffusion) and in the huggingface diffusers (https://github.com/huggingface/diffusers), they use the memory efficient attention from xformers (https://github.com/facebookresearch/xformers). Should we try to adopt the same on our diffusion models (or even transformers?)

[Warvito]

@ericspod xFormers depends on torch >= 1.12 (https://github.com/facebookresearch/xformers/blob/e117abd00ee6c92ade49a6fc5bb7d24ee095d7e8/requirements.txt#L3) would this be a problem in case we want to add it as an optional dependency?

[Warvito]

Recently diffuser improved their attention mechanism performance by using baddbmm(), it is something that we could adopt to our attention mechanisms
huggingface/diffusers#1203

[ericspod]

It's fine to add this an optional dependency. We can define networks or components to use xformers when it's present but use other code when not since MONAI still supports older versions of Pytorch.