Introduce model helper
This commit is contained in:
parent
a043703a3d
commit
fb15d0031e
@ -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' ]
|
||||
|
11
facefusion/model_helper.py
Normal file
11
facefusion/model_helper.py
Normal 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])
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user