From dc8ce0723d96d1e9ed056784ce799f5771955e48 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sat, 2 Mar 2024 09:12:19 -0500
Subject: [PATCH] improve xformers installer

---
 CHANGELOG.md | 1 +
 installer.py | 7 ++++---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4e80de52b..06f445599 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -53,6 +53,7 @@
   - improve model cpu offload compatibility
   - improve model sequential offload compatibility
   - improve bfloat16 compatibility
+  - improve xformers installer to match cuda version and install triton
   - fix extra networks refresh
   - fix sdp memory attention in backend original
   - fix autodetect sd21 models
diff --git a/installer.py b/installer.py
index 50e2c9db0..6ec35e663 100644
--- a/installer.py
+++ b/installer.py
@@ -425,9 +425,12 @@ def check_torch():
         log.info('nVidia CUDA toolkit detected: nvidia-smi present')
         if not args.use_xformers:
             torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/cu121')
+            xformers_package = os.environ.get('XFORMERS_PACKAGE', '--pre triton xformers --index-url https://download.pytorch.org/whl/cu121')
         else:
             torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/cu118')
-        xformers_package = os.environ.get('XFORMERS_PACKAGE', '--pre xformers' if opts.get('cross_attention_optimization', '') == 'xFormers' else 'none')
+            xformers_package = os.environ.get('XFORMERS_PACKAGE', '--pre triton xformers --index-url https://download.pytorch.org/whl/cu118')
+        if opts.get('cross_attention_optimization', '') != 'xFormers':
+            xformers_package = 'none'
         install('onnxruntime-gpu', 'onnxruntime-gpu', ignore=True)
     elif is_rocm_available():
         is_windows = platform.system() == 'Windows'
@@ -629,8 +632,6 @@ def check_torch():
             install(f'--no-deps {xformers_package}', ignore=True)
             import torch
             import xformers
-            if torch.__version__ != '2.0.1+cu118' and xformers.__version__ in ['0.0.22', '0.0.21', '0.0.20']:
-                log.warning(f'Likely incompatible torch with: xformers=={xformers.__version__} installed: torch=={torch.__version__} required: torch==2.1.0+cu118 - build xformers manually or downgrade torch')
         elif not args.experimental and not args.use_xformers:
             uninstall('xformers')
     except Exception as e: