HY-2012

Upload ASR and TTS models

a37072b verified 3 months ago

9.19 kB

/**************************************************************************************************
 * ZipVoice AXERA C++ Port
 *
 * EngineWrapper implementation
 **************************************************************************************************/

#include "EngineWrapper.hpp"
#include "utils/io.hpp"

#include <cstdlib>

static const char *strAlgoModelType[AX_ENGINE_MODEL_TYPE_BUTT] = {"3.6T", "7.2T", "10.8T"};

// NPU type enum
typedef enum axNPU_TYPE_E {
    AX_NPU_DEFAULT = 0,
    AX_STD_VNPU_1 = (1 << 0),
    AX_STD_VNPU_2 = (1 << 1),
    AX_STD_VNPU_3 = (1 << 2),
    AX_BL_VNPU_1  = (1 << 3),
    AX_BL_VNPU_2  = (1 << 4)
} AX_NPU_TYPE_E;

static AX_S32 CheckModelVNpu(const std::string &strModel,
                              const AX_ENGINE_MODEL_TYPE_T &eModelType,
                              const AX_S32 &nNpuType, AX_U32 &nNpuSet) {
    AX_ENGINE_NPU_ATTR_T stNpuAttr;
    memset(&stNpuAttr, 0x00, sizeof(stNpuAttr));

    auto ret = AX_ENGINE_GetVNPUAttr(&stNpuAttr);
    if (ret == 0) {
        if (stNpuAttr.eHardMode == AX_ENGINE_VIRTUAL_NPU_DISABLE) {
            nNpuSet = 0x01;
        } else if (stNpuAttr.eHardMode == AX_ENGINE_VIRTUAL_NPU_STD) {
            if (eModelType == AX_ENGINE_MODEL_TYPE1 || eModelType == AX_ENGINE_MODEL_TYPE2)
                return -1;
            if (nNpuType == 0) nNpuSet = 0x02;
            else {
                if (nNpuType & AX_STD_VNPU_1) nNpuSet |= 0x01;
                if (nNpuType & AX_STD_VNPU_2) nNpuSet |= 0x02;
                if (nNpuType & AX_STD_VNPU_3) nNpuSet |= 0x04;
            }
        } else if (stNpuAttr.eHardMode == AX_ENGINE_VIRTUAL_NPU_BIG_LITTLE) {
            if (eModelType == AX_ENGINE_MODEL_TYPE2) return -1;
            if (nNpuType == 0) {
                nNpuSet = (eModelType == AX_ENGINE_MODEL_TYPE1) ? 0x01 : 0x02;
            } else {
                if (eModelType == AX_ENGINE_MODEL_TYPE1) {
                    if (nNpuType & AX_BL_VNPU_2) return -1;
                    if (nNpuType & AX_BL_VNPU_1) nNpuSet |= 0x01;
                } else {
                    if (nNpuType & AX_BL_VNPU_1) nNpuSet |= 0x01;
                    if (nNpuType & AX_BL_VNPU_2) nNpuSet |= 0x02;
                }
            }
        }
    }
    return ret;
}

EngineWrapper::EngineWrapper()
    : m_hasInit(false), m_handle(nullptr), m_io_info(nullptr),
      m_input_num(0), m_output_num(0) {
    memset(&m_io, 0, sizeof(m_io));
}

EngineWrapper::~EngineWrapper() {
    Release();
}

