From e9ebced9234da30692577349e32d0793f33fbd27 Mon Sep 17 00:00:00 2001 From: tarepan Date: Sat, 2 Dec 2023 10:08:17 +0000 Subject: [PATCH] Refactor core load types/names --- .../synthesis_engine/core_wrapper.py | 152 +++++++++++------- 1 file changed, 91 insertions(+), 61 deletions(-) diff --git a/voicevox_engine/synthesis_engine/core_wrapper.py b/voicevox_engine/synthesis_engine/core_wrapper.py index 720c470f4..26d9e1450 100644 --- a/voicevox_engine/synthesis_engine/core_wrapper.py +++ b/voicevox_engine/synthesis_engine/core_wrapper.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from enum import Enum, auto from pathlib import Path -from typing import List, Optional +from typing import List, Literal import numpy as np @@ -59,88 +59,88 @@ class GPUType(Enum): @dataclass(frozen=True) -class CoreInfo: - name: str - platform: str - arch: str - core_type: str - gpu_type: GPUType +class _CoreInfo: + name: str # Coreファイル名 + platform: Literal["Windows", "Linux", "Darwin"] # 対応システム/OS + arch: Literal["x64", "x86", "armv7l", "aarch64", "universal"] # 対応アーキテクチャ + core_type: Literal["libtorch", "onnxruntime"] # `model_type` + gpu_type: GPUType # NONE | CUDA | DIRECT_ML # version 0.12 より前のコアの情報 -CORE_INFOS = [ +_CORE_INFOS = [ # Windows - CoreInfo( + _CoreInfo( name="core.dll", platform="Windows", arch="x64", core_type="libtorch", gpu_type=GPUType.CUDA, ), - CoreInfo( + _CoreInfo( name="core_cpu.dll", platform="Windows", arch="x64", core_type="libtorch", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="core_gpu_x64_nvidia.dll", platform="Windows", arch="x64", core_type="onnxruntime", gpu_type=GPUType.CUDA, ), - CoreInfo( + _CoreInfo( name="core_gpu_x64_directml.dll", platform="Windows", arch="x64", core_type="onnxruntime", gpu_type=GPUType.DIRECT_ML, ), - CoreInfo( + _CoreInfo( name="core_cpu_x64.dll", platform="Windows", arch="x64", core_type="onnxruntime", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="core_cpu_x86.dll", platform="Windows", arch="x86", core_type="onnxruntime", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="core_gpu_x86_directml.dll", platform="Windows", arch="x86", core_type="onnxruntime", gpu_type=GPUType.DIRECT_ML, ), - CoreInfo( + _CoreInfo( name="core_cpu_arm.dll", platform="Windows", arch="armv7l", core_type="onnxruntime", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="core_gpu_arm_directml.dll", platform="Windows", arch="armv7l", core_type="onnxruntime", gpu_type=GPUType.DIRECT_ML, ), - CoreInfo( + _CoreInfo( name="core_cpu_arm64.dll", platform="Windows", arch="aarch64", core_type="onnxruntime", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="core_gpu_arm64_directml.dll", platform="Windows", arch="aarch64", @@ -148,42 +148,42 @@ class CoreInfo: gpu_type=GPUType.DIRECT_ML, ), # Linux - CoreInfo( + _CoreInfo( name="libcore.so", platform="Linux", arch="x64", core_type="libtorch", gpu_type=GPUType.CUDA, ), - CoreInfo( + _CoreInfo( name="libcore_cpu.so", platform="Linux", arch="x64", core_type="libtorch", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="libcore_gpu_x64_nvidia.so", platform="Linux", arch="x64", core_type="onnxruntime", gpu_type=GPUType.CUDA, ), - CoreInfo( + _CoreInfo( name="libcore_cpu_x64.so", platform="Linux", arch="x64", core_type="onnxruntime", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="libcore_cpu_armhf.so", platform="Linux", arch="armv7l", core_type="onnxruntime", gpu_type=GPUType.NONE, ), - CoreInfo( + _CoreInfo( name="libcore_cpu_arm64.so", platform="Linux", arch="aarch64", @@ -191,7 +191,7 @@ class CoreInfo: gpu_type=GPUType.NONE, ), # macOS - CoreInfo( + _CoreInfo( name="libcore_cpu_universal2.dylib", platform="Darwin", arch="universal", @@ -204,17 +204,16 @@ class CoreInfo: # version 0.12 以降のコアの名前の辞書 # - version 0.12, 0.13 のコアの名前: core # - version 0.14 からのコアの名前: voicevox_core -CORENAME_DICT = { +_CORENAME_DICT = { "Windows": ("voicevox_core.dll", "core.dll"), "Linux": ("libvoicevox_core.so", "libcore.so"), "Darwin": ("libvoicevox_core.dylib", "libcore.dylib"), } -def find_version_0_12_core_or_later(core_dir: Path) -> Optional[str]: +def _find_version_0_12_core_or_later(core_dir: Path) -> str | None: """ - core_dir で指定したディレクトリにあるコアライブラリが Version 0.12 以降である場合、 - 見つかった共有ライブラリの名前を返す。 + `core_dir`直下に存在する コア Version 0.12 以降の共有ライブラリ名（None: 不在） Version 0.12 以降と判定する条件は、 @@ -227,20 +226,19 @@ def find_version_0_12_core_or_later(core_dir: Path) -> Optional[str]: if (core_dir / "metas.json").exists(): return None - for core_name in CORENAME_DICT[platform.system()]: + for core_name in _CORENAME_DICT[platform.system()]: if (core_dir / core_name).is_file(): return core_name return None -def get_arch_name() -> Optional[str]: +def _get_arch_name() -> Literal["x64", "x86", "aarch64", "armv7l"] | None: """ - platform.machine() が特定のアーキテクチャ上で複数パターンの文字列を返し得るので、 - 一意な文字列に変換する - サポート外のアーキテクチャである場合、None を返す + 実行中マシンのアーキテクチャ（None: サポート外アーキテクチャ） """ machine = platform.machine() + # 特定のアーキテクチャ上で複数パターンの文字列を返し得るので一意に変換 if machine == "x86_64" or machine == "x64" or machine == "AMD64": return "x64" elif machine == "i386" or machine == "x86": @@ -253,18 +251,33 @@ def get_arch_name() -> Optional[str]: return None -def get_core_name( - arch_name: str, +def _get_core_name( + arch_name: Literal["x64", "x86", "aarch64", "armv7l"], platform_name: str, - model_type: str, + model_type: Literal["libtorch", "onnxruntime"], gpu_type: GPUType, -) -> Optional[str]: +) -> str | None: + """ + 設定値を満たすCoreの名前（None: サポート外） + Parameters + ---------- + arch_name : Literal["x64", "x86", "aarch64", "armv7l"] + 実行中マシンのアーキテクチャ + platform_name : str + 実行中マシンのシステム名 + model_type: Literal["libtorch", "onnxruntime"] + gpu_type: GPUType + Returns + ------- + name : str | None + Core名（None: サポート外） + """ if platform_name == "Darwin": if gpu_type == GPUType.NONE and (arch_name == "x64" or arch_name == "aarch64"): arch_name = "universal" else: return None - for core_info in CORE_INFOS: + for core_info in _CORE_INFOS: if ( core_info.platform == platform_name and core_info.arch == arch_name @@ -275,27 +288,30 @@ def get_core_name( return None -def get_suitable_core_name( - model_type: str, +def _get_suitable_core_name( + model_type: Literal["libtorch", "onnxruntime"], gpu_type: GPUType, -) -> Optional[str]: - arch_name = get_arch_name() +) -> str | None: + """実行中マシン・引数設定値でサポートされるコアのファイル名（None: サポート外）""" + # 実行中マシンのアーキテクチャ・システム名 + arch_name = _get_arch_name() + platform_name = platform.system() if arch_name is None: return None - platform_name = platform.system() - return get_core_name(arch_name, platform_name, model_type, gpu_type) + return _get_core_name(arch_name, platform_name, model_type, gpu_type) -def check_core_type(core_dir: Path) -> Optional[str]: - # libtorch版はDirectML未対応なので、ここでは`gpu_type=GPUType.DIRECT_ML`は入れない +def _check_core_type(core_dir: Path) -> Literal["libtorch", "onnxruntime"] | None: + """`core_dir`直下に存在し実行中マシンで利用可能な Core の model_type（None: 利用可能 Core 無し）""" libtorch_core_names = [ - get_suitable_core_name("libtorch", gpu_type=GPUType.CUDA), - get_suitable_core_name("libtorch", gpu_type=GPUType.NONE), + _get_suitable_core_name("libtorch", gpu_type=GPUType.CUDA), + _get_suitable_core_name("libtorch", gpu_type=GPUType.NONE), + # ("libtorch", GPUType.DIRECT_ML): libtorch版はDirectML未対応 ] onnxruntime_core_names = [ - get_suitable_core_name("onnxruntime", gpu_type=GPUType.CUDA), - get_suitable_core_name("onnxruntime", gpu_type=GPUType.DIRECT_ML), - get_suitable_core_name("onnxruntime", gpu_type=GPUType.NONE), + _get_suitable_core_name("onnxruntime", gpu_type=GPUType.CUDA), + _get_suitable_core_name("onnxruntime", gpu_type=GPUType.DIRECT_ML), + _get_suitable_core_name("onnxruntime", gpu_type=GPUType.NONE), ] if any([(core_dir / name).is_file() for name in libtorch_core_names if name]): return "libtorch" @@ -306,7 +322,20 @@ def check_core_type(core_dir: Path) -> Optional[str]: def load_core(core_dir: Path, use_gpu: bool) -> CDLL: - core_name = find_version_0_12_core_or_later(core_dir) + """ + `core_dir` 直下に存在し実行中マシンでサポートされるコアDLLのロード + Parameters + ---------- + core_dir : Path + 直下にコア（共有ライブラリ）が存在するディレクトリ + use_gpu + Returns + ------- + core : CDLL + コアDLL + """ + # Core>=0.12 + core_name = _find_version_0_12_core_or_later(core_dir) if core_name: try: # NOTE: CDLL クラスのコンストラクタの引数 name には文字列を渡す必要がある。 @@ -315,29 +344,30 @@ def load_core(core_dir: Path, use_gpu: bool) -> CDLL: except OSError as err: raise RuntimeError(f"コアの読み込みに失敗しました：{err}") - model_type = check_core_type(core_dir) + # Core<0.12 + model_type = _check_core_type(core_dir) if model_type is None: raise RuntimeError("コアが見つかりません") if use_gpu or model_type == "onnxruntime": - core_name = get_suitable_core_name(model_type, gpu_type=GPUType.CUDA) + core_name = _get_suitable_core_name(model_type, gpu_type=GPUType.CUDA) if core_name: try: return CDLL(str((core_dir / core_name).resolve(strict=True))) except OSError: pass - core_name = get_suitable_core_name(model_type, gpu_type=GPUType.DIRECT_ML) + core_name = _get_suitable_core_name(model_type, gpu_type=GPUType.DIRECT_ML) if core_name: try: return CDLL(str((core_dir / core_name).resolve(strict=True))) except OSError: pass - core_name = get_suitable_core_name(model_type, gpu_type=GPUType.NONE) + core_name = _get_suitable_core_name(model_type, gpu_type=GPUType.NONE) if core_name: try: return CDLL(str((core_dir / core_name).resolve(strict=True))) except OSError as err: if model_type == "libtorch": - core_name = get_suitable_core_name(model_type, gpu_type=GPUType.CUDA) + core_name = _get_suitable_core_name(model_type, gpu_type=GPUType.CUDA) if core_name: try: return CDLL(str((core_dir / core_name).resolve(strict=True))) @@ -373,7 +403,7 @@ def __init__( self.exist_is_model_loaded = False is_version_0_12_core_or_later = ( - find_version_0_12_core_or_later(core_dir) is not None + _find_version_0_12_core_or_later(core_dir) is not None ) if is_version_0_12_core_or_later: model_type = "onnxruntime" @@ -384,7 +414,7 @@ def __init__( self.core.is_model_loaded.argtypes = (c_long,) self.core.is_model_loaded.restype = c_bool else: - model_type = check_core_type(core_dir) + model_type = _check_core_type(core_dir) assert model_type is not None if model_type == "onnxruntime":