diff --git a/playground/export_bert.py b/GPT_SoVITS/export_roberta_onnx.py similarity index 91% rename from playground/export_bert.py rename to GPT_SoVITS/export_roberta_onnx.py index b8d413c9..860d1af2 100644 --- a/playground/export_bert.py +++ b/GPT_SoVITS/export_roberta_onnx.py @@ -8,6 +8,8 @@ import argparse import os import shutil import numpy as np +import onnxsim +import onnx class CombinedBERTModel(nn.Module): """Wrapper class that combines BERT tokenizer preprocessing and model inference.""" @@ -58,10 +60,16 @@ def export_bert_to_onnx( output_names=['logits'], dynamic_axes={ 'input_ids': {0: 'batch_size', 1: 'sequence_length'}, - 'logits': {0: 'batch_size', 1: 'sequence_length'} + 'logits': {0: 'logits_length'} } ) - + # Load the ONNX model + model = onnx.load(onnx_path) + # Simplify the model + model_simplified, _ = onnxsim.simplify(model) + # Save the simplified model + onnx.save(model_simplified, onnx_path) + # Copy tokenizer.json if it exists tokenizer_cache_dir = combined_model.tokenizer.name_or_path if os.path.isdir(tokenizer_cache_dir): @@ -117,16 +125,20 @@ def test_model_equivalence(original_model, onnx_path: str, max_seq_length: int = # Create test input batch_size = 1 test_input_ids = torch.randint(0, original_model.tokenizer.vocab_size, (batch_size, max_seq_length)) - + input_ids = original_model.tokenizer.encode("原神,启动!", return_tensors="pt") + + # Get PyTorch output original_model.eval() with torch.no_grad(): - pytorch_output = original_model(test_input_ids).numpy() + pytorch_output = original_model(input_ids).numpy() # Get ONNX output ort_session = ort.InferenceSession(onnx_path) - onnx_output = ort_session.run(None, {"input_ids": test_input_ids.numpy()})[0] - + onnx_output = ort_session.run(None, {"input_ids": input_ids.numpy()})[0] + + print(f"PyTorch output shape: {pytorch_output.shape}") + print(f"ONNX output shape: {onnx_output.shape}") # Compare outputs max_diff = np.max(np.abs(pytorch_output - onnx_output)) mean_diff = np.mean(np.abs(pytorch_output - onnx_output)) diff --git a/GPT_SoVITS/onnx_export_v1v2.py b/GPT_SoVITS/onnx_export_v1v2.py index 700787c3..556174ed 100644 --- a/GPT_SoVITS/onnx_export_v1v2.py +++ b/GPT_SoVITS/onnx_export_v1v2.py @@ -462,6 +462,19 @@ def export(vits_path, gpt_path, project_name, voice_model_version, t2s_model_com convert_onnx_to_half(f"onnx/{project_name}/{project_name}_t2s_init_step.onnx") convert_onnx_to_half(f"onnx/{project_name}/{project_name}_t2s_stage_step.onnx") + configJson = { + "project_name": project_name, + "type": "GPTSoVits", + "version" : voice_model_version, + "bert_base_path": 'GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large', + "cnhuhbert_base_path": 'GPT_SoVITS/pretrained_models/chinese-hubert-base', + "t2s_weights_path": gpt_path, + "vits_weights_path": vits_path, + "half_precision": half_precision + } + with open(f"onnx/{project_name}/config.json", "w", encoding="utf-8") as f: + json.dump(configJson, f, ensure_ascii=False, indent=4) + if __name__ == "__main__": try: os.mkdir("onnx")