int EngineWrapper::Init(const char* strModelPath, uint32_t nNpuType) {
    // 1. Load model
    AX_BOOL bLoadModelUseCmm = AX_TRUE;
    AX_CHAR *pModelBufferVirAddr = nullptr;
    AX_U64 u64ModelBufferPhyAddr = 0;
    AX_U32 nModelBufferSize = 0;

    if (bLoadModelUseCmm) {
        if (!utils::read_file(strModelPath, (AX_VOID **)&pModelBufferVirAddr,
                               u64ModelBufferPhyAddr, nModelBufferSize)) {
            printf("Read model(%s) fail\n", strModelPath);
            return -1;
        }
    } else {
        std::vector<char> model_buffer;
        if (!utils::read_file(strModelPath, model_buffer)) {
            printf("Read model(%s) fail\n", strModelPath);
            return -1;
        }
        pModelBufferVirAddr = model_buffer.data();
        nModelBufferSize = model_buffer.size();
    }

    auto freeModelBuffer = [&]() {
        if (bLoadModelUseCmm) {
            if (u64ModelBufferPhyAddr != 0)
                AX_SYS_MemFree(u64ModelBufferPhyAddr, &pModelBufferVirAddr);
        }
    };

    // 1.1 Get Model Type
    AX_ENGINE_MODEL_TYPE_T eModelType = AX_ENGINE_MODEL_TYPE0;
    AX_S32 ret = AX_ENGINE_GetModelType(pModelBufferVirAddr, nModelBufferSize, &eModelType);
    if (0 != ret || eModelType >= AX_ENGINE_MODEL_TYPE_BUTT) {
        printf("%s AX_ENGINE_GetModelType fail ret=%x\n", strModelPath, ret);
        freeModelBuffer();
        return -1;
    }

    // 1.2 Check VNPU
    AX_U32 nNpuSet = 0;
    ret = CheckModelVNpu(strModelPath, eModelType, nNpuType, nNpuSet);
    if (0 != ret) {
        printf("CheckModelVNpu fail\n");
        freeModelBuffer();
        return -1;
    }

    // 2. Create handle
    AX_ENGINE_HANDLE handle = nullptr;
    ret = AX_ENGINE_CreateHandle(&handle, pModelBufferVirAddr, nModelBufferSize);
    freeModelBuffer();

    auto deinit_handle = [&handle]() {
        if (handle) { AX_ENGINE_DestroyHandle(handle); }
        return -1;
    };

    if (0 != ret || !handle) {
        printf("Create model(%s) handle fail\n", strModelPath);
        return deinit_handle();
    }

    // 3. Create context
    ret = AX_ENGINE_CreateContext(handle);
    if (0 != ret) return deinit_handle();

    // 4. Get IO info
    m_io_info = nullptr;
    ret = AX_ENGINE_GetIOInfo(handle, &m_io_info);
    if (0 != ret) return deinit_handle();

    m_input_num = m_io_info->nInputSize;
    m_output_num = m_io_info->nOutputSize;

    // Build name-to-index maps
    m_input_name_to_idx.clear();
    for (int i = 0; i < m_input_num; ++i) {
        m_input_name_to_idx[std::string(m_io_info->pInputs[i].pName)] = i;
    }
    m_output_name_to_idx.clear();
    for (int i = 0; i < m_output_num; ++i) {
        m_output_name_to_idx[std::string(m_io_info->pOutputs[i].pName)] = i;
    }

    // 5. Prepare IO buffers
    ret = utils::prepare_io("enc", m_io_info, m_io, utils::IO_BUFFER_STRATEGY_DEFAULT);
    if (0 != ret) {
        printf("prepare io failed!\n");
        utils::free_io(m_io);
        return deinit_handle();
    }

    m_handle = handle;
    m_hasInit = true;
    return 0;
}

int EngineWrapper::SetInput(void* pInput, int index) {
    if (!m_hasInit || index < 0 || index >= m_input_num) return -1;
    return utils::push_io_input(pInput, index, m_io);
}

int EngineWrapper::RunSync() {
    if (!m_hasInit) return -1;
    auto ret = AX_ENGINE_RunSync(m_handle, &m_io);
    if (0 != ret) {
        printf("AX_ENGINE_RunSync failed. ret=0x%x\n", ret);
    }
    return ret;
}

int EngineWrapper::GetOutput(void* pOutput, int index) {
    if (!m_hasInit || index < 0 || index >= m_output_num) return -1;
    return utils::push_io_output(pOutput, index, m_io);
}

int EngineWrapper::SetInputByName(const char* name, void* pInput) {
    auto it = m_input_name_to_idx.find(std::string(name));
    if (it == m_input_name_to_idx.end()) {
        printf("Input '%s' not found in model\n", name);
        return -1;
    }
    return SetInput(pInput, it->second);
}

int EngineWrapper::GetOutputByName(const char* name, void* pOutput) {
    auto it = m_output_name_to_idx.find(std::string(name));
    if (it == m_output_name_to_idx.end()) {
        printf("Output '%s' not found in model\n", name);
        return -1;
    }
    return GetOutput(pOutput, it->second);
}

int EngineWrapper::GetInputSize(int index) {
    if (index < 0 || index >= m_input_num) return -1;
    return m_io.pInputs[index].nSize;
}

int EngineWrapper::GetOutputSize(int index) {
    if (index < 0 || index >= m_output_num) return -1;
    return m_io.pOutputs[index].nSize;
}

int EngineWrapper::GetInputSizeByName(const char* name) {
    int idx = GetInputIndex(name);
    if (idx < 0) return -1;
    return GetInputSize(idx);
}

int EngineWrapper::GetOutputSizeByName(const char* name) {
    int idx = GetOutputIndex(name);
    if (idx < 0) return -1;
    return GetOutputSize(idx);
}

