Refactor and add one-click training

resample.py

@@ -65,6 +65,8 @@ if __name__ == "__main__":
                 twople = (spk_dir, filename, args)
                 tasks.append(twople)
 
+    if len(tasks) == 0:
+        raise ValueError(f"No wav files found in {args.in_dir}")
     for _ in tqdm(
         pool.imap_unordered(process, tasks), file=sys.stdout, total=len(tasks)
     ):

tools/subprocess_utils.py

@@ -0,0 +1,33 @@
+import subprocess
+import sys
+
+from .log import logger
+
+python = sys.executable
+
+
+def run_script_with_log(cmd: list[str]) -> tuple[bool, str]:
+    logger.info(f"Running: {' '.join(cmd)}")
+    result = subprocess.run(
+        [python] + cmd,
+        stdout=sys.stdout,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    if result.returncode != 0:
+        logger.error(f"Error: {' '.join(cmd)}")
+        print(result.stderr)
+        return False, result.stderr
+    elif result.stderr:
+        logger.warning(f"Warning: {' '.join(cmd)}")
+        print(result.stderr)
+        return True, result.stderr
+    logger.success(f"Success: {' '.join(cmd)}")
+    return True, ""
+
+
+def second_elem_of(original_function):
+    def inner_function(*args, **kwargs):
+        return original_function(*args, **kwargs)[1]
+
+    return inner_function

webui_train.py

@@ -1,22 +1,14 @@
 import json
 import os
 import shutil
-import subprocess
 import sys
+from multiprocessing import Pool, cpu_count
 
 import gradio as gr
 import yaml
 
-python = sys.executable
+from tools.log import logger
-
+from tools.subprocess_utils import run_script_with_log, second_elem_of
-
-def subprocess_wrapper(cmd):
-    return subprocess.run(
-        cmd,
-        stdout=sys.stdout,
-        stderr=subprocess.PIPE,
-        text=True,
-    )
 
 
 def get_path(model_name):
@@ -30,6 +22,7 @@ def get_path(model_name):
 
 
 def initialize(model_name, batch_size, epochs, 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):
         config = json.load(open(config_path, "r", encoding="utf-8"))
@@ -47,9 +40,11 @@ def initialize(model_name, batch_size, epochs, bf16_run):
     try:
         shutil.copytree(src="pretrained", dst=model_path)
     except FileExistsError:
-        return f"Error: モデルフォルダ {model_path} が既に存在します。問題なければ削除してください。"
+        logger.error(f"Step 1: {model_path} already exists.")
+        return False, f"Step1, Error: モデルフォルダ {model_path} が既に存在します。問題なければ削除してください。"
     except FileNotFoundError:
-        return "Error: pretrainedフォルダが見つかりません。"
+        logger.error("Step 1: `pretrained` folder not found.")
+        return False, "Step 1, Error: pretrainedフォルダが見つかりません。"
 
     with open(config_path, "w", encoding="utf-8") as f:
         json.dump(config, f, indent=2)
@@ -62,16 +57,17 @@ def initialize(model_name, batch_size, epochs, bf16_run):
     yml_data["dataset_path"] = dataset_path
     with open("config.yml", "w", encoding="utf-8") as f:
         yaml.dump(yml_data, f, allow_unicode=True)
-    return "Step 1: 初期設定が完了しました"
+    logger.success("Step 1: initialization finished.")
+    return True, "Step 1, Success: 初期設定が完了しました"
 
 
 def resample(model_name):
+    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")
-    result = subprocess_wrapper(
+    success, message = run_script_with_log(
         [
-            python,
             "resample.py",
             "--in_dir",
             in_dir,
@@ -81,14 +77,24 @@ def resample(model_name):
             "44100",
         ]
     )
-    if result.stderr:
+    if not success:
-        return f"{result.stderr}"
+        logger.error(f"Step 2: resampling failed.")
-    return "Step 2: 音声ファイルの前処理が完了しました"
+        return False, f"Step 2, Error: 音声ファイルの前処理に失敗しました:\n{message}"
+    elif message:
+        logger.warning(f"Step 2: resampling finished with stderr.")
+        return True, f"Step 2, Success: 音声ファイルの前処理が完了しました:\n{message}"
+    logger.success("Step 2: resampling finished.")
+    return True, "Step 2, Success: 音声ファイルの前処理が完了しました"
 
 
 def preprocess_text(model_name):
+    logger.info("Step 3: start preprocessing text...")
     dataset_path, lbl_path, train_path, val_path, config_path = get_path(model_name)
-    lines = open(lbl_path, "r", encoding="utf-8").readlines()
+    try:
+        lines = open(lbl_path, "r", encoding="utf-8").readlines()
+    except FileNotFoundError:
+        logger.error(f"Step 3: {lbl_path} not found.")
+        return False, f"Step 3, Error: 書き起こしファイル {lbl_path} が見つかりません。"
     with open(lbl_path, "w", encoding="utf-8") as f:
         for line in lines:
             path, spk, language, text = line.strip().split("|")
@@ -96,9 +102,8 @@ def preprocess_text(model_name):
                 "\\", "/"
             )
             f.writelines(f"{path}|{spk}|{language}|{text}\n")
-    result = subprocess_wrapper(
+    success, message = run_script_with_log(
         [
-            python,
             "preprocess_text.py",
             "--config-path",
             config_path,
@@ -110,25 +115,33 @@ def preprocess_text(model_name):
             val_path,
         ]
     )
-
+    if not success:
-    if result.stderr:
+        logger.error(f"Step 3: preprocessing text failed.")
-        return f"{result.stderr}"
+        return False, f"Step 3, Error: 書き起こしファイルの前処理に失敗しました:\n{message}"
-    return "Step 3: 書き起こしファイルの前処理が完了しました"
+    elif message:
+        logger.warning(f"Step 3: preprocessing text finished with stderr.")
+        return True, f"Step 3, Success: 書き起こしファイルの前処理が完了しました:\n{message}"
+    logger.success("Step 3: preprocessing text finished.")
+    return True, "Step 3, Success: 書き起こしファイルの前処理が完了しました"
 
 
 def bert_gen(model_name):
     _, _, _, _, config_path = get_path(model_name)
-    result = subprocess_wrapper([python, "bert_gen.py", "--config", config_path])
+    success, message = run_script_with_log(["bert_gen.py", "--config", config_path])
-    if result.stderr:
+    if not success:
-        return f"{result.stderr}"
+        logger.error(f"Step 4: bert_gen failed.")
-    return "Step 4: BERT特徴ファイルの生成が完了しました"
+        return False, f"Step 4, Error: BERT特徴ファイルの生成に失敗しました:\n{message}"
+    elif message:
+        logger.warning(f"Step 4: bert_gen finished with stderr.")
+        return True, f"Step 4, Success: BERT特徴ファイルの生成が完了しました:\n{message}"
+    logger.success("Step 4: bert_gen finished.")
+    return True, "Step 4, Success: BERT特徴ファイルの生成が完了しました"
 
 
 def style_gen(model_name, num_processes):
     _, _, _, _, config_path = get_path(model_name)
-    result = subprocess_wrapper(
+    success, message = run_script_with_log(
         [
-            python,
             "style_gen.py",
             "--config",
             config_path,
@@ -136,19 +149,49 @@ def style_gen(model_name, num_processes):
             str(num_processes),
         ]
     )
-    if result.stderr:
+    if not success:
-        return f"{result.stderr}"
+        logger.error(f"Step 5: style_gen failed.")
-    return "Step 5: スタイル特徴ファイルの生成が完了しました"
+        return False, f"Step 5, Error: スタイル特徴ファイルの生成に失敗しました:\n{message}"
+    elif message:
+        logger.warning(f"Step 5: style_gen finished with stderr.")
+        return True, f"Step 5, Success: スタイル特徴ファイルの生成が完了しました:\n{message}"
+    logger.success("Step 5: style_gen finished.")
+    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)
+    if not success:
+        return success, message
+    success, message = resample(model_name)
+    if not success:
+        return success, message
+    success, message = preprocess_text(model_name)
+    if not success:
+        return success, message
+    success, message = bert_gen(model_name)
+    if not success:
+        return success, message
+    success, message = style_gen(model_name, num_processes)
+    if not success:
+        return success, message
+    logger.success("Success: all preprocess finished.")
+    return True, "Success: 全ての前処理が完了しました。ターミナルを確認しておかしいところがないか確認するのをおすすめします。"
 
 
 def train(model_name):
     dataset_path, _, _, _, config_path = get_path(model_name)
-    result = subprocess_wrapper(
+    success, message = run_script_with_log(
-        [python, "train_ms.py", "--config", config_path, "--model", dataset_path]
+        ["train_ms.py", "--config", config_path, "--model", dataset_path]
     )
-    if result.stderr:
+    if not success:
-        return f"{result.stderr}"
+        logger.error(f"Final Step: train failed.")
-    return "Final Step: 学習が完了しました!"
+        return False, f"Final Step, Error: 学習に失敗しました:\n{message}"
+    elif message:
+        logger.warning(f"Final Step: train finished with stderr.")
+        return True, f"Final Step, Success: 学習が完了しました:\n{message}"
+    logger.success("Final Step: train finished.")
+    return True, "Final Step, Success: 学習が完了しました"
 
 
 initial_md = """
@@ -195,11 +238,9 @@ 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__":
@@ -211,66 +252,123 @@ if __name__ == "__main__":
             label="モデル名",
         )
         info = gr.Textbox(label="状況")
+        gr.Markdown("### 自動前処理")
         with gr.Row(variant="panel"):
-            with gr.Column(variant="panel", min_width=160):
+            batch_size = gr.Slider(
-                gr.Markdown(value="Step 1: 設定ファイルの生成")
+                label="バッチサイズ",
-                batch_size = gr.Slider(
+                info="VRAM 12GBで4くらい",
-                    label="バッチサイズ",
+                value=4,
-                    info="VRAM 12GBで4くらい",
+                minimum=1,
-                    value=4,
+                maximum=64,
-                    minimum=1,
+                step=1,
-                    maximum=64,
+            )
-                    step=1,
+            epochs = gr.Slider(
-                )
+                label="エポック数",
-                epochs = gr.Slider(
+                info="100もあれば十分そう",
-                    label="エポック数",
+                value=100,
-                    info="100もあれば十分そう",
+                minimum=1,
-                    value=100,
+                maximum=1000,
-                    minimum=1,
+                step=1,
-                    maximum=1000,
+            )
-                    step=1,
+            bf16_run = gr.Checkbox(
-                )
+                label="bfloat16を使う",
-                bf16_run = gr.Checkbox(
+                info="bfloat16を使うかどうか。新しめのグラボだと学習が早くなるかも、古いグラボだと動かないかも。",
-                    label="bfloat16を使う",
+                value=True,
-                    info="新しめのグラボだと学習が早くなるかも、古いグラボだと動かないかも",
+            )
-                    value=True,
+            num_processes = gr.Slider(
-                )
+                label="プロセス数",
-                generate_config_btn = gr.Button(value="実行", variant="primary")
+                info="前処理時の並列処理プロセス数、大きすぎるとフリーズするかも",
-            with gr.Column(variant="panel", min_width=160):
+                value=cpu_count() // 2,
-                gr.Markdown(value="Step 2: 音声ファイルの前処理")
+                minimum=1,
-                resample_btn = gr.Button(value="実行", variant="primary")
+                maximum=cpu_count(),
-            with gr.Column(variant="panel", min_width=160):
+                step=1,
-                gr.Markdown(value="Step 3: 書き起こしファイルの前処理")
+            )
-                preprocess_text_btn = gr.Button(value="実行", variant="primary")
+            preprocess_button = gr.Button(value="実行", variant="primary")
-            with gr.Column(variant="panel", min_width=160):
+        with gr.Accordion(open=False, label="手動前処理"):
-                gr.Markdown(value="Step 4: BERT特徴ファイルの生成")
+            with gr.Row(variant="panel"):
-                bert_gen_btn = gr.Button(value="実行", variant="primary")
+                with gr.Column(variant="panel", min_width=160):
-            with gr.Column(variant="panel", min_width=160):
+                    gr.Markdown(value="#### Step 1: 設定ファイルの生成")
-                gr.Markdown(value="Step 5: スタイル特徴ファイルの生成")
+                    batch_size_manual = gr.Slider(
-                num_processes = gr.Slider(
+                        label="バッチサイズ",
-                    label="プロセス数",
+                        value=4,
-                    value=4,
+                        minimum=1,
-                    minimum=1,
+                        maximum=64,
-                    maximum=16,
+                        step=1,
-                    step=1,
+                    )
-                )
+                    epochs_manual = gr.Slider(
-                style_gen_btn = gr.Button(value="実行", variant="primary")
+                        label="エポック数",
+                        value=100,
+                        minimum=1,
+                        maximum=1000,
+                        step=1,
+                    )
+                    bf16_run_manual = gr.Checkbox(
+                        label="bfloat16を使う",
+                        value=True,
+                    )
+                    generate_config_btn = gr.Button(value="実行", variant="primary")
+                with gr.Column(variant="panel", min_width=160):
+                    gr.Markdown(value="#### Step 2: 音声ファイルの前処理")
+                    num_processes_resample = gr.Slider(
+                        label="プロセス数",
+                        value=cpu_count() // 2,
+                        minimum=1,
+                        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,
+                    )
+                    bert_gen_btn = gr.Button(value="実行", variant="primary")
+                with gr.Column(variant="panel", min_width=160):
+                    gr.Markdown(value="#### Step 5: スタイル特徴ファイルの生成")
+                    num_processes_style = gr.Slider(
+                        label="プロセス数",
+                        value=cpu_count() // 2,
+                        minimum=1,
+                        maximum=cpu_count(),
+                        step=1,
+                    )
+                    style_gen_btn = gr.Button(value="実行", variant="primary")
         with gr.Row(variant="panel"):
             with gr.Column():
-                gr.Markdown(value="Final Step: 学習")
+                gr.Markdown("## 学習")
-                train_btn = gr.Button(value="学習", variant="primary")
+                train_btn = gr.Button(value="学習を開始する", variant="primary")
 
-        generate_config_btn.click(
+        preprocess_button.click(
-            initialize,
+            second_elem_of(preprocess_all),
-            inputs=[model_name, batch_size, epochs, bf16_run],
+            inputs=[model_name, batch_size, epochs, bf16_run, num_processes],
             outputs=[info],
         )
-        resample_btn.click(resample, inputs=[model_name], outputs=[info])
+        generate_config_btn.click(
-        preprocess_text_btn.click(preprocess_text, inputs=[model_name], outputs=[info])
+            second_elem_of(initialize),
-        bert_gen_btn.click(bert_gen, inputs=[model_name], outputs=[info])
+            inputs=[model_name, batch_size_manual, epochs_manual, bf16_run_manual],
-        style_gen_btn.click(
+            outputs=[info],
-            style_gen, inputs=[model_name, num_processes], outputs=[info]
         )
-        train_btn.click(train, inputs=[model_name], outputs=[info])
+        resample_btn.click(
+            second_elem_of(resample), inputs=[model_name], outputs=[info]
+        )
+        preprocess_text_btn.click(
+            second_elem_of(preprocess_text), inputs=[model_name], outputs=[info]
+        )
+        bert_gen_btn.click(
+            second_elem_of(bert_gen), inputs=[model_name], outputs=[info]
+        )
+        style_gen_btn.click(
+            second_elem_of(style_gen),
+            inputs=[model_name, num_processes_style],
+            outputs=[info],
+        )
+        train_btn.click(second_elem_of(train), inputs=[model_name], outputs=[info])
 
-    app.launch(share=False, inbrowser=True)
+    app.launch(inbrowser=True)