Improve WebUI and add loudness normalization

webui_train.py

@@ -2,7 +2,7 @@ import json
 import os
 import shutil
 import sys
-from multiprocessing import Pool, cpu_count
+from multiprocessing import cpu_count
 
 import gradio as gr
 import yaml
@@ -21,7 +21,7 @@ def get_path(model_name):
     return dataset_path, lbl_path, train_path, val_path, config_path
 
 
-def initialize(model_name, batch_size, epochs, bf16_run):
+def initialize(model_name, batch_size, epochs, save_every_steps, bf16_run):
     logger.info("Step 1: start initialization...")
     dataset_path, _, train_path, val_path, config_path = get_path(model_name)
     if os.path.isfile(config_path):
@@ -35,6 +35,7 @@ def initialize(model_name, batch_size, epochs, bf16_run):
     config["train"]["batch_size"] = batch_size
     config["train"]["epochs"] = epochs
     config["train"]["bf16_run"] = bf16_run
+    config["train"]["eval_interval"] = save_every_steps
 
     model_path = os.path.join(dataset_path, "models")
     try:
@@ -61,22 +62,25 @@ def initialize(model_name, batch_size, epochs, bf16_run):
     return True, "Step 1, Success: 初期設定が完了しました"
 
 
-def resample(model_name):
+def resample(model_name, normalize, num_processes):
     logger.info("Step 2: start resampling...")
     dataset_path, _, _, _, _ = get_path(model_name)
     in_dir = os.path.join(dataset_path, "raw")
     out_dir = os.path.join(dataset_path, "wavs")
-    success, message = run_script_with_log(
-        [
-            "resample.py",
-            "--in_dir",
-            in_dir,
-            "--out_dir",
-            out_dir,
-            "--sr",
-            "44100",
-        ]
-    )
+    cmd = [
+        "resample.py",
+        "--in_dir",
+        in_dir,
+        "--out_dir",
+        out_dir,
+        "--num_processes",
+        str(num_processes),
+        "--sr",
+        "44100",
+    ]
+    if normalize:
+        cmd.append("--normalize")
+    success, message = run_script_with_log(cmd)
     if not success:
         logger.error(f"Step 2: resampling failed.")
         return False, f"Step 2, Error: 音声ファイルの前処理に失敗しました:\n{message}"
@@ -126,8 +130,17 @@ def preprocess_text(model_name):
 
 
 def bert_gen(model_name):
+    logger.info("Step 4: start bert_gen...")
     _, _, _, _, config_path = get_path(model_name)
-    success, message = run_script_with_log(["bert_gen.py", "--config", config_path])
+    success, message = run_script_with_log(
+        [
+            "bert_gen.py",
+            "--config",
+            config_path,
+            # "--num_processes",  # bert_genは重いのでプロセス数いじらない
+            #  str(num_processes),
+        ]
+    )
     if not success:
         logger.error(f"Step 4: bert_gen failed.")
         return False, f"Step 4, Error: BERT特徴ファイルの生成に失敗しました:\n{message}"
@@ -139,6 +152,7 @@ def bert_gen(model_name):
 
 
 def style_gen(model_name, num_processes):
