Improve: Add option to generate AIVM/AIVMX files directly when running convert_onnx.py, add license information

convert_bert_onnx.py

@@ -1,5 +1,26 @@
 # Usage: .venv/bin/python convert_bert_onnx.py --language JP
-# ref: https://github.com/tuna2134/sbv2-api/blob/main/convert/convert_deberta.py
+# https://github.com/tuna2134/sbv2-api/blob/main/scripts/convert/convert_deberta.py を参考に実装した
+
+# MIT License
+# Copyright (c) 2024 tuna2134
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
 
 import time
 from argparse import ArgumentParser
@@ -216,9 +237,9 @@ if __name__ == "__main__":
     # モデルの入出力先ファイルパスを取得
     language = Languages(args.language)
     pretrained_model_name_or_path = DEFAULT_BERT_MODEL_PATHS[language]
-    onnx_temp_model_path = Path(pretrained_model_name_or_path) / f"model_temp.onnx"
+    onnx_temp_model_path = Path(pretrained_model_name_or_path) / "model_temp.onnx"
-    onnx_fp32_model_path = Path(pretrained_model_name_or_path) / f"model.onnx"
+    onnx_fp32_model_path = Path(pretrained_model_name_or_path) / "model.onnx"
-    onnx_fp16_model_path = Path(pretrained_model_name_or_path) / f"model_fp16.onnx"
+    onnx_fp16_model_path = Path(pretrained_model_name_or_path) / "model_fp16.onnx"
     tokenizer_json_path = Path(pretrained_model_name_or_path) / "tokenizer.json"
 
     print(Rule(characters="=", style=Style(color="blue")))
@@ -274,7 +295,7 @@ if __name__ == "__main__":
 
     # モデルを ONNX に変換
     print(Rule(characters="=", style=Style(color="blue")))
