Introduce model helper

This commit is contained in:
henryruhs 2024-11-16 22:52:26 +01:00
parent a043703a3d
commit fb15d0031e
3 changed files with 13 additions and 10 deletions

View File

@ -1,15 +1,13 @@
from functools import lru_cache
from time import sleep
from typing import List
import onnx
from onnxruntime import InferenceSession
from facefusion import process_manager, state_manager
from facefusion.app_context import detect_app_context
from facefusion.execution import create_execution_providers, has_execution_provider
from facefusion.thread_helper import thread_lock
from facefusion.typing import DownloadSet, ExecutionProviderKey, InferencePool, InferencePoolSet, ModelInitializer
from facefusion.typing import DownloadSet, ExecutionProviderKey, InferencePool, InferencePoolSet
INFERENCE_POOLS : InferencePoolSet =\
{
@ -61,12 +59,6 @@ def create_inference_session(model_path : str, execution_device_id : str, execut
return InferenceSession(model_path, providers = execution_providers)
@lru_cache(maxsize = None)
def get_static_model_initializer(model_path : str) -> ModelInitializer:
model = onnx.load(model_path)
return onnx.numpy_helper.to_array(model.graph.initializer[-1])
def resolve_execution_provider_keys(model_context : str) -> List[ExecutionProviderKey]:
if has_execution_provider('coreml') and (model_context.startswith('facefusion.processors.modules.age_modifier') or model_context.startswith('facefusion.processors.modules.frame_colorizer')):
return [ 'cpu' ]

View File

@ -0,0 +1,11 @@
from functools import lru_cache
import onnx
from facefusion.typing import ModelInitializer
@lru_cache(maxsize = None)
def get_static_model_initializer(model_path : str) -> ModelInitializer:
model = onnx.load(model_path)
return onnx.numpy_helper.to_array(model.graph.initializer[-1])

View File

@ -16,7 +16,7 @@ from facefusion.face_masker import create_occlusion_mask, create_region_mask, cr
from facefusion.face_selector import find_similar_faces, sort_and_filter_faces, sort_faces_by_order
from facefusion.face_store import get_reference_faces
from facefusion.filesystem import filter_image_paths, has_image, in_directory, is_image, is_video, resolve_relative_path, same_file_extension
from facefusion.inference_manager import get_static_model_initializer
from facefusion.model_helper import get_static_model_initializer
from facefusion.processors import choices as processors_choices
from facefusion.processors.pixel_boost import explode_pixel_boost, implode_pixel_boost
from facefusion.processors.typing import FaceSwapperInputs