Improve: Add option to generate AIVM/AIVMX files directly when running convert_onnx.py, add license information
This commit is contained in:
tsukumi
@@ -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_fp32_model_path = Path(pretrained_model_name_or_path) / f"model.onnx"
|
||||
onnx_fp16_model_path = Path(pretrained_model_name_or_path) / f"model_fp16.onnx"
|
||||
onnx_temp_model_path = Path(pretrained_model_name_or_path) / "model_temp.onnx"
|
||||
onnx_fp32_model_path = Path(pretrained_model_name_or_path) / "model.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"],
|
||||
|
||||
136
convert_onnx.py
136
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)
|
||||
# ref: https://github.com/tuna2134/sbv2-api/blob/main/convert/convert_model.py
|
||||
# .venv/bin/python convert_onnx.py --model model_assets/ (All models in the directory will be converted)
|
||||
# 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")))
|
||||
|
||||
@@ -106,9 +106,7 @@ def load_model(
|
||||
# 推論時に一番優先される ExecutionProvider の名前を取得
|
||||
assert len(onnx_providers) > 0
|
||||
first_provider_name = (
|
||||
onnx_providers[0]
|
||||
if type(onnx_providers[0]) is str
|
||||
else onnx_providers[0][0]
|
||||
onnx_providers[0] if type(onnx_providers[0]) is str else onnx_providers[0][0]
|
||||
)
|
||||
|
||||
# 推論セッションの設定
|
||||
|
||||
Reference in New Issue
Block a user