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

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

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

    # ONNX模型路径
    ONNX_MODEL = "mm_projector.onnx"
    # 输出RKNN模型路径
    RKNN_MODEL = "mm_projector.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=['last_hidden_state'],
                         input_size_list=[[1, 256, 3072]])
    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()