switch to has_arena_leak guard

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] = []

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):

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')