diff --git a/facefusion.py b/facefusion.py index 223d8534..4d265873 100755 --- a/facefusion.py +++ b/facefusion.py @@ -1,11 +1,10 @@ #!/usr/bin/env python3 -import os +from facefusion import environment -os.environ['OMP_NUM_THREADS'] = '1' +environment.setup() -from facefusion import conda, core +from facefusion import core if __name__ == '__main__': - conda.setup() core.cli() diff --git a/facefusion/conda.py b/facefusion/conda.py deleted file mode 100644 index ab8bd586..00000000 --- a/facefusion/conda.py +++ /dev/null @@ -1,41 +0,0 @@ -import os -import sys -from typing import List - -from facefusion.common_helper import is_linux, is_windows - - -def setup() -> None: - conda_prefix = os.getenv('CONDA_PREFIX') - conda_ready = os.getenv('CONDA_READY') - - if conda_prefix and not conda_ready: - if is_linux(): - python_id = 'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor) - library_paths : List[str] =\ - [ - os.path.join(conda_prefix, 'lib'), - os.path.join(conda_prefix, 'lib', python_id, 'site-packages', 'tensorrt_libs') - ] - library_paths = list(filter(os.path.exists, library_paths)) - - if library_paths: - if os.getenv('LD_LIBRARY_PATH'): - library_paths.append(os.getenv('LD_LIBRARY_PATH')) - os.environ['LD_LIBRARY_PATH'] = os.pathsep.join(library_paths) - os.environ['CONDA_READY'] = '1' - os.execv(sys.executable, [ sys.executable ] + sys.argv) - - if is_windows(): - library_paths =\ - [ - os.path.join(conda_prefix, 'Lib'), - os.path.join(conda_prefix, 'Lib', 'site-packages', 'tensorrt_libs') - ] - library_paths = list(filter(os.path.exists, library_paths)) - - if library_paths: - if os.getenv('PATH'): - library_paths.append(os.getenv('PATH')) - os.environ['PATH'] = os.pathsep.join(library_paths) - os.environ['CONDA_READY'] = '1' diff --git a/facefusion/environment.py b/facefusion/environment.py new file mode 100644 index 00000000..81728cc4 --- /dev/null +++ b/facefusion/environment.py @@ -0,0 +1,49 @@ +import os +import sys +from typing import List + +from facefusion.common_helper import is_linux, is_windows + + +def setup() -> None: + conda_prefix = os.getenv('CONDA_PREFIX') + environment_ready = os.getenv('ENVIRONMENT_READY') + + if not environment_ready: + if is_linux(): + os.environ['OMP_NUM_THREADS'] = '1' + os.environ['MALLOC_ARENA_MAX'] = '4' + os.environ['ENVIRONMENT_READY'] = '1' + + if conda_prefix: + python_id = 'python' + str(sys.version_info.major) + '.' + str(sys.version_info.minor) + library_paths : List[str] =\ + [ + os.path.join(conda_prefix, 'lib'), + os.path.join(conda_prefix, 'lib', python_id, 'site-packages', 'tensorrt_libs') + ] + library_paths = list(filter(os.path.exists, library_paths)) + + if library_paths: + if os.getenv('LD_LIBRARY_PATH'): + library_paths.append(os.getenv('LD_LIBRARY_PATH')) + os.environ['LD_LIBRARY_PATH'] = os.pathsep.join(library_paths) + + os.execv(sys.executable, [ sys.executable ] + sys.argv) + + if is_windows(): + os.environ['OMP_NUM_THREADS'] = '1' + os.environ['ENVIRONMENT_READY'] = '1' + + if conda_prefix: + library_paths =\ + [ + os.path.join(conda_prefix, 'Lib'), + os.path.join(conda_prefix, 'Lib', 'site-packages', 'tensorrt_libs') + ] + library_paths = list(filter(os.path.exists, library_paths)) + + if library_paths: + if os.getenv('PATH'): + library_paths.append(os.getenv('PATH')) + os.environ['PATH'] = os.pathsep.join(library_paths)