-    print(f"[bold cyan]Exporting ONNX model...[/bold cyan]")
+    print("[bold cyan]Exporting ONNX model...[/bold cyan]")
     print(Rule(characters="=", style=Style(color="blue")))
     export_start_time = time.time()
     torch.onnx.export(
@@ -304,7 +325,7 @@ if __name__ == "__main__":
 
     # ONNX モデルを最適化
     print(Rule(characters="=", style=Style(color="blue")))
-    print(f"[bold cyan]Optimizing ONNX model...[/bold cyan]")
+    print("[bold cyan]Optimizing ONNX model...[/bold cyan]")
     print(Rule(characters="=", style=Style(color="blue")))
     optimize_start_time = time.time()
     onnx_model = onnx.load(onnx_temp_model_path)
@@ -321,7 +342,7 @@ if __name__ == "__main__":
 
     # FP32 モデルの検証
     print(Rule(characters="=", style=Style(color="blue")))
-    print(f"[bold cyan]Validating FP32 model...[/bold cyan]")
+    print("[bold cyan]Validating FP32 model...[/bold cyan]")
     session = InferenceSession(
         str(onnx_fp32_model_path),
         providers=["CPUExecutionProvider"],
@@ -333,7 +354,7 @@ if __name__ == "__main__":
     if is_valid:
         # FP16 への変換
         print(Rule(characters="=", style=Style(color="blue")))
-        print(f"[bold cyan]Converting to FP16...[/bold cyan]")
+        print("[bold cyan]Converting to FP16...[/bold cyan]")
         print(Rule(characters="=", style=Style(color="blue")))
         fp16_start_time = time.time()
         fp16_model = float16_converter.convert_float_to_float16(
@@ -348,7 +369,7 @@ if __name__ == "__main__":
 
         # FP16 モデルの検証
         print(Rule(characters="=", style=Style(color="blue")))
-        print(f"[bold cyan]Validating FP16 model...[/bold cyan]")
+        print("[bold cyan]Validating FP16 model...[/bold cyan]")
         session = InferenceSession(
             str(onnx_fp16_model_path),
             providers=["CPUExecutionProvider"],

convert_onnx.py

@@ -1,11 +1,34 @@
 # Usage: .venv/bin/python convert_onnx.py --model model_assets/koharune-ami/koharune-ami.safetensors
-# Usage: .venv/bin/python convert_onnx.py --model model_assets/ (All models in the directory will be converted)
+#        .venv/bin/python convert_onnx.py --model model_assets/ (All models in the directory will be converted)
-# ref: https://github.com/tuna2134/sbv2-api/blob/main/convert/convert_model.py
+# https://github.com/tuna2134/sbv2-api/blob/main/scripts/convert/convert_model.py を参考に実装した
 
+# MIT License
+# Copyright (c) 2024 tuna2134
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import re
 import time
+import uuid
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import cast
+from typing import BinaryIO, cast
 
 import onnx
 import torch
@@ -28,6 +51,45 @@ from style_bert_vits2.models.models_jp_extra import (
 from style_bert_vits2.tts_model import TTSModel
 
 
+def generate_aivm_metadata(
+    hyper_parameters_file: BinaryIO,
+    style_vectors_file: BinaryIO,
+    model_file_name: str,
+    model_uuid: uuid.UUID,
+):
+    try:
+        import aivmlib
+        from aivmlib.schemas.aivm_manifest import ModelArchitecture
+    except ImportError:
+        raise ImportError(
+            "aivmlib is not installed. Please install it using `pip install aivmlib`."
+        )
+
+    # AIVM メタデータを生成
+    metadata = aivmlib.generate_aivm_metadata(
+        # 実際に JP-Extra かどうかはハイパーパラメータの値を元に自動判定されるので、ここでは JP-Extra を指定
+        ModelArchitecture.StyleBertVITS2JPExtra,
+        hyper_parameters_file,
+        style_vectors_file,
+    )
+
+    # モデルファイル名からエポック数とステップ数を抽出
+    epoch_match = re.search(r"e(\d{2,})", model_file_name)  # "e" の後ろに2桁以上の数字
+    step_match = re.search(r"s(\d{2,})", model_file_name)  # "s" の後ろに2桁以上の数字
+
+    # エポック数を設定
+    if epoch_match:
+        metadata.manifest.training_epochs = int(epoch_match.group(1))
+    # ステップ数を設定
+    if step_match:
+        metadata.manifest.training_steps = int(step_match.group(1))
+
+    # UUID を設定
+    metadata.manifest.uuid = model_uuid
+
+    return metadata
+
+
 if __name__ == "__main__":
     start_time = time.time()
     parser = ArgumentParser()
@@ -39,6 +101,16 @@ if __name__ == "__main__":
         action="store_true",
         help="Already converted models will be overwritten",
     )
+    parser.add_argument(
+        "--aivm",
+        action="store_true",
+        help="Generate AIVM file from Safetensors model",
+    )
+    parser.add_argument(
+        "--aivmx",
+        action="store_true",
+        help="Generate AIVMX file from ONNX model",
+    )
     args = parser.parse_args()
 
     # --model に指定されたパスがディレクトリの時、配下にある全ての .safetensors ファイルを対象に変換する
@@ -58,6 +130,8 @@ if __name__ == "__main__":
         onnx_optimized_model_path = model_path.parent / f"{model_path.stem}.onnx"
         config_path = model_path.parent / "config.json"
         style_vec_path = model_path.parent / "style_vectors.npy"
+        aivm_path = model_path.parent / f"{model_path.stem}.aivm"
+        aivmx_path = model_path.parent / f"{model_path.stem}.aivmx"
         assert model_path.exists(), "Model file does not exist"
         assert config_path.exists(), "Config file does not exist"
         assert style_vec_path.exists(), "Style vector file does not exist"
@@ -77,8 +151,9 @@ if __name__ == "__main__":
                 "[bold]If you want to overwrite it, use the --force-convert option.[/bold]"
             )
             print(Rule(characters="=", style=Style(color="blue")))
-            continue
 
+        # ONNX モデルが存在しない場合、ONNX モデルを生成
+        else:
             # PyTorch モデルを読み込む
             device = "cpu"
             tts_model = TTSModel(
@@ -88,7 +163,10 @@ if __name__ == "__main__":
                 device=device,
             )
             tts_model.load()
+            if DEFAULT_STYLE in tts_model.style2id:
                 style_id = tts_model.style2id[DEFAULT_STYLE]
+            else:
+                style_id = 0  # 通常デフォルトスタイルのインデックスは 0 になる
             assert tts_model.net_g is not None, "Model is not loaded"
 
             # 音声合成に必要な BERT 特徴量・音素列・アクセント列・言語 ID を取得
@@ -162,7 +240,7 @@ if __name__ == "__main__":
                 # モデルを ONNX に変換
                 print(Rule(characters="=", style=Style(color="blue")))
                 print(
-                f"[bold cyan]Exporting ONNX model... (Architecture: JP-Extra)[/bold cyan]"
+                    "[bold cyan]Exporting ONNX model... (Architecture: JP-Extra)[/bold cyan]"
                 )
                 print(Rule(characters="=", style=Style(color="blue")))
                 export_start_time = time.time()
@@ -253,7 +331,7 @@ if __name__ == "__main__":
                 # モデルを ONNX に変換
                 print(Rule(characters="=", style=Style(color="blue")))
                 print(
-                f"[bold cyan]Exporting ONNX model... (Architecture: Non-JP-Extra)[/bold cyan]"
+                    "[bold cyan]Exporting ONNX model... (Architecture: Non-JP-Extra)[/bold cyan]"
                 )
                 print(Rule(characters="=", style=Style(color="blue")))
                 export_start_time = time.time()
@@ -310,7 +388,7 @@ if __name__ == "__main__":
 
             # ONNX モデルを最適化
             print(Rule(characters="=", style=Style(color="blue")))
-        print(f"[bold cyan]Optimizing ONNX model...[/bold cyan]")
+            print("[bold cyan]Optimizing ONNX model...[/bold cyan]")
             print(Rule(characters="=", style=Style(color="blue")))
             optimize_start_time = time.time()
             onnx_model = onnx.load(onnx_temp_model_path)
@@ -329,3 +407,47 @@ if __name__ == "__main__":
             print("[bold cyan]Optimized model info:[/bold cyan]")
             model_info.print_simplifying_info(onnx_model, simplified_onnx_model)
             print(Rule(characters="=", style=Style(color="blue")))
+
+        # AIVM/AIVMX ファイルを生成
+        if args.aivm or args.aivmx:
+            try:
+                import aivmlib
+            except ImportError:
+                raise ImportError(
+                    "aivmlib is not installed. Please install it using `pip install aivmlib`."
+                )
+
+            # 共通の UUID を生成
+            model_uuid = uuid.uuid4()
+
+            # AIVM メタデータを生成
+            with config_path.open("rb") as hyper_parameters_file:
+                with style_vec_path.open("rb") as style_vectors_file:
+                    aivm_metadata = generate_aivm_metadata(
+                        hyper_parameters_file,
+                        style_vectors_file,
+                        model_path.name,
+                        model_uuid,
+                    )
+
+            # AIVM ファイルを生成
+            if args.aivm and (not aivm_path.exists() or args.force_convert):
+                print("[bold cyan]Generating AIVM file...[/bold cyan]")
+                print(Rule(characters="=", style=Style(color="blue")))
+                with model_path.open("rb") as safetensors_file:
+                    new_aivm_file_content = aivmlib.write_aivm_metadata(safetensors_file, aivm_metadata)  # fmt: skip
+                    with aivm_path.open("wb") as f:
+                        f.write(new_aivm_file_content)
+                print(f"[bold green]Generated AIVM file: {aivm_path}[/bold green]")
+                print(Rule(characters="=", style=Style(color="blue")))
+
+            # AIVMX ファイルを生成
+            if args.aivmx and (not aivmx_path.exists() or args.force_convert):
+                print("[bold cyan]Generating AIVMX file...[/bold cyan]")
+                print(Rule(characters="=", style=Style(color="blue")))
+                with onnx_optimized_model_path.open("rb") as onnx_file:
+                    new_aivmx_file_content = aivmlib.write_aivmx_metadata(onnx_file, aivm_metadata)  # fmt: skip
+                    with aivmx_path.open("wb") as f:
+                        f.write(new_aivmx_file_content)
+                print(f"[bold green]Generated AIVMX file: {aivmx_path}[/bold green]")
+                print(Rule(characters="=", style=Style(color="blue")))

style_bert_vits2/nlp/onnx_bert_models.py

@@ -106,9 +106,7 @@ def load_model(
     # 推論時に一番優先される ExecutionProvider の名前を取得
     assert len(onnx_providers) > 0
     first_provider_name = (
-        onnx_providers[0]
+        onnx_providers[0] if type(onnx_providers[0]) is str else onnx_providers[0][0]
-        if type(onnx_providers[0]) is str
-        else onnx_providers[0][0]
     )
 
     # 推論セッションの設定