File size: 3,696 Bytes
3914b35 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import os
import shutil
import signal
import subprocess
import sys
from argparse import ArgumentParser, HelpFormatter
from facefusion import metadata, wording
from facefusion.common_helper import is_linux, is_windows
ONNXRUNTIME_SET =\
{
'default': ('onnxruntime', '1.21.0')
}
if is_windows() or is_linux():
ONNXRUNTIME_SET['cuda'] = ('onnxruntime-gpu', '1.21.0')
ONNXRUNTIME_SET['openvino'] = ('onnxruntime-openvino', '1.20.0')
if is_windows():
ONNXRUNTIME_SET['directml'] = ('onnxruntime-directml', '1.17.3')
if is_linux():
ONNXRUNTIME_SET['rocm'] = ('onnxruntime-rocm', '1.19.0')
def cli() -> None:
signal.signal(signal.SIGINT, lambda signal_number, frame: sys.exit(0))
program = ArgumentParser(formatter_class = lambda prog: HelpFormatter(prog, max_help_position = 50))
program.add_argument('--onnxruntime', help = wording.get('help.install_dependency').format(dependency = 'onnxruntime'), choices = ONNXRUNTIME_SET.keys(), required = True)
program.add_argument('--skip-conda', help = wording.get('help.skip_conda'), action = 'store_true')
program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version')
run(program)
def run(program : ArgumentParser) -> None:
args = program.parse_args()
has_conda = 'CONDA_PREFIX' in os.environ
onnxruntime_name, onnxruntime_version = ONNXRUNTIME_SET.get(args.onnxruntime)
if not args.skip_conda and not has_conda:
sys.stdout.write(wording.get('conda_not_activated') + os.linesep)
sys.exit(1)
with open('requirements.txt') as file:
for line in file.readlines():
__line__ = line.strip()
if not __line__.startswith('onnxruntime'):
subprocess.call([ shutil.which('pip'), 'install', line, '--force-reinstall' ])
if args.onnxruntime == 'rocm':
python_id = 'cp' + str(sys.version_info.major) + str(sys.version_info.minor)
if python_id in [ 'cp310', 'cp312' ]:
wheel_name = 'onnxruntime_rocm-' + onnxruntime_version + '-' + python_id + '-' + python_id + '-linux_x86_64.whl'
wheel_url = 'https://repo.radeon.com/rocm/manylinux/rocm-rel-6.3.4/' + wheel_name
subprocess.call([ shutil.which('pip'), 'install', wheel_url, '--force-reinstall' ])
else:
subprocess.call([ shutil.which('pip'), 'install', onnxruntime_name + '==' + onnxruntime_version, '--force-reinstall' ])
if args.onnxruntime == 'cuda' and has_conda:
library_paths = []
if is_linux():
if os.getenv('LD_LIBRARY_PATH'):
library_paths = os.getenv('LD_LIBRARY_PATH').split(os.pathsep)
python_id = 'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor)
library_paths.extend(
[
os.path.join(os.getenv('CONDA_PREFIX'), 'lib'),
os.path.join(os.getenv('CONDA_PREFIX'), 'lib', python_id, 'site-packages', 'tensorrt_libs')
])
library_paths = list(dict.fromkeys([ library_path for library_path in library_paths if os.path.exists(library_path) ]))
subprocess.call([ shutil.which('conda'), 'env', 'config', 'vars', 'set', 'LD_LIBRARY_PATH=' + os.pathsep.join(library_paths) ])
if is_windows():
if os.getenv('PATH'):
library_paths = os.getenv('PATH').split(os.pathsep)
library_paths.extend(
[
os.path.join(os.getenv('CONDA_PREFIX'), 'Lib'),
os.path.join(os.getenv('CONDA_PREFIX'), 'Lib', 'site-packages', 'tensorrt_libs')
])
library_paths = list(dict.fromkeys([ library_path for library_path in library_paths if os.path.exists(library_path) ]))
subprocess.call([ shutil.which('conda'), 'env', 'config', 'vars', 'set', 'PATH=' + os.pathsep.join(library_paths) ])
if args.onnxruntime in [ 'directml', 'rocm' ]:
subprocess.call([ shutil.which('pip'), 'install', 'numpy==1.26.4', '--force-reinstall' ])
|