From 2d6d4edbbdb3c6d7013df1db9369634355a75846 Mon Sep 17 00:00:00 2001 From: Suraj Patil Date: Tue, 22 Nov 2022 13:37:17 +0100 Subject: [PATCH] use memory_efficient_attention by default (#1354) * use memory_efficient_attention by default * Update src/diffusers/models/attention.py Co-authored-by: Pedro Cuenca --- src/diffusers/models/attention.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 69522f76b0..7e11bde273 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +import warnings from dataclasses import dataclass from typing import Optional @@ -396,6 +397,16 @@ class BasicTransformerBlock(nn.Module): self.norm2 = nn.LayerNorm(dim) self.norm3 = nn.LayerNorm(dim) + # if xformers is installed try to use memory_efficient_attention by default + if is_xformers_available(): + try: + self._set_use_memory_efficient_attention_xformers(True) + except Exception as e: + warnings.warn( + "Could not enable memory efficient attention. Make sure xformers is installed" + f" correctly and a GPU is available: {e}" + ) + def _set_attention_slice(self, slice_size): self.attn1._slice_size = slice_size self.attn2._slice_size = slice_size