Fix: In the “Merge” tab of the Web UI, re-create the TTSModelHolder using the passed TTSModelHolder instance variable so that ONNX models do not get mixed up

This commit is contained in:
tsukumi
2024-09-23 08:13:16 +09:00
parent b7cd512978
commit 245bfe54ae
2 changed files with 13 additions and 2 deletions

View File

@@ -1003,6 +1003,11 @@ def method_change(x: str):
def create_merge_app(model_holder: TTSModelHolder) -> gr.Blocks: def create_merge_app(model_holder: TTSModelHolder) -> gr.Blocks:
# ONNX モデルが混じらないよう、渡された TTSModelHolder のインスタンス変数を使って TTSModelHolder を作り直す
model_holder = TTSModelHolder(
model_holder.root_dir, model_holder.device, model_holder.onnx_providers, ignore_onnx=True
)
model_names = model_holder.model_names model_names = model_holder.model_names
if len(model_names) == 0: if len(model_names) == 0:
logger.error( logger.error(
@@ -1527,7 +1532,7 @@ def create_merge_app(model_holder: TTSModelHolder) -> gr.Blocks:
if __name__ == "__main__": if __name__ == "__main__":
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
model_holder = TTSModelHolder( model_holder = TTSModelHolder(
assets_root, device, torch_device_to_onnx_providers(device) assets_root, device, torch_device_to_onnx_providers(device), ignore_onnx=True
) )
app = create_merge_app(model_holder) app = create_merge_app(model_holder)
app.launch(inbrowser=True) app.launch(inbrowser=True)

View File

@@ -467,6 +467,7 @@ class TTSModelHolder:
model_root_dir: Path, model_root_dir: Path,
device: str, device: str,
onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]], onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]],
ignore_onnx: bool = False,
) -> None: ) -> None:
""" """
Style-Bert-VITS2 の音声合成モデルを管理するクラスを初期化する。 Style-Bert-VITS2 の音声合成モデルを管理するクラスを初期化する。
@@ -488,11 +489,13 @@ class TTSModelHolder:
model_root_dir (Path): 音声合成モデルが配置されているディレクトリのパス model_root_dir (Path): 音声合成モデルが配置されているディレクトリのパス
device (str): PyTorch 推論での音声合成時に利用するデバイス (cpu, cuda, mps など) device (str): PyTorch 推論での音声合成時に利用するデバイス (cpu, cuda, mps など)
onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など) onnx_providers (list[str]): ONNX 推論で利用する ExecutionProvider (CPUExecutionProvider, CUDAExecutionProvider など)
ignore_onnx (bool, optional): ONNX モデルを除外するかどうか. Defaults to False.
""" """
self.root_dir: Path = model_root_dir self.root_dir: Path = model_root_dir
self.device: str = device self.device: str = device
self.onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = onnx_providers # fmt: skip self.onnx_providers: Sequence[Union[str, tuple[str, dict[str, Any]]]] = onnx_providers # fmt: skip
self.ignore_onnx: bool = ignore_onnx
self.model_files_dict: dict[str, list[Path]] = {} self.model_files_dict: dict[str, list[Path]] = {}
self.current_model: Optional[TTSModel] = None self.current_model: Optional[TTSModel] = None
self.model_names: list[str] = [] self.model_names: list[str] = []
@@ -513,11 +516,14 @@ class TTSModelHolder:
for model_dir in model_dirs: for model_dir in model_dirs:
if model_dir.name.startswith("."): if model_dir.name.startswith("."):
continue continue
suffixes = [".pth", ".pt", ".safetensors"]
if self.ignore_onnx is False:
suffixes.append(".onnx")
model_files = sorted( model_files = sorted(
[ [
f f
for f in model_dir.iterdir() for f in model_dir.iterdir()
if f.suffix in [".pth", ".pt", ".safetensors", ".onnx"] if f.suffix in suffixes
] ]
) )
if len(model_files) == 0: if len(model_files) == 0: