From d8dd3989de59c34674a3209fdedae9d621abab25 Mon Sep 17 00:00:00 2001
From: Wang Qiang <37444407+XDUWQ@users.noreply.github.com>
Date: Tue, 11 Jul 2023 16:28:35 +0800
Subject: [PATCH] Xformers accelerate memory efficient attention (#362)

* xformers accelerate memory efficient attention

* xformers

* precommit

---------

Co-authored-by: XDUWQ <yijing.wq@alibaba-inc.com>
---
 .../stable_diffusion/stable_diffusion.py           | 14 ++++++++++++++
 modelscope/utils/error.py                          |  6 ++++++
 modelscope/utils/import_utils.py                   |  1 +
 3 files changed, 21 insertions(+)

diff --git a/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py b/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py
index 88cb4969..6b829485 100644
--- a/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py
+++ b/modelscope/models/multi_modal/stable_diffusion/stable_diffusion.py
@@ -6,6 +6,7 @@ from typing import Callable, List, Optional, Union
 import torch
 import torch.nn.functional as F
 from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
+from packaging import version
 from transformers import CLIPTextModel, CLIPTokenizer
 
 from modelscope.metainfo import Models
@@ -34,6 +35,7 @@ class StableDiffusion(TorchModel):
         """
         super().__init__(model_dir, *args, **kwargs)
         revision = kwargs.pop('revision', None)
+        xformers_enable = kwargs.pop('xformers_enable', False)
         self.lora_tune = kwargs.pop('lora_tune', False)
         self.dreambooth_tune = kwargs.pop('dreambooth_tune', False)
 
@@ -66,6 +68,18 @@ class StableDiffusion(TorchModel):
                 self.unet.requires_grad_(False)
             self.unet = self.unet.to(self.device)
 
+        # xformers accelerate memory efficient attention
+        if xformers_enable:
+            import xformers
+
+            xformers_version = version.parse(xformers.__version__)
+            if xformers_version == version.parse('0.0.16'):
+                logger.warn(
+                    'xFormers 0.0.16 cannot be used for training in some GPUs. '
+                    'If you observe problems during training, please update xFormers to at least 0.0.17.'
+                )
+            self.unet.enable_xformers_memory_efficient_attention()
+
     def tokenize_caption(self, captions):
         """ Convert caption text to token data.
 
diff --git a/modelscope/utils/error.py b/modelscope/utils/error.py
index 841662c0..8259c7ce 100644
--- a/modelscope/utils/error.py
+++ b/modelscope/utils/error.py
@@ -168,3 +168,9 @@ TAMING_IMPORT_ERROR = """
 {0} requires the timm library but it was not found in your environment. You can install it with pip:
 `pip install taming-transformers-rom1504`
 """
+
+# docstyle-ignore
+XFORMERS_IMPORT_ERROR = """
+{0} requires the timm library but it was not found in your environment. You can install it with pip:
+`pip install xformers>=0.0.17`
+"""
diff --git a/modelscope/utils/import_utils.py b/modelscope/utils/import_utils.py
index 3e8be2e1..f2fc7e37 100644
--- a/modelscope/utils/import_utils.py
+++ b/modelscope/utils/import_utils.py
@@ -306,6 +306,7 @@ REQUIREMENTS_MAAPING = OrderedDict([
     ('mpi4py', (is_package_available('mpi4py'), MPI4PY_IMPORT_ERROR)),
     ('open_clip', (is_package_available('open_clip'), OPENCLIP_IMPORT_ERROR)),
     ('taming', (is_package_available('taming'), TAMING_IMPORT_ERROR)),
+    ('xformers', (is_package_available('xformers'), XFORMERS_IMPORT_ERROR)),
 ])
 
 SYSTEM_PACKAGE = set(['os', 'sys', 'typing'])