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

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 = "dit_step.onnx"
    # 输出RKNN模型路径
    RKNN_MODEL = "dit_step.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=['x', 'mu', 't', 'cond', 'dt'],
                         input_size_list=[[1, 64, 2], [1, 1024], [1], [1, 64, 2], [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.accuracy_analysis(inputs=["dumps_dit_step/x.npy", "dumps_dit_step/mu.npy", "dumps_dit_step/t.npy", "dumps_dit_step/cond.npy", "dumps_dit_step/dt.npy"], target="rk3588", device_id=None)

    rknn.release()

if __name__ == '__main__':
    main()