From d4a728ccaa8fecf59481cc92e00f51801e4da038 Mon Sep 17 00:00:00 2001 From: Hameer Abbasi Date: Sun, 8 Oct 2023 14:57:34 +0000 Subject: [PATCH] Install right PyTorch version if xformers is enabled. --- installer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/installer.py b/installer.py index 36e783fec..96ed6276c 100644 --- a/installer.py +++ b/installer.py @@ -360,8 +360,10 @@ 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', '--pre xformers<0.0.24' if opts.get('cross_attention_optimization', '') == 'xFormers' else 'none') + xformers_enabled = opts.get('cross_attention_optimization', '') == 'xFormers' + cuda_version = "118" if xformers_enabled else "121" + torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/cu{cuda_version}') + xformers_package = os.environ.get('XFORMERS_PACKAGE', '--pre xformers<0.0.24' if xformers_enabled 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')