Xformers accelerate memory efficient attention (#362)

* xformers accelerate memory efficient attention

* xformers

* precommit

---------

Co-authored-by: XDUWQ <yijing.wq@alibaba-inc.com>

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.
 

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`
+"""

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'])