Cuda 12 installer (#362)

* Make cuda nightly (12) the default

* Better keep legacy cuda just in case

* Use CUDA and ROCM versions

* Remove MacOS options from installer (CoreML include in default package)
This commit is contained in:
Henry Ruhs
2024-02-01 10:14:31 +01:00
committed by GitHub
parent 0d00c74bc1
commit d44b38103d
2 changed files with 16 additions and 15 deletions
+15 -14
View File
@@ -16,22 +16,21 @@ TORCH : Dict[str, str] =\
}
ONNXRUNTIMES : Dict[str, Tuple[str, str]] =\
{
'default': ('onnxruntime', '1.16.3')
'default': ('onnxruntime', '1.17.0')
}
if platform.system().lower() == 'linux' or platform.system().lower() == 'windows':
TORCH['cuda'] = 'cu118'
TORCH['cuda-nightly'] = 'cu121'
ONNXRUNTIMES['cuda'] = ('onnxruntime-gpu', '1.16.3')
ONNXRUNTIMES['cuda-nightly'] = ('onnxruntime-gpu', '1.17.0')
TORCH['cuda-12.1'] = 'cu121'
TORCH['cuda-11.8'] = 'cu118'
ONNXRUNTIMES['cuda-12.1'] = ('onnxruntime-gpu', '1.17.0')
ONNXRUNTIMES['cuda-11.8'] = ('onnxruntime-gpu', '1.17.0')
ONNXRUNTIMES['openvino'] = ('onnxruntime-openvino', '1.16.0')
if platform.system().lower() == 'linux':
TORCH['rocm'] = 'rocm5.6'
ONNXRUNTIMES['rocm'] = ('onnxruntime-rocm', '1.16.3')
if platform.system().lower() == 'darwin':
ONNXRUNTIMES['coreml-legacy'] = ('onnxruntime-coreml', '1.13.1')
ONNXRUNTIMES['coreml-silicon'] = ('onnxruntime-silicon', '1.16.0')
TORCH['rocm-5.4.2'] = 'rocm5.4.2'
TORCH['rocm-5.6'] = 'rocm5.6'
ONNXRUNTIMES['rocm-5.4.2'] = ('onnxruntime-rocm', '1.16.3')
ONNXRUNTIMES['rocm-5.6'] = ('onnxruntime-rocm', '1.16.3')
if platform.system().lower() == 'windows':
ONNXRUNTIMES['directml'] = ('onnxruntime-directml', '1.16.3')
ONNXRUNTIMES['directml'] = ('onnxruntime-directml', '1.17.0')
def cli() -> None:
@@ -72,9 +71,11 @@ def run(program : ArgumentParser) -> None:
subprocess.call([ 'pip', 'install', '-r', 'requirements.txt', '--force-reinstall' ])
else:
subprocess.call([ 'pip', 'install', '-r', 'requirements.txt', '--extra-index-url', 'https://download.pytorch.org/whl/' + torch_wheel, '--force-reinstall' ])
if onnxruntime == 'rocm':
if onnxruntime == 'rocm-5.4.2' or onnxruntime == 'rocm-5.6':
if python_id in [ 'cp39', 'cp310', 'cp311' ]:
wheel_name = 'onnxruntime_training-' + onnxruntime_version + '+rocm56-' + python_id + '-' + python_id + '-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'
rocm_version = onnxruntime.replace('-', '')
rocm_version = rocm_version.replace('.', '')
wheel_name = 'onnxruntime_training-' + onnxruntime_version + '+' + rocm_version + '-' + python_id + '-' + python_id + '-manylinux_2_17_x86_64.manylinux2014_x86_64.whl'
wheel_path = os.path.join(tempfile.gettempdir(), wheel_name)
wheel_url = 'https://download.onnxruntime.ai/' + wheel_name
subprocess.call([ 'curl', '--silent', '--location', '--continue-at', '-', '--output', wheel_path, wheel_url ])
@@ -83,7 +84,7 @@ def run(program : ArgumentParser) -> None:
os.remove(wheel_path)
else:
subprocess.call([ 'pip', 'uninstall', 'onnxruntime', onnxruntime_name, '-y', '-q' ])
if onnxruntime == 'cuda-nightly':
if onnxruntime == 'cuda-12.1':
subprocess.call([ 'pip', 'install', onnxruntime_name + '==' + onnxruntime_version, '--extra-index-url', 'https://pkgs.dev.azure.com/onnxruntime/onnxruntime/_packaging/onnxruntime-cuda-12/pypi/simple', '--force-reinstall' ])
else:
subprocess.call([ 'pip', 'install', onnxruntime_name + '==' + onnxruntime_version, '--force-reinstall' ])
+1 -1
View File
@@ -3,7 +3,7 @@ filetype==1.2.0
gradio==3.50.2
numpy==1.26.2
onnx==1.15.0
onnxruntime==1.16.3
onnxruntime==1.17.0
opencv-python==4.8.1.78
psutil==5.9.6
realesrgan==0.3.0