tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
tsukumi
2025-03-26 20:01:17 +09:00
parent 810ca43615
commit 2aa503233a
3 changed files with 396 additions and 255 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
convert_bert_onnx.py
View File

@@ -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"],

608
convert_onnx.py
View File

@@ -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,255 +151,303 @@ if __name__ == "__main__":
"[bold]If you want to overwrite it, use the --force-convert option.[/bold]"
)
print(Rule(characters="=", style=Style(color="blue")))
continue
# PyTorch モデルを読み込む
device = "cpu"
tts_model = TTSModel(
model_path=model_path,
config_path=config_path,
style_vec_path=style_vec_path,
device=device,
)
tts_model.load()
style_id = tts_model.style2id[DEFAULT_STYLE]
assert tts_model.net_g is not None, "Model is not loaded"
# 音声合成に必要な BERT 特徴量・音素列・アクセント列・言語 ID を取得
# JP-Extra モデルアーキテクチャの場合、bert (中国語の BERT 特徴量) や en_bert (英語の BERT 特徴量) は
# torch.zeros() で適当に埋められており、推論には ja_bert (日本語の BERT 特徴量) のみが使用される
bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
"今日はいい天気ですね。",
Languages.JP,
tts_model.hyper_parameters,
device,
assist_text=None,
assist_text_weight=DEFAULT_ASSIST_TEXT_WEIGHT,
given_phone=None,
given_tone=None,
)
# スタイルベクトルを取得
style_vector = tts_model.get_style_vector(style_id, DEFAULT_STYLE_WEIGHT)
# モデルの入力を作成
x_tst = phones.to(device).unsqueeze(0)
tones = tones.to(device).unsqueeze(0)
lang_ids = lang_ids.to(device).unsqueeze(0)
bert = bert.to(device).unsqueeze(0)
ja_bert = ja_bert.to(device).unsqueeze(0)
en_bert = en_bert.to(device).unsqueeze(0)
x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
style_vec_tensor = torch.from_numpy(style_vector).to(device).unsqueeze(0)
sid = 0
sid_tensor = torch.LongTensor([sid]).to(device)
length_scale = torch.tensor(1.0)
sdp_ratio = torch.tensor(0.0)
noise_scale = torch.tensor(0.667)
noise_scale_w = torch.tensor(0.8)
# JP-Extra モデルアーキテクチャ向けの ONNX 変換ロジック
if isinstance(tts_model.net_g, SynthesizerTrnJPExtra):
# SynthesizerTrnJPExtra の forward メソッドをオーバーライド
def forward_jp_extra(
x: torch.Tensor,
x_lengths: torch.Tensor,
sid: torch.Tensor,
tone: torch.Tensor,
language: torch.Tensor,
bert: torch.Tensor,
style_vec: torch.Tensor,
length_scale: float = 1.0,
sdp_ratio: float = 0.0,
noise_scale: float = 0.667,
noise_scale_w: float = 0.8,
) -> tuple[
torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]
]:
return cast(SynthesizerTrnJPExtra, tts_model.net_g).infer(
x,
x_lengths,
sid,
tone,
language,
bert,
style_vec,
length_scale=length_scale,
sdp_ratio=sdp_ratio,
noise_scale=noise_scale,
noise_scale_w=noise_scale_w,
)
tts_model.net_g.forward = forward_jp_extra # type: ignore
# モデルを ONNX に変換
print(Rule(characters="=", style=Style(color="blue")))
print(
f"[bold cyan]Exporting ONNX model... (Architecture: JP-Extra)[/bold cyan]"
)
print(Rule(characters="=", style=Style(color="blue")))
export_start_time = time.time()
torch.onnx.export(
model=tts_model.net_g,
args=(
x_tst,
x_tst_lengths,
sid_tensor,
tones,
lang_ids,
ja_bert,
style_vec_tensor,
length_scale,
sdp_ratio,
noise_scale,
noise_scale_w,
),
f=str(onnx_temp_model_path),
verbose=False,
input_names=[
"x_tst",
"x_tst_lengths",
"sid",
"tones",
"language",
"bert",
"style_vec",
"length_scale",
"sdp_ratio",
"noise_scale",
"noise_scale_w",
],
output_names=["output"],
dynamic_axes={
"x_tst": {0: "batch_size", 1: "x_tst_max_length"},
"x_tst_lengths": {0: "batch_size"},
"sid": {0: "batch_size"},
"tones": {0: "batch_size", 1: "x_tst_max_length"},
"language": {0: "batch_size", 1: "x_tst_max_length"},
"bert": {0: "batch_size", 2: "x_tst_max_length"},
"style_vec": {0: "batch_size"},
},
)
print(
f"[bold green]ONNX model exported to {onnx_temp_model_path} ({time.time() - export_start_time:.2f}s)[/bold green]"
)
# 非 JP-Extra モデルアーキテクチャ向けの ONNX 変換ロジック
# ONNX モデルが存在しない場合、ONNX モデルを生成
else:
# PyTorch モデルを読み込む
device = "cpu"
tts_model = TTSModel(
model_path=model_path,
config_path=config_path,
style_vec_path=style_vec_path,
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"
# SynthesizerTrn の forward メソッドをオーバーライド
def forward_non_jp_extra(
x: torch.Tensor,
x_lengths: torch.Tensor,
sid: torch.Tensor,
tone: torch.Tensor,
language: torch.Tensor,
bert: torch.Tensor,
ja_bert: torch.Tensor,
en_bert: torch.Tensor,
style_vec: torch.Tensor,
length_scale: float = 1.0,
sdp_ratio: float = 0.0,
noise_scale: float = 0.667,
noise_scale_w: float = 0.8,
) -> tuple[
torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]
]:
return cast(SynthesizerTrn, tts_model.net_g).infer(
x,
x_lengths,
sid,
tone,
language,
bert,
ja_bert,
en_bert,
style_vec,
length_scale=length_scale,
sdp_ratio=sdp_ratio,
noise_scale=noise_scale,
noise_scale_w=noise_scale_w,
# 音声合成に必要な BERT 特徴量・音素列・アクセント列・言語 ID を取得
# JP-Extra モデルアーキテクチャの場合、bert (中国語の BERT 特徴量) や en_bert (英語の BERT 特徴量) は
# torch.zeros() で適当に埋められており、推論には ja_bert (日本語の BERT 特徴量) のみが使用される
bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
"今日はいい天気ですね。",
Languages.JP,
tts_model.hyper_parameters,
device,
assist_text=None,
assist_text_weight=DEFAULT_ASSIST_TEXT_WEIGHT,
given_phone=None,
given_tone=None,
)
# スタイルベクトルを取得
style_vector = tts_model.get_style_vector(style_id, DEFAULT_STYLE_WEIGHT)
# モデルの入力を作成
x_tst = phones.to(device).unsqueeze(0)
tones = tones.to(device).unsqueeze(0)
lang_ids = lang_ids.to(device).unsqueeze(0)
bert = bert.to(device).unsqueeze(0)
ja_bert = ja_bert.to(device).unsqueeze(0)
en_bert = en_bert.to(device).unsqueeze(0)
x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
style_vec_tensor = torch.from_numpy(style_vector).to(device).unsqueeze(0)
sid = 0
sid_tensor = torch.LongTensor([sid]).to(device)
length_scale = torch.tensor(1.0)
sdp_ratio = torch.tensor(0.0)
noise_scale = torch.tensor(0.667)
noise_scale_w = torch.tensor(0.8)
# JP-Extra モデルアーキテクチャ向けの ONNX 変換ロジック
if isinstance(tts_model.net_g, SynthesizerTrnJPExtra):
# SynthesizerTrnJPExtra の forward メソッドをオーバーライド
def forward_jp_extra(
x: torch.Tensor,
x_lengths: torch.Tensor,
sid: torch.Tensor,
tone: torch.Tensor,
language: torch.Tensor,
bert: torch.Tensor,
style_vec: torch.Tensor,
length_scale: float = 1.0,
sdp_ratio: float = 0.0,
noise_scale: float = 0.667,
noise_scale_w: float = 0.8,
) -> tuple[
torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]
]:
return cast(SynthesizerTrnJPExtra, tts_model.net_g).infer(
x,
x_lengths,
sid,
tone,
language,
bert,
style_vec,
length_scale=length_scale,
sdp_ratio=sdp_ratio,
noise_scale=noise_scale,
noise_scale_w=noise_scale_w,
)
tts_model.net_g.forward = forward_jp_extra # type: ignore
# モデルを ONNX に変換
print(Rule(characters="=", style=Style(color="blue")))
print(
"[bold cyan]Exporting ONNX model... (Architecture: JP-Extra)[/bold cyan]"
)
print(Rule(characters="=", style=Style(color="blue")))
export_start_time = time.time()
torch.onnx.export(
model=tts_model.net_g,
args=(
x_tst,
x_tst_lengths,
sid_tensor,
tones,
lang_ids,
ja_bert,
style_vec_tensor,
length_scale,
sdp_ratio,
noise_scale,
noise_scale_w,
),
f=str(onnx_temp_model_path),
verbose=False,
input_names=[
"x_tst",
"x_tst_lengths",
"sid",
"tones",
"language",
"bert",
"style_vec",
"length_scale",
"sdp_ratio",
"noise_scale",
"noise_scale_w",
],
output_names=["output"],
dynamic_axes={
"x_tst": {0: "batch_size", 1: "x_tst_max_length"},
"x_tst_lengths": {0: "batch_size"},
"sid": {0: "batch_size"},
"tones": {0: "batch_size", 1: "x_tst_max_length"},
"language": {0: "batch_size", 1: "x_tst_max_length"},
"bert": {0: "batch_size", 2: "x_tst_max_length"},
"style_vec": {0: "batch_size"},
},
)
print(
f"[bold green]ONNX model exported to {onnx_temp_model_path} ({time.time() - export_start_time:.2f}s)[/bold green]"
)
tts_model.net_g.forward = forward_non_jp_extra # type: ignore
# 非 JP-Extra モデルアーキテクチャ向けの ONNX 変換ロジック
else:
# モデルを ONNX に変換
# SynthesizerTrn の forward メソッドをオーバーライド
def forward_non_jp_extra(
x: torch.Tensor,
x_lengths: torch.Tensor,
sid: torch.Tensor,
tone: torch.Tensor,
language: torch.Tensor,
bert: torch.Tensor,
ja_bert: torch.Tensor,
en_bert: torch.Tensor,
style_vec: torch.Tensor,
length_scale: float = 1.0,
sdp_ratio: float = 0.0,
noise_scale: float = 0.667,
noise_scale_w: float = 0.8,
) -> tuple[
torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]
]:
return cast(SynthesizerTrn, tts_model.net_g).infer(
x,
x_lengths,
sid,
tone,
language,
bert,
ja_bert,
en_bert,
style_vec,
length_scale=length_scale,
sdp_ratio=sdp_ratio,
noise_scale=noise_scale,
noise_scale_w=noise_scale_w,
)
tts_model.net_g.forward = forward_non_jp_extra # type: ignore
# モデルを ONNX に変換
print(Rule(characters="=", style=Style(color="blue")))
print(
"[bold cyan]Exporting ONNX model... (Architecture: Non-JP-Extra)[/bold cyan]"
)
print(Rule(characters="=", style=Style(color="blue")))
export_start_time = time.time()
torch.onnx.export(
model=tts_model.net_g,
args=(
x_tst,
x_tst_lengths,
sid_tensor,
tones,
lang_ids,
bert,
ja_bert,
en_bert,
style_vec_tensor,
length_scale,
sdp_ratio,
noise_scale,
noise_scale_w,
),
f=str(onnx_temp_model_path),
verbose=False,
input_names=[
"x_tst",
"x_tst_lengths",
"sid",
"tones",
"language",
"bert",
"ja_bert",
"en_bert",
"style_vec",
"length_scale",
"sdp_ratio",
"noise_scale",
"noise_scale_w",
],
output_names=["output"],
dynamic_axes={
"x_tst": {0: "batch_size", 1: "x_tst_max_length"},
"x_tst_lengths": {0: "batch_size"},
"sid": {0: "batch_size"},
"tones": {0: "batch_size", 1: "x_tst_max_length"},
"language": {0: "batch_size", 1: "x_tst_max_length"},
"bert": {0: "batch_size", 2: "x_tst_max_length"},
"ja_bert": {0: "batch_size", 2: "x_tst_max_length"},
"en_bert": {0: "batch_size", 2: "x_tst_max_length"},
"style_vec": {0: "batch_size"},
},
)
print(
f"[bold green]ONNX model exported to {onnx_temp_model_path} ({time.time() - export_start_time:.2f}s)[/bold green]"
)
# ONNX モデルを最適化
print(Rule(characters="=", style=Style(color="blue")))
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)
simplified_onnx_model, check = simplify(onnx_model)
onnx.save(simplified_onnx_model, onnx_optimized_model_path)
print(
f"[bold cyan]Exporting ONNX model... (Architecture: Non-JP-Extra)[/bold cyan]"
)
print(Rule(characters="=", style=Style(color="blue")))
export_start_time = time.time()
torch.onnx.export(
model=tts_model.net_g,
args=(
x_tst,
x_tst_lengths,
sid_tensor,
tones,
lang_ids,
bert,
ja_bert,
en_bert,
style_vec_tensor,
length_scale,
sdp_ratio,
noise_scale,
noise_scale_w,
),
f=str(onnx_temp_model_path),
verbose=False,
input_names=[
"x_tst",
"x_tst_lengths",
"sid",
"tones",
"language",
"bert",
"ja_bert",
"en_bert",
"style_vec",
"length_scale",
"sdp_ratio",
"noise_scale",
"noise_scale_w",
],
output_names=["output"],
dynamic_axes={
"x_tst": {0: "batch_size", 1: "x_tst_max_length"},
"x_tst_lengths": {0: "batch_size"},
"sid": {0: "batch_size"},
"tones": {0: "batch_size", 1: "x_tst_max_length"},
"language": {0: "batch_size", 1: "x_tst_max_length"},
"bert": {0: "batch_size", 2: "x_tst_max_length"},
"ja_bert": {0: "batch_size", 2: "x_tst_max_length"},
"en_bert": {0: "batch_size", 2: "x_tst_max_length"},
"style_vec": {0: "batch_size"},
},
f"[bold green]ONNX model optimized and saved to {onnx_optimized_model_path} ({time.time() - optimize_start_time:.2f}s)[/bold green]"
)
print(
f"[bold green]ONNX model exported to {onnx_temp_model_path} ({time.time() - export_start_time:.2f}s)[/bold green]"
f"[bold]Total Time: {time.time() - start_time:.2f}s / "
f"Size: {onnx_temp_model_path.stat().st_size / 1000 / 1000:.2f}MB -> "
f"{onnx_optimized_model_path.stat().st_size / 1000 / 1000:.2f}MB[/bold]"
)
onnx_temp_model_path.unlink()
print(Rule(characters="=", style=Style(color="blue")))
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")))
# ONNX モデルを最適化
print(Rule(characters="=", style=Style(color="blue")))
print(f"[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)
simplified_onnx_model, check = simplify(onnx_model)
onnx.save(simplified_onnx_model, onnx_optimized_model_path)
print(
f"[bold green]ONNX model optimized and saved to {onnx_optimized_model_path} ({time.time() - optimize_start_time:.2f}s)[/bold green]"
)
print(
f"[bold]Total Time: {time.time() - start_time:.2f}s / "
f"Size: {onnx_temp_model_path.stat().st_size / 1000 / 1000:.2f}MB -> "
f"{onnx_optimized_model_path.stat().st_size / 1000 / 1000:.2f}MB[/bold]"
)
onnx_temp_model_path.unlink()
print(Rule(characters="=", style=Style(color="blue")))
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")))

4
style_bert_vits2/nlp/onnx_bert_models.py
View File

@@ -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]
)
# 推論セッションの設定