diff --git a/installer.py b/installer.py index 71c200f29..b128e4fed 100644 --- a/installer.py +++ b/installer.py @@ -361,8 +361,8 @@ def check_torch(): pass 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') + torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/cu121') + 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') @@ -405,11 +405,14 @@ def check_torch(): 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 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}') + torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm{rocm_ver}') + elif rocm_ver in {"5.5", "5.6"}: + torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --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') + # ROCm 5.5 is oldest for PyTorch 2.1 + torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm5.5') 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 @@ -431,7 +434,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') @@ -695,7 +698,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: @@ -703,7 +706,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