diff --git a/facefusion/libraries/aom.py b/facefusion/libraries/aom.py new file mode 100644 index 00000000..2fff393c --- /dev/null +++ b/facefusion/libraries/aom.py @@ -0,0 +1,122 @@ +import ctypes +from functools import lru_cache +from typing import Optional + +from facefusion.common_helper import is_linux, is_macos, is_windows +from facefusion.download import conditional_download_hashes, conditional_download_sources, resolve_download_url_by_provider +from facefusion.filesystem import resolve_relative_path +from facefusion.types import LibrarySet + + +@lru_cache +def create_static_library_set() -> Optional[LibrarySet]: + if is_linux(): + return\ + { + 'hashes': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libaom.hash'), + 'path': resolve_relative_path('../.libraries/libaom.hash') + } + }, + 'sources': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'linux/libaom.so'), + 'path': resolve_relative_path('../.libraries/libaom.so') + } + } + } + if is_macos(): + return\ + { + 'hashes': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libaom.hash'), + 'path': resolve_relative_path('../.libraries/libaom.hash') + } + }, + 'sources': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'macos/libaom.dylib'), + 'path': resolve_relative_path('../.libraries/libaom.dylib') + } + } + } + if is_windows(): + return\ + { + 'hashes': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/aom.hash'), + 'path': resolve_relative_path('../.libraries/aom.hash') + } + }, + 'sources': + { + 'aom': + { + 'url': resolve_download_url_by_provider('huggingface', 'libraries-4.0.0', 'windows/aom.dll'), + 'path': resolve_relative_path('../.libraries/aom.dll') + } + } + } + + return None + + +def pre_check() -> bool: + library_hash_set = create_static_library_set().get('hashes') + library_source_set = create_static_library_set().get('sources') + + return conditional_download_hashes(library_hash_set) and conditional_download_sources(library_source_set) + + +@lru_cache +def create_static_library() -> Optional[ctypes.CDLL]: + library_path = create_static_library_set().get('sources').get('aom').get('path') + + if library_path: + if is_windows(): + library = ctypes.CDLL(library_path, winmode = 0) + else: + library = ctypes.CDLL(library_path) + + if library: + return init_ctypes(library) + + return None + + +def init_ctypes(library : ctypes.CDLL) -> ctypes.CDLL: + library.aom_codec_enc_config_default.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint ] + library.aom_codec_enc_config_default.restype = ctypes.c_int + + library.aom_codec_enc_init_ver.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_long, ctypes.c_int ] + library.aom_codec_enc_init_ver.restype = ctypes.c_int + + library.aom_codec_encode.argtypes = [ ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int64, ctypes.c_ulong, ctypes.c_long, ctypes.c_ulong ] + library.aom_codec_encode.restype = ctypes.c_int + + library.aom_codec_get_cx_data.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p) ] + library.aom_codec_get_cx_data.restype = ctypes.c_void_p + + library.aom_codec_destroy.argtypes = [ ctypes.c_void_p ] + library.aom_codec_destroy.restype = ctypes.c_int + + library.aom_img_wrap.argtypes = [ ctypes.c_void_p, ctypes.c_int, ctypes.c_uint, ctypes.c_uint, ctypes.c_uint, ctypes.c_void_p ] + library.aom_img_wrap.restype = ctypes.c_void_p + + library.aom_codec_control.argtypes = [ ctypes.c_void_p, ctypes.c_int, ctypes.c_int ] + library.aom_codec_control.restype = ctypes.c_int + + return library diff --git a/tests/test_library_aom.py b/tests/test_library_aom.py new file mode 100644 index 00000000..f43370f3 --- /dev/null +++ b/tests/test_library_aom.py @@ -0,0 +1,17 @@ +import ctypes + +import pytest + +from facefusion import state_manager +from facefusion.libraries import aom as aom_module + + +@pytest.fixture(scope = 'module', autouse = True) +def before_all() -> None: + state_manager.init_item('download_providers', [ 'github', 'huggingface' ]) + + aom_module.pre_check() + + +def test_create_static_library() -> None: + assert isinstance(aom_module.create_static_library(), ctypes.CDLL)