File size: 5,156 Bytes
773f76d | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | import importlib
import random
from functools import lru_cache
from time import sleep, time
from typing import List
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.exit_helper import fatal_exit
from facefusion.filesystem import get_file_name, is_file
from facefusion.time_helper import calculate_end_time
from facefusion.types import DownloadSet, ExecutionProvider, InferencePool, InferencePoolSet, InferenceProvider
INFERENCE_POOL_SET : InferencePoolSet =\
{
'cli': {},
'ui': {}
}
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')
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 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):
INFERENCE_POOL_SET['ui'][inference_context] = INFERENCE_POOL_SET.get('cli').get(inference_context)
if not INFERENCE_POOL_SET.get(app_context).get(inference_context):
inference_providers = resolve_static_inference_providers(module_name, execution_device_id)
INFERENCE_POOL_SET[app_context][inference_context] = create_inference_pool(model_source_set, inference_providers)
current_inference_context = get_inference_context(module_name, model_names, random.choice(execution_device_ids), execution_providers)
return INFERENCE_POOL_SET.get(app_context).get(current_inference_context)
def create_inference_pool(model_source_set : DownloadSet, inference_providers : List[InferenceProvider]) -> InferencePool:
inference_pool : InferencePool = {}
for model_name in model_source_set.keys():
model_path = model_source_set.get(model_name).get('path')
if is_file(model_path):
inference_pool[model_name] = create_inference_session(model_path, inference_providers)
return inference_pool
def clear_inference_pool(module_name : str, model_names : List[str]) -> None:
execution_device_ids = state_manager.get_item('execution_device_ids')
execution_providers = state_manager.get_item('execution_providers')
app_context = detect_app_context()
if is_windows() and has_execution_provider('directml'):
INFERENCE_POOL_SET[app_context].clear()
for execution_device_id in execution_device_ids:
inference_context = get_inference_context(module_name, model_names, execution_device_id, execution_providers)
if INFERENCE_POOL_SET.get(app_context).get(inference_context):
del INFERENCE_POOL_SET[app_context][inference_context]
def create_inference_session(model_path : str, inference_providers : List[InferenceProvider]) -> InferenceSession:
model_file_name = get_file_name(model_path)
start_time = time()
try:
inference_session = InferenceSession(model_path, providers = inference_providers)
logger.debug(translator.get('loading_model_succeeded').format(model_name = model_file_name, seconds = calculate_end_time(start_time)), __name__)
return inference_session
except Exception:
logger.error(translator.get('loading_model_failed').format(model_name = model_file_name), __name__)
fatal_exit(1)
def get_inference_context(module_name : str, model_names : List[str], execution_device_id : int, execution_providers : List[ExecutionProvider]) -> str:
inference_context = '.'.join([ module_name ] + model_names + [ str(execution_device_id) ] + list(execution_providers))
return inference_context
@lru_cache()
def resolve_static_inference_providers(module_name : str, execution_device_id : int) -> List[InferenceProvider]:
module = importlib.import_module(module_name)
execution_providers = state_manager.get_item('execution_providers')
if hasattr(module, 'override_inference_providers'):
override_inference_providers = getattr(module, 'override_inference_providers')()
if override_inference_providers:
return override_inference_providers
if hasattr(module, 'adjust_inference_providers'):
adjust_inference_providers = getattr(module, 'adjust_inference_providers')()
if adjust_inference_providers:
inference_providers = create_inference_providers(execution_device_id, execution_providers)
for adjust_inference_provider in adjust_inference_providers:
for inference_provider in inference_providers:
if inference_provider[0] == adjust_inference_provider[0] and inference_provider[1]:
inference_provider[1].update(adjust_inference_provider[1])
return inference_providers
return create_inference_providers(execution_device_id, execution_providers)
|