int EngineWrapper::GetInputIndex(const char* name) {
    auto it = m_input_name_to_idx.find(std::string(name));
    return (it != m_input_name_to_idx.end()) ? it->second : -1;
}

int EngineWrapper::GetOutputIndex(const char* name) {
    auto it = m_output_name_to_idx.find(std::string(name));
    return (it != m_output_name_to_idx.end()) ? it->second : -1;
}

const char* EngineWrapper::GetInputName(int index) const {
    if (index < 0 || index >= m_input_num) return nullptr;
    return m_io_info->pInputs[index].pName;
}

const char* EngineWrapper::GetOutputName(int index) const {
    if (index < 0 || index >= m_output_num) return nullptr;
    return m_io_info->pOutputs[index].pName;
}

static const char* dtype_str(AX_ENGINE_DATA_TYPE_T t) {
    switch (t) {
        case AX_ENGINE_DT_FLOAT32: return "float32";
        case AX_ENGINE_DT_FLOAT64: return "float64";
        case AX_ENGINE_DT_SINT8:   return "sint8";
        case AX_ENGINE_DT_UINT8:   return "uint8";
        case AX_ENGINE_DT_SINT16:  return "sint16";
        case AX_ENGINE_DT_UINT16:  return "uint16";
        case AX_ENGINE_DT_SINT32:  return "sint32";
        case AX_ENGINE_DT_UINT32:  return "uint32";
        default: return "unknown";
    }
}

const char* EngineWrapper::GetInputDtypeStr(int index) const {
    if (index < 0 || index >= m_input_num) return "?";
    return dtype_str(m_io_info->pInputs[index].eDataType);
}

const char* EngineWrapper::GetOutputDtypeStr(int index) const {
    if (index < 0 || index >= m_output_num) return "?";
    return dtype_str(m_io_info->pOutputs[index].eDataType);
}

