#!/usr/bin/env python3
# ztu_somemodelruntime_rknn2: F5_Transformer_opset19

from rknn.api import RKNN
import os
import numpy as np

def main():
    # 创建RKNN实例
    rknn = RKNN(verbose=True)

    # ONNX模型路径
    ONNX_MODEL = "F5_Transformer_opset19.onnx"
    # 输出RKNN模型路径
    RKNN_MODEL = "F5_Transformer_opset19.rknn"

    # 配置参数
    print("--> Config model")
    ret = rknn.config(target_platform="rk3588",
                      dynamic_input=None,
                      disable_rules=[])
    if ret != 0:
        print('Config model failed!')
        exit(ret)

    # 加载ONNX模型
    print("--> Loading model")
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         inputs=['noise', 'rope_cos_q', 'rope_sin_q', 'rope_cos_k', 'rope_sin_k', 'cat_mel_text', 'cat_mel_text_drop', 'time_step.1'],
                         input_size_list=[[1, 1536, 100], [2, 16, 1536, 64], [2, 16, 1536, 64], [2, 16, 64, 1536], [2, 16, 64, 1536], [1, 1536, 612], [1, 1536, 612], [1]])
    if ret != 0:
        print('Load model failed!')
        exit(ret)

    # 构建模型
    print("--> Building model")
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)

    # 导出RKNN模型
    print("--> Export RKNN model")
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)

    print(f'Done! The converted RKNN model has been saved to: ' + RKNN_MODEL)
    rknn.release()

if __name__ == '__main__':
    main()