Ver 2.7.0 maybe

gradio_tabs/convert_onnx.py

@@ -0,0 +1,101 @@
+from pathlib import Path
+
+import gradio as gr
+
+from style_bert_vits2.constants import GRADIO_THEME
+from style_bert_vits2.logging import logger
+from style_bert_vits2.tts_model import NullModelParam, TTSModelHolder
+from style_bert_vits2.utils.subprocess import run_script_with_log
+
+
+def call_convert_onnx(
+    model: str,
+):
+    if model == "":
+        return "Error: モデル名を入力してください。"
+    logger.info("Start converting model to onnx...")
+    cmd = [
+        "convert_onnx.py",
+        "--model",
+        model,
+    ]
+    success, message = run_script_with_log(cmd, ignore_warning=True)
+    if not success:
+        return f"Error: {message}"
+    return "ONNX変換が完了しました。"
+
+
+initial_md = """
+safetensors形式のモデルをONNX形式に変換します。
+このONNXモデルは、[AIVM Generator](https://aivm-generator.aivis-project.com/) 等でさらにAIVM形式・AIVMX形式に変換して[AivisSpeech](https://aivis-project.com/)で利用できます。
+
+変換には5分以上ほどの時間がかかります。進捗状況はターミナルのログを参照してください。
+"""
+
+
+def create_onnx_app(model_holder: TTSModelHolder) -> gr.Blocks:
+    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_id = 0
+    initial_pth_files = get_model_files(model_names[initial_id])
+
+    with gr.Blocks(theme=GRADIO_THEME) as app:
+        gr.Markdown(initial_md)
+        with gr.Row():
+            with gr.Column():
+                model_name = gr.Dropdown(
+                    label="モデル一覧",
+                    choices=model_names,
+                    value=model_names[initial_id],
+                )
+                model_path = gr.Dropdown(
+                    label="モデルファイル",
+                    choices=initial_pth_files,
+                    value=initial_pth_files[0],
+                )
+            refresh_button = gr.Button("更新")
+        convert_button = gr.Button("ONNX形式に変換", variant="primary")
+        info = gr.Textbox(label="情報")
+
+        model_name.change(
+            model_holder.update_model_files_for_gradio,
+            inputs=[model_name],
+            outputs=[model_path],
+        )
+
+        def refresh_fn() -> tuple[gr.Dropdown, gr.Dropdown]:
+            names, files, _ = model_holder.update_model_names_for_gradio()
+            return names, files
+
+        refresh_button.click(
+            refresh_fn,
+            outputs=[model_name, model_path],
+        )
+        convert_button.click(
+            call_convert_onnx,
+            inputs=[model_path],
+            outputs=[info],
+        )
+
+    return app
+
+
+if __name__ == "__main__":
+    from config import get_path_config
+
+    path_config = get_path_config()
+    assets_root = path_config.assets_root
+    model_holder = TTSModelHolder(assets_root, "cpu", "", ignore_onnx=True)
+    app = create_onnx_app(model_holder)
+    app.launch(inbrowser=True)