Feat: open tensorboard button

docs/CHANGELOG.md

@@ -7,9 +7,11 @@
 - 学習の際のメモリ使用量を削減しようと頑張った
 
 ### バグ修正や改善
+- 学習WebUIからTensorboardのログを見れるように
 - 音声合成（やそのAPI）において、同時に別の話者が選択され音声合成がリクエストされた場合に発生するエラーを修正
 - モデルマージ時に、そのレシピを`recipe.json`ファイルへ保存するように変更
 - 「改行で分けて生成」がより感情が乗る旨の明記等、軽微な説明文の改善
+- 「`ーーそれは面白い`」や「`なるほど。ーーーそういうことか。`」等、長音記号の前が母音でない場合、長音記号`ー`でなくダッシュ`―`の勘違いだと思われるので、ダッシュ記号として処理するように変更
 
 ## v2.0.1 (2024-02-05)
 

webui_train.py

@@ -2,6 +2,11 @@ import argparse
 import json
 import os
 import shutil
+import subprocess
+import socket
+import sys
+import time
+import webbrowser
 from datetime import datetime
 from multiprocessing import cpu_count
 
@@ -10,9 +15,11 @@ import yaml
 
 from common.constants import LATEST_VERSION
 from common.log import logger
+from common.stdout_wrapper import SAFE_STDOUT
 from common.subprocess_utils import run_script_with_log, second_elem_of
 
 logger_handler = None
+tensorboard_executed = False
 
 # Get path settings
 with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
@@ -316,6 +323,46 @@ def train(model_name, skip_style=False, use_jp_extra=True, speedup=False):
     return True, "Success: 学習が完了しました"
 
 
+def wait_for_tensorboard(port=6006, timeout=10):
+    start_time = time.time()
+    while True:
+        try:
+            with socket.create_connection(("localhost", port), timeout=1):
+                return True  # ポートが開いている場合
+        except OSError:
+            pass  # ポートがまだ開いていない場合
+
+        if time.time() - start_time > timeout:
+            return False  # タイムアウト
+
+        time.sleep(0.1)
+
+
+def run_tensorboard(model_name):
+    global tensorboard_executed
+    if not tensorboard_executed:
+        python = sys.executable
+        tensorboard_cmd = [
+            python,
+            "-m",
+            "tensorboard.main",
+            "--logdir",
+            f"Data/{model_name}/models",
+        ]
+        subprocess.Popen(
+            tensorboard_cmd,
+            stdout=SAFE_STDOUT,  # type: ignore
+            stderr=SAFE_STDOUT,  # type: ignore
+        )
+        yield gr.Button("起動中…")
+        if wait_for_tensorboard():
+            tensorboard_executed = True
+        else:
+            logger.error("Tensorboard did not start in the expected time.")
+    webbrowser.open("http://localhost:6006")
+    yield gr.Button("Tensorboardを開く")
+
+
 initial_md = f"""
 # Style-Bert-VITS2 ver {LATEST_VERSION} 学習用WebUI
 
@@ -369,7 +416,7 @@ english_teacher.wav|Mary|EN|How are you? I'm fine, thank you, and you?
 """
 
 if __name__ == "__main__":
-    with gr.Blocks(theme="NoCrypt/miku") as app:
+    with gr.Blocks(theme="NoCrypt/miku").queue() as app:
         gr.Markdown(initial_md)
         with gr.Accordion(label="データの前準備", open=False):
             gr.Markdown(prepare_md)
@@ -548,7 +595,7 @@ if __name__ == "__main__":
                     style_gen_btn = gr.Button(value="実行", variant="primary")
                     info_style = gr.Textbox(label="状況")
         gr.Markdown("## 学習")
-        with gr.Row(variant="panel"):
+        with gr.Row():
             skip_style = gr.Checkbox(
                 label="スタイルファイルの生成をスキップする",
                 info="学習再開の場合の場合はチェックしてください",
@@ -564,7 +611,8 @@ if __name__ == "__main__":
                 visible=False,  # Experimental
             )
             train_btn = gr.Button(value="学習を開始する", variant="primary")
-            info_train = gr.Textbox(label="状況")
+            tensorboard_btn = gr.Button(value="Tensorboardを開く")
+        info_train = gr.Textbox(label="状況")
 
         preprocess_button.click(
             second_elem_of(preprocess_all),
@@ -635,6 +683,10 @@ if __name__ == "__main__":
             inputs=[model_name, skip_style, use_jp_extra_train, speedup],
             outputs=[info_train],
         )
+        tensorboard_btn.click(
+            run_tensorboard, inputs=[model_name], outputs=[tensorboard_btn]
+        )
+
         use_jp_extra.change(
             lambda x: gr.Checkbox(value=x),
             inputs=[use_jp_extra],