init

.gitignore

@@ -0,0 +1,44 @@
+__pycache__/
+venv/
+.venv/
+dist/
+.coverage
+.ipynb_checkpoints/
+.ruff_cache/
+
+/*.yml
+!/default_config.yml
+# /bert/*/*.bin
+# /bert/*/*.h5
+# /bert/*/*.model
+# /bert/*/*.safetensors
+# /bert/*/*.msgpack
+
+/configs/paths.yml
+
+/pretrained/*.safetensors
+/pretrained/*.pth
+
+/pretrained_jp_extra/*.safetensors
+/pretrained_jp_extra/*.pth
+
+/slm/*/*.bin
+
+/scripts/test/
+/scripts/lib/
+/scripts/Style-Bert-VITS2/
+/scripts/sbv2/
+*.zip
+*.csv
+*.bak
+/mos_results/
+
+safetensors.ipynb
+*.wav
+/static/
+
+# pyopenjtalk's dictionary
+*.dic
+
+playground.ipynb
+playgrounds/

app.py

@@ -0,0 +1,253 @@
+import datetime
+from pathlib import Path
+import gradio as gr
+import random
+from style_bert_vits2.constants import (
+    DEFAULT_LENGTH,
+    DEFAULT_LINE_SPLIT,
+    DEFAULT_NOISE,
+    DEFAULT_NOISEW,
+    DEFAULT_SPLIT_INTERVAL,
+)
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models.infer import InvalidToneError
+from style_bert_vits2.nlp.japanese import pyopenjtalk_worker as pyopenjtalk
+from style_bert_vits2.tts_model import TTSModelHolder
+
+
+pyopenjtalk.initialize_worker()
+
+example_file = "chupa_examples.txt"
+
+initial_text = (
+    "ちゅぱ、ちゅるる、ぢゅ、んく、れーれゅれろれろれろ、じゅぽぽぽぽぽ……ちゅううう！"
+)
+
+with open(example_file, "r", encoding="utf-8") as f:
+    examples = f.read().splitlines()
+
+
+def get_random_text() -> str:
+    return random.choice(examples)
+
+
+initial_md = """
+# チュパ音合成デモ
+
+2024-07-07: initial ver
+"""
+
+
+def make_interactive():
+    return gr.update(interactive=True, value="音声合成")
+
+
+def make_non_interactive():
+    return gr.update(interactive=False, value="音声合成（モデルをロードしてください）")
+
+
+def gr_util(item):
+    if item == "プリセットから選ぶ":
+        return (gr.update(visible=True), gr.Audio(visible=False, value=None))
+    else:
+        return (gr.update(visible=False), gr.update(visible=True))
+
+
+def create_inference_app(model_holder: TTSModelHolder) -> gr.Blocks:
+    def tts_fn(
+        model_name,
+        model_path,
+        text,
+        language,
+        sdp_ratio,
+        noise_scale,
+        noise_scale_w,
+        length_scale,
+        line_split,
+        split_interval,
+        speaker,
+    ):
+        model_holder.get_model(model_name, model_path)
+        assert model_holder.current_model is not None
+
+        speaker_id = model_holder.current_model.spk2id[speaker]
+
+        start_time = datetime.datetime.now()
+
+        try:
+            sr, audio = model_holder.current_model.infer(
+                text=text,
+                language=language,
+                sdp_ratio=sdp_ratio,
+                noise=noise_scale,
+                noise_w=noise_scale_w,
+                length=length_scale,
+                line_split=line_split,
+                split_interval=split_interval,
+                speaker_id=speaker_id,
+            )
+        except InvalidToneError as e:
+            logger.error(f"Tone error: {e}")
+            return f"Error: アクセント指定が不正です:\n{e}", None
+        except ValueError as e:
+            logger.error(f"Value error: {e}")
+            return f"Error: {e}", None
+
+        end_time = datetime.datetime.now()
+        duration = (end_time - start_time).total_seconds()
+
+        message = f"Success, time: {duration} seconds."
+        return message, (sr, audio)
+
+    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:
+        logger.error(
+            f"モデルが見つかりませんでした。{model_holder.root_dir}にモデルを置いてください。"
+        )
+        with gr.Blocks() as app:
+            gr.Markdown(
+                f"Error: モデルが見つかりませんでした。{model_holder.root_dir}にモデルを置いてください。"
+            )
+        return app
+
+    initial_pth_files = get_model_files(model_names[0])
+    model = model_holder.get_model(model_names[0], initial_pth_files[0])
+    speakers = list(model.spk2id.keys())
+
+    with gr.Blocks(theme="ParityError/Anime") as app:
+        gr.Markdown(initial_md)
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        model_name = gr.Dropdown(
+                            label="モデル一覧",
+                            choices=model_names,
+                            value=model_names[0],
+                        )
+                        model_path = gr.Dropdown(
+                            label="モデルファイル",
+                            choices=initial_pth_files,
+                            value=initial_pth_files[0],
+                        )
+                    refresh_button = gr.Button("更新", scale=1, visible=False)
+                    load_button = gr.Button("ロード", scale=1, variant="primary")
+                with gr.Row():
+                    text_input = gr.TextArea(
+                        label="テキスト", value=initial_text, scale=3
+                    )
+                    random_button = gr.Button("例から選ぶ 🎲", scale=1)
+                    random_button.click(get_random_text, outputs=[text_input])
+                with gr.Row():
+                    length_scale = gr.Slider(
+                        minimum=0.1,
+                        maximum=2,
+                        value=DEFAULT_LENGTH,
+                        step=0.1,
+                        label="生成音声の長さ（Length）",
+                    )
+                    sdp_ratio = gr.Slider(
+                        minimum=0,
+                        maximum=1,
+                        value=1,
+                        step=0.1,
+                        label="SDP Ratio",
+                    )
+                line_split = gr.Checkbox(
+                    label="改行で分けて生成（分けたほうが感情が乗ります）",
+                    value=DEFAULT_LINE_SPLIT,
+                    visible=False,
+                )
+                split_interval = gr.Slider(
+                    minimum=0.0,
+                    maximum=2,
+                    value=DEFAULT_SPLIT_INTERVAL,
+                    step=0.1,
+                    label="改行ごとに挟む無音の長さ（秒）",
+                )
+                line_split.change(
+                    lambda x: (gr.Slider(visible=x)),
+                    inputs=[line_split],
+                    outputs=[split_interval],
+                )
+                language = gr.Dropdown(
+                    choices=["JP"], value="JP", label="Language", visible=False
+                )
+                speaker = gr.Dropdown(label="話者", choices=speakers, value=speakers[0])
+                with gr.Accordion(label="詳細設定", open=True):
+                    noise_scale = gr.Slider(
+                        minimum=0.1,
+                        maximum=2,
+                        value=DEFAULT_NOISE,
+                        step=0.1,
+                        label="Noise",
+                    )
+                    noise_scale_w = gr.Slider(
+                        minimum=0.1,
+                        maximum=2,
+                        value=DEFAULT_NOISEW,
+                        step=0.1,
+                        label="Noise_W",
+                    )
+            with gr.Column():
+                tts_button = gr.Button("音声合成", variant="primary")
+                text_output = gr.Textbox(label="情報")
+                audio_output = gr.Audio(label="結果")
+
+        tts_button.click(
+            tts_fn,
+            inputs=[
+                model_name,
+                model_path,
+                text_input,
+                language,
+                sdp_ratio,
+                noise_scale,
+                noise_scale_w,
+                length_scale,
+                line_split,
+                split_interval,
+                speaker,
+            ],
+            outputs=[text_output, audio_output],
+        )
+
+        model_name.change(
+            model_holder.update_model_files_for_gradio,
+            inputs=[model_name],
+            outputs=[model_path],
+        )
+
+        model_path.change(make_non_interactive, outputs=[tts_button])
+
+        refresh_button.click(
+            model_holder.update_model_names_for_gradio,
+            outputs=[model_name, model_path, tts_button],
+        )
+        style = gr.Dropdown(label="スタイル", choices=[], visible=False)
+
+        load_button.click(
+            model_holder.get_model_for_gradio,
+            inputs=[model_name, model_path],
+            outputs=[style, tts_button, speaker],
+        )
+
+    return app
+
+
+if __name__ == "__main__":
+    import torch
+
+    from style_bert_vits2.constants import Languages
+    from style_bert_vits2.nlp import bert_models
+
+    bert_models.load_model(Languages.JP)
+    bert_models.load_tokenizer(Languages.JP)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model_holder = TTSModelHolder(Path("model_assets"), device)
+    app = create_inference_app(model_holder)
+    app.launch(inbrowser=True)

bert/bert_models.json

@@ -0,0 +1,14 @@
+{
+    "deberta-v2-large-japanese-char-wwm": {
+        "repo_id": "ku-nlp/deberta-v2-large-japanese-char-wwm",
+        "files": ["pytorch_model.bin"]
+    },
+    "chinese-roberta-wwm-ext-large": {
+        "repo_id": "hfl/chinese-roberta-wwm-ext-large",
+        "files": ["pytorch_model.bin"]
+    },
+    "deberta-v3-large": {
+        "repo_id": "microsoft/deberta-v3-large",
+        "files": ["spm.model", "pytorch_model.bin"]
+    }
+}

bert/chinese-roberta-wwm-ext-large/.gitattributes

@@ -0,0 +1,9 @@
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tar.gz filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text

bert/chinese-roberta-wwm-ext-large/README.md

@@ -0,0 +1,57 @@
+---
+language:
+- zh
+tags:
+- bert
+license: "apache-2.0"
+---
+
+# Please use 'Bert' related functions to load this model!
+
+## Chinese BERT with Whole Word Masking
+For further accelerating Chinese natural language processing, we provide **Chinese pre-trained BERT with Whole Word Masking**.
+
+**[Pre-Training with Whole Word Masking for Chinese BERT](https://arxiv.org/abs/1906.08101)**
+Yiming Cui, Wanxiang Che, Ting Liu, Bing Qin, Ziqing Yang, Shijin Wang, Guoping Hu
+
+This repository is developed based on：https://github.com/google-research/bert
+
+You may also interested in,
+- Chinese BERT series: https://github.com/ymcui/Chinese-BERT-wwm
+- Chinese MacBERT: https://github.com/ymcui/MacBERT
+- Chinese ELECTRA: https://github.com/ymcui/Chinese-ELECTRA
+- Chinese XLNet: https://github.com/ymcui/Chinese-XLNet
+- Knowledge Distillation Toolkit - TextBrewer: https://github.com/airaria/TextBrewer
+
+More resources by HFL: https://github.com/ymcui/HFL-Anthology
+
+## Citation
+If you find the technical report or resource is useful, please cite the following technical report in your paper.
+- Primary: https://arxiv.org/abs/2004.13922
+```
+@inproceedings{cui-etal-2020-revisiting,
+    title = "Revisiting Pre-Trained Models for {C}hinese Natural Language Processing",
+    author = "Cui, Yiming  and
+      Che, Wanxiang  and
+      Liu, Ting  and
+      Qin, Bing  and
+      Wang, Shijin  and
+      Hu, Guoping",
+    booktitle = "Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings",
+    month = nov,
+    year = "2020",
+    address = "Online",
+    publisher = "Association for Computational Linguistics",
+    url = "https://www.aclweb.org/anthology/2020.findings-emnlp.58",
+    pages = "657--668",
+}
+```
+- Secondary: https://arxiv.org/abs/1906.08101
+```
+@article{chinese-bert-wwm,
+  title={Pre-Training with Whole Word Masking for Chinese BERT},
+  author={Cui, Yiming and Che, Wanxiang and Liu, Ting and Qin, Bing and Yang, Ziqing and Wang, Shijin and Hu, Guoping},
+  journal={arXiv preprint arXiv:1906.08101},
+  year={2019}
+ }
+```

bert/chinese-roberta-wwm-ext-large/added_tokens.json

@@ -0,0 +1 @@
+{}

bert/chinese-roberta-wwm-ext-large/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "directionality": "bidi",
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "pooler_fc_size": 768,
+  "pooler_num_attention_heads": 12,
+  "pooler_num_fc_layers": 3,
+  "pooler_size_per_head": 128,
+  "pooler_type": "first_token_transform",
+  "type_vocab_size": 2,
+  "vocab_size": 21128
+}

bert/chinese-roberta-wwm-ext-large/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ac62d49144d770c5ca9a5d1d3039c4995665a080febe63198189857c6bd11cd
+size 1306484351

bert/chinese-roberta-wwm-ext-large/special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

bert/chinese-roberta-wwm-ext-large/tokenizer_config.json

@@ -0,0 +1 @@
+{"init_inputs": []}

bert/deberta-v2-large-japanese-char-wwm/.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

bert/deberta-v2-large-japanese-char-wwm/README.md