int EngineWrapper::Release() {
    if (m_handle) {
        utils::free_io(m_io);
        AX_ENGINE_DestroyHandle(m_handle);
        m_handle = nullptr;
    }
    m_hasInit = false;
    m_input_name_to_idx.clear();
    m_output_name_to_idx.clear();
    return 0;
}

	/**************************************************************************************************
	* ZipVoice AXERA C++ Port
	*
	* EngineWrapper implementation
	**************************************************************************************************/

	#include "EngineWrapper.hpp"
	#include "utils/io.hpp"

	#include <cstdlib>

	static const char *strAlgoModelType[AX_ENGINE_MODEL_TYPE_BUTT] = {"3.6T", "7.2T", "10.8T"};

	// NPU type enum
	typedef enum axNPU_TYPE_E {
	AX_NPU_DEFAULT = 0,
	AX_STD_VNPU_1 = (1 << 0),
	AX_STD_VNPU_2 = (1 << 1),
	AX_STD_VNPU_3 = (1 << 2),
	AX_BL_VNPU_1 = (1 << 3),
	AX_BL_VNPU_2 = (1 << 4)
	} AX_NPU_TYPE_E;

	static AX_S32 CheckModelVNpu(const std::string &strModel,
	const AX_ENGINE_MODEL_TYPE_T &eModelType,
	const AX_S32 &nNpuType, AX_U32 &nNpuSet) {
	AX_ENGINE_NPU_ATTR_T stNpuAttr;
	memset(&stNpuAttr, 0x00, sizeof(stNpuAttr));

	auto ret = AX_ENGINE_GetVNPUAttr(&stNpuAttr);
	if (ret == 0) {
	if (stNpuAttr.eHardMode == AX_ENGINE_VIRTUAL_NPU_DISABLE) {
	nNpuSet = 0x01;
	} else if (stNpuAttr.eHardMode == AX_ENGINE_VIRTUAL_NPU_STD) {
	if (eModelType == AX_ENGINE_MODEL_TYPE1 \|\| eModelType == AX_ENGINE_MODEL_TYPE2)
	return -1;
	if (nNpuType == 0) nNpuSet = 0x02;
	else {
	if (nNpuType & AX_STD_VNPU_1) nNpuSet \|= 0x01;
	if (nNpuType & AX_STD_VNPU_2) nNpuSet \|= 0x02;
	if (nNpuType & AX_STD_VNPU_3) nNpuSet \|= 0x04;
	}
	} else if (stNpuAttr.eHardMode == AX_ENGINE_VIRTUAL_NPU_BIG_LITTLE) {
	if (eModelType == AX_ENGINE_MODEL_TYPE2) return -1;
	if (nNpuType == 0) {
	nNpuSet = (eModelType == AX_ENGINE_MODEL_TYPE1) ? 0x01 : 0x02;
	} else {
	if (eModelType == AX_ENGINE_MODEL_TYPE1) {
	if (nNpuType & AX_BL_VNPU_2) return -1;
	if (nNpuType & AX_BL_VNPU_1) nNpuSet \|= 0x01;
	} else {
	if (nNpuType & AX_BL_VNPU_1) nNpuSet \|= 0x01;
	if (nNpuType & AX_BL_VNPU_2) nNpuSet \|= 0x02;
	}
	}
	}
	}
	return ret;
	}

	EngineWrapper::EngineWrapper()
	: m_hasInit(false), m_handle(nullptr), m_io_info(nullptr),
	m_input_num(0), m_output_num(0) {
	memset(&m_io, 0, sizeof(m_io));
	}

	EngineWrapper::~EngineWrapper() {
	Release();
	}

	int EngineWrapper::Init(const char* strModelPath, uint32_t nNpuType) {
	// 1. Load model
	AX_BOOL bLoadModelUseCmm = AX_TRUE;
	AX_CHAR *pModelBufferVirAddr = nullptr;
	AX_U64 u64ModelBufferPhyAddr = 0;
	AX_U32 nModelBufferSize = 0;

	if (bLoadModelUseCmm) {
	if (!utils::read_file(strModelPath, (AX_VOID **)&pModelBufferVirAddr,
	u64ModelBufferPhyAddr, nModelBufferSize)) {
	printf("Read model(%s) fail\n", strModelPath);
	return -1;
	}
	} else {
	std::vector<char> model_buffer;
	if (!utils::read_file(strModelPath, model_buffer)) {
	printf("Read model(%s) fail\n", strModelPath);
	return -1;
	}
	pModelBufferVirAddr = model_buffer.data();
	nModelBufferSize = model_buffer.size();
	}

	auto freeModelBuffer = [&]() {
	if (bLoadModelUseCmm) {
	if (u64ModelBufferPhyAddr != 0)
	AX_SYS_MemFree(u64ModelBufferPhyAddr, &pModelBufferVirAddr);
	}
	};

	// 1.1 Get Model Type
	AX_ENGINE_MODEL_TYPE_T eModelType = AX_ENGINE_MODEL_TYPE0;
	AX_S32 ret = AX_ENGINE_GetModelType(pModelBufferVirAddr, nModelBufferSize, &eModelType);
	if (0 != ret \|\| eModelType >= AX_ENGINE_MODEL_TYPE_BUTT) {
	printf("%s AX_ENGINE_GetModelType fail ret=%x\n", strModelPath, ret);
	freeModelBuffer();
	return -1;
	}

	// 1.2 Check VNPU
	AX_U32 nNpuSet = 0;
	ret = CheckModelVNpu(strModelPath, eModelType, nNpuType, nNpuSet);
	if (0 != ret) {
	printf("CheckModelVNpu fail\n");
	freeModelBuffer();
	return -1;
	}

	// 2. Create handle
	AX_ENGINE_HANDLE handle = nullptr;
	ret = AX_ENGINE_CreateHandle(&handle, pModelBufferVirAddr, nModelBufferSize);
	freeModelBuffer();

	auto deinit_handle = [&handle]() {
	if (handle) { AX_ENGINE_DestroyHandle(handle); }
	return -1;
	};

	if (0 != ret \|\| !handle) {
	printf("Create model(%s) handle fail\n", strModelPath);
	return deinit_handle();
	}

	// 3. Create context
	ret = AX_ENGINE_CreateContext(handle);
	if (0 != ret) return deinit_handle();

	// 4. Get IO info
	m_io_info = nullptr;
	ret = AX_ENGINE_GetIOInfo(handle, &m_io_info);
	if (0 != ret) return deinit_handle();

	m_input_num = m_io_info->nInputSize;
	m_output_num = m_io_info->nOutputSize;

	// Build name-to-index maps
	m_input_name_to_idx.clear();
	for (int i = 0; i < m_input_num; ++i) {
	m_input_name_to_idx[std::string(m_io_info->pInputs[i].pName)] = i;
	}
	m_output_name_to_idx.clear();
	for (int i = 0; i < m_output_num; ++i) {
	m_output_name_to_idx[std::string(m_io_info->pOutputs[i].pName)] = i;
	}

	// 5. Prepare IO buffers
	ret = utils::prepare_io("enc", m_io_info, m_io, utils::IO_BUFFER_STRATEGY_DEFAULT);
	if (0 != ret) {
	printf("prepare io failed!\n");
	utils::free_io(m_io);
	return deinit_handle();
	}

	m_handle = handle;
	m_hasInit = true;
	return 0;
	}

	int EngineWrapper::SetInput(void* pInput, int index) {
	if (!m_hasInit \|\| index < 0 \|\| index >= m_input_num) return -1;
	return utils::push_io_input(pInput, index, m_io);
	}

	int EngineWrapper::RunSync() {
	if (!m_hasInit) return -1;
	auto ret = AX_ENGINE_RunSync(m_handle, &m_io);
	if (0 != ret) {
	printf("AX_ENGINE_RunSync failed. ret=0x%x\n", ret);
	}
	return ret;
	}

	int EngineWrapper::GetOutput(void* pOutput, int index) {
	if (!m_hasInit \|\| index < 0 \|\| index >= m_output_num) return -1;
	return utils::push_io_output(pOutput, index, m_io);
	}

	int EngineWrapper::SetInputByName(const char* name, void* pInput) {
	auto it = m_input_name_to_idx.find(std::string(name));
	if (it == m_input_name_to_idx.end()) {
	printf("Input '%s' not found in model\n", name);
	return -1;
	}
	return SetInput(pInput, it->second);
	}

	int EngineWrapper::GetOutputByName(const char* name, void* pOutput) {
	auto it = m_output_name_to_idx.find(std::string(name));
	if (it == m_output_name_to_idx.end()) {
	printf("Output '%s' not found in model\n", name);
	return -1;
	}
	return GetOutput(pOutput, it->second);
	}

	int EngineWrapper::GetInputSize(int index) {
	if (index < 0 \|\| index >= m_input_num) return -1;
	return m_io.pInputs[index].nSize;
	}

	int EngineWrapper::GetOutputSize(int index) {
	if (index < 0 \|\| index >= m_output_num) return -1;
	return m_io.pOutputs[index].nSize;
	}

	int EngineWrapper::GetInputSizeByName(const char* name) {
	int idx = GetInputIndex(name);
	if (idx < 0) return -1;
	return GetInputSize(idx);
	}

	int EngineWrapper::GetOutputSizeByName(const char* name) {
	int idx = GetOutputIndex(name);
	if (idx < 0) return -1;
	return GetOutputSize(idx);
	}

	int EngineWrapper::GetInputIndex(const char* name) {
	auto it = m_input_name_to_idx.find(std::string(name));
	return (it != m_input_name_to_idx.end()) ? it->second : -1;
	}

	int EngineWrapper::GetOutputIndex(const char* name) {
	auto it = m_output_name_to_idx.find(std::string(name));
	return (it != m_output_name_to_idx.end()) ? it->second : -1;
	}

	const char* EngineWrapper::GetInputName(int index) const {
	if (index < 0 \|\| index >= m_input_num) return nullptr;
	return m_io_info->pInputs[index].pName;
	}

	const char* EngineWrapper::GetOutputName(int index) const {
	if (index < 0 \|\| index >= m_output_num) return nullptr;
	return m_io_info->pOutputs[index].pName;
	}

	static const char* dtype_str(AX_ENGINE_DATA_TYPE_T t) {
	switch (t) {
	case AX_ENGINE_DT_FLOAT32: return "float32";
	case AX_ENGINE_DT_FLOAT64: return "float64";
	case AX_ENGINE_DT_SINT8: return "sint8";
	case AX_ENGINE_DT_UINT8: return "uint8";
	case AX_ENGINE_DT_SINT16: return "sint16";
	case AX_ENGINE_DT_UINT16: return "uint16";
	case AX_ENGINE_DT_SINT32: return "sint32";
	case AX_ENGINE_DT_UINT32: return "uint32";
	default: return "unknown";
	}
	}

	const char* EngineWrapper::GetInputDtypeStr(int index) const {
	if (index < 0 \|\| index >= m_input_num) return "?";
	return dtype_str(m_io_info->pInputs[index].eDataType);
	}

	const char* EngineWrapper::GetOutputDtypeStr(int index) const {
	if (index < 0 \|\| index >= m_output_num) return "?";
	return dtype_str(m_io_info->pOutputs[index].eDataType);
	}

	int EngineWrapper::Release() {
	if (m_handle) {
	utils::free_io(m_io);
	AX_ENGINE_DestroyHandle(m_handle);
	m_handle = nullptr;
	}
	m_hasInit = false;
	m_input_name_to_idx.clear();
	m_output_name_to_idx.clear();
	return 0;
	}