diff --git a/server_editor.py b/server_editor.py index aa85ef0..c6f856e 100644 --- a/server_editor.py +++ b/server_editor.py @@ -52,7 +52,7 @@ from style_bert_vits2.nlp.japanese.user_dict import ( rewrite_word, update_dict, ) -from style_bert_vits2.tts_model import TTSModelHolder +from style_bert_vits2.tts_model import TTSModelHolder, TTSModelInfo # ---フロントエンド部分に関する処理--- @@ -250,7 +250,7 @@ async def normalize(item: TextRequest): return normalize_text(item.text) -@router.get("/models_info") +@router.get("/models_info", response_model=list[TTSModelInfo]) def models_info(): return model_holder.models_info diff --git a/style_bert_vits2/tts_model.py b/style_bert_vits2/tts_model.py index fe4dcb5..527b4ab 100644 --- a/style_bert_vits2/tts_model.py +++ b/style_bert_vits2/tts_model.py @@ -1,6 +1,6 @@ import warnings from pathlib import Path -from typing import Any, Optional, Union, TypedDict +from typing import Any, Optional, Union import gradio as gr import numpy as np @@ -8,6 +8,7 @@ import pyannote.audio import torch from gradio.processing_utils import convert_to_16_bit_wav from numpy.typing import NDArray +from pydantic import BaseModel from style_bert_vits2.constants import ( DEFAULT_ASSIST_TEXT_WEIGHT, @@ -263,7 +264,7 @@ class TTSModel: return (self.hyper_parameters.data.sampling_rate, audio) -class TTSModelInfo(TypedDict): +class TTSModelInfo(BaseModel): name: str files: list[str] styles: list[str] @@ -341,12 +342,12 @@ class TTSModelHolder: styles = list(style2id.keys()) spk2id: dict[str, int] = hyper_parameters.data.spk2id speakers = list(spk2id.keys()) - self.models_info.append({ - "name": model_dir.name, - "files": [str(f) for f in model_files], - "styles": styles, - "speakers": speakers, - }) + self.models_info.append(TTSModelInfo( + name = model_dir.name, + files = [str(f) for f in model_files], + styles = styles, + speakers = speakers, + )) def get_model(self, model_name: str, model_path_str: str) -> TTSModel: diff --git a/tests/test_main.py b/tests/test_main.py index c5d786f..0c0fe77 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -13,12 +13,12 @@ def synthesize(device: str = 'cpu'): # jvnv-F2-jp モデルを探す for model_info in model_holder.models_info: - if model_info['name'] == 'jvnv-F2-jp': + if model_info.name == 'jvnv-F2-jp': # すべてのスタイルに対して音声合成を実行 - for style in model_info['styles']: + for style in model_info.styles: # 音声合成を実行 - model = model_holder.get_model(model_info['name'], model_info['files'][0]) + model = model_holder.get_model(model_info.name, model_info.files[0]) model.load() sample_rate, audio_data = model.infer( "あらゆる現実を、すべて自分のほうへねじ曲げたのだ。",