@@ -0,0 +1,89 @@
+---
+language: ja
+license: cc-by-sa-4.0
+library_name: transformers
+tags:
+  - deberta
+  - deberta-v2
+  - fill-mask
+  - character
+  - wwm
+datasets:
+  - wikipedia
+  - cc100
+  - oscar
+metrics:
+  - accuracy
+mask_token: "[MASK]"
+widget:
+    - text: "京都大学で自然言語処理を[MASK][MASK]する。"
+---
+
+# Model Card for Japanese character-level DeBERTa V2 large
+
+## Model description
+
+This is a Japanese DeBERTa V2 large model pre-trained on Japanese Wikipedia, the Japanese portion of CC-100, and the Japanese portion of OSCAR.
+This model is trained with character-level tokenization and whole word masking.
+
+## How to use
+
+You can use this model for masked language modeling as follows:
+
+```python
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+tokenizer = AutoTokenizer.from_pretrained('ku-nlp/deberta-v2-large-japanese-char-wwm')
+model = AutoModelForMaskedLM.from_pretrained('ku-nlp/deberta-v2-large-japanese-char-wwm')
+
+sentence = '京都大学で自然言語処理を[MASK][MASK]する。'
+encoding = tokenizer(sentence, return_tensors='pt')
+...
+```
+
+You can also fine-tune this model on downstream tasks.
+
+## Tokenization
+
+There is no need to tokenize texts in advance, and you can give raw texts to the tokenizer.
+The texts are tokenized into character-level tokens by [sentencepiece](https://github.com/google/sentencepiece).
+
+## Training data
+
+We used the following corpora for pre-training:
+
+- Japanese Wikipedia (as of 20221020, 3.2GB, 27M sentences, 1.3M documents)
+- Japanese portion of CC-100 (85GB, 619M sentences, 66M documents)
+- Japanese portion of OSCAR (54GB, 326M sentences, 25M documents)
+
+Note that we filtered out documents annotated with "header", "footer", or "noisy" tags in OSCAR.
+Also note that Japanese Wikipedia was duplicated 10 times to make the total size of the corpus comparable to that of CC-100 and OSCAR. As a result, the total size of the training data is 171GB.
+
+## Training procedure
+
+We first segmented texts in the corpora into words using [Juman++ 2.0.0-rc3](https://github.com/ku-nlp/jumanpp/releases/tag/v2.0.0-rc3) for whole word masking.
+Then, we built a sentencepiece model with 22,012 tokens including all characters that appear in the training corpus.
+
+We tokenized raw corpora into character-level subwords using the sentencepiece model and trained the Japanese DeBERTa model using [transformers](https://github.com/huggingface/transformers) library.
+The training took 26 days using 16 NVIDIA A100-SXM4-40GB GPUs.
+
+The following hyperparameters were used during pre-training:
+
+- learning_rate: 1e-4
+- per_device_train_batch_size: 26
+- distributed_type: multi-GPU
+- num_devices: 16
+- gradient_accumulation_steps: 8
+- total_train_batch_size: 3,328
+- max_seq_length: 512
+- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-06
+- lr_scheduler_type: linear schedule with warmup (lr = 0 at 300k steps)
+- training_steps: 260,000
+- warmup_steps: 10,000
+
+The accuracy of the trained model on the masked language modeling task was 0.795.
+The evaluation set consists of 5,000 randomly sampled documents from each of the training corpora.
+
+## Acknowledgments
+
+This work was supported by Joint Usage/Research Center for Interdisciplinary Large-scale Information Infrastructures (JHPCN) through General Collaboration Project no. jh221004, "Developing a Platform for Constructing and Sharing of Large-Scale Japanese Language Models".
+For training models, we used the mdx: a platform for the data-driven future.

bert/deberta-v2-large-japanese-char-wwm/config.json

@@ -0,0 +1,37 @@
+{
+  "architectures": [
+    "DebertaV2ForMaskedLM"
+  ],
+  "attention_head_size": 64,
+  "attention_probs_dropout_prob": 0.1,
+  "conv_act": "gelu",
+  "conv_kernel_size": 3,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-07,
+  "max_position_embeddings": 512,
+  "max_relative_positions": -1,
+  "model_type": "deberta-v2",
+  "norm_rel_ebd": "layer_norm",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "pad_token_id": 0,
+  "pooler_dropout": 0,
+  "pooler_hidden_act": "gelu",
+  "pooler_hidden_size": 1024,
+  "pos_att_type": [
+    "p2c",
+    "c2p"
+  ],
+  "position_biased_input": false,
+  "position_buckets": 256,
+  "relative_attention": true,
+  "share_att_key": true,
+  "torch_dtype": "float16",
+  "transformers_version": "4.25.1",
+  "type_vocab_size": 0,
+  "vocab_size": 22012
+}

bert/deberta-v2-large-japanese-char-wwm/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf0dab8ad87bd7c22e85ec71e04f2240804fda6d33196157d6b5923af6ea1201
+size 1318456639

bert/deberta-v2-large-japanese-char-wwm/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

bert/deberta-v2-large-japanese-char-wwm/tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "do_subword_tokenize": true,
+  "do_word_tokenize": true,
+  "jumanpp_kwargs": null,
+  "mask_token": "[MASK]",
+  "mecab_kwargs": null,
+  "model_max_length": 1000000000000000019884624838656,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "special_tokens_map_file": null,
+  "subword_tokenizer_type": "character",
+  "sudachi_kwargs": null,
+  "tokenizer_class": "BertJapaneseTokenizer",
+  "unk_token": "[UNK]",
+  "word_tokenizer_type": "basic"
+}

bert/deberta-v3-large/.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

bert/deberta-v3-large/README.md

@@ -0,0 +1,93 @@
+---
+language: en
+tags:
+  - deberta
+  - deberta-v3
+  - fill-mask
+thumbnail: https://huggingface.co/front/thumbnails/microsoft.png
+license: mit
+---
+
+## DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing
+
+[DeBERTa](https://arxiv.org/abs/2006.03654) improves the BERT and RoBERTa models using disentangled attention and enhanced mask decoder. With those two improvements, DeBERTa out perform RoBERTa on a majority of NLU tasks with 80GB training data.
+
+In [DeBERTa V3](https://arxiv.org/abs/2111.09543), we further improved the efficiency of DeBERTa using ELECTRA-Style pre-training with Gradient Disentangled Embedding Sharing. Compared to DeBERTa,  our V3 version significantly improves the model performance on downstream tasks.  You can find more technique details about the new model from our [paper](https://arxiv.org/abs/2111.09543).
+
+Please check the [official repository](https://github.com/microsoft/DeBERTa) for more implementation details and updates.
+
+The DeBERTa V3 large model comes with 24 layers and a hidden size of 1024. It has 304M backbone parameters  with a vocabulary containing 128K tokens which introduces 131M parameters in the Embedding layer.  This model was trained using the 160GB data as DeBERTa V2.
+
+
+#### Fine-tuning on NLU tasks
+
+We present the dev results on SQuAD 2.0 and MNLI tasks.
+
+| Model             |Vocabulary(K)|Backbone #Params(M)| SQuAD 2.0(F1/EM) | MNLI-m/mm(ACC)|
+|-------------------|----------|-------------------|-----------|----------|
+| RoBERTa-large     |50     |304                | 89.4/86.5 | 90.2   |
+| XLNet-large       |32     |-                  | 90.6/87.9 | 90.8   |
+| DeBERTa-large     |50     |-                  | 90.7/88.0 | 91.3   |
+| **DeBERTa-v3-large**|128|304                  |  **91.5/89.0**| **91.8/91.9**|
+
+
+#### Fine-tuning with HF transformers
+
+```bash
+#!/bin/bash
+
+cd transformers/examples/pytorch/text-classification/
+
+pip install datasets
+export TASK_NAME=mnli
+
+output_dir="ds_results"
+
+num_gpus=8
+
+batch_size=8
+
+python -m torch.distributed.launch --nproc_per_node=${num_gpus} \
+  run_glue.py \
+  --model_name_or_path microsoft/deberta-v3-large \
+  --task_name $TASK_NAME \
+  --do_train \
+  --do_eval \
+  --evaluation_strategy steps \
+  --max_seq_length 256 \
+  --warmup_steps 50 \
+  --per_device_train_batch_size ${batch_size} \
+  --learning_rate 6e-6 \
+  --num_train_epochs 2 \
+  --output_dir $output_dir \
+  --overwrite_output_dir \
+  --logging_steps 1000 \
+  --logging_dir $output_dir
+
+```
+
+### Citation
+
+If you find DeBERTa useful for your work, please cite the following papers:
+
+``` latex
+@misc{he2021debertav3,
+      title={DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing},
+      author={Pengcheng He and Jianfeng Gao and Weizhu Chen},
+      year={2021},
+      eprint={2111.09543},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```
+
+``` latex
+@inproceedings{
+he2021deberta,
+title={DEBERTA: DECODING-ENHANCED BERT WITH DISENTANGLED ATTENTION},
+author={Pengcheng He and Xiaodong Liu and Jianfeng Gao and Weizhu Chen},
+booktitle={International Conference on Learning Representations},
+year={2021},
+url={https://openreview.net/forum?id=XPZIaotutsD}
+}
+```

bert/deberta-v3-large/config.json

@@ -0,0 +1,22 @@
+{
+		"model_type": "deberta-v2",
+		"attention_probs_dropout_prob": 0.1,
+		"hidden_act": "gelu",
+		"hidden_dropout_prob": 0.1,
+		"hidden_size": 1024,
+		"initializer_range": 0.02,
+		"intermediate_size": 4096,
+		"max_position_embeddings": 512,
+		"relative_attention": true,
+		"position_buckets": 256,
+		"norm_rel_ebd": "layer_norm",
+		"share_att_key": true,
+		"pos_att_type": "p2c|c2p",
+		"layer_norm_eps": 1e-7,
+		"max_relative_positions": -1,
+		"position_biased_input": false,
+		"num_attention_heads": 16,
+		"num_hidden_layers": 24,
+		"type_vocab_size": 0,
+		"vocab_size": 128100
+}

bert/deberta-v3-large/generator_config.json

@@ -0,0 +1,22 @@
+{
+	"model_type": "deberta-v2",
+	"attention_probs_dropout_prob": 0.1,
+	"hidden_act": "gelu",
+	"hidden_dropout_prob": 0.1,
+	"hidden_size": 1024,
+	"initializer_range": 0.02,
+	"intermediate_size": 4096,
+	"max_position_embeddings": 512,
+	"relative_attention": true,
+	"position_buckets": 256,
+	"norm_rel_ebd": "layer_norm",
+	"share_att_key": true,
+	"pos_att_type": "p2c|c2p",
+	"layer_norm_eps": 1e-7,
+	"max_relative_positions": -1,
+	"position_biased_input": false,
+	"num_attention_heads": 16,
+	"num_hidden_layers": 12,
+	"type_vocab_size": 0,
+	"vocab_size": 128100
+}

bert/deberta-v3-large/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd5b5d93e2db101aaf281df0ea1216c07ad73620ff59c5b42dccac4bf2eef5b5
+size 873673253

bert/deberta-v3-large/pytorch_model.bin.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd5b5d93e2db101aaf281df0ea1216c07ad73620ff59c5b42dccac4bf2eef5b5
+size 873673253

bert/deberta-v3-large/spm.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c679fbf93643d19aab7ee10c0b99e460bdbc02fedf34b92b05af343b4af586fd
+size 2464616

bert/deberta-v3-large/tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "do_lower_case": false,
+  "vocab_type": "spm"
+}

model_assets/chupa_1/chupa_1spk_e1000_s194312.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8af08fae399f64bbc506a4accaf6c56b0d294def6435235dbe60755728784d8c
+size 251150980

model_assets/chupa_1/config.json

@@ -0,0 +1,87 @@
+{
+  "model_name": "chupa_1spk",
+  "train": {
+    "log_interval": 50,
+    "eval_interval": 1000,
+    "seed": 42,
+    "epochs": 1000,
+    "learning_rate": 0.0001,
+    "betas": [0.8, 0.99],
+    "eps": 1e-9,
+    "batch_size": 2,
+    "bf16_run": false,
+    "fp16_run": false,
+    "lr_decay": 0.99996,
+    "segment_size": 16384,
+    "init_lr_ratio": 1,
+    "warmup_epochs": 0,
+    "c_mel": 45,
+    "c_kl": 1.0,
+    "c_commit": 100,
+    "skip_optimizer": false,
+    "freeze_ZH_bert": false,
+    "freeze_JP_bert": false,
+    "freeze_EN_bert": false,
+    "freeze_emo": false,
+    "freeze_style": false,
+    "freeze_decoder": false
+  },
+  "data": {
+    "use_jp_extra": true,
+    "training_files": "Data/chupa_1/train.list",
+    "validation_files": "Data/chupa_1/val.list",
+    "max_wav_value": 32768.0,
+    "sampling_rate": 44100,
+    "filter_length": 2048,
+    "hop_length": 512,
+    "win_length": 2048,
+    "n_mel_channels": 128,
+    "mel_fmin": 0.0,
+    "mel_fmax": null,
+    "add_blank": true,
+    "n_speakers": 1,
+    "spk2id": {
+      "1": 0
+    },
+    "cleaned_text": true,
+    "num_styles": 1,
+    "style2id": {
+      "Neutral": 0
+    }
+  },
+  "model": {
+    "use_spk_conditioned_encoder": true,
+    "use_noise_scaled_mas": true,
+    "use_mel_posterior_encoder": false,
+    "use_duration_discriminator": false,
+    "use_wavlm_discriminator": true,
+    "inter_channels": 192,
+    "hidden_channels": 192,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 6,
+    "kernel_size": 3,
+    "p_dropout": 0.1,
+    "resblock": "1",
+    "resblock_kernel_sizes": [3, 7, 11],
+    "resblock_dilation_sizes": [
+      [1, 3, 5],
+      [1, 3, 5],
+      [1, 3, 5]
+    ],
+    "upsample_rates": [8, 8, 2, 2, 2],
+    "upsample_initial_channel": 512,
+    "upsample_kernel_sizes": [16, 16, 8, 2, 2],
+    "n_layers_q": 3,
+    "use_spectral_norm": false,
+    "gin_channels": 512,
+    "slm": {
+      "model": "./slm/wavlm-base-plus",
+      "sr": 16000,
+      "hidden": 768,
+      "nlayers": 13,
+      "initial_channel": 64
+    }
+  },
+  "version": "2.6.0-JP-Extra"
+}

model_assets/chupa_1/style_vectors.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fd42ba186b887c87b57fa66f5781f3fdf4382504d971d5338288d50b8b40461
+size 1152

requirements.txt

@@ -0,0 +1,23 @@
+cmudict
+cn2an
+# faster-whisper==0.10.1
+g2p_en
+GPUtil
+gradio
+jieba
+# librosa==0.9.2
+loguru
+num2words
+numpy<2
+# protobuf==4.25
+psutil
+# punctuators
+pyannote.audio>=3.1.0
+# pyloudnorm
+pyopenjtalk-dict
+pypinyin
+pyworld-prebuilt
+# stable_ts
+# tensorboard
+torch
+transformers

style_bert_vits2/.editorconfig

@@ -0,0 +1,15 @@
+root = true
+
+[*]
+charset = utf-8
+end_of_line = lf
+insert_final_newline = true
+indent_size = 4
+indent_style = space
+trim_trailing_whitespace = true
+
+[*.md]
+trim_trailing_whitespace = false
+
+[*.yml]
+indent_size = 2

style_bert_vits2/constants.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+
+from style_bert_vits2.utils.strenum import StrEnum
+
+
+# Style-Bert-VITS2 のバージョン
+VERSION = "2.6.0"
+
+# Style-Bert-VITS2 のベースディレクトリ
+BASE_DIR = Path(__file__).parent.parent
+
+
+# 利用可能な言語
+## JP-Extra モデル利用時は JP 以外の言語の音声合成はできない
+class Languages(StrEnum):
+    JP = "JP"
+    EN = "EN"
+    ZH = "ZH"
+
+
+# 言語ごとのデフォルトの BERT トークナイザーのパス
+DEFAULT_BERT_TOKENIZER_PATHS = {
+    Languages.JP: BASE_DIR / "bert" / "deberta-v2-large-japanese-char-wwm",
+    Languages.EN: BASE_DIR / "bert" / "deberta-v3-large",
+    Languages.ZH: BASE_DIR / "bert" / "chinese-roberta-wwm-ext-large",
+}
+
+# デフォルトのユーザー辞書ディレクトリ
+## style_bert_vits2.nlp.japanese.user_dict モジュールのデフォルト値として利用される
+## ライブラリとしての利用などで外部のユーザー辞書を指定したい場合は、user_dict 以下の各関数の実行時、引数に辞書データファイルのパスを指定する
+DEFAULT_USER_DICT_DIR = BASE_DIR / "dict_data"
+
+# デフォルトの推論パラメータ
+DEFAULT_STYLE = "Neutral"
+DEFAULT_STYLE_WEIGHT = 1.0
+DEFAULT_SDP_RATIO = 0.2
+DEFAULT_NOISE = 0.6
+DEFAULT_NOISEW = 0.8
+DEFAULT_LENGTH = 1.0
+DEFAULT_LINE_SPLIT = True
+DEFAULT_SPLIT_INTERVAL = 0.5
+DEFAULT_ASSIST_TEXT_WEIGHT = 0.7
+DEFAULT_ASSIST_TEXT_WEIGHT = 1.0
+
+# Gradio のテーマ
+## Built-in theme: "default", "base", "monochrome", "soft", "glass"
+## See https://huggingface.co/spaces/gradio/theme-gallery for more themes
+GRADIO_THEME = "NoCrypt/miku"

style_bert_vits2/logging.py

@@ -0,0 +1,15 @@
+from loguru import logger
+
+from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT
+
+
+# Remove all default handlers
+logger.remove()
+
+# Add a new handler
+logger.add(
+    SAFE_STDOUT,
+    format="<g>{time:MM-DD HH:mm:ss}</g> |<lvl>{level:^8}</lvl>| {file}:{line} | {message}",
+    backtrace=True,
+    diagnose=True,
+)

style_bert_vits2/models/attentions.py

@@ -0,0 +1,491 @@
+import math
+from typing import Any, Optional
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from style_bert_vits2.models import commons
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+@torch.jit.script  # type: ignore
+def fused_add_tanh_sigmoid_multiply(
+    input_a: torch.Tensor, input_b: torch.Tensor, n_channels: list[int]
+) -> torch.Tensor:
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int = 1,
+        p_dropout: float = 0.0,
+        window_size: int = 4,
+        isflow: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        # if isflow:
+        #  cond_layer = torch.nn.Conv1d(256, 2*hidden_channels*n_layers, 1)
+        #  self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1)
+        #  self.cond_layer = weight_norm(cond_layer, name='weight')
+        #  self.gin_channels = 256
+        self.cond_layer_idx = self.n_layers
+        if "gin_channels" in kwargs:
+            self.gin_channels = kwargs["gin_channels"]
+            if self.gin_channels != 0:
+                self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
+                # vits2 says 3rd block, so idx is 2 by default
+                self.cond_layer_idx = (
+                    kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2
+                )
+                # logger.debug(self.gin_channels, self.cond_layer_idx)
+                assert (
+                    self.cond_layer_idx < self.n_layers
+                ), "cond_layer_idx should be less than n_layers"
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(
+        self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            if i == self.cond_layer_idx and g is not None:
+                g = self.spk_emb_linear(g.transpose(1, 2))
+                assert g is not None
+                g = g.transpose(1, 2)
+                x = x + g
+                x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int = 1,
+        p_dropout: float = 0.0,
+        proximal_bias: bool = False,
+        proximal_init: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        h: torch.Tensor,
+        h_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        x: decoder input
+        h: encoder output
+        """
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: int,
+        n_heads: int,
+        p_dropout: float = 0.0,
+        window_size: Optional[int] = None,
+        heads_share: bool = True,
+        block_length: Optional[int] = None,
+        proximal_bias: bool = False,
+        proximal_init: bool = False,
+    ) -> None:
+        super().__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                assert self.conv_k.bias is not None
+                assert self.conv_q.bias is not None
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def forward(
+        self, x: torch.Tensor, c: torch.Tensor, attn_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        return output, p_attn
+
+    def _matmul_with_relative_values(
+        self, x: torch.Tensor, y: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        x: [b, h, l, m]
+        y: [h or 1, m, d]
+        ret: [b, h, l, d]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(
+        self, x: torch.Tensor, y: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        x: [b, h, l, d]
+        y: [h or 1, m, d]
+        ret: [b, h, l, m]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
+
+    def _get_relative_embeddings(
+        self, relative_embeddings: torch.Tensor, length: int
+    ) -> torch.Tensor:
+        assert self.window_size is not None
+        2 * self.window_size + 1  # type: ignore
+        # Pad first before slice to avoid using cond ops.
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: [b, h, l, 2*l-1]
+        ret: [b, h, l, l]
+        """
+        batch, heads, length, _ = x.size()
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(
+            x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: [b, h, l, l]
+        ret: [b, h, l, 2*l-1]
+        """
+        batch, heads, length, _ = x.size()
+        # pad along column
+        x = F.pad(
+            x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
+        )
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length: int) -> torch.Tensor:
+        """Bias for self-attention to encourage attention to close positions.
+        Args:
+          length: an integer scalar.
+        Returns:
+          a Tensor with shape [1, 1, length, length]
+        """
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float = 0.0,
+        activation: Optional[str] = None,
+        causal: bool = False,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x: torch.Tensor) -> torch.Tensor:
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x: torch.Tensor) -> torch.Tensor:
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x

style_bert_vits2/models/commons.py

@@ -0,0 +1,223 @@
+"""
+以下に記述されている関数のコメントはリファクタリング時に GPT-4 に生成させたもので、
+コードと完全に一致している保証はない。あくまで参考程度とすること。
+"""
+
+from typing import Any, Optional, Union
+
+import torch
+from torch.nn import functional as F
+
+
+def init_weights(m: torch.nn.Module, mean: float = 0.0, std: float = 0.01) -> None:
+    """
+    モジュールの重みを初期化する
+
+    Args:
+        m (torch.nn.Module): 重みを初期化する対象のモジュール
+        mean (float): 正規分布の平均
+        std (float): 正規分布の標準偏差
+    """
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    """
+    カーネルサイズと膨張率からパディングの大きさを計算する
+
+    Args:
+        kernel_size (int): カーネルのサイズ
+        dilation (int): 膨張率
+
+    Returns:
+        int: 計算されたパディングの大きさ
+    """
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape: list[list[Any]]) -> list[Any]:
+    """
+    パディングの形状を変換する
+
+    Args:
+        pad_shape (list[list[Any]]): 変換前のパディングの形状
+
+    Returns:
+        list[Any]: 変換後のパディングの形状
+    """
+    layer = pad_shape[::-1]
+    new_pad_shape = [item for sublist in layer for item in sublist]
+    return new_pad_shape
+
+
+def intersperse(lst: list[Any], item: Any) -> list[Any]:
+    """
+    リストの要素の間に特定のアイテムを挿入する
+
+    Args:
+        lst (list[Any]): 元のリスト
+        item (Any): 挿入するアイテム
+
+    Returns:
+        list[Any]: 新しいリスト
+    """
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+
+def slice_segments(
+    x: torch.Tensor, ids_str: torch.Tensor, segment_size: int = 4
+) -> torch.Tensor:
+    """
+    テンソルからセグメントをスライスする
+
+    Args:
+        x (torch.Tensor): 入力テンソル
+        ids_str (torch.Tensor): スライスを開始するインデックス
+        segment_size (int, optional): スライスのサイズ (デフォルト: 4)
+
+    Returns:
+        torch.Tensor: スライスされたセグメント
+    """
+    gather_indices = ids_str.view(x.size(0), 1, 1).repeat(
+        1, x.size(1), 1
+    ) + torch.arange(segment_size, device=x.device)
+    return torch.gather(x, 2, gather_indices)
+
+
+def rand_slice_segments(
+    x: torch.Tensor, x_lengths: Optional[torch.Tensor] = None, segment_size: int = 4
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    ランダムなセグメントをスライスする
+
+    Args:
+        x (torch.Tensor): 入力テンソル
+        x_lengths (Optional[torch.Tensor], optional): 各バッチの長さ (デフォルト: None)
+        segment_size (int, optional): スライスのサイズ (デフォルト: 4)
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]: スライスされたセグメントと開始インデックス
+    """
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t  # type: ignore
+    ids_str_max = torch.clamp(x_lengths - segment_size + 1, min=0)  # type: ignore
+    ids_str = (torch.rand([b], device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def subsequent_mask(length: int) -> torch.Tensor:
+    """
+    後続のマスクを生成する
+
+    Args:
+        length (int): マスクのサイズ
+
+    Returns:
+        torch.Tensor: 生成されたマスク
+    """
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script  # type: ignore
+def fused_add_tanh_sigmoid_multiply(
+    input_a: torch.Tensor, input_b: torch.Tensor, n_channels: torch.Tensor
+) -> torch.Tensor:
+    """
+    加算、tanh、sigmoid の活性化関数を組み合わせた演算を行う
+
+    Args:
+        input_a (torch.Tensor): 入力テンソル A
+        input_b (torch.Tensor): 入力テンソル B
+        n_channels (torch.Tensor): チャネル数
+
+    Returns:
+        torch.Tensor: 演算結果
+    """
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def sequence_mask(
+    length: torch.Tensor, max_length: Optional[int] = None
+) -> torch.Tensor:
+    """
+    シーケンスマスクを生成する
+
+    Args:
+        length (torch.Tensor): 各シーケンスの長さ
+        max_length (Optional[int]): 最大のシーケンス長さ。指定されていない場合は length の最大値を使用
+
+    Returns:
+        torch.Tensor: 生成されたシーケンスマスク
+    """
+    if max_length is None:
+        max_length = length.max()  # type: ignore
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)  # type: ignore
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """
+    パスを生成する
+
+    Args:
+        duration (torch.Tensor): 各時間ステップの持続時間
+        mask (torch.Tensor): マスクテンソル
+
+    Returns:
+        torch.Tensor: 生成されたパス
+    """
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(
+    parameters: Union[torch.Tensor, list[torch.Tensor]],
+    clip_value: Optional[float],
+    norm_type: float = 2.0,
+) -> float:
+    """
+    勾配の値をクリップする
+
+    Args:
+        parameters (Union[torch.Tensor, list[torch.Tensor]]): クリップするパラメータ
+        clip_value (Optional[float]): クリップする値。None の場合はクリップしない
+        norm_type (float): ノルムの種類
+
+    Returns:
+        float: 総ノルム
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0.0
+    for p in parameters:
+        assert p.grad is not None
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm

style_bert_vits2/models/hyper_parameters.py

@@ -0,0 +1,129 @@
+"""
+Style-Bert-VITS2 モデルのハイパーパラメータを表す Pydantic モデル。
+デフォルト値は configs/config_jp_extra.json 内の定義と概ね同一で、
+万が一ロードした config.json に存在しないキーがあった際のフェイルセーフとして適用される。
+"""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from pydantic import BaseModel, ConfigDict
+
+
+class HyperParametersTrain(BaseModel):
+    log_interval: int = 200
+    eval_interval: int = 1000
+    seed: int = 42
+    epochs: int = 1000
+    learning_rate: float = 0.0001
+    betas: tuple[float, float] = (0.8, 0.99)
+    eps: float = 1e-9
+    batch_size: int = 2
+    bf16_run: bool = False
+    fp16_run: bool = False
+    lr_decay: float = 0.99996
+    segment_size: int = 16384
+    init_lr_ratio: int = 1
+    warmup_epochs: int = 0
+    c_mel: int = 45
+    c_kl: float = 1.0
+    c_commit: int = 100
+    skip_optimizer: bool = False
+    freeze_ZH_bert: bool = False
+    freeze_JP_bert: bool = False
+    freeze_EN_bert: bool = False
+    freeze_emo: bool = False
+    freeze_style: bool = False
+    freeze_decoder: bool = False
+
+
+class HyperParametersData(BaseModel):
+    use_jp_extra: bool = True
+    training_files: str = "Data/Dummy/train.list"
+    validation_files: str = "Data/Dummy/val.list"
+    max_wav_value: float = 32768.0
+    sampling_rate: int = 44100
+    filter_length: int = 2048
+    hop_length: int = 512
+    win_length: int = 2048
+    n_mel_channels: int = 128
+    mel_fmin: float = 0.0
+    mel_fmax: Optional[float] = None
+    add_blank: bool = True
+    n_speakers: int = 1
+    cleaned_text: bool = True
+    spk2id: dict[str, int] = {
+        "Dummy": 0,
+    }
+    num_styles: int = 1
+    style2id: dict[str, int] = {
+        "Neutral": 0,
+    }
+
+
+class HyperParametersModelSLM(BaseModel):
+    model: str = "./slm/wavlm-base-plus"
+    sr: int = 16000
+    hidden: int = 768
+    nlayers: int = 13
+    initial_channel: int = 64
+
+
+class HyperParametersModel(BaseModel):
+    use_spk_conditioned_encoder: bool = True
+    use_noise_scaled_mas: bool = True
+    use_mel_posterior_encoder: bool = False
+    use_duration_discriminator: bool = False
+    use_wavlm_discriminator: bool = True
+    inter_channels: int = 192
+    hidden_channels: int = 192
+    filter_channels: int = 768
+    n_heads: int = 2
+    n_layers: int = 6
+    kernel_size: int = 3
+    p_dropout: float = 0.1
+    resblock: str = "1"
+    resblock_kernel_sizes: list[int] = [3, 7, 11]
+    resblock_dilation_sizes: list[list[int]] = [
+        [1, 3, 5],
+        [1, 3, 5],
+        [1, 3, 5],
+    ]
+    upsample_rates: list[int] = [8, 8, 2, 2, 2]
+    upsample_initial_channel: int = 512
+    upsample_kernel_sizes: list[int] = [16, 16, 8, 2, 2]
+    n_layers_q: int = 3
+    use_spectral_norm: bool = False
+    gin_channels: int = 512
+    slm: HyperParametersModelSLM = HyperParametersModelSLM()
+
+
+class HyperParameters(BaseModel):
+    model_name: str = "Dummy"
+    version: str = "2.0-JP-Extra"
+    train: HyperParametersTrain = HyperParametersTrain()
+    data: HyperParametersData = HyperParametersData()
+    model: HyperParametersModel = HyperParametersModel()
+
+    # 以下は学習時にのみ動的に設定されるパラメータ (通常 config.json には存在しない)
+    model_dir: Optional[str] = None
+    speedup: bool = False
+    repo_id: Optional[str] = None
+
+    # model_ 以下を Pydantic の保護対象から除外する
+    model_config = ConfigDict(protected_namespaces=())
+
+    @staticmethod
+    def load_from_json(json_path: Union[str, Path]) -> "HyperParameters":
+        """
+        与えられた JSON ファイルからハイパーパラメータを読み込む。
+
+        Args:
+            json_path (Union[str, Path]): JSON ファイルのパス
+
+        Returns:
+            HyperParameters: ハイパーパラメータ
+        """
+
+        with open(json_path, encoding="utf-8") as f:
+            return HyperParameters.model_validate_json(f.read())

style_bert_vits2/models/infer.py

@@ -0,0 +1,308 @@
+from typing import Any, Optional, Union, cast
+
+import torch
+from numpy.typing import NDArray
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models import commons, utils
+from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.models.models import SynthesizerTrn
+from style_bert_vits2.models.models_jp_extra import (
+    SynthesizerTrn as SynthesizerTrnJPExtra,
+)
+from style_bert_vits2.nlp import (
+    clean_text,
+    cleaned_text_to_sequence,
+    extract_bert_feature,
+)
+from style_bert_vits2.nlp.symbols import SYMBOLS
+
+
+def get_net_g(model_path: str, version: str, device: str, hps: HyperParameters):
+    if version.endswith("JP-Extra"):
+        logger.info("Using JP-Extra model")
+        net_g = SynthesizerTrnJPExtra(
+            n_vocab=len(SYMBOLS),
+            spec_channels=hps.data.filter_length // 2 + 1,
+            segment_size=hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            # hps.model 以下のすべての値を引数に渡す
+            use_spk_conditioned_encoder=hps.model.use_spk_conditioned_encoder,
+            use_noise_scaled_mas=hps.model.use_noise_scaled_mas,
+            use_mel_posterior_encoder=hps.model.use_mel_posterior_encoder,
+            use_duration_discriminator=hps.model.use_duration_discriminator,
+            use_wavlm_discriminator=hps.model.use_wavlm_discriminator,
+            inter_channels=hps.model.inter_channels,
+            hidden_channels=hps.model.hidden_channels,
+            filter_channels=hps.model.filter_channels,
+            n_heads=hps.model.n_heads,
+            n_layers=hps.model.n_layers,
+            kernel_size=hps.model.kernel_size,
+            p_dropout=hps.model.p_dropout,
+            resblock=hps.model.resblock,
+            resblock_kernel_sizes=hps.model.resblock_kernel_sizes,
+            resblock_dilation_sizes=hps.model.resblock_dilation_sizes,
+            upsample_rates=hps.model.upsample_rates,
+            upsample_initial_channel=hps.model.upsample_initial_channel,
+            upsample_kernel_sizes=hps.model.upsample_kernel_sizes,
+            n_layers_q=hps.model.n_layers_q,
+            use_spectral_norm=hps.model.use_spectral_norm,
+            gin_channels=hps.model.gin_channels,
+            slm=hps.model.slm,
+        ).to(device)
+    else:
+        logger.info("Using normal model")
+        net_g = SynthesizerTrn(
+            n_vocab=len(SYMBOLS),
+            spec_channels=hps.data.filter_length // 2 + 1,
+            segment_size=hps.train.segment_size // hps.data.hop_length,
+            n_speakers=hps.data.n_speakers,
+            # hps.model 以下のすべての値を引数に渡す
+            use_spk_conditioned_encoder=hps.model.use_spk_conditioned_encoder,
+            use_noise_scaled_mas=hps.model.use_noise_scaled_mas,
+            use_mel_posterior_encoder=hps.model.use_mel_posterior_encoder,
+            use_duration_discriminator=hps.model.use_duration_discriminator,
+            use_wavlm_discriminator=hps.model.use_wavlm_discriminator,
+            inter_channels=hps.model.inter_channels,
+            hidden_channels=hps.model.hidden_channels,
+            filter_channels=hps.model.filter_channels,
+            n_heads=hps.model.n_heads,
+            n_layers=hps.model.n_layers,
+            kernel_size=hps.model.kernel_size,
+            p_dropout=hps.model.p_dropout,
+            resblock=hps.model.resblock,
+            resblock_kernel_sizes=hps.model.resblock_kernel_sizes,
+            resblock_dilation_sizes=hps.model.resblock_dilation_sizes,
+            upsample_rates=hps.model.upsample_rates,
+            upsample_initial_channel=hps.model.upsample_initial_channel,
+            upsample_kernel_sizes=hps.model.upsample_kernel_sizes,
+            n_layers_q=hps.model.n_layers_q,
+            use_spectral_norm=hps.model.use_spectral_norm,
+            gin_channels=hps.model.gin_channels,
+            slm=hps.model.slm,
+        ).to(device)
+    net_g.state_dict()
+    _ = net_g.eval()
+    if model_path.endswith(".pth") or model_path.endswith(".pt"):
+        _ = utils.checkpoints.load_checkpoint(
+            model_path, net_g, None, skip_optimizer=True
+        )
+    elif model_path.endswith(".safetensors"):
+        _ = utils.safetensors.load_safetensors(model_path, net_g, True)
+    else:
+        raise ValueError(f"Unknown model format: {model_path}")
+    return net_g
+
+
+def get_text(
+    text: str,
+    language_str: Languages,
+    hps: HyperParameters,
+    device: str,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+    given_phone: Optional[list[str]] = None,
+    given_tone: Optional[list[int]] = None,
+):
+    use_jp_extra = hps.version.endswith("JP-Extra")
+    # 推論時のみ呼び出されるので、raise_yomi_error は False に設定
+    norm_text, phone, tone, word2ph = clean_text(
+        text,
+        language_str,
+        use_jp_extra=use_jp_extra,
+        raise_yomi_error=False,
+    )
+    # phone と tone の両方が与えられた場合はそれを使う
+    if given_phone is not None and given_tone is not None:
+        # 指定された phone と指定された tone 両方の長さが一致していなければならない
+        if len(given_phone) != len(given_tone):
+            raise InvalidPhoneError(
+                f"Length of given_phone ({len(given_phone)}) != length of given_tone ({len(given_tone)})"
+            )
+        # 与えられた音素数と pyopenjtalk で生成した読みの音素数が一致しない
+        if len(given_phone) != sum(word2ph):
+            # 日本語の場合、len(given_phone) と sum(word2ph) が一致するように word2ph を適切に調整する
+            # 他の言語は word2ph の調整方法が思いつかないのでエラー
+            if language_str == Languages.JP:
+                from style_bert_vits2.nlp.japanese.g2p import adjust_word2ph
+
+                word2ph = adjust_word2ph(word2ph, phone, given_phone)
+                # 上記処理により word2ph の合計が given_phone の長さと一致するはず
+                # それでも一致しない場合、大半は読み上げテキストと given_phone が著しく乖離していて調整し切れなかったことを意味する
+                if len(given_phone) != sum(word2ph):
+                    raise InvalidPhoneError(
+                        f"Length of given_phone ({len(given_phone)}) != sum of word2ph ({sum(word2ph)})"
+                    )
+            else:
+                raise InvalidPhoneError(
+                    f"Length of given_phone ({len(given_phone)}) != sum of word2ph ({sum(word2ph)})"
+                )
+        phone = given_phone
+        # 生成あるいは指定された phone と指定された tone 両方の長さが一致していなければならない
+        if len(phone) != len(given_tone):
+            raise InvalidToneError(
+                f"Length of phone ({len(phone)}) != length of given_tone ({len(given_tone)})"
+            )
+        tone = given_tone
+    # tone だけが与えられた場合は clean_text() で生成した phone と合わせて使う
+    elif given_tone is not None:
+        # 生成した phone と指定された tone 両方の長さが一致していなければならない
+        if len(phone) != len(given_tone):
+            raise InvalidToneError(
+                f"Length of phone ({len(phone)}) != length of given_tone ({len(given_tone)})"
+            )
+        tone = given_tone
+    phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)
+
+    if hps.data.add_blank:
+        phone = commons.intersperse(phone, 0)
+        tone = commons.intersperse(tone, 0)
+        language = commons.intersperse(language, 0)
+        for i in range(len(word2ph)):
+            word2ph[i] = word2ph[i] * 2
+        word2ph[0] += 1
+    bert_ori = extract_bert_feature(
+        norm_text,
+        word2ph,
+        language_str,
+        device,
+        assist_text,
+        assist_text_weight,
+    )
+    del word2ph
+    assert bert_ori.shape[-1] == len(phone), phone
+
+    if language_str == Languages.ZH:
+        bert = bert_ori
+        ja_bert = torch.zeros(1024, len(phone))
+        en_bert = torch.zeros(1024, len(phone))
+    elif language_str == Languages.JP:
+        bert = torch.zeros(1024, len(phone))
+        ja_bert = bert_ori
+        en_bert = torch.zeros(1024, len(phone))
+    elif language_str == Languages.EN:
+        bert = torch.zeros(1024, len(phone))
+        ja_bert = torch.zeros(1024, len(phone))
+        en_bert = bert_ori
+    else:
+        raise ValueError("language_str should be ZH, JP or EN")
+
+    assert bert.shape[-1] == len(
+        phone
+    ), f"Bert seq len {bert.shape[-1]} != {len(phone)}"
+
+    phone = torch.LongTensor(phone)
+    tone = torch.LongTensor(tone)
+    language = torch.LongTensor(language)
+    return bert, ja_bert, en_bert, phone, tone, language
+
+
+def infer(
+    text: str,
+    style_vec: NDArray[Any],
+    sdp_ratio: float,
+    noise_scale: float,
+    noise_scale_w: float,
+    length_scale: float,
+    sid: int,  # In the original Bert-VITS2, its speaker_name: str, but here it's id
+    language: Languages,
+    hps: HyperParameters,
+    net_g: Union[SynthesizerTrn, SynthesizerTrnJPExtra],
+    device: str,
+    skip_start: bool = False,
+    skip_end: bool = False,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+    given_phone: Optional[list[str]] = None,
+    given_tone: Optional[list[int]] = None,
+):
+    is_jp_extra = hps.version.endswith("JP-Extra")
+    bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
+        text,
+        language,
+        hps,
+        device,
+        assist_text=assist_text,
+        assist_text_weight=assist_text_weight,
+        given_phone=given_phone,
+        given_tone=given_tone,
+    )
+    if skip_start:
+        phones = phones[3:]
+        tones = tones[3:]
+        lang_ids = lang_ids[3:]
+        bert = bert[:, 3:]
+        ja_bert = ja_bert[:, 3:]
+        en_bert = en_bert[:, 3:]
+    if skip_end:
+        phones = phones[:-2]
+        tones = tones[:-2]
+        lang_ids = lang_ids[:-2]
+        bert = bert[:, :-2]
+        ja_bert = ja_bert[:, :-2]
+        en_bert = en_bert[:, :-2]
+    with torch.no_grad():
+        x_tst = phones.to(device).unsqueeze(0)
+        tones = tones.to(device).unsqueeze(0)
+        lang_ids = lang_ids.to(device).unsqueeze(0)
+        bert = bert.to(device).unsqueeze(0)
+        ja_bert = ja_bert.to(device).unsqueeze(0)
+        en_bert = en_bert.to(device).unsqueeze(0)
+        x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
+        style_vec_tensor = torch.from_numpy(style_vec).to(device).unsqueeze(0)
+        del phones
+        sid_tensor = torch.LongTensor([sid]).to(device)
+        if is_jp_extra:
+            output = cast(SynthesizerTrnJPExtra, net_g).infer(
+                x_tst,
+                x_tst_lengths,
+                sid_tensor,
+                tones,
+                lang_ids,
+                ja_bert,
+                style_vec=style_vec_tensor,
+                sdp_ratio=sdp_ratio,
+                noise_scale=noise_scale,
+                noise_scale_w=noise_scale_w,
+                length_scale=length_scale,
+            )
+        else:
+            output = cast(SynthesizerTrn, net_g).infer(
+                x_tst,
+                x_tst_lengths,
+                sid_tensor,
+                tones,
+                lang_ids,
+                bert,
+                ja_bert,
+                en_bert,
+                style_vec=style_vec_tensor,
+                sdp_ratio=sdp_ratio,
+                noise_scale=noise_scale,
+                noise_scale_w=noise_scale_w,
+                length_scale=length_scale,
+            )
+        audio = output[0][0, 0].data.cpu().float().numpy()
+        del (
+            x_tst,
+            tones,
+            lang_ids,
+            bert,
+            x_tst_lengths,
+            sid_tensor,
+            ja_bert,
+            en_bert,
+            style_vec,
+        )  # , emo
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return audio
+
+
+class InvalidPhoneError(ValueError):
+    pass
+
+
+class InvalidToneError(ValueError):
+    pass

style_bert_vits2/models/models.py

@@ -0,0 +1,1102 @@
+import math
+from typing import Any, Optional
+
+import torch
+from torch import nn
+from torch.nn import Conv1d, Conv2d, ConvTranspose1d
+from torch.nn import functional as F
+from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
+
+from style_bert_vits2.models import attentions, commons, modules, monotonic_alignment
+from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
+
+
+class DurationDiscriminator(nn.Module):  # vits2
+    def __init__(
+        self,
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.LayerNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.LayerNorm(filter_channels)
+        self.dur_proj = nn.Conv1d(1, filter_channels, 1)
+
+        self.pre_out_conv_1 = nn.Conv1d(
+            2 * filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.pre_out_norm_1 = modules.LayerNorm(filter_channels)
+        self.pre_out_conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.pre_out_norm_2 = modules.LayerNorm(filter_channels)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+        self.output_layer = nn.Sequential(nn.Linear(filter_channels, 1), nn.Sigmoid())
+
+    def forward_probability(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        dur: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        dur = self.dur_proj(dur)
+        x = torch.cat([x, dur], dim=1)
+        x = self.pre_out_conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.pre_out_norm_1(x)
+        x = self.drop(x)
+        x = self.pre_out_conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.pre_out_norm_2(x)
+        x = self.drop(x)
+        x = x * x_mask
+        x = x.transpose(1, 2)
+        output_prob = self.output_layer(x)
+        return output_prob
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        dur_r: torch.Tensor,
+        dur_hat: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> list[torch.Tensor]:
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+
+        output_probs = []
+        for dur in [dur_r, dur_hat]:
+            output_prob = self.forward_probability(x, x_mask, dur, g)
+            output_probs.append(output_prob)
+
+        return output_probs
+
+
+class TransformerCouplingBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+        share_parameter: bool = False,
+    ) -> None:
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+
+        self.wn = (
+            # attentions.FFT(
+            #     hidden_channels,
+            #     filter_channels,
+            #     n_heads,
+            #     n_layers,
+            #     kernel_size,
+            #     p_dropout,
+            #     isflow=True,
+            #     gin_channels=self.gin_channels,
+            # )
+            None
+            if share_parameter
+            else None
+        )
+
+        for i in range(n_flows):
+            self.flows.append(
+                modules.TransformerCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    n_layers,
+                    n_heads,
+                    p_dropout,
+                    filter_channels,
+                    mean_only=True,
+                    wn_sharing_parameter=self.wn,
+                    gin_channels=self.gin_channels,
+                )
+            )
+            self.flows.append(modules.Flip())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class StochasticDurationPredictor(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        filter_channels = in_channels  # it needs to be removed from future version.
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.log_flow = modules.Log()
+        self.flows = nn.ModuleList()
+        self.flows.append(modules.ElementwiseAffine(2))
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.flows.append(modules.Flip())
+
+        self.post_pre = nn.Conv1d(1, filter_channels, 1)
+        self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.post_convs = modules.DDSConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(modules.ElementwiseAffine(2))
+        for i in range(4):
+            self.post_flows.append(
+                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.post_flows.append(modules.Flip())
+
+        self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+        self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.convs = modules.DDSConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        w: Optional[torch.Tensor] = None,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+        noise_scale: float = 1.0,
+    ) -> torch.Tensor:
+        x = torch.detach(x)
+        x = self.pre(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.convs(x, x_mask)
+        x = self.proj(x) * x_mask
+
+        if not reverse:
+            flows = self.flows
+            assert w is not None
+
+            logdet_tot_q = 0
+            h_w = self.post_pre(w)
+            h_w = self.post_convs(h_w, x_mask)
+            h_w = self.post_proj(h_w) * x_mask
+            e_q = (
+                torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype)
+                * x_mask
+            )
+            z_q = e_q
+            for flow in self.post_flows:
+                z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+                logdet_tot_q += logdet_q
+            z_u, z1 = torch.split(z_q, [1, 1], 1)
+            u = torch.sigmoid(z_u) * x_mask
+            z0 = (w - u) * x_mask
+            logdet_tot_q += torch.sum(
+                (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2]
+            )
+            logq = (
+                torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2])
+                - logdet_tot_q
+            )
+
+            logdet_tot = 0
+            z0, logdet = self.log_flow(z0, x_mask)
+            logdet_tot += logdet
+            z = torch.cat([z0, z1], 1)
+            for flow in flows:
+                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+                logdet_tot = logdet_tot + logdet
+            nll = (
+                torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2])
+                - logdet_tot
+            )
+            return nll + logq  # [b]
+        else:
+            flows = list(reversed(self.flows))
+            flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
+            z = (
+                torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype)
+                * noise_scale
+            )
+            for flow in flows:
+                z = flow(z, x_mask, g=x, reverse=reverse)
+            z0, z1 = torch.split(z, [1, 1], 1)
+            logw = z0
+            return logw
+
+
+class DurationPredictor(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.LayerNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.LayerNorm(filter_channels)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+    def forward(
+        self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        n_vocab: int,
+        out_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        n_speakers: int,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+        self.emb = nn.Embedding(len(SYMBOLS), hidden_channels)
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+        self.tone_emb = nn.Embedding(NUM_TONES, hidden_channels)
+        nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
+        self.language_emb = nn.Embedding(NUM_LANGUAGES, hidden_channels)
+        nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
+        self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
+        self.ja_bert_proj = nn.Conv1d(1024, hidden_channels, 1)
+        self.en_bert_proj = nn.Conv1d(1024, hidden_channels, 1)
+        self.style_proj = nn.Linear(256, hidden_channels)
+
+        self.encoder = attentions.Encoder(
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        tone: torch.Tensor,
+        language: torch.Tensor,
+        bert: torch.Tensor,
+        ja_bert: torch.Tensor,
+        en_bert: torch.Tensor,
+        style_vec: torch.Tensor,
+        sid: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        bert_emb = self.bert_proj(bert).transpose(1, 2)
+        ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
+        en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2)
+        style_emb = self.style_proj(style_vec.unsqueeze(1))
+
+        x = (
+            self.emb(x)
+            + self.tone_emb(tone)
+            + self.language_emb(language)
+            + bert_emb
+            + ja_bert_emb
+            + en_bert_emb
+            + style_emb
+        ) * math.sqrt(
+            self.hidden_channels
+        )  # [b, t, h]
+        x = torch.transpose(x, 1, -1)  # [b, h, t]
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+
+        x = self.encoder(x * x_mask, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return x, m, logs, x_mask
+
+
+class ResidualCouplingBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ResidualCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    n_layers,
+                    gin_channels=gin_channels,
+                    mean_only=True,
+                )
+            )
+            self.flows.append(modules.Flip())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class PosteriorEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+
+class Generator(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel: int,
+        resblock_str: str,
+        resblock_kernel_sizes: list[int],
+        resblock_dilation_sizes: list[list[int]],
+        upsample_rates: list[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list[int],
+        gin_channels: int = 0,
+    ) -> None:
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ResBlock1 if resblock_str == "1" else modules.ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        ch = None
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))  # type: ignore
+
+        assert ch is not None
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(commons.init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(
+        self, x: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            assert xs is not None
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self) -> None:
+        print("Removing weight norm...")
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(
+        self,
+        period: int,
+        kernel_size: int = 5,
+        stride: int = 3,
+        use_spectral_norm: bool = False,
+    ) -> None:
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        self.use_spectral_norm = use_spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(
+                    Conv2d(
+                        1,
+                        32,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        32,
+                        128,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        128,
+                        512,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        512,
+                        1024,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        1024,
+                        1024,
+                        (kernel_size, 1),
+                        1,
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+            ]
+        )
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for layer in self.convs:
+            x = layer(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class DiscriminatorS(torch.nn.Module):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
+        super(DiscriminatorS, self).__init__()
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+                norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+                norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+                norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+                norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+                norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+
+        for layer in self.convs:
+            x = layer(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
+        super(MultiPeriodDiscriminator, self).__init__()
+        periods = [2, 3, 5, 7, 11]
+
+        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+        discs = discs + [
+            DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods
+        ]
+        self.discriminators = nn.ModuleList(discs)
+
+    def forward(
+        self,
+        y: torch.Tensor,
+        y_hat: torch.Tensor,
+    ) -> tuple[
+        list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]
+    ]:
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class ReferenceEncoder(nn.Module):
+    """
+    inputs --- [N, Ty/r, n_mels*r]  mels
+    outputs --- [N, ref_enc_gru_size]
+    """
+
+    def __init__(self, spec_channels: int, gin_channels: int = 0) -> None:
+        super().__init__()
+        self.spec_channels = spec_channels
+        ref_enc_filters = [32, 32, 64, 64, 128, 128]
+        K = len(ref_enc_filters)
+        filters = [1] + ref_enc_filters
+        convs = [
+            weight_norm(
+                nn.Conv2d(
+                    in_channels=filters[i],
+                    out_channels=filters[i + 1],
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    padding=(1, 1),
+                )
+            )
+            for i in range(K)
+        ]
+        self.convs = nn.ModuleList(convs)
+        # self.wns = nn.ModuleList([weight_norm(num_features=ref_enc_filters[i]) for i in range(K)])
+
+        out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
+        self.gru = nn.GRU(
+            input_size=ref_enc_filters[-1] * out_channels,
+            hidden_size=256 // 2,
+            batch_first=True,
+        )
+        self.proj = nn.Linear(128, gin_channels)
+
+    def forward(
+        self, inputs: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        N = inputs.size(0)
+        out = inputs.view(N, 1, -1, self.spec_channels)  # [N, 1, Ty, n_freqs]
+        for conv in self.convs:
+            out = conv(out)
+            # out = wn(out)
+            out = F.relu(out)  # [N, 128, Ty//2^K, n_mels//2^K]
+
+        out = out.transpose(1, 2)  # [N, Ty//2^K, 128, n_mels//2^K]
+        T = out.size(1)
+        N = out.size(0)
+        out = out.contiguous().view(N, T, -1)  # [N, Ty//2^K, 128*n_mels//2^K]
+
+        self.gru.flatten_parameters()
+        memory, out = self.gru(out)  # out --- [1, N, 128]
+
+        return self.proj(out.squeeze(0))
+
+    def calculate_channels(
+        self, L: int, kernel_size: int, stride: int, pad: int, n_convs: int
+    ) -> int:
+        for i in range(n_convs):
+            L = (L - kernel_size + 2 * pad) // stride + 1
+        return L
+
+
+class SynthesizerTrn(nn.Module):
+    """
+    Synthesizer for Training
+    """
+
+    def __init__(
+        self,
+        n_vocab: int,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        resblock: str,
+        resblock_kernel_sizes: list[int],
+        resblock_dilation_sizes: list[list[int]],
+        upsample_rates: list[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list[int],
+        n_speakers: int = 256,
+        gin_channels: int = 256,
+        use_sdp: bool = True,
+        n_flow_layer: int = 4,
+        n_layers_trans_flow: int = 4,
+        flow_share_parameter: bool = False,
+        use_transformer_flow: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.n_speakers = n_speakers
+        self.gin_channels = gin_channels
+        self.n_layers_trans_flow = n_layers_trans_flow
+        self.use_spk_conditioned_encoder = kwargs.get(
+            "use_spk_conditioned_encoder", True
+        )
+        self.use_sdp = use_sdp
+        self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
+        self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
+        self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
+        self.current_mas_noise_scale = self.mas_noise_scale_initial
+        if self.use_spk_conditioned_encoder and gin_channels > 0:
+            self.enc_gin_channels = gin_channels
+        self.enc_p = TextEncoder(
+            n_vocab,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            self.n_speakers,
+            gin_channels=self.enc_gin_channels,
+        )
+        self.dec = Generator(
+            inter_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        )
+        self.enc_q = PosteriorEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+        if use_transformer_flow:
+            self.flow = TransformerCouplingBlock(
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers_trans_flow,
+                5,
+                p_dropout,
+                n_flow_layer,
+                gin_channels=gin_channels,
+                share_parameter=flow_share_parameter,
+            )
+        else:
+            self.flow = ResidualCouplingBlock(
+                inter_channels,
+                hidden_channels,
+                5,
+                1,
+                n_flow_layer,
+                gin_channels=gin_channels,
+            )
+        self.sdp = StochasticDurationPredictor(
+            hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
+        )
+        self.dp = DurationPredictor(
+            hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
+        )
+
+        if n_speakers >= 1:
+            self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        else:
+            self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        y: torch.Tensor,
+        y_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        tone: torch.Tensor,
+        language: torch.Tensor,
+        bert: torch.Tensor,
+        ja_bert: torch.Tensor,
+        en_bert: torch.Tensor,
+        style_vec: torch.Tensor,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        tuple[torch.Tensor, ...],
+        tuple[torch.Tensor, ...],
+    ]:
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, ja_bert, en_bert, style_vec, sid, g=g
+        )
+        z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
+        z_p = self.flow(z, y_mask, g=g)
+
+        with torch.no_grad():
+            # negative cross-entropy
+            s_p_sq_r = torch.exp(-2 * logs_p)  # [b, d, t]
+            neg_cent1 = torch.sum(
+                -0.5 * math.log(2 * math.pi) - logs_p, [1], keepdim=True
+            )  # [b, 1, t_s]
+            neg_cent2 = torch.matmul(
+                -0.5 * (z_p**2).transpose(1, 2), s_p_sq_r
+            )  # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
+            neg_cent3 = torch.matmul(
+                z_p.transpose(1, 2), (m_p * s_p_sq_r)
+            )  # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
+            neg_cent4 = torch.sum(
+                -0.5 * (m_p**2) * s_p_sq_r, [1], keepdim=True
+            )  # [b, 1, t_s]
+            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+            if self.use_noise_scaled_mas:
+                epsilon = (
+                    torch.std(neg_cent)
+                    * torch.randn_like(neg_cent)
+                    * self.current_mas_noise_scale
+                )
+                neg_cent = neg_cent + epsilon
+
+            attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+            attn = (
+                monotonic_alignment.maximum_path(neg_cent, attn_mask.squeeze(1))
+                .unsqueeze(1)
+                .detach()
+            )
+
+        w = attn.sum(2)
+
+        l_length_sdp = self.sdp(x, x_mask, w, g=g)
+        l_length_sdp = l_length_sdp / torch.sum(x_mask)
+
+        logw_ = torch.log(w + 1e-6) * x_mask
+        logw = self.dp(x, x_mask, g=g)
+        # logw_sdp = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=1.0)
+        l_length_dp = torch.sum((logw - logw_) ** 2, [1, 2]) / torch.sum(
+            x_mask
+        )  # for averaging
+        # l_length_sdp += torch.sum((logw_sdp - logw_) ** 2, [1, 2]) / torch.sum(x_mask)
+
+        l_length = l_length_dp + l_length_sdp
+
+        # expand prior
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2)
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2)
+
+        z_slice, ids_slice = commons.rand_slice_segments(
+            z, y_lengths, self.segment_size
+        )
+        o = self.dec(z_slice, g=g)
+        return (
+            o,
+            l_length,
+            attn,
+            ids_slice,
+            x_mask,
+            y_mask,
+            (z, z_p, m_p, logs_p, m_q, logs_q),
+            (x, logw, logw_),
+        )
+
+    def infer(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        tone: torch.Tensor,
+        language: torch.Tensor,
+        bert: torch.Tensor,
+        ja_bert: torch.Tensor,
+        en_bert: torch.Tensor,
+        style_vec: torch.Tensor,
+        noise_scale: float = 0.667,
+        length_scale: float = 1.0,
+        noise_scale_w: float = 0.8,
+        max_len: Optional[int] = None,
+        sdp_ratio: float = 0.0,
+        y: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]]:
+        # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)
+        # g = self.gst(y)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            assert y is not None
+            g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, ja_bert, en_bert, style_vec, sid, g=g
+        )
+        logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
+            sdp_ratio
+        ) + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
+        w = torch.exp(logw) * x_mask * length_scale
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(
+            x_mask.dtype
+        )
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.generate_path(w_ceil, attn_mask)
+
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(
+            1, 2
+        )  # [b, t', t], [b, t, d] -> [b, d, t']
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(
+            1, 2
+        )  # [b, t', t], [b, t, d] -> [b, d, t']
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        o = self.dec((z * y_mask)[:, :, :max_len], g=g)
+        return o, attn, y_mask, (z, z_p, m_p, logs_p)

style_bert_vits2/models/models_jp_extra.py

@@ -0,0 +1,1157 @@
+import math
+from typing import Any, Optional
+
+import torch
+from torch import nn
+from torch.nn import Conv1d, Conv2d, ConvTranspose1d
+from torch.nn import functional as F
+from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
+
+from style_bert_vits2.models import attentions, commons, modules, monotonic_alignment
+from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
+
+
+class DurationDiscriminator(nn.Module):  # vits2
+    def __init__(
+        self,
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.LayerNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.LayerNorm(filter_channels)
+        self.dur_proj = nn.Conv1d(1, filter_channels, 1)
+
+        self.LSTM = nn.LSTM(
+            2 * filter_channels, filter_channels, batch_first=True, bidirectional=True
+        )
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+        self.output_layer = nn.Sequential(
+            nn.Linear(2 * filter_channels, 1), nn.Sigmoid()
+        )
+
+    def forward_probability(self, x: torch.Tensor, dur: torch.Tensor) -> torch.Tensor:
+        dur = self.dur_proj(dur)
+        x = torch.cat([x, dur], dim=1)
+        x = x.transpose(1, 2)
+        x, _ = self.LSTM(x)
+        output_prob = self.output_layer(x)
+        return output_prob
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        dur_r: torch.Tensor,
+        dur_hat: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> list[torch.Tensor]:
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+
+        output_probs = []
+        for dur in [dur_r, dur_hat]:
+            output_prob = self.forward_probability(x, dur)
+            output_probs.append(output_prob)
+
+        return output_probs
+
+
+class TransformerCouplingBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+        share_parameter: bool = False,
+    ) -> None:
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+
+        self.wn = (
+            # attentions.FFT(
+            #     hidden_channels,
+            #     filter_channels,
+            #     n_heads,
+            #     n_layers,
+            #     kernel_size,
+            #     p_dropout,
+            #     isflow=True,
+            #     gin_channels=self.gin_channels,
+            # )
+            None
+            if share_parameter
+            else None
+        )
+
+        for i in range(n_flows):
+            self.flows.append(
+                modules.TransformerCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    n_layers,
+                    n_heads,
+                    p_dropout,
+                    filter_channels,
+                    mean_only=True,
+                    wn_sharing_parameter=self.wn,
+                    gin_channels=self.gin_channels,
+                )
+            )
+            self.flows.append(modules.Flip())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class StochasticDurationPredictor(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        filter_channels = in_channels  # it needs to be removed from future version.
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.log_flow = modules.Log()
+        self.flows = nn.ModuleList()
+        self.flows.append(modules.ElementwiseAffine(2))
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.flows.append(modules.Flip())
+
+        self.post_pre = nn.Conv1d(1, filter_channels, 1)
+        self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.post_convs = modules.DDSConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(modules.ElementwiseAffine(2))
+        for i in range(4):
+            self.post_flows.append(
+                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.post_flows.append(modules.Flip())
+
+        self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+        self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.convs = modules.DDSConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        w: Optional[torch.Tensor] = None,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+        noise_scale: float = 1.0,
+    ) -> torch.Tensor:
+        x = torch.detach(x)
+        x = self.pre(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.convs(x, x_mask)
+        x = self.proj(x) * x_mask
+
+        if not reverse:
+            flows = self.flows
+            assert w is not None
+
+            logdet_tot_q = 0
+            h_w = self.post_pre(w)
+            h_w = self.post_convs(h_w, x_mask)
+            h_w = self.post_proj(h_w) * x_mask
+            e_q = (
+                torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype)
+                * x_mask
+            )
+            z_q = e_q
+            for flow in self.post_flows:
+                z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+                logdet_tot_q += logdet_q
+            z_u, z1 = torch.split(z_q, [1, 1], 1)
+            u = torch.sigmoid(z_u) * x_mask
+            z0 = (w - u) * x_mask
+            logdet_tot_q += torch.sum(
+                (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2]
+            )
+            logq = (
+                torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2])
+                - logdet_tot_q
+            )
+
+            logdet_tot = 0
+            z0, logdet = self.log_flow(z0, x_mask)
+            logdet_tot += logdet
+            z = torch.cat([z0, z1], 1)
+            for flow in flows:
+                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+                logdet_tot = logdet_tot + logdet
+            nll = (
+                torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2])
+                - logdet_tot
+            )
+            return nll + logq  # [b]
+        else:
+            flows = list(reversed(self.flows))
+            flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
+            z = (
+                torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype)
+                * noise_scale
+            )
+            for flow in flows:
+                z = flow(z, x_mask, g=x, reverse=reverse)
+            z0, z1 = torch.split(z, [1, 1], 1)
+            logw = z0
+            return logw
+
+
+class DurationPredictor(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.LayerNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.LayerNorm(filter_channels)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+    def forward(
+        self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+
+
+class Bottleneck(nn.Sequential):
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        c_fc1 = nn.Linear(in_dim, hidden_dim, bias=False)
+        c_fc2 = nn.Linear(in_dim, hidden_dim, bias=False)
+        super().__init__(c_fc1, c_fc2)
+
+
+class Block(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(in_dim)
+        self.mlp = MLP(in_dim, hidden_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.mlp(self.norm(x))
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.c_fc1 = nn.Linear(in_dim, hidden_dim, bias=False)
+        self.c_fc2 = nn.Linear(in_dim, hidden_dim, bias=False)
+        self.c_proj = nn.Linear(hidden_dim, in_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
+        x = self.c_proj(x)
+        return x
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        n_vocab: int,
+        out_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+        self.emb = nn.Embedding(len(SYMBOLS), hidden_channels)
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+        self.tone_emb = nn.Embedding(NUM_TONES, hidden_channels)
+        nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
+        self.language_emb = nn.Embedding(NUM_LANGUAGES, hidden_channels)
+        nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
+        self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
+
+        # Remove emo_vq since it's not working well.
+        self.style_proj = nn.Linear(256, hidden_channels)
+
+        self.encoder = attentions.Encoder(
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        tone: torch.Tensor,
+        language: torch.Tensor,
+        bert: torch.Tensor,
+        style_vec: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        bert_emb = self.bert_proj(bert).transpose(1, 2)
+        style_emb = self.style_proj(style_vec.unsqueeze(1))
+        x = (
+            self.emb(x)
+            + self.tone_emb(tone)
+            + self.language_emb(language)
+            + bert_emb
+            + style_emb
+        ) * math.sqrt(
+            self.hidden_channels
+        )  # [b, t, h]
+        x = torch.transpose(x, 1, -1)  # [b, h, t]
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+
+        x = self.encoder(x * x_mask, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return x, m, logs, x_mask
+
+
+class ResidualCouplingBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        n_flows: int = 4,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ResidualCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    n_layers,
+                    gin_channels=gin_channels,
+                    mean_only=True,
+                )
+            )
+            self.flows.append(modules.Flip())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class PosteriorEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        gin_channels: int = 0,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+
+class Generator(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel: int,
+        resblock_str: str,
+        resblock_kernel_sizes: list[int],
+        resblock_dilation_sizes: list[list[int]],
+        upsample_rates: list[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list[int],
+        gin_channels: int = 0,
+    ) -> None:
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ResBlock1 if resblock_str == "1" else modules.ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        ch = None
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))  # type: ignore
+
+        assert ch is not None
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(commons.init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(
+        self, x: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            assert xs is not None
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self) -> None:
+        print("Removing weight norm...")
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(
+        self,
+        period: int,
+        kernel_size: int = 5,
+        stride: int = 3,
+        use_spectral_norm: bool = False,
+    ) -> None:
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        self.use_spectral_norm = use_spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(
+                    Conv2d(
+                        1,
+                        32,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        32,
+                        128,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        128,
+                        512,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        512,
+                        1024,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+                norm_f(
+                    Conv2d(
+                        1024,
+                        1024,
+                        (kernel_size, 1),
+                        1,
+                        padding=(commons.get_padding(kernel_size, 1), 0),
+                    )
+                ),
+            ]
+        )
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for layer in self.convs:
+            x = layer(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class DiscriminatorS(torch.nn.Module):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
+        super(DiscriminatorS, self).__init__()
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+                norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+                norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+                norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+                norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+                norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+
+        for layer in self.convs:
+            x = layer(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self, use_spectral_norm: bool = False) -> None:
+        super(MultiPeriodDiscriminator, self).__init__()
+        periods = [2, 3, 5, 7, 11]
+
+        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+        discs = discs + [
+            DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods
+        ]
+        self.discriminators = nn.ModuleList(discs)
+
+    def forward(
+        self,
+        y: torch.Tensor,
+        y_hat: torch.Tensor,
+    ) -> tuple[
+        list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]
+    ]:
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class WavLMDiscriminator(nn.Module):
+    """docstring for Discriminator."""
+
+    def __init__(
+        self,
+        slm_hidden: int = 768,
+        slm_layers: int = 13,
+        initial_channel: int = 64,
+        use_spectral_norm: bool = False,
+    ) -> None:
+        super(WavLMDiscriminator, self).__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.pre = norm_f(
+            Conv1d(slm_hidden * slm_layers, initial_channel, 1, 1, padding=0)
+        )
+
+        self.convs = nn.ModuleList(
+            [
+                norm_f(
+                    nn.Conv1d(
+                        initial_channel, initial_channel * 2, kernel_size=5, padding=2
+                    )
+                ),
+                norm_f(
+                    nn.Conv1d(
+                        initial_channel * 2,
+                        initial_channel * 4,
+                        kernel_size=5,
+                        padding=2,
+                    )
+                ),
+                norm_f(
+                    nn.Conv1d(initial_channel * 4, initial_channel * 4, 5, 1, padding=2)
+                ),
+            ]
+        )
+
+        self.conv_post = norm_f(Conv1d(initial_channel * 4, 1, 3, 1, padding=1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.pre(x)
+
+        fmap = []
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x
+
+
+class ReferenceEncoder(nn.Module):
+    """
+    inputs --- [N, Ty/r, n_mels*r]  mels
+    outputs --- [N, ref_enc_gru_size]
+    """
+
+    def __init__(self, spec_channels: int, gin_channels: int = 0) -> None:
+        super().__init__()
+        self.spec_channels = spec_channels
+        ref_enc_filters = [32, 32, 64, 64, 128, 128]
+        K = len(ref_enc_filters)
+        filters = [1] + ref_enc_filters
+        convs = [
+            weight_norm(
+                nn.Conv2d(
+                    in_channels=filters[i],
+                    out_channels=filters[i + 1],
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    padding=(1, 1),
+                )
+            )
+            for i in range(K)
+        ]
+        self.convs = nn.ModuleList(convs)
+        # self.wns = nn.ModuleList([weight_norm(num_features=ref_enc_filters[i]) for i in range(K)])
+
+        out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
+        self.gru = nn.GRU(
+            input_size=ref_enc_filters[-1] * out_channels,
+            hidden_size=256 // 2,
+            batch_first=True,
+        )
+        self.proj = nn.Linear(128, gin_channels)
+
+    def forward(
+        self, inputs: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        N = inputs.size(0)
+        out = inputs.view(N, 1, -1, self.spec_channels)  # [N, 1, Ty, n_freqs]
+        for conv in self.convs:
+            out = conv(out)
+            # out = wn(out)
+            out = F.relu(out)  # [N, 128, Ty//2^K, n_mels//2^K]
+
+        out = out.transpose(1, 2)  # [N, Ty//2^K, 128, n_mels//2^K]
+        T = out.size(1)
+        N = out.size(0)
+        out = out.contiguous().view(N, T, -1)  # [N, Ty//2^K, 128*n_mels//2^K]
+
+        self.gru.flatten_parameters()
+        memory, out = self.gru(out)  # out --- [1, N, 128]
+
+        return self.proj(out.squeeze(0))
+
+    def calculate_channels(
+        self, L: int, kernel_size: int, stride: int, pad: int, n_convs: int
+    ) -> int:
+        for i in range(n_convs):
+            L = (L - kernel_size + 2 * pad) // stride + 1
+        return L
+
+
+class SynthesizerTrn(nn.Module):
+    """
+    Synthesizer for Training
+    """
+
+    def __init__(
+        self,
+        n_vocab: int,
+        spec_channels: int,
+        segment_size: int,
+        inter_channels: int,
+        hidden_channels: int,
+        filter_channels: int,
+        n_heads: int,
+        n_layers: int,
+        kernel_size: int,
+        p_dropout: float,
+        resblock: str,
+        resblock_kernel_sizes: list[int],
+        resblock_dilation_sizes: list[list[int]],
+        upsample_rates: list[int],
+        upsample_initial_channel: int,
+        upsample_kernel_sizes: list[int],
+        n_speakers: int = 256,
+        gin_channels: int = 256,
+        use_sdp: bool = True,
+        n_flow_layer: int = 4,
+        n_layers_trans_flow: int = 6,
+        flow_share_parameter: bool = False,
+        use_transformer_flow: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.n_speakers = n_speakers
+        self.gin_channels = gin_channels
+        self.n_layers_trans_flow = n_layers_trans_flow
+        self.use_spk_conditioned_encoder = kwargs.get(
+            "use_spk_conditioned_encoder", True
+        )
+        self.use_sdp = use_sdp
+        self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
+        self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
+        self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
+        self.current_mas_noise_scale = self.mas_noise_scale_initial
+        if self.use_spk_conditioned_encoder and gin_channels > 0:
+            self.enc_gin_channels = gin_channels
+        self.enc_p = TextEncoder(
+            n_vocab,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.enc_gin_channels,
+        )
+        self.dec = Generator(
+            inter_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        )
+        self.enc_q = PosteriorEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+        if use_transformer_flow:
+            self.flow = TransformerCouplingBlock(
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers_trans_flow,
+                5,
+                p_dropout,
+                n_flow_layer,
+                gin_channels=gin_channels,
+                share_parameter=flow_share_parameter,
+            )
+        else:
+            self.flow = ResidualCouplingBlock(
+                inter_channels,
+                hidden_channels,
+                5,
+                1,
+                n_flow_layer,
+                gin_channels=gin_channels,
+            )
+        self.sdp = StochasticDurationPredictor(
+            hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
+        )
+        self.dp = DurationPredictor(
+            hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
+        )
+
+        if n_speakers >= 1:
+            self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        else:
+            self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        y: torch.Tensor,
+        y_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        tone: torch.Tensor,
+        language: torch.Tensor,
+        bert: torch.Tensor,
+        style_vec: torch.Tensor,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        tuple[torch.Tensor, ...],
+        tuple[torch.Tensor, ...],
+    ]:
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, style_vec, g=g
+        )
+        z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
+        z_p = self.flow(z, y_mask, g=g)
+
+        with torch.no_grad():
+            # negative cross-entropy
+            s_p_sq_r = torch.exp(-2 * logs_p)  # [b, d, t]
+            neg_cent1 = torch.sum(
+                -0.5 * math.log(2 * math.pi) - logs_p, [1], keepdim=True
+            )  # [b, 1, t_s]
+            neg_cent2 = torch.matmul(
+                -0.5 * (z_p**2).transpose(1, 2), s_p_sq_r
+            )  # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
+            neg_cent3 = torch.matmul(
+                z_p.transpose(1, 2), (m_p * s_p_sq_r)
+            )  # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
+            neg_cent4 = torch.sum(
+                -0.5 * (m_p**2) * s_p_sq_r, [1], keepdim=True
+            )  # [b, 1, t_s]
+            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+            if self.use_noise_scaled_mas:
+                epsilon = (
+                    torch.std(neg_cent)
+                    * torch.randn_like(neg_cent)
+                    * self.current_mas_noise_scale
+                )
+                neg_cent = neg_cent + epsilon
+
+            attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+            attn = (
+                monotonic_alignment.maximum_path(neg_cent, attn_mask.squeeze(1))
+                .unsqueeze(1)
+                .detach()
+            )
+
+        w = attn.sum(2)
+
+        l_length_sdp = self.sdp(x, x_mask, w, g=g)
+        l_length_sdp = l_length_sdp / torch.sum(x_mask)
+
+        logw_ = torch.log(w + 1e-6) * x_mask
+        logw = self.dp(x, x_mask, g=g)
+        # logw_sdp = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=1.0)
+        l_length_dp = torch.sum((logw - logw_) ** 2, [1, 2]) / torch.sum(
+            x_mask
+        )  # for averaging
+        # l_length_sdp += torch.sum((logw_sdp - logw_) ** 2, [1, 2]) / torch.sum(x_mask)
+
+        l_length = l_length_dp + l_length_sdp
+
+        # expand prior
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2)
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2)
+
+        z_slice, ids_slice = commons.rand_slice_segments(
+            z, y_lengths, self.segment_size
+        )
+        o = self.dec(z_slice, g=g)
+        return (
+            o,
+            l_length,
+            attn,
+            ids_slice,
+            x_mask,
+            y_mask,
+            (z, z_p, m_p, logs_p, m_q, logs_q),  # type: ignore
+            (x, logw, logw_),  # , logw_sdp),
+            g,
+        )
+
+    def infer(
+        self,
+        x: torch.Tensor,
+        x_lengths: torch.Tensor,
+        sid: torch.Tensor,
+        tone: torch.Tensor,
+        language: torch.Tensor,
+        bert: torch.Tensor,
+        style_vec: torch.Tensor,
+        noise_scale: float = 0.667,
+        length_scale: float = 1.0,
+        noise_scale_w: float = 0.8,
+        max_len: Optional[int] = None,
+        sdp_ratio: float = 0.0,
+        y: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]]:
+        # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)
+        # g = self.gst(y)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            assert y is not None
+            g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, style_vec, g=g
+        )
+        logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
+            sdp_ratio
+        ) + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
+        w = torch.exp(logw) * x_mask * length_scale
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(
+            x_mask.dtype
+        )
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.generate_path(w_ceil, attn_mask)
+
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(
+            1, 2
+        )  # [b, t', t], [b, t, d] -> [b, d, t']
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(
+            1, 2
+        )  # [b, t', t], [b, t, d] -> [b, d, t']
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        o = self.dec((z * y_mask)[:, :, :max_len], g=g)
+        return o, attn, y_mask, (z, z_p, m_p, logs_p)

style_bert_vits2/models/modules.py

@@ -0,0 +1,642 @@
+import math
+from typing import Any, Optional, Union
+
+import torch
+from torch import nn
+from torch.nn import Conv1d
+from torch.nn import functional as F
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+from style_bert_vits2.models import commons
+from style_bert_vits2.models.attentions import Encoder
+from style_bert_vits2.models.transforms import piecewise_rational_quadratic_transform
+
+
+LRELU_SLOPE = 0.1
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        n_layers: int,
+        p_dropout: float,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(
+                in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+            )
+        )
+        self.norm_layers.append(LayerNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    hidden_channels,
+                    hidden_channels,
+                    kernel_size,
+                    padding=kernel_size // 2,
+                )
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        assert self.proj.bias is not None
+        self.proj.bias.data.zero_()
+
+    def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DDSConv(nn.Module):
+    """
+    Dialted and Depth-Separable Convolution
+    """
+
+    def __init__(
+        self, channels: int, kernel_size: int, n_layers: int, p_dropout: float = 0.0
+    ) -> None:
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(LayerNorm(channels))
+            self.norms_2.append(LayerNorm(channels))
+
+    def forward(
+        self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
+
+
+class WN(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        gin_channels: int = 0,
+        p_dropout: float = 0,
+    ) -> None:
+        super(WN, self).__init__()
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self) -> None:
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple[int, int, int] = (1, 3, 5),
+    ) -> None:
+        super(ResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=commons.get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=commons.get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=commons.get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(commons.init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=commons.get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=commons.get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=commons.get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(commons.init_weights)
+
+    def forward(
+        self, x: torch.Tensor, x_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(
+        self, channels: int, kernel_size: int = 3, dilation: tuple[int, int] = (1, 3)
+    ) -> None:
+        super(ResBlock2, self).__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=commons.get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=commons.get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(commons.init_weights)
+
+    def forward(
+        self, x: torch.Tensor, x_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class Log(nn.Module):
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        reverse: bool = False,
+        **kwargs: Any,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
+
+class Flip(nn.Module):
+    def forward(
+        self,
+        x: torch.Tensor,
+        *args: Any,
+        reverse: bool = False,
+        **kwargs: Any,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class ElementwiseAffine(nn.Module):
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        reverse: bool = False,
+        **kwargs: Any,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
+
+
+class ResidualCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        n_layers: int,
+        p_dropout: float = 0,
+        gin_channels: int = 0,
+        mean_only: bool = False,
+    ) -> None:
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        assert self.post.bias is not None
+        self.post.bias.data.zero_()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+
+class ConvFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        filter_channels: int,
+        kernel_size: int,
+        n_layers: int,
+        num_bins: int = 10,
+        tail_bound: float = 5.0,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        assert self.proj.bias is not None
+        self.proj.bias.data.zero_()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)  # [b, cx?, t] -> [b, c, t, ?]
+
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
+            self.filter_channels
+        )
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x
+
+
+class TransformerCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        n_layers: int,
+        n_heads: int,
+        p_dropout: float = 0,
+        filter_channels: int = 0,
+        mean_only: bool = False,
+        wn_sharing_parameter: Optional[nn.Module] = None,
+        gin_channels: int = 0,
+    ) -> None:
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = (
+            Encoder(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=gin_channels,
+            )
+            if wn_sharing_parameter is None
+            else wn_sharing_parameter
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        assert self.post.bias is not None
+        self.post.bias.data.zero_()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        g: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x

style_bert_vits2/models/monotonic_alignment.py

@@ -0,0 +1,89 @@
+"""
+以下に記述されている関数のコメントはリファクタリング時に GPT-4 に生成させたもので、
+コードと完全に一致している保証はない。あくまで参考程度とすること。
+"""
+
+from typing import Any
+
+import numba
+import torch
+from numpy import float32, int32, zeros
+
+
+def maximum_path(neg_cent: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """
+    与えられた負の中心とマスクを使用して最大パスを計算する
+
+    Args:
+        neg_cent (torch.Tensor): 負の中心を表すテンソル
+        mask (torch.Tensor): マスクを表すテンソル
+
+    Returns:
+        Tensor: 計算された最大パスを表すテンソル
+    """
+
+    device = neg_cent.device
+    dtype = neg_cent.dtype
+    neg_cent = neg_cent.data.cpu().numpy().astype(float32)
+    path = zeros(neg_cent.shape, dtype=int32)
+
+    t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
+    t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
+    __maximum_path_jit(path, neg_cent, t_t_max, t_s_max)
+
+    return torch.from_numpy(path).to(device=device, dtype=dtype)
+
+
+@numba.jit(
+    numba.void(
+        numba.int32[:, :, ::1],
+        numba.float32[:, :, ::1],
+        numba.int32[::1],
+        numba.int32[::1],
+    ),
+    nopython=True,
+    nogil=True,
+)  # type: ignore
+def __maximum_path_jit(paths: Any, values: Any, t_ys: Any, t_xs: Any) -> None:
+    """
+    与えられたパス、値、およびターゲットの y と x 座標を使用して JIT で最大パスを計算する
+
+    Args:
+        paths: 計算されたパスを格納するための整数型の 3 次元配列
+        values: 値を格納するための浮動小数点型の 3 次元配列
+        t_ys: ターゲットの y 座標を格納するための整数型の 1 次元配列
+        t_xs: ターゲットの x 座標を格納するための整数型の 1 次元配列
+    """
+
+    b = paths.shape[0]
+    max_neg_val = -1e9
+    for i in range(int(b)):
+        path = paths[i]
+        value = values[i]
+        t_y = t_ys[i]
+        t_x = t_xs[i]
+
+        v_prev = v_cur = 0.0
+        index = t_x - 1
+
+        for y in range(t_y):
+            for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+                if x == y:
+                    v_cur = max_neg_val
+                else:
+                    v_cur = value[y - 1, x]
+                if x == 0:
+                    if y == 0:
+                        v_prev = 0.0
+                    else:
+                        v_prev = max_neg_val
+                else:
+                    v_prev = value[y - 1, x - 1]
+                value[y, x] += max(v_prev, v_cur)
+
+        for y in range(t_y - 1, -1, -1):
+            path[y, index] = 1
+            if index != 0 and (
+                index == y or value[y - 1, index] < value[y - 1, index - 1]
+            ):
+                index = index - 1

style_bert_vits2/models/transforms.py

@@ -0,0 +1,215 @@
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(
+    inputs: torch.Tensor,
+    unnormalized_widths: torch.Tensor,
+    unnormalized_heights: torch.Tensor,
+    unnormalized_derivatives: torch.Tensor,
+    inverse: bool = False,
+    tails: Optional[str] = None,
+    tail_bound: float = 1.0,
+    min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative: float = DEFAULT_MIN_DERIVATIVE,
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs,  # type: ignore
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(
+    bin_locations: torch.Tensor, inputs: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unconstrained_rational_quadratic_spline(
+    inputs: torch.Tensor,
+    unnormalized_widths: torch.Tensor,
+    unnormalized_heights: torch.Tensor,
+    unnormalized_derivatives: torch.Tensor,
+    inverse: bool = False,
+    tails: str = "linear",
+    tail_bound: float = 1.0,
+    min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative: float = DEFAULT_MIN_DERIVATIVE,
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError(f"{tails} tails are not implemented.")
+
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def rational_quadratic_spline(
+    inputs: torch.Tensor,
+    unnormalized_widths: torch.Tensor,
+    unnormalized_heights: torch.Tensor,
+    unnormalized_derivatives: torch.Tensor,
+    inverse: bool = False,
+    left: float = 0.0,
+    right: float = 1.0,
+    bottom: float = 0.0,
+    top: float = 1.0,
+    min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative: float = DEFAULT_MIN_DERIVATIVE,
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

style_bert_vits2/models/utils/__init__.py

@@ -0,0 +1,264 @@
+import glob
+import logging
+import os
+import re
+import subprocess
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+
+from style_bert_vits2.logging import logger
+from style_bert_vits2.models.utils import checkpoints  # type: ignore
+from style_bert_vits2.models.utils import safetensors  # type: ignore
+
+
+if TYPE_CHECKING:
+    # tensorboard はライブラリとしてインストールされている場合は依存関係に含まれないため、型チェック時のみインポートする
+    from torch.utils.tensorboard import SummaryWriter
+
+
+__is_matplotlib_imported = False
+
+
+def summarize(
+    writer: "SummaryWriter",
+    global_step: int,
+    scalars: dict[str, float] = {},
+    histograms: dict[str, Any] = {},
+    images: dict[str, Any] = {},
+    audios: dict[str, Any] = {},
+    audio_sampling_rate: int = 22050,
+) -> None:
+    """
+    指定されたデータを TensorBoard にまとめて追加する
+
+    Args:
+        writer (SummaryWriter): TensorBoard への書き込みを行うオブジェクト
+        global_step (int): グローバルステップ数
+        scalars (dict[str, float]): スカラー値の辞書
+        histograms (dict[str, Any]): ヒストグラムの辞書
+        images (dict[str, Any]): 画像データの辞書
+        audios (dict[str, Any]): 音声データの辞書
+        audio_sampling_rate (int): 音声データのサンプリングレート
+    """
+    for k, v in scalars.items():
+        writer.add_scalar(k, v, global_step)
+    for k, v in histograms.items():
+        writer.add_histogram(k, v, global_step)
+    for k, v in images.items():
+        writer.add_image(k, v, global_step, dataformats="HWC")
+    for k, v in audios.items():
+        writer.add_audio(k, v, global_step, audio_sampling_rate)
+
+
+def is_resuming(dir_path: Union[str, Path]) -> bool:
+    """
+    指定されたディレクトリパスに再開可能なモデルが存在するかどうかを返す
+
+    Args:
+        dir_path: チェックするディレクトリのパス
+
+    Returns:
+        bool: 再開可能なモデルが存在するかどうか
+    """
+    # JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する
+    g_list = glob.glob(os.path.join(dir_path, "G_*.pth"))
+    # d_list = glob.glob(os.path.join(dir_path, "D_*.pth"))
+    # dur_list = glob.glob(os.path.join(dir_path, "DUR_*.pth"))
+    return len(g_list) > 0
+
+
+def plot_spectrogram_to_numpy(spectrogram: NDArray[Any]) -> NDArray[Any]:
+    """
+    指定されたスペクトログラムを画像データに変換する
+
+    Args:
+        spectrogram (NDArray[Any]): スペクトログラム
+
+    Returns:
+        NDArray[Any]: 画像データ
+    """
+
+    global __is_matplotlib_imported
+    if not __is_matplotlib_imported:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        __is_matplotlib_imported = True
+        mpl_logger = logging.getLogger("matplotlib")
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(10, 2))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+    plt.colorbar(im, ax=ax)
+    plt.xlabel("Frames")
+    plt.ylabel("Channels")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")  # type: ignore
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def plot_alignment_to_numpy(
+    alignment: NDArray[Any], info: Optional[str] = None
+) -> NDArray[Any]:
+    """
+    指定されたアライメントを画像データに変換する
+
+    Args:
+        alignment (NDArray[Any]): アライメント
+        info (Optional[str]): 画像に追加する情報
+
+    Returns:
+        NDArray[Any]: 画像データ
+    """
+
+    global __is_matplotlib_imported
+    if not __is_matplotlib_imported:
+        import matplotlib
+
+        matplotlib.use("Agg")
+        __is_matplotlib_imported = True
+        mpl_logger = logging.getLogger("matplotlib")
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+
+    fig, ax = plt.subplots(figsize=(6, 4))
+    im = ax.imshow(
+        alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
+    )
+    fig.colorbar(im, ax=ax)
+    xlabel = "Decoder timestep"
+    if info is not None:
+        xlabel += "\n\n" + info
+    plt.xlabel(xlabel)
+    plt.ylabel("Encoder timestep")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")  # type: ignore
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def load_wav_to_torch(full_path: Union[str, Path]) -> tuple[torch.FloatTensor, int]:
+    """
+    指定された音声ファイルを読み込み、PyTorch のテンソルに変換して返す
+
+    Args:
+        full_path (Union[str, Path]): 音声ファイルのパス
+
+    Returns:
+        tuple[torch.FloatTensor, int]: 音声データのテンソルとサンプリングレート
+    """
+
+    # この関数は学習時以外使われないため、ライブラリとしての style_bert_vits2 が
+    # 重たい scipy に依存しないように遅延 import する
+    try:
+        from scipy.io.wavfile import read
+    except ImportError:
+        raise ImportError("scipy is required to load wav file")
+
+    sampling_rate, data = read(full_path)
+    return torch.FloatTensor(data.astype(np.float32)), sampling_rate
+
+
+def load_filepaths_and_text(
+    filename: Union[str, Path], split: str = "|"
+) -> list[list[str]]:
+    """
+    指定されたファイルからファイルパスとテキストを読み込む
+
+    Args:
+        filename (Union[str, Path]): ファイルのパス
+        split (str): ファイルの区切り文字 (デフォルト: "|")
+
+    Returns:
+        list[list[str]]: ファイルパスとテキストのリスト
+    """
+
+    with open(filename, encoding="utf-8") as f:
+        filepaths_and_text = [line.strip().split(split) for line in f]
+    return filepaths_and_text
+
+
+def get_logger(
+    model_dir_path: Union[str, Path], filename: str = "train.log"
+) -> logging.Logger:
+    """
+    ロガーを取得する
+
+    Args:
+        model_dir_path (Union[str, Path]): ログを保存するディレクトリのパス
+        filename (str): ログファイルの名前 (デフォルト: "train.log")
+
+    Returns:
+        logging.Logger: ロガー
+    """
+
+    global logger
+    logger = logging.getLogger(os.path.basename(model_dir_path))
+    logger.setLevel(logging.DEBUG)
+
+    formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
+    if not os.path.exists(model_dir_path):
+        os.makedirs(model_dir_path)
+    h = logging.FileHandler(os.path.join(model_dir_path, filename))
+    h.setLevel(logging.DEBUG)
+    h.setFormatter(formatter)
+    logger.addHandler(h)
+    return logger
+
+
+def get_steps(model_path: Union[str, Path]) -> Optional[int]:
+    """
+    モデルのパスからイテレーション回数を取得する
+
+    Args:
+        model_path (Union[str, Path]): モデルのパス
+
+    Returns:
+        Optional[int]: イテレーション回数
+    """
+
+    matches = re.findall(r"\d+", model_path)  # type: ignore
+    return matches[-1] if matches else None
+
+
+def check_git_hash(model_dir_path: Union[str, Path]) -> None:
+    """
+    モデルのディレクトリに .git ディレクトリが存在する場合、ハッシュ値を比較する
+
+    Args:
+        model_dir_path (Union[str, Path]): モデルのディレクトリのパス
+    """
+
+    source_dir = os.path.dirname(os.path.realpath(__file__))
+    if not os.path.exists(os.path.join(source_dir, ".git")):
+        logger.warning(
+            f"{source_dir} is not a git repository, therefore hash value comparison will be ignored."
+        )
+        return
+
+    cur_hash = subprocess.getoutput("git rev-parse HEAD")
+
+    path = os.path.join(model_dir_path, "githash")
+    if os.path.exists(path):
+        with open(path, encoding="utf-8") as f:
+            saved_hash = f.read()
+        if saved_hash != cur_hash:
+            logger.warning(
+                f"git hash values are different. {saved_hash[:8]}(saved) != {cur_hash[:8]}(current)"
+            )
+    else:
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(cur_hash)

style_bert_vits2/models/utils/checkpoints.py

@@ -0,0 +1,202 @@
+import glob
+import os
+import re
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import torch
+
+from style_bert_vits2.logging import logger
+
+
+def load_checkpoint(
+    checkpoint_path: Union[str, Path],
+    model: torch.nn.Module,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    skip_optimizer: bool = False,
+    for_infer: bool = False,
+) -> tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]:
+    """
+    指定されたパスからチェックポイントを読み込み、モデルとオプティマイザーを更新する。
+
+    Args:
+        checkpoint_path (Union[str, Path]): チェックポイントファイルのパス
+        model (torch.nn.Module): 更新するモデル
+        optimizer (Optional[torch.optim.Optimizer]): 更新するオプティマイザー。None の場合は更新しない
+        skip_optimizer (bool): オプティマイザーの更新をスキップするかどうかのフラグ
+        for_infer (bool): 推論用に読み込むかどうかのフラグ
+
+    Returns:
+        tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]: 更新されたモデルとオプティマイザー、学習率、イテレーション回数
+    """
+
+    assert os.path.isfile(checkpoint_path)
+    checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
+    iteration = checkpoint_dict["iteration"]
+    learning_rate = checkpoint_dict["learning_rate"]
+    logger.info(
+        f"Loading model and optimizer at iteration {iteration} from {checkpoint_path}"
+    )
+    if (
+        optimizer is not None
+        and not skip_optimizer
+        and checkpoint_dict["optimizer"] is not None
+    ):
+        optimizer.load_state_dict(checkpoint_dict["optimizer"])
+    elif optimizer is None and not skip_optimizer:
+        # else:      Disable this line if Infer and resume checkpoint,then enable the line upper
+        new_opt_dict = optimizer.state_dict()  # type: ignore
+        new_opt_dict_params = new_opt_dict["param_groups"][0]["params"]
+        new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"]
+        new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params
+        optimizer.load_state_dict(new_opt_dict)  # type: ignore
+
+    saved_state_dict = checkpoint_dict["model"]
+    if hasattr(model, "module"):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        try:
+            # assert "emb_g" not in k
+            new_state_dict[k] = saved_state_dict[k]
+            assert saved_state_dict[k].shape == v.shape, (
+                saved_state_dict[k].shape,
+                v.shape,
+            )
+        except:
+            # For upgrading from the old version
+            if "ja_bert_proj" in k:
+                v = torch.zeros_like(v)
+                logger.warning(
+                    f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
+                )
+            elif "enc_q" in k and for_infer:
+                continue
+            else:
+                logger.error(f"{k} is not in the checkpoint {checkpoint_path}")
+
+            new_state_dict[k] = v
+
+    if hasattr(model, "module"):
+        model.module.load_state_dict(new_state_dict, strict=False)
+    else:
+        model.load_state_dict(new_state_dict, strict=False)
+
+    logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
+
+    return model, optimizer, learning_rate, iteration
+
+
+def save_checkpoint(
+    model: torch.nn.Module,
+    optimizer: Union[torch.optim.Optimizer, torch.optim.AdamW],
+    learning_rate: float,
+    iteration: int,
+    checkpoint_path: Union[str, Path],
+) -> None:
+    """
+    モデルとオプティマイザーの状態を指定されたパスに保存する。
+
+    Args:
+        model (torch.nn.Module): 保存するモデル
+        optimizer (Union[torch.optim.Optimizer, torch.optim.AdamW]): 保存するオプティマイザー
+        learning_rate (float): 学習率
+        iteration (int): イテレーション回数
+        checkpoint_path (Union[str, Path]): 保存先のパス
+    """
+    logger.info(
+        f"Saving model and optimizer state at iteration {iteration} to {checkpoint_path}"
+    )
+    if hasattr(model, "module"):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+    torch.save(
+        {
+            "model": state_dict,
+            "iteration": iteration,
+            "optimizer": optimizer.state_dict(),
+            "learning_rate": learning_rate,
+        },
+        checkpoint_path,
+    )
+
+
+def clean_checkpoints(
+    model_dir_path: Union[str, Path] = "logs/44k/",
+    n_ckpts_to_keep: int = 2,
+    sort_by_time: bool = True,
+) -> None:
+    """
+    指定されたディレクトリから古いチェックポイントを削除して空き容量を確保する
+
+    Args:
+        model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス
+        n_ckpts_to_keep (int): 保持するチェックポイントの数（G_0.pth と D_0.pth を除く）
+        sort_by_time (bool): True の場合、時間順に削除。False の場合、名前順に削除
+    """
+
+    ckpts_files = [
+        f
+        for f in os.listdir(model_dir_path)
+        if os.path.isfile(os.path.join(model_dir_path, f))
+    ]
+
+    def name_key(_f: str) -> int:
+        return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))  # type: ignore
+
+    def time_key(_f: str) -> float:
+        return os.path.getmtime(os.path.join(model_dir_path, _f))
+
+    sort_key = time_key if sort_by_time else name_key
+
+    def x_sorted(_x: str) -> list[str]:
+        return sorted(
+            [f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")],
+            key=sort_key,
+        )
+
+    to_del = [
+        os.path.join(model_dir_path, fn)
+        for fn in (
+            x_sorted("G_")[:-n_ckpts_to_keep]
+            + x_sorted("D_")[:-n_ckpts_to_keep]
+            + x_sorted("WD_")[:-n_ckpts_to_keep]
+            + x_sorted("DUR_")[:-n_ckpts_to_keep]
+        )
+    ]
+
+    def del_info(fn: str) -> None:
+        return logger.info(f"Free up space by deleting ckpt {fn}")
+
+    def del_routine(x: str) -> list[Any]:
+        return [os.remove(x), del_info(x)]
+
+    [del_routine(fn) for fn in to_del]
+
+
+def get_latest_checkpoint_path(
+    model_dir_path: Union[str, Path], regex: str = "G_*.pth"
+) -> str:
+    """
+    指定されたディレクトリから最新のチェックポイントのパスを取得する
+
+    Args:
+        model_dir_path (Union[str, Path]): モデルが保存されているディレクトリのパス
+        regex (str): チェックポイントのファイル名の正規表現
+
+    Returns:
+        str: 最新のチェックポイントのパス
+    """
+
+    f_list = glob.glob(os.path.join(str(model_dir_path), regex))
+    f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
+    try:
+        x = f_list[-1]
+    except IndexError:
+        raise ValueError(f"No checkpoint found in {model_dir_path} with regex {regex}")
+
+    return x

style_bert_vits2/models/utils/safetensors.py

@@ -0,0 +1,91 @@
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from style_bert_vits2.logging import logger
+
+
+def load_safetensors(
+    checkpoint_path: Union[str, Path],
+    model: torch.nn.Module,
+    for_infer: bool = False,
+) -> tuple[torch.nn.Module, Optional[int]]:
+    """
+    指定されたパスから safetensors モデルを読み込み、モデルとイテレーションを返す。
+
+    Args:
+        checkpoint_path (Union[str, Path]): モデルのチェックポイントファイルのパス
+        model (torch.nn.Module): 読み込む対象のモデル
+        for_infer (bool): 推論用に読み込むかどうかのフラグ
+
+    Returns:
+        tuple[torch.nn.Module, Optional[int]]: 読み込まれたモデルとイテレーション回数（存在する場合）
+    """
+
+    tensors: dict[str, Any] = {}
+    iteration: Optional[int] = None
+    with safe_open(str(checkpoint_path), framework="pt", device="cpu") as f:  # type: ignore
+        for key in f.keys():
+            if key == "iteration":
+                iteration = f.get_tensor(key).item()
+            tensors[key] = f.get_tensor(key)
+    if hasattr(model, "module"):
+        result = model.module.load_state_dict(tensors, strict=False)
+    else:
+        result = model.load_state_dict(tensors, strict=False)
+    for key in result.missing_keys:
+        if key.startswith("enc_q") and for_infer:
+            continue
+        logger.warning(f"Missing key: {key}")
+    for key in result.unexpected_keys:
+        if key == "iteration":
+            continue
+        logger.warning(f"Unexpected key: {key}")
+    if iteration is None:
+        logger.info(f"Loaded '{checkpoint_path}'")
+    else:
+        logger.info(f"Loaded '{checkpoint_path}' (iteration {iteration})")
+
+    return model, iteration
+
+
+def save_safetensors(
+    model: torch.nn.Module,
+    iteration: int,
+    checkpoint_path: Union[str, Path],
+    is_half: bool = False,
+    for_infer: bool = False,
+) -> None:
+    """
+    モデルを safetensors 形式で保存する。
+
+    Args:
+        model (torch.nn.Module): 保存するモデル
+        iteration (int): イテレーション回数
+        checkpoint_path (Union[str, Path]): 保存先のパス
+        is_half (bool): モデルを半精度で保存するかどうかのフラグ
+        for_infer (bool): 推論用に保存するかどうかのフラグ
+    """
+
+    if hasattr(model, "module"):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+    keys = []
+    for k in state_dict:
+        if "enc_q" in k and for_infer:
+            continue
+        keys.append(k)
+
+    new_dict = (
+        {k: state_dict[k].half() for k in keys}
+        if is_half
+        else {k: state_dict[k] for k in keys}
+    )
+    new_dict["iteration"] = torch.LongTensor([iteration])
+    logger.info(f"Saved safetensors to {checkpoint_path}")
+
+    save_file(new_dict, checkpoint_path)

style_bert_vits2/nlp/__init__.py

@@ -0,0 +1,120 @@
+from typing import TYPE_CHECKING, Optional
+
+from style_bert_vits2.constants import Languages
+from style_bert_vits2.nlp.symbols import (
+    LANGUAGE_ID_MAP,
+    LANGUAGE_TONE_START_MAP,
+    SYMBOLS,
+)
+
+
+# __init__.py は配下のモジュールをインポートした時点で実行される
+# PyTorch のインポートは重いので、型チェック時以外はインポートしない
+if TYPE_CHECKING:
+    import torch
+
+
+__symbol_to_id = {s: i for i, s in enumerate(SYMBOLS)}
+
+
+def extract_bert_feature(
+    text: str,
+    word2ph: list[int],
+    language: Languages,
+    device: str,
+    assist_text: Optional[str] = None,
+    assist_text_weight: float = 0.7,
+) -> "torch.Tensor":
+    """
+    テキストから BERT の特徴量を抽出する
+
+    Args:
+        text (str): テキスト
+        word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
+        language (Languages): テキストの言語
+        device (str): 推論に利用するデバイス
+        assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
+        assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
+
+    Returns:
+        torch.Tensor: BERT の特徴量
+    """
+
+    if language == Languages.JP:
+        from style_bert_vits2.nlp.japanese.bert_feature import extract_bert_feature
+    elif language == Languages.EN:
+        from style_bert_vits2.nlp.english.bert_feature import extract_bert_feature
+    elif language == Languages.ZH:
+        from style_bert_vits2.nlp.chinese.bert_feature import extract_bert_feature
+    else:
+        raise ValueError(f"Language {language} not supported")
+
+    return extract_bert_feature(text, word2ph, device, assist_text, assist_text_weight)
+
+
+def clean_text(
+    text: str,
+    language: Languages,
+    use_jp_extra: bool = True,
+    raise_yomi_error: bool = False,
+) -> tuple[str, list[str], list[int], list[int]]:
+    """
+    テキストをクリーニングし、音素に変換する
+
+    Args:
+        text (str): クリーニングするテキスト
+        language (Languages): テキストの言語
+        use_jp_extra (bool, optional): テキストが日本語の場合に JP-Extra モデルを利用するかどうか。Defaults to True.
+        raise_yomi_error (bool, optional): False の場合、読めない文字が消えたような扱いとして処理される。Defaults to False.
+
+    Returns:
+        tuple[str, list[str], list[int], list[int]]: クリーニングされたテキストと、音素・アクセント・元のテキストの各文字に音素が何個割り当てられるかのリスト
+    """
+
+    # Changed to import inside if condition to avoid unnecessary import
+    if language == Languages.JP:
+        from style_bert_vits2.nlp.japanese.g2p import g2p
+        from style_bert_vits2.nlp.japanese.normalizer import normalize_text
+
+        norm_text = normalize_text(text)
+        phones, tones, word2ph = g2p(norm_text, use_jp_extra, raise_yomi_error)
+    elif language == Languages.EN:
+        from style_bert_vits2.nlp.english.g2p import g2p
+        from style_bert_vits2.nlp.english.normalizer import normalize_text
+
+        norm_text = normalize_text(text)
+        phones, tones, word2ph = g2p(norm_text)
+    elif language == Languages.ZH:
+        from style_bert_vits2.nlp.chinese.g2p import g2p
+        from style_bert_vits2.nlp.chinese.normalizer import normalize_text
+
+        norm_text = normalize_text(text)
+        phones, tones, word2ph = g2p(norm_text)
+    else:
+        raise ValueError(f"Language {language} not supported")
+
+    return norm_text, phones, tones, word2ph
+
+
+def cleaned_text_to_sequence(
+    cleaned_phones: list[str], tones: list[int], language: Languages
+) -> tuple[list[int], list[int], list[int]]:
+    """
+    音素リスト・アクセントリスト・言語を、テキスト内の対応する ID に変換する
+
+    Args:
+        cleaned_phones (list[str]): clean_text() でクリーニングされた音素のリスト
+        tones (list[int]): 各音素のアクセント
+        language (Languages): テキストの言語
+
+    Returns:
+        tuple[list[int], list[int], list[int]]: List of integers corresponding to the symbols in the text
+    """
+
+    phones = [__symbol_to_id[symbol] for symbol in cleaned_phones]
+    tone_start = LANGUAGE_TONE_START_MAP[language]
+    tones = [i + tone_start for i in tones]
+    lang_id = LANGUAGE_ID_MAP[language]
+    lang_ids = [lang_id for i in phones]
+
+    return phones, tones, lang_ids