diff --git a/scripts/convert/README.md b/scripts/convert/README.md index 3dbf86d..60dc1c1 100644 --- a/scripts/convert/README.md +++ b/scripts/convert/README.md @@ -25,6 +25,24 @@ python convert_model.py --style_file "ここにstyle_vectors.npyの場所" --config_file "同様にconfig.json場所" --model_file "同様に.safetensorsで終わるファイルの場所" ``` +TensorRT用のモデルを生成する場合は `--tensorrt` を追加します。 + +```sh +python convert_model.py --tensorrt --style_file "style_vectors.npy" --config_file "config.json" --model_file "model.safetensors" +``` + +TensorRTモードでは、非対応Opを含む確率的 Duration Predictor の代わりに決定論的 Duration Predictor を使用します。また、`RandomNormalLike` はTensorRT対応の要素演算による決定論的ノイズへ置換されます。 + +### trtexecで検証 + +変換したONNXからTensorRTエンジンをビルドし、実データ型のダミー入力で1回推論します。 + +```sh +uv run python test_tensorrt.py ../../models/model_名前.onnx --fp16 +``` + +`trtexec` がPATHにない場合は `--trtexec /path/to/trtexec` または環境変数 `TRTEXEC` で指定します。デフォルトの動的長プロファイルは1～256、最適値とテスト値は32です。必要に応じて `--min-length`、`--opt-length`、`--max-length`、`--test-length` を指定できます。追加の `trtexec` 引数は `--` の後ろへ渡します。 + 4. `models/名前.sbv2`というファイルが出力されます。GUI版のモデルファイルに入れてあげたら使えます。 ## Deberta変換 diff --git a/scripts/convert/convert_model.py b/scripts/convert/convert_model.py index 689809b..d054f83 100644 --- a/scripts/convert/convert_model.py +++ b/scripts/convert/convert_model.py @@ -4,6 +4,7 @@ from io import BytesIO from style_bert_vits2.nlp import bert_models from style_bert_vits2.constants import Languages from style_bert_vits2.models.infer import get_net_g, get_text +from style_bert_vits2.models import commons from style_bert_vits2.models.hyper_parameters import HyperParameters import torch from style_bert_vits2.constants import ( @@ -23,6 +24,11 @@ parser = ArgumentParser() parser.add_argument("--style_file", required=True) parser.add_argument("--config_file", required=True) parser.add_argument("--model_file", required=True) +parser.add_argument( + "--tensorrt", + action="store_true", + help="Export a TensorRT-compatible deterministic inference graph.", +) args = parser.parse_args() style_file = args.style_file config_file = args.config_file @@ -94,20 +100,70 @@ model = get_net_g( ) -def forward(x, x_len, sid, tone, lang, bert, style, length_scale, sdp_ratio, noise_scale, noise_scale_w): - return model.infer( - x, - x_len, - sid, - tone, - lang, - bert, - style, - sdp_ratio=sdp_ratio, - length_scale=length_scale, - noise_scale=noise_scale, - noise_scale_w=noise_scale_w, +def _tensorrt_noise_like(value, phase): + """Generate deterministic noise using TensorRT-supported elementwise ops.""" + indices = torch.arange( + value.numel(), dtype=value.dtype, device=value.device + ).reshape_as(value) + # A sum of independent-looking sinusoids has approximately normal variance. + return ( + torch.sin(indices * 12.9898 + phase) + + torch.sin(indices * 78.233 + phase * 1.37) + + torch.sin(indices * 37.719 + phase * 2.11) + ) * 0.816496580927726 + + +def forward( + x, + x_len, + sid, + tone, + lang, + bert, + style, + length_scale, + sdp_ratio, + noise_scale, + noise_scale_w, +): + if not args.tensorrt: + return model.infer( + x, + x_len, + sid, + tone, + lang, + bert, + style, + sdp_ratio=sdp_ratio, + length_scale=length_scale, + noise_scale=noise_scale, + noise_scale_w=noise_scale_w, + ) + + # The stochastic duration predictor contains NonZero and ScatterND, and + # torch.randn_like exports as RandomNormalLike. Those operators cannot be + # compiled by TensorRT. Use the deterministic duration predictor and create + # latent noise from elementwise operations instead. + g = model.emb_g(sid).unsqueeze(-1) + hidden, m_p, logs_p, x_mask = model.enc_p( + x, x_len, tone, lang, bert, style, g=g ) + logw = model.dp(hidden, x_mask, g=g) + w = torch.exp(logw) * x_mask * length_scale + w_ceil = torch.ceil(w) + y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long() + y_mask = commons.sequence_mask(y_lengths, None).unsqueeze(1).to(x_mask.dtype) + attn_mask = x_mask.unsqueeze(2) * y_mask.unsqueeze(-1) + attn = commons.generate_path(w_ceil, attn_mask) + m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) + logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) + # Keep both legacy controls as part of the deterministic phase so the ONNX + # input contract remains compatible with existing runtimes. + phase = noise_scale_w + sdp_ratio + z_p = m_p + _tensorrt_noise_like(m_p, phase) * torch.exp(logs_p) * noise_scale + z = model.flow(z_p, y_mask, g=g, reverse=True) + return model.dec(z * y_mask, g=g) model.forward = forward @@ -128,6 +184,9 @@ torch.onnx.export( torch.tensor(0.8), ), f"../../models/model_{out_name}.onnx", + # Preserve the established Style-Bert-VITS2 ONNX graph. The dynamo exporter + # (the default in recent PyTorch) cannot lower the complete inference path. + dynamo=False, verbose=True, dynamic_axes={ "x_tst": {0: "batch_size", 1: "x_tst_max_length"}, @@ -153,7 +212,20 @@ torch.onnx.export( ], output_names=["output"], ) -os.system(f"onnxsim ../../models/model_{out_name}.onnx ../../models/model_{out_name}.onnx") +onnx_path = f"../../models/model_{out_name}.onnx" +os.system(f"onnxsim {onnx_path} {onnx_path}") + +if args.tensorrt: + import onnx + + unsupported_ops = {"NonZero", "RandomNormalLike", "ScatterND"} + graph_ops = {node.op_type for node in onnx.load(onnx_path).graph.node} + remaining_ops = unsupported_ops & graph_ops + if remaining_ops: + raise RuntimeError( + "TensorRT-incompatible operators remain after export: " + + ", ".join(sorted(remaining_ops)) + ) onnxfile = open(f"../../models/model_{out_name}.onnx", "rb").read() stylefile = open(f"../../models/style_vectors_{out_name}.json", "rb").read() version = bytes("1", "utf8")