+    logger.info("Step 5: start style_gen...")
     _, _, _, _, config_path = get_path(model_name)
     success, message = run_script_with_log(
         [
@@ -159,23 +173,29 @@ def style_gen(model_name, num_processes):
     return True, "Step 5, Success: スタイル特徴ファイルの生成が完了しました"
 
 
-def preprocess_all(model_name, batch_size, epochs, bf16_run, num_processes):
-    success, message = initialize(model_name, batch_size, epochs, bf16_run)
+def preprocess_all(
+    model_name, batch_size, epochs, save_every_steps, bf16_run, num_processes, normalize
+):
+    if model_name == "":
+        return False, "Error: モデル名を入力してください"
+    success, message = initialize(
+        model_name, batch_size, epochs, save_every_steps, bf16_run
+    )
     if not success:
-        return success, message
-    success, message = resample(model_name)
+        return False, message
+    success, message = resample(model_name, normalize, num_processes)
     if not success:
-        return success, message
+        return False, message
     success, message = preprocess_text(model_name)
     if not success:
-        return success, message
-    success, message = bert_gen(model_name)
+        return False, message
+    success, message = bert_gen(model_name)  # bert_genは重いのでプロセス数いじらない
     if not success:
-        return success, message
+        return False, message
     success, message = style_gen(model_name, num_processes)
     if not success:
-        return success, message
-    logger.success("Success: all preprocess finished.")
+        return False, message
+    logger.success("Success: All preprocess finished!")
     return True, "Success: 全ての前処理が完了しました。ターミナルを確認しておかしいところがないか確認するのをおすすめします。"
 
 
@@ -199,16 +219,20 @@ initial_md = """
 
 ## 使い方
 
-- データを準備して、各ステップを順に実行してください。進捗状況等はターミナルに表示されます。
+- データを準備して、モデル名を入力して、必要なら設定を調整してから、「自動前処理を実行」ボタンを押してください。進捗状況等はターミナルに表示されます。
 
-- 途中から学習を再開する場合は、モデル名を入力してFinal Stepだけ実行すればよいです。
+- 各ステップごとに実行する場合は「手動前処理」を使ってください（基本的には自動でいいはず）。
+
+- 前処理が終わったら、「学習を開始する」ボタンを押すと学習が開始されます。
+
+- 途中から学習を再開する場合は、モデル名を入力してから「学習を開始する」を押せばよいです。
 
 注意: 音声合成で使うには、スタイルベクトルファイル`style_vectors.npy`を作る必要があります。これは、`Style.bat`を実行してそこで作成してください。
 動作は軽いはずなので、学習中でも実行でき、何度でも繰り返して試せます。
 """
 
 prepare_md = """
-まず音声データ（wavファイルで1ファイルが2-15秒程度の、長すぎず短すぎない発話のものをいくつか）と、書き起こしテキストを用意してください。
+まず音声データ（wavファイルで1ファイルが2-12秒程度の、長すぎず短すぎない発話のものをいくつか）と、書き起こしテキストを用意してください。
 
 それを次のように配置します。
 ```
@@ -237,56 +261,64 @@ english_teacher.wav|Mary|EN|How are you? I'm fine, thank you, and you?
 日本語話者の単一話者データセットでも構いません。
 """
 
-css = """
-div.panel {
-    justify-content: space-between;
-}
-"""
-
 if __name__ == "__main__":
-    with gr.Blocks(theme="NoCrypt/miku", css=css) as app:
+    with gr.Blocks(theme="NoCrypt/miku") as app:
         gr.Markdown(initial_md)
         with gr.Accordion(label="データの前準備", open=False):
             gr.Markdown(prepare_md)
         model_name = gr.Textbox(
             label="モデル名",
         )
-        info = gr.Textbox(label="状況")
         gr.Markdown("### 自動前処理")
         with gr.Row(variant="panel"):
-            batch_size = gr.Slider(
-                label="バッチサイズ",
-                info="VRAM 12GBで4くらい",
-                value=4,
-                minimum=1,
-                maximum=64,
-                step=1,
-            )
-            epochs = gr.Slider(
-                label="エポック数",
-                info="100もあれば十分そう",
-                value=100,
-                minimum=1,
-                maximum=1000,
-                step=1,
-            )
-            bf16_run = gr.Checkbox(
-                label="bfloat16を使う",
-                info="bfloat16を使うかどうか。新しめのグラボだと学習が早くなるかも、古いグラボだと動かないかも。",
-                value=True,
-            )
-            num_processes = gr.Slider(
-                label="プロセス数",
-                info="前処理時の並列処理プロセス数、大きすぎるとフリーズするかも",
-                value=cpu_count() // 2,
-                minimum=1,
-                maximum=cpu_count(),
-                step=1,
-            )
-            preprocess_button = gr.Button(value="実行", variant="primary")
+            with gr.Column():
+                batch_size = gr.Slider(
+                    label="バッチサイズ",
+                    info="VRAM 12GBで4くらい",
+                    value=4,
+                    minimum=1,
+                    maximum=64,
+                    step=1,
+                )
+                epochs = gr.Slider(
+                    label="エポック数",
+                    info="100もあれば十分そう",
+                    value=100,
+                    minimum=1,
+                    maximum=1000,
+                    step=1,
+                )
+                save_every_steps = gr.Slider(
+                    label="何ステップごとに結果を保存するか",
+                    info="エポック数とは違うことに注意",
+                    value=1000,
+                    minimum=100,
+                    maximum=10000,
+                    step=100,
+                )
+                bf16_run = gr.Checkbox(
+                    label="bfloat16を使う",
+                    info="bfloat16を使うかどうか。新しめのグラボだと学習が早くなるかも、古いグラボだと動かないかも。",
+                    value=True,
+                )
+                num_processes = gr.Slider(
+                    label="プロセス数",
+                    info="前処理時の並列処理プロセス数、大きすぎるとフリーズするかも",
+                    value=cpu_count() // 2,
+                    minimum=1,
+                    maximum=cpu_count(),
+                    step=1,
+                )
+                normalize = gr.Checkbox(
+                    label="音声の音量を正規化する",
+                    value=True,
+                )
+            with gr.Column():
+                preprocess_button = gr.Button(value="自動前処理を実行", variant="primary")
+                info_all = gr.Textbox(label="状況")
         with gr.Accordion(open=False, label="手動前処理"):
             with gr.Row(variant="panel"):
-                with gr.Column(variant="panel", min_width=160):
+                with gr.Column():
                     gr.Markdown(value="#### Step 1: 設定ファイルの生成")
                     batch_size_manual = gr.Slider(
                         label="バッチサイズ",
@@ -302,12 +334,22 @@ if __name__ == "__main__":
                         maximum=1000,
                         step=1,
                     )
+                    save_every_steps_manual = gr.Slider(
+                        label="何ステップごとに結果を保存するか",
+                        value=1000,
+                        minimum=100,
+                        maximum=10000,
+                        step=100,
+                    )
                     bf16_run_manual = gr.Checkbox(
                         label="bfloat16を使う",
                         value=True,
                     )
+                with gr.Column():
                     generate_config_btn = gr.Button(value="実行", variant="primary")
-                with gr.Column(variant="panel", min_width=160):
+                    info_init = gr.Textbox(label="状況")
+            with gr.Row(variant="panel"):
+                with gr.Column():
                     gr.Markdown(value="#### Step 2: 音声ファイルの前処理")
                     num_processes_resample = gr.Slider(
                         label="プロセス数",
@@ -316,21 +358,35 @@ if __name__ == "__main__":
                         maximum=cpu_count(),
                         step=1,
                     )
-                    resample_btn = gr.Button(value="実行", variant="primary")
-                with gr.Column(variant="panel", min_width=160):
-                    gr.Markdown(value="#### Step 3: 書き起こしファイルの前処理")
-                    preprocess_text_btn = gr.Button(value="実行", variant="primary")
-                with gr.Column(variant="panel", min_width=160):
-                    gr.Markdown(value="#### Step 4: BERT特徴ファイルの生成")
-                    num_processes_bert = gr.Slider(
-                        label="プロセス数",
-                        value=cpu_count() // 2,
-                        minimum=1,
-                        maximum=cpu_count(),
-                        step=1,
+                    normalize_resample = gr.Checkbox(
+                        label="音声の音量を正規化する",
+                        value=True,
                     )
+                with gr.Column():
+                    resample_btn = gr.Button(value="実行", variant="primary")
+                    info_resample = gr.Textbox(label="状況")
+            with gr.Row(variant="panel"):
+                with gr.Column():
+                    gr.Markdown(value="#### Step 3: 書き起こしファイルの前処理")
+                with gr.Column():
+                    preprocess_text_btn = gr.Button(value="実行", variant="primary")
+                    info_preprocess_text = gr.Textbox(label="状況")
+            with gr.Row(variant="panel"):
+                with gr.Column():
+                    gr.Markdown(value="#### Step 4: BERT特徴ファイルの生成")
+                    # num_processes_bert = gr.Slider(
+                    #     label="プロセス数",
+                    #     value=cpu_count() // 2,
+                    #     minimum=1,
+                    #     maximum=cpu_count(),
+                    #     step=1,
+                    # )
+                    # bert_genは重いようで、初期の4がよいみたい
+                with gr.Column():
                     bert_gen_btn = gr.Button(value="実行", variant="primary")
-                with gr.Column(variant="panel", min_width=160):
+                    info_bert = gr.Textbox(label="状況")
+            with gr.Row(variant="panel"):
+                with gr.Column():
                     gr.Markdown(value="#### Step 5: スタイル特徴ファイルの生成")
                     num_processes_style = gr.Slider(
                         label="プロセス数",
@@ -339,36 +395,60 @@ if __name__ == "__main__":
                         maximum=cpu_count(),
                         step=1,
                     )
+                with gr.Column():
                     style_gen_btn = gr.Button(value="実行", variant="primary")
+                    info_style = gr.Textbox(label="状況")
+        gr.Markdown("## 学習")
         with gr.Row(variant="panel"):
-            with gr.Column():
-                gr.Markdown("## 学習")
-                train_btn = gr.Button(value="学習を開始する", variant="primary")
+            train_btn = gr.Button(value="学習を開始する", variant="primary")
+            info_train = gr.Textbox(label="状況")
 
         preprocess_button.click(
             second_elem_of(preprocess_all),
-            inputs=[model_name, batch_size, epochs, bf16_run, num_processes],
-            outputs=[info],
+            inputs=[
+                model_name,
+                batch_size,
+                epochs,
+                save_every_steps,
+                bf16_run,
+                num_processes,
+                normalize,
+            ],
+            outputs=[info_all],
         )
         generate_config_btn.click(
             second_elem_of(initialize),
-            inputs=[model_name, batch_size_manual, epochs_manual, bf16_run_manual],
-            outputs=[info],
+            inputs=[
+                model_name,
+                batch_size_manual,
+                epochs_manual,
+                save_every_steps_manual,
+                bf16_run_manual,
+            ],
+            outputs=[info_init],
         )
         resample_btn.click(
-            second_elem_of(resample), inputs=[model_name], outputs=[info]
+            second_elem_of(resample),
+            inputs=[model_name, normalize_resample, num_processes_resample],
+            outputs=[info_resample],
         )
         preprocess_text_btn.click(
-            second_elem_of(preprocess_text), inputs=[model_name], outputs=[info]
+            second_elem_of(preprocess_text),
+            inputs=[model_name],
+            outputs=[info_preprocess_text],
         )
         bert_gen_btn.click(
-            second_elem_of(bert_gen), inputs=[model_name], outputs=[info]
+            second_elem_of(bert_gen),
+            inputs=[model_name],
+            outputs=[info_bert],
         )
         style_gen_btn.click(
             second_elem_of(style_gen),
             inputs=[model_name, num_processes_style],
-            outputs=[info],
+            outputs=[info_style],
+        )
+        train_btn.click(
+            second_elem_of(train), inputs=[model_name], outputs=[info_train]
         )
-        train_btn.click(second_elem_of(train), inputs=[model_name], outputs=[info])
 
     app.launch(inbrowser=True)