diff --git a/.gitignore b/.gitignore index f239417..8625c49 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,4 @@ safetensors.ipynb *.dic playground.ipynb +playgrounds/ diff --git a/gradio_tabs/inference.py b/gradio_tabs/inference.py index f5af097..83f6b75 100644 --- a/gradio_tabs/inference.py +++ b/gradio_tabs/inference.py @@ -1,6 +1,6 @@ import datetime import json -from typing import Optional +from typing import Any, Optional, Union import gradio as gr @@ -186,6 +186,10 @@ style_md = f""" - どのくらいに強さがいいかはモデルやスタイルによって異なるようです。 - 音声ファイルを入力する場合は、学習データと似た声音の話者（特に同じ性別）でないとよい効果が出ないかもしれません。 """ +voice_keys = ["dec"] +voice_pitch_keys = ["flow"] +speech_style_keys = ["enc_p"] +tempo_keys = ["sdp", "dp"] def make_interactive(): @@ -203,6 +207,35 @@ def gr_util(item): return (gr.update(visible=False), gr.update(visible=True)) +null_models_frame = 0 + + +def change_null_model_row( + null_model_index: int, + null_model_name: str, + null_model_path: str, + null_voice_weights: float, + null_voice_pitch_weights: float, + null_speech_style_weights: float, + null_tempo_weights: float, + null_models: dict[int, dict[str, Any]], +): + null_models[null_model_index] = { + "name": null_model_name, + "path": null_model_path, + "weight": null_voice_weights, + "pitch": null_voice_pitch_weights, + "style": null_speech_style_weights, + "tempo": null_tempo_weights, + } + if len(null_models) > null_models_frame: + keys_to_keep = list(range(null_models_frame)) + result = {k: null_models[k] for k in keys_to_keep} + else: + result = null_models + return result, True + + def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks: def tts_fn( model_name, @@ -226,9 +259,12 @@ def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks: speaker, pitch_scale, intonation_scale, + null_models: dict[int, dict[str, Union[str, float]]], + force_reload_model: bool, ): model_holder.get_model(model_name, model_path) assert model_holder.current_model is not None + logger.debug(f"Null models setting: {null_models}") wrong_tone_message = "" kata_tone: Optional[list[tuple[str, int]]] = None @@ -283,6 +319,8 @@ def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks: speaker_id=speaker_id, pitch_scale=pitch_scale, intonation_scale=intonation_scale, + null_model_params=null_models, + force_reload_model=force_reload_model, ) except InvalidToneError as e: logger.error(f"Tone error: {e}") @@ -304,7 +342,10 @@ def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks: message = f"Success, time: {duration} seconds." if wrong_tone_message != "": message = wrong_tone_message + "\n" + message - return message, (sr, audio), kata_tone_json_str + return message, (sr, audio), kata_tone_json_str, False + + def get_model_files(model_name: str): + return [str(f) for f in model_holder.model_files_dict[model_name]] model_names = model_holder.model_names if len(model_names) == 0: @@ -317,13 +358,13 @@ def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks: ) return app initial_id = 0 - initial_pth_files = [ - str(f) for f in model_holder.model_files_dict[model_names[initial_id]] - ] + initial_pth_files = get_model_files(model_names[initial_id]) with gr.Blocks(theme=GRADIO_THEME) as app: gr.Markdown(initial_md) gr.Markdown(terms_of_use_md) + null_models = gr.State({}) + force_reload_model = gr.State(False) with gr.Accordion(label="使い方", open=False): gr.Markdown(how_to_md) with gr.Row(): @@ -437,6 +478,166 @@ def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks: inputs=[use_assist_text], outputs=[assist_text, assist_text_weight], ) + with gr.Accordion(label="ヌルモデル", open=False): + with gr.Row(): + null_models_count = gr.Number( + label="ヌルモデルの数", value=0, step=1 + ) + with gr.Column(variant="panel"): + + @gr.render(inputs=[null_models_count]) + def render_null_models( + null_models_count: int, + ): + global null_models_frame + null_models_frame = null_models_count + for i in range(null_models_count): + with gr.Row(): + null_model_index = gr.Number( + value=i, + key=f"null_model_index_{i}", + visible=False, + ) + null_model_name = gr.Dropdown( + label="モデル一覧", + choices=model_names, + key=f"null_model_name_{i}", + value=model_names[initial_id], + ) + null_model_path = gr.Dropdown( + label="モデルファイル", + key=f"null_model_path_{i}", + # FIXME: 再レンダー時に選択肢が消えるのでどうにかしたい + # 現在は再レンダーでvalueは保存されるが選択肢は保存されないので選択肢が空になる + # そのときに選択肢にない値となるので、それを許す + allow_custom_value=True, + ) + null_voice_weights = gr.Slider( + minimum=0, + maximum=1, + value=1, + step=0.1, + key=f"null_voice_weights_{i}", + label="声質", + ) + null_voice_pitch_weights = gr.Slider( + minimum=0, + maximum=1, + value=1, + step=0.1, + key=f"null_voice_pitch_weights_{i}", + label="声の高さ", + ) + null_speech_style_weights = gr.Slider( + minimum=0, + maximum=1, + value=1, + step=0.1, + key=f"null_speech_style_weights_{i}", + label="話し方", + ) + null_tempo_weights = gr.Slider( + minimum=0, + maximum=1, + value=1, + step=0.1, + key=f"null_tempo_weights_{i}", + label="テンポ", + ) + + null_model_name.change( + model_holder.update_model_files_for_gradio, + inputs=[null_model_name], + outputs=[null_model_path], + ) + null_model_path.change( + make_non_interactive, outputs=[tts_button] + ) + # 愚直すぎるのでもう少しなんとかしたい + null_model_path.change( + change_null_model_row, + inputs=[ + null_model_index, + null_model_name, + null_model_path, + null_voice_weights, + null_voice_pitch_weights, + null_speech_style_weights, + null_tempo_weights, + null_models, + ], + outputs=[null_models, force_reload_model], + ) + null_voice_weights.change( + change_null_model_row, + inputs=[ + null_model_index, + null_model_name, + null_model_path, + null_voice_weights, + null_voice_pitch_weights, + null_speech_style_weights, + null_tempo_weights, + null_models, + ], + outputs=[null_models, force_reload_model], + ) + null_voice_pitch_weights.change( + change_null_model_row, + inputs=[ + null_model_index, + null_model_name, + null_model_path, + null_voice_weights, + null_voice_pitch_weights, + null_speech_style_weights, + null_tempo_weights, + null_models, + ], + outputs=[null_models, force_reload_model], + ) + null_speech_style_weights.change( + change_null_model_row, + inputs=[ + null_model_index, + null_model_name, + null_model_path, + null_voice_weights, + null_voice_pitch_weights, + null_speech_style_weights, + null_tempo_weights, + null_models, + ], + outputs=[null_models, force_reload_model], + ) + null_tempo_weights.change( + change_null_model_row, + inputs=[ + null_model_index, + null_model_name, + null_model_path, + null_voice_weights, + null_voice_pitch_weights, + null_speech_style_weights, + null_tempo_weights, + null_models, + ], + outputs=[null_models, force_reload_model], + ) + + add_btn = gr.Button("ヌルモデルを増やす") + del_btn = gr.Button("ヌルモデルを減らす") + add_btn.click( + lambda x: x + 1, + inputs=[null_models_count], + outputs=[null_models_count], + ) + del_btn.click( + lambda x: x - 1 if x > 0 else 0, + inputs=[null_models_count], + outputs=[null_models_count], + ) + with gr.Column(): with gr.Accordion("スタイルについて詳細", open=False): gr.Markdown(style_md) @@ -494,8 +695,10 @@ def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks: speaker, pitch_scale, intonation_scale, + null_models, + force_reload_model, ], - outputs=[text_output, audio_output, tone], + outputs=[text_output, audio_output, tone, force_reload_model], ) model_name.change( diff --git a/gradio_tabs/merge.py b/gradio_tabs/merge.py index d55b04f..1bd3d13 100644 --- a/gradio_tabs/merge.py +++ b/gradio_tabs/merge.py @@ -105,6 +105,8 @@ def merge_style_usual( new_config = config_a.copy() new_config["data"]["num_styles"] = len(new_style2id) new_config["data"]["style2id"] = new_style2id + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} new_config["model_name"] = output_name save_config(new_config, output_name) @@ -160,6 +162,8 @@ def merge_style_add_diff( new_config = config_a.copy() new_config["data"]["num_styles"] = len(new_style2id) new_config["data"]["style2id"] = new_style2id + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} new_config["model_name"] = output_name save_config(new_config, output_name) @@ -219,6 +223,8 @@ def merge_style_weighted_sum( new_config = config_a.copy() new_config["data"]["num_styles"] = len(new_style2id) new_config["data"]["style2id"] = new_style2id + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} new_config["model_name"] = output_name save_config(new_config, output_name) @@ -268,6 +274,8 @@ def merge_style_add_null( new_config = config_a.copy() new_config["data"]["num_styles"] = len(new_style2id) new_config["data"]["style2id"] = new_style2id + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} new_config["model_name"] = output_name save_config(new_config, output_name) @@ -362,6 +370,8 @@ def merge_models_usual( new_config["model_name"] = output_name new_config["data"]["num_styles"] = 1 new_config["data"]["style2id"] = {DEFAULT_STYLE: 0} + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} save_config(new_config, output_name) neutral_vector_a = style_vectors_a[0] @@ -444,6 +454,8 @@ def merge_models_add_diff( new_config["model_name"] = output_name new_config["data"]["num_styles"] = 1 new_config["data"]["style2id"] = {DEFAULT_STYLE: 0} + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} with open(assets_root / output_name / "config.json", "w", encoding="utf-8") as f: json.dump(new_config, f, indent=2, ensure_ascii=False) @@ -519,6 +531,8 @@ def merge_models_weighted_sum( new_config["model_name"] = output_name new_config["data"]["num_styles"] = 1 new_config["data"]["style2id"] = {DEFAULT_STYLE: 0} + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} with open(assets_root / output_name / "config.json", "w", encoding="utf-8") as f: json.dump(new_config, f, indent=2, ensure_ascii=False) @@ -595,6 +609,8 @@ def merge_models_add_null( new_config["model_name"] = output_name new_config["data"]["num_styles"] = 1 new_config["data"]["style2id"] = {DEFAULT_STYLE: 0} + if new_config["data"]["n_speakers"] == 1: + new_config["data"]["spk2id"] = { output_name : 0} with open(assets_root / output_name / "config.json", "w", encoding="utf-8") as f: json.dump(new_config, f, indent=2, ensure_ascii=False) diff --git a/style_bert_vits2/tts_model.py b/style_bert_vits2/tts_model.py index 1e57167..2994575 100644 --- a/style_bert_vits2/tts_model.py +++ b/style_bert_vits2/tts_model.py @@ -108,8 +108,9 @@ class TTSModel: ) self.style_vector_inference: Optional[Any] = None - # net_g は PyTorch 推論時のみ遅延初期化される + # net_g / null_model_params は PyTorch 推論時のみ遅延初期化される self.net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra, None] = None + self.null_model_params: dict[int, dict[str, Union[float, str]]] = {} # onnx_session は ONNX 推論時のみ遅延初期化される self.onnx_session: Optional[onnxruntime.InferenceSession] = None @@ -135,6 +136,44 @@ class TTSModel: f'Model loaded successfully from {self.model_path} to "{self.device}" device ({time.time() - start_time:.2f}s)' ) + # ここからはヌルモデルのロード用パラメータが指定されている場合のみ + if len(self.null_model_params.keys()) == 0: + return + + # 推論対象のモデルの重みとヌルモデルの重みをマージ + for null_model_info in self.null_model_params.values(): + logger.info(f"Adding null model: {null_model_info['path']}...") + null_model_add = get_net_g( + model_path=str(null_model_info["path"]), + version=self.hyper_parameters.version, + device=self.device, + hps=self.hyper_parameters, + ) + # 愚直。もっと上手い方法ありそう + params = zip(self.net_g.dec.parameters(), null_model_add.dec.parameters()) + for v in params: + v[0].data.add_(v[1].data, alpha=float(null_model_info["weight"])) + params = zip( + self.net_g.flow.parameters(), null_model_add.flow.parameters() + ) + for v in params: + v[0].data.add_(v[1].data, alpha=float(null_model_info["pitch"])) + + params = zip( + self.net_g.enc_p.parameters(), null_model_add.enc_p.parameters() + ) + for v in params: + v[0].data.add_(v[1].data, alpha=float(null_model_info["style"])) + # テンポは sdp と dp 二つあるからとりあえずどっちも足す + params = zip(self.net_g.sdp.parameters(), null_model_add.sdp.parameters()) + for v in params: + v[0].data.add_(v[1].data, alpha=float(null_model_info["tempo"])) + params = zip(self.net_g.dp.parameters(), null_model_add.dp.parameters()) + for v in params: + v[0].data.add_(v[1].data, alpha=float(null_model_info["tempo"])) + + logger.info(f"Null models merged successfully ({time.time() - start_time:.2f}s)") + # ONNX 推論時 else: sess_options = onnxruntime.SessionOptions() @@ -289,6 +328,8 @@ class TTSModel: given_tone: Optional[list[int]] = None, pitch_scale: float = 1.0, intonation_scale: float = 1.0, + null_model_params: dict[int, dict[str, Union[str, float]]] = {}, + force_reload_model: bool = False, ) -> tuple[int, NDArray[Any]]: """ テキストから音声を合成する。 @@ -313,7 +354,8 @@ class TTSModel: given_tone (Optional[list[int]], optional): アクセントのトーンのリスト. Defaults to None. pitch_scale (float, optional): ピッチの高さ (1.0 から変更すると若干音質が低下する). Defaults to 1.0. intonation_scale (float, optional): 抑揚の平均からの変化幅 (1.0 から変更すると若干音質が低下する). Defaults to 1.0. - + null_model_params (dict[int, dict[str, Union[str, float]]], optional): 推論時に使用するヌルモデルの名前、適用割合の dict が入った dict 。ONNX 推論では無視される。 + force_reload_model (bool, optional): モデルを強制的に再ロードするかどうか. Defaults to False. Returns: tuple[int, NDArray[Any]]: サンプリングレートと音声データ (16bit PCM) """ @@ -344,6 +386,15 @@ class TTSModel: from style_bert_vits2.models.infer import infer + if null_model_params is not {}: + self.null_model_params = null_model_params + else: + self.null_model_params = {} + + # force_reload_model が True のとき、メモリ上に保持されているモデルを破棄する + if force_reload_model is True: + self.net_g = None + # モデルがロードされていない場合はロードする if self.net_g is None: self.load() @@ -402,6 +453,10 @@ class TTSModel: else: from style_bert_vits2.models.infer_onnx import infer_onnx + # force_reload_model が True のとき、メモリ上に保持されているモデルを破棄する + if force_reload_model is True: + self.onnx_session = None + # モデルがロードされていない場合はロードする if self.onnx_session is None: self.load() diff --git a/train_ms.py b/train_ms.py index 1adee4e..7244bb8 100644 --- a/train_ms.py +++ b/train_ms.py @@ -13,6 +13,7 @@ from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm +from transformers.trainer_pt_utils import DistributedLengthGroupedSampler # logging.getLogger("numba").setLevel(logging.WARNING) import default_style @@ -242,15 +243,23 @@ def run(): # prefetch_factor=6, ) else: + train_sampler = DistributedLengthGroupedSampler( + dataset=train_dataset, + batch_size=hps.train.batch_size, + num_replicas=n_gpus, + rank=rank, + lengths=train_dataset.lengths, + drop_last=True, + ) train_loader = DataLoader( train_dataset, # メモリ消費量を減らそうとnum_workersを1にしてみる # num_workers=min(config.train_ms_config.num_workers, os.cpu_count() // 2), num_workers=1, - shuffle=True, + # shuffle=True, pin_memory=True, collate_fn=collate_fn, - # batch_sampler=train_sampler, + sampler=train_sampler, batch_size=hps.train.batch_size, persistent_workers=True, # これもメモリ消費量を減らそうとしてコメントアウト diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py index e5c5bd1..8c33858 100644 --- a/train_ms_jp_extra.py +++ b/train_ms_jp_extra.py @@ -13,6 +13,7 @@ from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm +from transformers.trainer_pt_utils import DistributedLengthGroupedSampler # logging.getLogger("numba").setLevel(logging.WARNING) import default_style @@ -243,20 +244,32 @@ def run(): # prefetch_factor=6, ) else: + train_sampler = DistributedLengthGroupedSampler( + dataset=train_dataset, + batch_size=hps.train.batch_size, + num_replicas=n_gpus, + rank=rank, + lengths=train_dataset.lengths, + drop_last=True, + ) train_loader = DataLoader( train_dataset, # メモリ消費量を減らそうとnum_workersを1にしてみる # num_workers=min(config.train_ms_config.num_workers, os.cpu_count() // 2), num_workers=1, - shuffle=True, + # shuffle=True, pin_memory=True, collate_fn=collate_fn, - # batch_sampler=train_sampler, + sampler=train_sampler, batch_size=hps.train.batch_size, persistent_workers=True, # これもメモリ消費量を減らそうとしてコメントアウト # prefetch_factor=6, ) + logger.info("Using DistributedLengthGroupedSampler for training.") + logger.debug(f"len(train_dataset): {len(train_dataset)}") + logger.debug(f"len(train_loader): {len(train_loader)}") + eval_dataset = None eval_loader = None if rank == 0 and not args.speedup: