From 765aeb94fd008e4971e015d9af57a8d6a8e1a7e3 Mon Sep 17 00:00:00 2001
From: Hameer Abbasi <hameerabbasi@yahoo.com>
Date: Sun, 8 Oct 2023 14:44:35 +0000
Subject: [PATCH] Update to PyTorch 2.1.0 with pre-release xformers.

---
 installer.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/installer.py b/installer.py
index 4da612567..36e783fec 100644
--- a/installer.py
+++ b/installer.py
@@ -361,7 +361,7 @@ def check_torch():
     elif allow_cuda and (shutil.which('nvidia-smi') is not None or os.path.exists(os.path.join(os.environ.get('SystemRoot') or r'C:\Windows', 'System32', 'nvidia-smi.exe'))):
         log.info('nVidia CUDA toolkit detected')
         torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/cu121')
-        xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.22' if opts.get('cross_attention_optimization', '') == 'xFormers' else 'none')
+        xformers_package = os.environ.get('XFORMERS_PACKAGE', '--pre xformers<0.0.24' if opts.get('cross_attention_optimization', '') == 'xFormers' else 'none')
     elif allow_rocm and (shutil.which('rocminfo') is not None or os.path.exists('/opt/rocm/bin/rocminfo') or os.path.exists('/dev/kfd')):
         log.info('AMD ROCm toolkit detected')
         os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512')
@@ -397,18 +397,18 @@ def check_torch():
                 log.debug(f'HSA_OVERRIDE_GFX_VERSION auto config is skipped for {gpu}')
         try:
             command = subprocess.run('hipconfig --version', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-            arr = command.stdout.decode(encoding="utf8", errors="ignore").split('.')
-            if len(arr) >= 2:
-                rocm_ver = f'{arr[0]}.{arr[1]}'
+            rocm_ver_tuple = tuple(int(v) for v in command.stdout.decode(encoding="utf8", errors="ignore").split('.'))
+            if len(rocm_ver_tuple) >= 2:
+                rocm_ver = f'{rocm_ver_tuple[0]}.{rocm_ver_tuple[1]}'
             log.debug(f'ROCm version detected: {rocm_ver}')
         except Exception as e:
             log.debug(f'ROCm hipconfig failed: {e}')
             rocm_ver = None
-        if rocm_ver in ['5.5', '5.6', '5.7']:
+        if rocm_ver_tuple in {(5, 7)}:
             # install torch nightly via torchvision to avoid wasting bandwidth when torchvision depends on torch from yesterday
             torch_command = os.environ.get('TORCH_COMMAND', f'torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm{rocm_ver}')
         else:
-            torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/rocm5.4.2')
+            torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm{rocm_ver}')
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
     elif allow_ipex and (args.use_ipex or shutil.which('sycl-ls') is not None or shutil.which('sycl-ls.exe') is not None or os.environ.get('ONEAPI_ROOT') is not None or os.path.exists('/opt/intel/oneapi') or os.path.exists("C:/Program Files (x86)/Intel/oneAPI") or os.path.exists("C:/oneAPI")):
         args.use_ipex = True # pylint: disable=attribute-defined-outside-init
@@ -430,7 +430,7 @@ def check_torch():
     else:
         machine = platform.machine()
         if sys.platform == 'darwin':
-            torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1 torchvision==0.15.2')
+            torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
         elif allow_directml and args.use_directml and ('arm' not in machine and 'aarch' not in machine):
             log.info('Using DirectML Backend')
             torch_command = os.environ.get('TORCH_COMMAND', 'torch-directml')
@@ -700,7 +700,7 @@ def ensure_base_requirements():
     except ImportError:
         install('setuptools', 'setuptools')
     try:
-        import setuptools # pylint: disable=unused-import
+        import setuptools # pylint: disable=unused-import # noqa: F811
     except ImportError:
         pass
     try:
@@ -708,7 +708,7 @@ def ensure_base_requirements():
     except ImportError:
         install('rich', 'rich')
     try:
-        import rich # pylint: disable=unused-import
+        import rich # pylint: disable=unused-import # noqa: F811
     except ImportError:
         pass