diff --git a/facefusion/execution.py b/facefusion/execution.py index 6a1d2283..061866e9 100644 --- a/facefusion/execution.py +++ b/facefusion/execution.py @@ -3,7 +3,7 @@ import shutil import subprocess import xml.etree.ElementTree as ElementTree from functools import lru_cache -from typing import List, Optional +from typing import List, Optional, Tuple import onnxruntime @@ -18,6 +18,16 @@ def has_execution_provider(execution_provider : ExecutionProvider) -> bool: return execution_provider in get_available_execution_providers() +@lru_cache() +def get_onnxruntime_version() -> Tuple[int, int, int]: + version_split = onnxruntime.__version__.split('.') + major_version = int(version_split[0]) + minor_version = int(version_split[1]) + patch_version = int(version_split[2].split('+')[0]) + + return major_version, minor_version, patch_version + + def get_available_execution_providers() -> List[ExecutionProvider]: inference_session_providers = onnxruntime.get_available_providers() available_execution_providers : List[ExecutionProvider] = [] diff --git a/facefusion/inference_manager.py b/facefusion/inference_manager.py index 8d5a3330..ff5bbb03 100644 --- a/facefusion/inference_manager.py +++ b/facefusion/inference_manager.py @@ -9,7 +9,7 @@ from onnxruntime import InferenceSession from facefusion import logger, process_manager, state_manager, translator from facefusion.app_context import detect_app_context from facefusion.common_helper import is_windows -from facefusion.execution import create_inference_providers, has_execution_provider +from facefusion.execution import create_inference_providers, get_onnxruntime_version, has_execution_provider from facefusion.exit_helper import fatal_exit from facefusion.filesystem import get_file_name, is_file from facefusion.time_helper import calculate_end_time @@ -25,14 +25,16 @@ INFERENCE_POOL_SET : InferencePoolSet =\ def get_inference_pool(module_name : str, model_names : List[str], model_source_set : DownloadSet) -> InferencePool: while process_manager.is_checking(): sleep(0.5) + execution_device_ids = state_manager.get_item('execution_device_ids') execution_providers = state_manager.get_item('execution_providers') + has_arena_leak = has_execution_provider('cuda') and get_onnxruntime_version() > (1, 24, 4) app_context = detect_app_context() for execution_device_id in execution_device_ids: inference_context = get_inference_context(module_name, model_names, execution_device_id, execution_providers) - if state_manager.get_item('video_memory_strategy') == 'tolerant': + if not has_arena_leak: if app_context == 'cli' and INFERENCE_POOL_SET.get('ui').get(inference_context): INFERENCE_POOL_SET['cli'][inference_context] = INFERENCE_POOL_SET.get('ui').get(inference_context) if app_context == 'ui' and INFERENCE_POOL_SET.get('cli').get(inference_context): diff --git a/tests/test_inference_manager.py b/tests/test_inference_manager.py index dd69d669..611bc625 100644 --- a/tests/test_inference_manager.py +++ b/tests/test_inference_manager.py @@ -20,26 +20,27 @@ def test_get_inference_pool() -> None: model_names = [ 'nsfw_1', 'nsfw_2', 'nsfw_3' ] _, model_source_set = content_analyser.collect_model_downloads() - state_manager.init_item('video_memory_strategy', 'strict') + with patch('facefusion.inference_manager.has_execution_provider', return_value = True): + with patch('facefusion.inference_manager.get_onnxruntime_version', return_value = (1, 26, 0)): - with patch('facefusion.inference_manager.detect_app_context', return_value = 'cli'): - cli_inference_pool = get_inference_pool('facefusion.content_analyser', model_names, model_source_set) + with patch('facefusion.inference_manager.detect_app_context', return_value = 'cli'): + cli_inference_pool = get_inference_pool('facefusion.content_analyser', model_names, model_source_set) - assert isinstance(cli_inference_pool.get('nsfw_1'), InferenceSession) + assert isinstance(cli_inference_pool.get('nsfw_1'), InferenceSession) - with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): - ui_inference_pool = get_inference_pool('facefusion.content_analyser', model_names, model_source_set) + with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): + ui_inference_pool = get_inference_pool('facefusion.content_analyser', model_names, model_source_set) - assert isinstance(ui_inference_pool.get('nsfw_1'), InferenceSession) + assert isinstance(ui_inference_pool.get('nsfw_1'), InferenceSession) - assert not (cli_inference_pool.get('nsfw_1') is ui_inference_pool.get('nsfw_1')) + assert not (cli_inference_pool.get('nsfw_1') is ui_inference_pool.get('nsfw_1')) - state_manager.init_item('video_memory_strategy', 'tolerant') + with patch('facefusion.inference_manager.get_onnxruntime_version', return_value = (1, 24, 4)): - with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): - ui_inference_pool = get_inference_pool('facefusion.content_analyser', model_names, model_source_set) + with patch('facefusion.inference_manager.detect_app_context', return_value = 'ui'): + ui_inference_pool = get_inference_pool('facefusion.content_analyser', model_names, model_source_set) - assert isinstance(ui_inference_pool.get('nsfw_1'), InferenceSession) + assert isinstance(ui_inference_pool.get('nsfw_1'), InferenceSession) assert cli_inference_pool.get('nsfw_1') is ui_inference_pool.get('nsfw_1')