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

import faulthandler
faulthandler.enable()
import os
os.chdir(os.path.dirname(os.path.abspath(__file__)))
from rknn.api import RKNN
import numpy as np

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

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

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

    # 加载ONNX模型
    print("--> Loading model")
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         inputs=['audio_wave'],
                         input_size_list=[[1, 1, 40960]])
    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.accuracy_analysis(inputs=["dumps_audio_vae_encode/audio_wave.npy"], target="rk3588", device_id=None)

    rknn.release()

if __name__ == '__main__':
    main()