simplify xfrmers options: --xformers to enable and that's it

4 files changed, 16 insertions, 10 deletions

diff --git a/launch.py b/launch.py
index a592e1ba..61f62096 100644
--- a/launch.py
+++ b/launch.py
@@ -32,7 +32,7 @@ def extract_arg(args, name):
 

 

 args, skip_torch_cuda_test = extract_arg(args, '--skip-torch-cuda-test')

-args, xformers = extract_arg(args, '--xformers')

+xformers = '--xformers' in args

 

 

 def repo_dir(name):

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 5d93f7f6..91e98c16 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -22,7 +22,7 @@ def apply_optimizations():
     undo_optimizations()

 

     ldm.modules.diffusionmodules.model.nonlinearity = silu

-    if not cmd_opts.disable_opt_xformers_attention and not (cmd_opts.opt_split_attention or torch.version.hip) and shared.xformers_available:

+    if cmd_opts.xformers and shared.xformers_available and not torch.version.hip:

         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.xformers_attention_forward

         ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.xformers_attnblock_forward

     elif cmd_opts.opt_split_attention_v1:

diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py
index 05023b6f..d23d733b 100644
--- a/modules/sd_hijack_optimizations.py
+++ b/modules/sd_hijack_optimizations.py
@@ -1,4 +1,7 @@
 import math

+import sys

+import traceback

+

 import torch

 from torch import einsum

 

@@ -7,13 +10,16 @@ from einops import rearrange
 

 from modules import shared

 

-try:

-    import xformers.ops

-    import functorch

-    xformers._is_functorch_available = True

-    shared.xformers_available = True

-except Exception:

-    print('Cannot find xformers, defaulting to split attention. Try adding --xformers commandline argument to your webui-user file if you wish to install it.')

+if shared.cmd_opts.xformers:

+    try:

+        import xformers.ops

+        import functorch

+        xformers._is_functorch_available = True

+        shared.xformers_available = True

+    except Exception:

+        print("Cannot import xformers", file=sys.stderr)

+        print(traceback.format_exc(), file=sys.stderr)

+

 

 # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion

 def split_cross_attention_forward_v1(self, x, context=None, mask=None):

diff --git a/modules/shared.py b/modules/shared.py
index d68df751..02cb2722 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -43,7 +43,7 @@ parser.add_argument("--realesrgan-models-path", type=str, help="Path to director
 parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=os.path.join(models_path, 'ScuNET'))

 parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(models_path, 'SwinIR'))

 parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(models_path, 'LDSR'))

-parser.add_argument("--disable-opt-xformers-attention", action='store_true', help="force-disables xformers attention optimization")

+parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")

 parser.add_argument("--opt-split-attention", action='store_true', help="force-enables cross-attention layer optimization. By default, it's on for torch.cuda and off for other torch devices.")

 parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")

 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")