Upload 34 files

.gitattributes

@@ -33,3 +33,18 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+audio_vae_decode.rknn filter=lfs diff=lfs merge=lfs -text
+audio_vae_encode.rknn filter=lfs diff=lfs merge=lfs -text
+base_lm.rkllm filter=lfs diff=lfs merge=lfs -text
+basic_ref_zh.wav filter=lfs diff=lfs merge=lfs -text
+dit_step.rknn filter=lfs diff=lfs merge=lfs -text
+fsq_layer.rknn filter=lfs diff=lfs merge=lfs -text
+librkllmrt.so filter=lfs diff=lfs merge=lfs -text
+lm_to_dit_proj.rknn filter=lfs diff=lfs merge=lfs -text
+locenc_1.rknn filter=lfs diff=lfs merge=lfs -text
+locenc_64.rknn filter=lfs diff=lfs merge=lfs -text
+model_structure.jpg filter=lfs diff=lfs merge=lfs -text
+res_to_dit_proj.rknn filter=lfs diff=lfs merge=lfs -text
+residual_lm.rkllm filter=lfs diff=lfs merge=lfs -text
+rknn_output.wav filter=lfs diff=lfs merge=lfs -text
+stop_head.rknn filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,349 @@
----
-license: agpl-3.0
----
+---
+license: agpl-3.0
+language:
+- en
+- zh
+base_model:
+- openbmb/VoxCPM-0.5B
+pipeline_tag: text-to-speech
+tags:
+- rknn
+- rkllm
+- text-to-speech
+- speech
+- speech generation
+- voice cloning
+---
+
+# VoxCPM-0.5B-RKNN2
+
+### (English README see below)
+
+VoxCPM 是一种创新的无分词器文本转语音（TTS）系统，重新定义了语音合成的真实感。通过在连续空间中建模语音，它克服了离散标记化的局限，并实现了两项核心能力：上下文感知的语音生成和逼真的零样本语音克隆。
+不同于将语音转换为离散标记的主流方法，VoxCPM 采用端到端的扩散自回归架构，直接从文本生成连续的语音表示。它基于 MiniCPM-4 主干构建，通过分层语言建模和 FSQ 约束实现了隐式的语义-声学解耦，极大地提升了表现力和生成稳定性。
+
+![模型架构](model_structure.jpg)
+
+
+- 推理速度(RKNN2)：RK3588上RTF约8（生成10s音频需要推理80s）
+- 大致内存占用(RKNN2)：约3.3GB
+
+## 使用方法
+
+1. 克隆项目到本地
+
+2. 安装依赖
+
+```bash
+pip install "numpy<2" scipy soundfile tqdm transformers sentencepiece rknn-toolkit-lite2
+```
+
+3. 运行
+
+```bash
+python onnx_infer-rknn2.py --onnx-dir . --tokenizer-dir . --base-hf-dir . --residual-hf-dir . --text "哇, 这个模型居然在RK3588这个辣鸡SoC上也能完美运行!" --prompt-audio basic_ref_zh.wav --prompt-text "对，这就是我，万人敬仰的太乙真人。" --output rknn_output.wav --cfg-value 2.0 --inference-timesteps 10 --seed 1234
+```
+
+可选参数：
+- `--text`: 要生成的文本
+- `--prompt-audio`: 参考音频路径（用于语音克隆）
+- `--prompt-text`: 参考音频对应的文本（使用参考音频时必填）
+- `--cfg-value`: CFG引导强度，默认2.0
+- `--inference-timesteps`: 扩散步数，默认10
+- `--seed`: 随机种子
+- `--output`: 输出音频路径
+
+## 运行效果
+
+
+```log
+> python onnx_infer-rknn2.py --onnx-dir . --tokenizer-dir . --base-hf-dir . --residual-hf-dir . --text "哇, 这个模型居然在RK3588这个辣鸡SoC上也能完美运行!" --prompt-audio basic_ref_zh.wav --prompt-text "对，这就是我，万人敬仰的太乙真人。" --output rknn_output.wav --cfg-value 2.0 --inference-timesteps 10 --seed 1234
+
+I rkllm: rkllm-runtime version: 1.2.3, rknpu driver version: 0.9.8, platform: RK3588
+I rkllm: loading rkllm model from ./base_lm.rkllm
+I rkllm: rkllm-toolkit version: 1.2.3, max_context_limit: 4096, npu_core_num: 1, target_platform: RK3588, model_dtype: FP16
+I rkllm: Enabled cpus: [4, 5, 6, 7]
+I rkllm: Enabled cpus num: 4
+I rkllm: rkllm-runtime version: 1.2.3, rknpu driver version: 0.9.8, platform: RK3588
+I rkllm: loading rkllm model from ./residual_lm.rkllm
+I rkllm: rkllm-toolkit version: 1.2.2, max_context_limit: 4096, npu_core_num: 3, target_platform: RK3588, model_dtype: FP16
+I rkllm: Enabled cpus: [4, 5, 6, 7]
+I rkllm: Enabled cpus num: 4
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:26.264] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:26.264] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:26.265] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:26.404] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:26.537] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:26.537] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:26.537] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:26.616] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:26.795] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:26.795] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:26.795] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.020] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.194] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.194] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.194] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.317] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.431] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.431] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.431] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (@2025-04-03T08:26:16)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: dynamic_shape
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.547] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.547] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.547] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.549] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.728] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.728] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.728] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.819] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.937] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.937] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.937] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.940] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:28.058] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:28.058] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:28.058] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:28.060] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+[time] vae_encode_0: 1601.56 ms
+[time] vae_encode_38400: 1605.46 ms
+[time] vae_encode_76800: 1591.07 ms
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.
+[time] locenc_0: 819.49 ms
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.
+[time] locenc_64: 818.33 ms
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.
+[time] locenc_128: 819.09 ms
+[time] base_lm initial: 579.08 ms
+[time] fsq_init_0: 2.54 ms
+[time] fsq_init_64: 1.86 ms
+[time] fsq_init_128: 1.79 ms
+[time] residual_lm initial: 139.10 ms
+gen_loop:   0%|                                                                          | 0/2000 [00:00<?, ?it/s][time] lm_to_dit: 0.82 ms
+[time] res_to_dit: 0.56 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 33.32it/s]
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.4it/s]
+[time] locenc_step: 16.32 ms
+gen_loop:   0%|                                                                  | 1/2000 [00:00<14:30,  2.30it/s][time] lm_to_dit: 0.57 ms
+[time] res_to_dit: 0.44 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 33.10it/s]
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.1it/s]
+[time] locenc_step: 15.84 ms
+gen_loop:   0%|                                                                  | 2/2000 [00:00<14:27,  2.30it/s][time] lm_to_dit: 0.56 ms
+[time] res_to_dit: 0.50 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 31.93it/s]
+
+...
+
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.5it/s]
+[time] locenc_step: 15.88 ms
+gen_loop:   6%|███▉                                                            | 123/2000 [00:53<13:35,  2.30it/s][time] lm_to_dit: 0.57 ms
+[time] res_to_dit: 0.49 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 32.94it/s]
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.6it/s]
+[time] locenc_step: 15.84 ms
+gen_loop:   6%|███▉                                                            | 123/2000 [00:54<13:44,  2.28it/s]
+[time] vae_decode_0: 1044.00 ms
+[time] vae_decode_60: 1018.03 ms
+[time] vae_decode_120: 1020.72 ms
+[time] vae_decode_180: 1021.19 ms
+[time] vae_decode_240: 1006.85 ms
+Saved: rknn_output.wav
+```
+
+## 模型转换
+
+#### 懒得写了，待补充
+
+## 已知问题
+
+- 某些情况下语音生成可能陷入死循环，原项目似乎有检测死循环的机制，但我这里没有实现。
+- 由于RKNN工具链的内部问题，locenc模型没有办法在一个模型里配置两种输入长度的两组shape，因此只能单独转换两个模型。
+- 由于RKLLM工具链/运行时的内部问题，两个LLM的输出张量的数值都只有正确结果的四分之一，手动乘4之后可以得到正确结果。
+- 由于RKNN工具链目前不支持非4维输入模型多batch使用多NPU核的数据并行推理，脚本中CFG是分两次单独进行的，速度较慢。
+
+
+## 参考
+- [openbmb/VoxCPM-0.5B](https://huggingface.co/openbmb/VoxCPM-0.5B)
+- [0seba/VoxCPMANE](https://github.com/0seba/VoxCPMANE)
+- [bluryar/VoxCPM-ONNX](https://github.com/bluryar/VoxCPM-ONNX)
+
+# English README
+
+VoxCPM is an innovative tokenizer-free Text-to-Speech (TTS) system that redefines realism in speech synthesis. By modeling speech in continuous space, it overcomes the limitations of discrete tokenization and achieves two core capabilities: context-aware speech generation and realistic zero-shot voice cloning.
+
+Unlike mainstream approaches that convert speech into discrete tokens, VoxCPM adopts an end-to-end diffusion autoregressive architecture that directly generates continuous speech representations from text. Built on the MiniCPM-4 backbone, it achieves implicit semantic-acoustic decoupling through hierarchical language modeling and FSQ constraints, greatly enhancing expressiveness and generation stability.
+
+- Inference speed (RKNN2): RTF approximately 8 on RK3588 (80s inference time to generate 10s audio)
+- Approximate memory usage (RKNN2): ~3.3GB
+
+## Usage
+
+1. Clone the project locally
+
+2. Install dependencies
+
+```bash
+pip install "numpy<2" scipy soundfile tqdm transformers sentencepiece rknn-toolkit-lite2
+```
+
+3. Run
+
+```bash
+python onnx_infer-rknn2.py --onnx-dir . --tokenizer-dir . --base-hf-dir . --residual-hf-dir . --text "Wow, this model actually runs perfectly on the RK3588 SoC!" --prompt-audio basic_ref_zh.wav --prompt-text "对，这就是我，万人敬仰的太乙真人。" --output rknn_output.wav --cfg-value 2.0 --inference-timesteps 10 --seed 1234
+```
+
+Optional parameters:
+- `--text`: Text to generate
+- `--prompt-audio`: Reference audio path (for voice cloning)
+- `--prompt-text`: Text corresponding to the reference audio (required when using reference audio)
+- `--cfg-value`: CFG guidance strength, default 2.0
+- `--inference-timesteps`: Number of diffusion steps, default 10
+- `--seed`: Random seed
+- `--output`: Output audio path
+
+## Performance
+
+
+```log
+> python onnx_infer-rknn2.py --onnx-dir . --tokenizer-dir . --base-hf-dir . --residual-hf-dir . --text "哇, 这个模型居然在RK3588这个辣鸡SoC上也能完美运行!" --prompt-audio basic_ref_zh.wav --prompt-text "对，这就是我，万人敬仰的太乙真人。" --output rknn_output.wav --cfg-value 2.0 --inference-timesteps 10 --seed 1234
+
+I rkllm: rkllm-runtime version: 1.2.3, rknpu driver version: 0.9.8, platform: RK3588
+I rkllm: loading rkllm model from ./base_lm.rkllm
+I rkllm: rkllm-toolkit version: 1.2.3, max_context_limit: 4096, npu_core_num: 1, target_platform: RK3588, model_dtype: FP16
+I rkllm: Enabled cpus: [4, 5, 6, 7]
+I rkllm: Enabled cpus num: 4
+I rkllm: rkllm-runtime version: 1.2.3, rknpu driver version: 0.9.8, platform: RK3588
+I rkllm: loading rkllm model from ./residual_lm.rkllm
+I rkllm: rkllm-toolkit version: 1.2.2, max_context_limit: 4096, npu_core_num: 3, target_platform: RK3588, model_dtype: FP16
+I rkllm: Enabled cpus: [4, 5, 6, 7]
+I rkllm: Enabled cpus num: 4
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:26.264] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:26.264] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:26.265] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:26.404] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:26.537] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:26.537] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:26.537] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:26.616] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:26.795] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:26.795] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:26.795] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.020] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.194] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.194] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.194] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.317] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.431] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.431] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.431] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (@2025-04-03T08:26:16)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: dynamic_shape
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.547] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.547] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.547] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.549] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.728] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.728] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.728] RKNN Model Information, version: 6, toolkit version: 2.3.2(compiler version: 2.3.2 (e045de294f@2025-04-07T19:48:25)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.819] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:27.937] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:27.937] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:27.937] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:27.940] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+W rknn-toolkit-lite2 version: 2.3.2
+I RKNN: [18:58:28.058] RKNN Runtime Information, librknnrt version: 2.3.2 (429f97ae6b@2025-04-09T09:09:27)
+I RKNN: [18:58:28.058] RKNN Driver Information, version: 0.9.8
+I RKNN: [18:58:28.058] RKNN Model Information, version: 6, toolkit version: 2.3.0(compiler version: 2.3.0 (@2024-11-07T08:11:34)), target: RKNPU v2, target platform: rk3588, framework name: ONNX, framework layout: NCHW, model inference type: static_shape
+W RKNN: [18:58:28.060] query RKNN_QUERY_INPUT_DYNAMIC_RANGE error, rknn model is static shape type, please export rknn with dynamic_shapes
+W Query dynamic range failed. Ret code: RKNN_ERR_MODEL_INVALID. (If it is a static shape RKNN model, please ignore the above warning message.)
+[time] vae_encode_0: 1601.56 ms
+[time] vae_encode_38400: 1605.46 ms
+[time] vae_encode_76800: 1591.07 ms
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.
+[time] locenc_0: 819.49 ms
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.
+[time] locenc_64: 818.33 ms
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.
+[time] locenc_128: 819.09 ms
+[time] base_lm initial: 579.08 ms
+[time] fsq_init_0: 2.54 ms
+[time] fsq_init_64: 1.86 ms
+[time] fsq_init_128: 1.79 ms
+[time] residual_lm initial: 139.10 ms
+gen_loop:   0%|                                                                          | 0/2000 [00:00<?, ?it/s][time] lm_to_dit: 0.82 ms
+[time] res_to_dit: 0.56 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 33.32it/s]
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.4it/s]
+[time] locenc_step: 16.32 ms
+gen_loop:   0%|                                                                  | 1/2000 [00:00<14:30,  2.30it/s][time] lm_to_dit: 0.57 ms
+[time] res_to_dit: 0.44 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 33.10it/s]
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.1it/s]
+[time] locenc_step: 15.84 ms
+gen_loop:   0%|                                                                  | 2/2000 [00:00<14:27,  2.30it/s][time] lm_to_dit: 0.56 ms
+[time] res_to_dit: 0.50 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 31.93it/s]
+
+...
+
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.5it/s]
+[time] locenc_step: 15.88 ms
+gen_loop:   6%|███▉                                                            | 123/2000 [00:53<13:35,  2.30it/s][time] lm_to_dit: 0.57 ms
+[time] res_to_dit: 0.49 ms
+100%|█████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 32.94it/s]
+W The input[0] need NHWC data format, but NCHW set, the data format and data buffer will be changed to NHWC.6it/s]
+[time] locenc_step: 15.84 ms
+gen_loop:   6%|███▉                                                            | 123/2000 [00:54<13:44,  2.28it/s]
+[time] vae_decode_0: 1044.00 ms
+[time] vae_decode_60: 1018.03 ms
+[time] vae_decode_120: 1020.72 ms
+[time] vae_decode_180: 1021.19 ms
+[time] vae_decode_240: 1006.85 ms
+Saved: rknn_output.wav
+```
+
+## Model Conversion
+
+#### TODO: Documentation to be added
+
+## Known Issues
+
+- In some cases, speech generation may fall into an infinite loop. The original project seems to have a mechanism to detect infinite loops, but it is not implemented here.
+- Due to internal issues with the RKNN toolchain, the locenc model cannot configure two sets of shapes for two different input lengths in a single model, so two separate models must be converted.
+- Due to internal issues with the RKLLM toolchain/runtime, the output tensor values of both LLMs are only one-quarter of the correct result. Multiplying by 4 manually yields the correct result.
+- Since the RKNN toolchain currently does not support data-parallel inference using multiple NPU cores for non-4D input models with multiple batches, CFG in the script is performed separately in two passes, which is relatively slow.
+
+## References
+- [openbmb/VoxCPM-0.5B](https://huggingface.co/openbmb/VoxCPM-0.5B)
+- [0seba/VoxCPMANE](https://github.com/0seba/VoxCPMANE)
+- [bluryar/VoxCPM-ONNX](https://github.com/bluryar/VoxCPM-ONNX)

audio_vae_decode.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ec4181f54e0ec9e3de9d4380beb79855f3fdb45b1562bcb89d403e2f5f130f5
+size 58597553

audio_vae_encode.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14f0cadac33b0bd5c722edd09df758eb3f4204738af112c584d4f4292ac781f5
+size 102299829

base_lm.rkllm

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d75d343623ac22402175e1aaa935c243338524a19e9cb1dae4955e643df15964
+size 1028092444

basic_ref_zh.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96724a113240d1f82c6ded1334122f0176b96c9226ccd3c919e625bcfd2a3ede
+size 324558

config.json

@@ -0,0 +1,52 @@
+{
+    "architecture": "voxcpm",
+    "lm_config": {
+        "bos_token_id": 1,
+        "eos_token_id": 2,
+        "hidden_size": 1024,
+        "intermediate_size": 4096,
+        "max_position_embeddings": 32768,
+        "num_attention_heads": 16,
+        "num_hidden_layers": 24,
+        "num_key_value_heads": 2,
+        "rms_norm_eps": 1e-05,
+        "rope_theta": 10000,
+        "rope_scaling": {
+            "type": "longrope", 
+            "long_factor": [1.0004360675811768, 1.0668443441390991, 1.1631425619125366, 1.3025742769241333, 1.5040205717086792, 1.7941505908966064, 2.2101221084594727, 2.802666664123535, 3.6389970779418945, 4.804192543029785, 6.39855432510376, 8.527148246765137, 11.277542114257812, 14.684998512268066, 18.69317054748535, 23.13019371032715, 27.72362518310547, 32.1606559753418, 36.168827056884766, 39.57627868652344, 42.32667541503906, 44.45526885986328, 46.04962921142578, 47.21482849121094, 48.05115509033203, 48.64370346069336, 49.05967712402344, 49.34980392456055, 49.551246643066406, 49.69068145751953, 49.78697967529297, 49.85338592529297],
+            "short_factor": [1.0004360675811768, 1.0668443441390991, 1.1631425619125366, 1.3025742769241333, 1.5040205717086792, 1.7941505908966064, 2.2101221084594727, 2.802666664123535, 3.6389970779418945, 4.804192543029785, 6.39855432510376, 8.527148246765137, 11.277542114257812, 14.684998512268066, 18.69317054748535, 23.13019371032715, 27.72362518310547, 32.1606559753418, 36.168827056884766, 39.57627868652344, 42.32667541503906, 44.45526885986328, 46.04962921142578, 47.21482849121094, 48.05115509033203, 48.64370346069336, 49.05967712402344, 49.34980392456055, 49.551246643066406, 49.69068145751953, 49.78697967529297, 49.85338592529297],
+            "original_max_position_embeddings": 32768
+        },    
+        "vocab_size": 73448,
+        "scale_emb": 12,
+        "dim_model_base": 256,
+        "scale_depth": 1.4,
+        "use_mup": false
+    },
+    "patch_size": 2,
+    "feat_dim": 64,
+    "scalar_quantization_latent_dim": 256,
+    "scalar_quantization_scale": 9,
+    "residual_lm_num_layers": 6,
+    "encoder_config": {
+        "hidden_dim": 1024,
+        "ffn_dim": 4096,
+        "num_heads": 16,
+        "num_layers": 4
+    },
+    "dit_config": {
+        "hidden_dim": 1024,
+        "ffn_dim": 4096,
+        "num_heads": 16,
+        "num_layers": 4,
+        "cfm_config": {
+            "sigma_min": 1e-06,
+            "solver": "euler",
+            "t_scheduler": "log-norm",
+            "inference_cfg_rate": 2.0
+        }
+    },
+    "max_length": 4096,
+    "device": "cuda",
+    "dtype": "bfloat16"
+}

convert_audio_vae_decode.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# ztu_somemodelruntime_rknn2: audio_vae_decode
+
+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_decode.onnx"
+    # 输出RKNN模型路径
+    RKNN_MODEL = "audio_vae_decode.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=['latent'],
+                         input_size_list=[[1, 64, 64]])
+    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_decode/latent.npy"], target="rk3588", device_id=None)
+
+    rknn.release()
+
+if __name__ == '__main__':
+    main()

convert_audio_vae_encode.py

@@ -0,0 +1,59 @@
+#!/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()

convert_dit_step.py

@@ -0,0 +1,59 @@
+#!/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()

convert_fsq_layer.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# ztu_somemodelruntime_rknn2: fsq_layer
+
+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 = "fsq_layer.onnx"
+    # 输出RKNN模型路径
+    RKNN_MODEL = "fsq_layer.rknn"
+
+    # 配置参数
+    print("--> Config model")
+    ret = rknn.config(target_platform="rk3588",
+                      dynamic_input=[[[1, 64, 1024]], [[1, 1, 1024]]])
+    if ret != 0:
+        print('Config model failed!')
+        exit(ret)
+
+    # 加载ONNX模型
+    print("--> Loading model")
+    ret = rknn.load_onnx(model=ONNX_MODEL,
+                         inputs=None,
+                         input_size_list=None)
+    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_fsq_layer/hidden.npy"], target="rk3588", device_id=None)
+
+    rknn.release()
+
+if __name__ == '__main__':
+    main()

convert_lm_to_dit_proj.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# ztu_somemodelruntime_rknn2: lm_to_dit_proj
+
+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 = "lm_to_dit_proj.onnx"
+    # 输出RKNN模型路径
+    RKNN_MODEL = "lm_to_dit_proj.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=['input'],
+                         input_size_list=[[1, 1024]])
+    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_lm_to_dit_proj/input.npy"], target="rk3588", device_id=None)
+
+    rknn.release()
+
+if __name__ == '__main__':
+    main()

convert_locenc.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+# ztu_somemodelruntime_rknn2: locenc
+
+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 = "locenc.onnx"
+    # 输出RKNN模型路径
+    RKNN_MODEL = "locenc_64.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'],
+                         input_size_list=[[1, 64, 2, 64]])
+    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_locenc/x.npy"], target="rk3588", device_id=None)
+
+    rknn.release()
+
+def main2():
+    # 创建RKNN实例
+    rknn = RKNN(verbose=True)
+
+    # ONNX模型路径
+    ONNX_MODEL = "locenc.onnx"
+    # 输出RKNN模型路径
+    RKNN_MODEL = "locenc_1.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'],
+                         input_size_list=[[1, 1, 2, 64]])
+    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_locenc/x.npy"], target="rk3588", device_id=None)
+
+    rknn.release()
+
+if __name__ == '__main__':
+    main()
+    main2()

convert_res_to_dit_proj.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# ztu_somemodelruntime_rknn2: res_to_dit_proj
+
+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 = "res_to_dit_proj.onnx"
+    # 输出RKNN模型路径
+    RKNN_MODEL = "res_to_dit_proj.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=['input'],
+                         input_size_list=[[1, 1024]])
+    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_res_to_dit_proj/input.npy"], target="rk3588", device_id=None)
+
+    rknn.release()
+
+if __name__ == '__main__':
+    main()

convert_stop_head.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# ztu_somemodelruntime_rknn2: stop_head
+
+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 = "stop_head.onnx"
+    # 输出RKNN模型路径
+    RKNN_MODEL = "stop_head.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=['hidden'],
+                         input_size_list=[[1, 1024]])
+    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_stop_head/hidden.npy"], target="rk3588", device_id=None)
+
+    rknn.release()
+
+if __name__ == '__main__':
+    main()

convert_vox_minicpm_to_hf.py

@@ -0,0 +1,115 @@
+import argparse
+import json
+import os
+import sys
+import torch
+import math
+
+
+def load_vox_configs(vox_config_path: str) -> tuple[dict, dict]:
+    """Return (base_lm_cfg, residual_cfg)."""
+    with open(vox_config_path, "r") as f:
+        data = json.load(f)
+
+    base = data["lm_config"]
+    rope = base.get("rope_scaling")
+    if rope:
+        rope = dict(rope)
+        # Vox config uses "type", transformers expects "rope_type"
+        if "type" in rope and "rope_type" not in rope:
+            rope["rope_type"] = rope.pop("type")
+        base["rope_scaling"] = rope
+
+    residual = dict(base)
+    residual["num_hidden_layers"] = data.get("residual_lm_num_layers", residual["num_hidden_layers"])
+    # keep vocab_size for easier loading; Vox sets 0 because inputs_embeds are provided
+    residual.setdefault("vocab_size", base.get("vocab_size"))
+
+    # Align transformers residual scaling with Vox (no scaling when use_mup=False)
+    if not base.get("use_mup", True):
+        base["scale_depth"] = math.sqrt(base["num_hidden_layers"])
+        residual["scale_depth"] = math.sqrt(residual["num_hidden_layers"])
+    return base, residual
+
+
+def build_hf_config(lm_cfg: dict, minicpm_dir: str):
+    sys.path.insert(0, minicpm_dir)
+    from configuration_minicpm import MiniCPMConfig
+
+    return MiniCPMConfig(**lm_cfg)
+
+
+def convert_state_dict(vox_state_path: str, lm_prefix: str) -> dict:
+    raw = torch.load(vox_state_path, map_location="cpu")
+    sd = raw["state_dict"] if isinstance(raw, dict) and "state_dict" in raw else raw
+
+    out = {}
+    prefix = f"{lm_prefix}."
+    for k, v in sd.items():
+        if not k.startswith(prefix):
+            continue
+        new_k = "model." + k[len(prefix) :]
+        out[new_k] = v
+
+    # Tie lm_head to embeddings for MiniCPMForCausalLM
+    if "model.embed_tokens.weight" in out:
+        out["lm_head.weight"] = out["model.embed_tokens.weight"]
+    return out
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert VoxCPM MiniCPM weights to transformers format")
+    parser.add_argument(
+        "--vox-config",
+        default="VoxCPM-0.5B/config.json",
+        help="Path to VoxCPM config.json (used to read lm_config)",
+    )
+    parser.add_argument(
+        "--vox-state",
+        default="VoxCPM-0.5B/pytorch_model.bin",
+        help="Path to VoxCPM checkpoint containing base_lm weights",
+    )
+    parser.add_argument(
+        "--minicpm-dir",
+        default="MiniCPM4-0.5B",
+        help="Path to local MiniCPM4-0.5B directory (provides configuration_minicpm.py)",
+    )
+    parser.add_argument(
+        "--out-dir",
+        default="converted-minicpm-hf",
+        help="Output directory for base LM transformers-style checkpoint",
+    )
+    parser.add_argument(
+        "--out-residual-dir",
+        default="converted-minicpm-residual-hf",
+        help="Output directory for residual LM checkpoint",
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    os.makedirs(args.out_residual_dir, exist_ok=True)
+
+    base_cfg, residual_cfg = load_vox_configs(args.vox_config)
+
+    hf_config = build_hf_config(base_cfg, args.minicpm_dir)
+    hf_config.save_pretrained(args.out_dir)
+
+    print("Loaded Vox lm_config and wrote transformers config to", args.out_dir)
+
+    hf_state = convert_state_dict(args.vox_state, lm_prefix="base_lm")
+    out_path = os.path.join(args.out_dir, "pytorch_model.bin")
+    torch.save(hf_state, out_path)
+    print("Saved base LM weights to", out_path)
+
+    residual_hf_config = build_hf_config(residual_cfg, args.minicpm_dir)
+    residual_hf_config.save_pretrained(args.out_residual_dir)
+    residual_state = convert_state_dict(args.vox_state, lm_prefix="residual_lm")
+    residual_out_path = os.path.join(args.out_residual_dir, "pytorch_model.bin")
+    torch.save(residual_state, residual_out_path)
+    print("Saved residual LM weights to", residual_out_path)
+
+    print("Load with MiniCPMForCausalLM.from_pretrained(...) or MiniCPMModel.from_pretrained(...).")
+
+
+if __name__ == "__main__":
+    main()

dit_step.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b7881815a088cfa90b06c804d4fd8a997930a00903facd84836592b01cef4ad
+size 137852823

embed_tokens.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a13e0da6a95e7eb2562468f65e8d64dc412f628057cd9387aaa68cbd4a4cc8f
+size 300843136

export_onnx.py

@@ -0,0 +1,290 @@
+import argparse
+import os
+import torch
+from torch import nn
+
+from voxcpm.model.voxcpm import VoxCPMModel
+
+
+def remove_weight_norm(module: nn.Module):
+    """Strip weight_norm wrappers for cleaner ONNX graphs."""
+    for name, child in module.named_children():
+        remove_weight_norm(child)
+        if isinstance(child, (nn.Conv1d, nn.ConvTranspose1d)):
+            try:
+                torch.nn.utils.remove_weight_norm(child)
+            except ValueError:
+                # not wrapped, skip
+                pass
+
+
+class VAEEncodeWrapper(nn.Module):
+    def __init__(self, audio_vae: nn.Module):
+        super().__init__()
+        self.audio_vae = audio_vae
+
+    def forward(self, audio_wave: torch.Tensor):
+        return self.audio_vae.encode(audio_wave, self.audio_vae.sample_rate)
+
+
+class VAEDecodeWrapper(nn.Module):
+    def __init__(self, audio_vae: nn.Module):
+        super().__init__()
+        self.audio_vae = audio_vae
+
+    def forward(self, latent: torch.Tensor):
+        return self.audio_vae.decode(latent)
+
+
+class LocEncWrapper(nn.Module):
+    def __init__(self, locenc: nn.Module):
+        super().__init__()
+        self.locenc = locenc
+
+    def forward(self, x: torch.Tensor):
+        # x: [B, T, P, D]
+        return self.locenc(x)
+
+
+class LocEncLmWrapper(nn.Module):
+    """LocEnc with enc_to_lm projection fused in a single graph."""
+
+    def __init__(self, locenc: nn.Module, proj: nn.Module):
+        super().__init__()
+        self.locenc = locenc
+        self.proj = proj
+
+    def forward(self, x: torch.Tensor):
+        # x: [B, T, P, D]
+        hidden = self.locenc(x)
+        return self.proj(hidden)
+
+
+class FSQWrapper(nn.Module):
+    def __init__(self, fsq: nn.Module):
+        super().__init__()
+        self.fsq = fsq
+
+    def forward(self, hidden: torch.Tensor):
+        return self.fsq(hidden)
+
+
+class StopHeadWrapper(nn.Module):
+    def __init__(self, stop_proj: nn.Linear, stop_actn: nn.Module, stop_head: nn.Linear):
+        super().__init__()
+        self.stop_proj = stop_proj
+        self.stop_actn = stop_actn
+        self.stop_head = stop_head
+
+    def forward(self, hidden: torch.Tensor):
+        hidden = self.stop_proj(hidden)
+        hidden = self.stop_actn(hidden)
+        return self.stop_head(hidden)
+
+
+class CFMWrapper(nn.Module):
+    """
+    Wrapper for one diffusion step block.
+
+    Note: the number of diffusion steps (n_timesteps) is fixed at export time.
+    """
+
+    def __init__(self, cfm: nn.Module, patch_size: int, n_timesteps: int, cfg_value: float):
+        super().__init__()
+        self.cfm = cfm
+        self.patch_size = patch_size
+        self.n_timesteps = n_timesteps
+        self.cfg_value = cfg_value
+
+    def forward(self, mu: torch.Tensor, cond: torch.Tensor):
+        # mu: [B, H_dit], cond: [B, D_feat, P]
+        return self.cfm(
+            mu=mu,
+            n_timesteps=self.n_timesteps,
+            patch_size=self.patch_size,
+            cond=cond,
+            cfg_value=self.cfg_value,
+        )
+
+
+class DiTStepWrapper(nn.Module):
+    """
+    Wrapper for a single VoxCPMLocDiT forward (one diffusion score estimation step).
+    Inputs match VoxCPMLocDiT.forward: x, mu, t, cond, dt.
+    """
+
+    def __init__(self, dit: nn.Module):
+        super().__init__()
+        self.dit = dit
+
+    def forward(self, x: torch.Tensor, mu: torch.Tensor, t: torch.Tensor, cond: torch.Tensor, dt: torch.Tensor):
+        return self.dit(x, mu, t, cond, dt)
+
+
+def export(model: nn.Module, inputs, path: str, dynamic_axes: dict, opset: int):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    torch.onnx.export(
+        model,
+        inputs,
+        path,
+        opset_version=opset,
+        do_constant_folding=True,
+        input_names=list(dynamic_axes.keys()),
+        output_names=["output"],
+        dynamic_axes=dynamic_axes,
+    )
+    print(f"Saved: {path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Export VoxCPM submodules to ONNX (LLM excluded).")
+    parser.add_argument("--model-dir", required=True, help="Path to VoxCPM model directory (config/weights).")
+    parser.add_argument("--out-dir", default="onnx_exports", help="Output directory for ONNX files.")
+    parser.add_argument("--opset", type=int, default=18, help="ONNX opset version.")
+    parser.add_argument("--audio-samples", type=int, default=1280, help="Dummy audio length for encoder export.")
+    parser.add_argument("--latent-steps", type=int, default=6, help="Dummy latent steps for decoder export.")
+    parser.add_argument("--seq-len", type=int, default=4, help="Dummy sequence length for LocEnc/FSQ export.")
+    parser.add_argument("--dit-step-t", type=float, default=0.5, help="Dummy diffusion time for DiT step export.")
+    parser.add_argument("--force-fp32", action="store_true", help="Force submodules to float32 for ONNX export.")
+    parser.add_argument("--dump-embeddings", action="store_true", help="Dump base_lm.embed_tokens weights to npy.")
+    args = parser.parse_args()
+
+    device = torch.device("cpu")
+    # Load full model once, then peel submodules; keep optimize disabled.
+    full_model = VoxCPMModel.from_local(args.model_dir, optimize=False).to(device).eval()
+    if args.force_fp32 or full_model.config.dtype != "float32":
+        full_model.config.dtype = "float32"
+        full_model = full_model.to(torch.float32)
+        full_model.audio_vae = full_model.audio_vae.to(torch.float32)
+    remove_weight_norm(full_model)
+
+    # Audio VAE encode
+    vae_enc = VAEEncodeWrapper(full_model.audio_vae).to(device).eval()
+    dummy_audio = torch.randn(1, 1, args.audio_samples, device=device)
+    export(
+        vae_enc,
+        dummy_audio,
+        os.path.join(args.out_dir, "audio_vae_encode.onnx"),
+        dynamic_axes={"audio_wave": {0: "batch", 2: "samples"}},
+        opset=args.opset,
+    )
+
+    # Audio VAE decode
+    vae_dec = VAEDecodeWrapper(full_model.audio_vae).to(device).eval()
+    dummy_latent = torch.randn(1, full_model.audio_vae.latent_dim, args.latent_steps, device=device)
+    export(
+        vae_dec,
+        dummy_latent,
+        os.path.join(args.out_dir, "audio_vae_decode.onnx"),
+        dynamic_axes={"latent": {0: "batch", 2: "latent_steps"}},
+        opset=args.opset,
+    )
+
+    # LocEnc with enc_to_lm projection fused
+    locenc = LocEncLmWrapper(full_model.feat_encoder, full_model.enc_to_lm_proj).to(device).eval()
+    dummy_seq = torch.randn(1, args.seq_len, full_model.patch_size, full_model.feat_dim, device=device)
+    export(
+        locenc,
+        dummy_seq,
+        os.path.join(args.out_dir, "locenc.onnx"),
+        dynamic_axes={"x": {0: "batch", 1: "seq_len"}},
+        opset=args.opset,
+    )
+
+    # FSQ layer
+    fsq = FSQWrapper(full_model.fsq_layer).to(device).eval()
+    hidden_size = full_model.config.lm_config.hidden_size
+    dummy_hidden = torch.randn(1, args.seq_len, hidden_size, device=device)
+    export(
+        fsq,
+        dummy_hidden,
+        os.path.join(args.out_dir, "fsq_layer.onnx"),
+        dynamic_axes={"hidden": {0: "batch", 1: "seq_len"}},
+        opset=args.opset,
+    )
+
+    # Stop head
+    stop = StopHeadWrapper(full_model.stop_proj, full_model.stop_actn, full_model.stop_head).to(device).eval()
+    dummy_stop_inp = torch.randn(1, hidden_size, device=device)
+    export(
+        stop,
+        dummy_stop_inp,
+        os.path.join(args.out_dir, "stop_head.onnx"),
+        dynamic_axes={"hidden": {0: "batch"}},
+        opset=args.opset,
+    )
+
+    # Projection layers
+    # export(
+    #     full_model.enc_to_lm_proj,
+    #     dummy_hidden,
+    #     os.path.join(args.out_dir, "enc_to_lm_proj.onnx"),
+    #     dynamic_axes={"input": {0: "batch", 1: "seq_len"}},
+    #     opset=args.opset,
+    # )
+    lm_hidden = torch.randn(1, full_model.config.lm_config.hidden_size, device=device)
+    export(
+        full_model.lm_to_dit_proj,
+        lm_hidden,
+        os.path.join(args.out_dir, "lm_to_dit_proj.onnx"),
+        dynamic_axes={"input": {0: "batch"}},
+        opset=args.opset,
+    )
+    export(
+        full_model.res_to_dit_proj,
+        lm_hidden,
+        os.path.join(args.out_dir, "res_to_dit_proj.onnx"),
+        dynamic_axes={"input": {0: "batch"}},
+        opset=args.opset,
+    )
+
+    # VoxCPMLocDiT single step (score function)
+    dit_step = DiTStepWrapper(full_model.feat_decoder.estimator).to(device).eval()
+    dummy_x = torch.randn(1, full_model.feat_dim, full_model.patch_size, device=device)
+    dummy_mu = torch.randn(1, full_model.config.dit_config.hidden_dim, device=device)
+    dummy_t = torch.full((1,), args.dit_step_t, device=device)
+    dummy_dt = torch.full((1,), 0.0, device=device)
+    dummy_cond = torch.randn(1, full_model.feat_dim, full_model.patch_size, device=device)
+    export(
+        dit_step,
+        (dummy_x, dummy_mu, dummy_t, dummy_cond, dummy_dt),
+        os.path.join(args.out_dir, "dit_step.onnx"),
+        dynamic_axes={
+            "x": {0: "batch"},
+            "mu": {0: "batch"},
+            "t": {0: "batch"},
+            "cond": {0: "batch"},
+            "dt": {0: "batch"},
+        },
+        opset=args.opset,
+    )
+
+    # # UnifiedCFM + VoxCPMLocDiT (single-step sampler unrolled with fixed n_timesteps)
+    # cfm = CFMWrapper(
+    #     full_model.feat_decoder,
+    #     patch_size=full_model.patch_size,
+    #     n_timesteps=args.cfm_steps,
+    #     cfg_value=args.cfg_value,
+    # ).to(device).eval()
+    # dummy_mu = torch.randn(1, full_model.config.dit_config.hidden_dim, device=device)
+    # dummy_cond = torch.randn(1, full_model.feat_dim, full_model.patch_size, device=device)
+    # export(
+    #     cfm,
+    #     (dummy_mu, dummy_cond),
+    #     os.path.join(args.out_dir, "cfm_step.onnx"),
+    #     dynamic_axes={"mu": {0: "batch"}, "cond": {0: "batch"}},
+    #     opset=args.opset,
+    # )
+
+    if args.dump_embeddings and hasattr(full_model.base_lm, "embed_tokens"):
+        import numpy as np
+        emb = full_model.base_lm.embed_tokens.weight.detach().cpu().numpy()
+        os.makedirs(args.out_dir, exist_ok=True)
+        np.save(os.path.join(args.out_dir, "embed_tokens.npy"), emb)
+        print(f"Saved: {os.path.join(args.out_dir, 'embed_tokens.npy')}")
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()

fsq_layer.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfcb8153c4bdc1e9ba992302cf4339d47bdd258a620ef4048d2c002bec8b2867
+size 1282675

librkllmrt.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbcf28a8666b9fbf7361d6aad892b957920f6ea92400c074899b48f4c5b2c96f
+size 7543744

lm_to_dit_proj.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0d5c01b463bd6ca197b4cb6cdd890d5ba0438ea8dda7cd4e88723f5c68d5ef5
+size 2139756

locenc_1.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:482c78a7b9e31bab61886fde2d0463c6314e4fbfb26356316f7c6b2223716a29
+size 131148908

locenc_64.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:780d8e0db5a009595930585c3f2c7999900e7ba941d6ee29156dc256ed1ad909
+size 179347310

onnx_infer-rknn2.py

@@ -0,0 +1,580 @@
+import argparse
+import os
+import math
+import random
+import numpy as np
+import time
+from tqdm import tqdm
+import soundfile as sf
+from scipy import signal
+
+from rkllm_binding import *
+
+from transformers import AutoTokenizer
+
+import ztu_somemodelruntime_rknnlite2 as ort
+
+def mask_multichar_chinese_tokens(tokenizer):
+    # Pre-compute multi-character tokens (length >= 2, pure Chinese characters)
+    multichar_tokens = {
+        token for token in tokenizer.vocab.keys() 
+        if len(token) >= 2 and all("\u4e00" <= c <= "\u9fff" for c in token)
+    }
+
+    class CharTokenizerWrapper:
+        """Wrapper class for tokenizers that handles multi-character Chinese tokens.
+        
+        This wrapper automatically splits multi-character Chinese tokens into
+        individual characters while preserving the original tokenizer's interface.
+        """
+        
+        def __init__(self, base_tokenizer) -> None:
+            """Initialize the wrapper with a base tokenizer.
+            
+            Args:
+                base_tokenizer: The tokenizer to wrap
+            """
+            self.tokenizer = base_tokenizer
+            self.multichar_tokens = multichar_tokens
+
+        def tokenize(self, text: str, **kwargs):
+            """Tokenize text and split multi-character Chinese tokens into single characters.
+            
+            Args:
+                text: Input text to tokenize
+                **kwargs: Additional arguments passed to the base tokenizer
+                
+            Returns:
+                List of processed tokens with multi-character Chinese tokens split
+                
+            Example:
+                >>> wrapper = CharTokenizerWrapper(tokenizer)
+                >>> tokens = wrapper.tokenize("你好世界")
+                >>> # Returns ["你", "好", "世", "界"] instead of ["你好", "世界"]
+            """
+            if not isinstance(text, str):
+                raise TypeError(f"Expected string input, got {type(text)}")
+                
+            tokens = self.tokenizer.tokenize(text, **kwargs)
+            processed = []
+            
+            for token in tokens:
+                # Remove possible subword prefix
+                clean_token = token.replace("▁", "")
+
+                if clean_token in self.multichar_tokens:
+                    # Split multi-character token into single characters
+                    chars = list(clean_token)
+                    processed.extend(chars)
+                else:
+                    processed.append(token)
+                    
+            return processed
+
+        def __call__(self, text: str, **kwargs):
+            """Call the tokenizer and return token IDs.
+            
+            This method provides the same interface as the original tokenizer
+            but with multi-character Chinese token handling.
+            
+            Args:
+                text: Input text to tokenize
+                **kwargs: Additional arguments passed to the base tokenizer
+                
+            Returns:
+                List of token IDs
+                
+            Raises:
+                TypeError: If input is not a string
+                ValueError: If tokenization fails
+            """
+            try:
+                tokens = self.tokenize(text, **kwargs)
+                result = self.tokenizer.convert_tokens_to_ids(tokens)
+                return result
+            except Exception as e:
+                raise ValueError(f"Tokenization failed: {str(e)}") from e
+
+    return CharTokenizerWrapper(tokenizer)
+
+
+def load_rknn(path: str, providers):
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"ONNX file not found: {path}")
+    return ort.InferenceSession(path, providers=providers)
+
+def ensure_numpy(arr, dtype=None):
+    np_arr = np.asarray(arr)
+    if dtype is not None:
+        np_arr = np_arr.astype(dtype, copy=False)
+    return np_arr
+
+
+def run_ort(session: ort.InferenceSession, inputs: dict, name: str = None):
+    start = time.perf_counter()
+    ort_inputs = {k: ensure_numpy(v) for k, v in inputs.items()}
+    outputs = session.run(None, ort_inputs)  # noqa: SLF001
+    if name:
+        elapsed_ms = (time.perf_counter() - start) * 1000
+        print(f"[time] {name}: {elapsed_ms:.2f} ms")
+    return np.asarray(outputs[0])
+
+
+def run_locenc_fixed_len(x: np.ndarray, sess: ort.InferenceSession, max_len: int, name: str):
+    """Run locenc in fixed-length chunks to avoid dynamic seq len requirements."""
+    b, t, p, d = x.shape
+    outputs = []
+    for start in range(0, t, max_len):
+        end = min(start + max_len, t)
+        chunk = x[:, start:end, :, :]
+        if chunk.shape[1] < max_len:
+            pad = np.zeros((b, max_len - chunk.shape[1], p, d), dtype=chunk.dtype)
+            chunk = np.concatenate([chunk, pad], axis=1)
+        out = run_ort(sess, {"x": chunk.astype(np.float32)}, name=f"{name}_{start}")
+        outputs.append(out[:, : end - start, :])
+    return np.concatenate(outputs, axis=1)
+
+
+def run_fsq_fixed_len(hidden: np.ndarray, sess: ort.InferenceSession, max_len: int, name: str):
+    """Run FSQ in fixed-length chunks to avoid dynamic seq len requirements."""
+    b, t, h = hidden.shape
+    outputs = []
+    for start in range(0, t, max_len):
+        end = min(start + max_len, t)
+        chunk = hidden[:, start:end, :]
+        if chunk.shape[1] < max_len:
+            pad = np.zeros((b, max_len - chunk.shape[1], h), dtype=chunk.dtype)
+            chunk = np.concatenate([chunk, pad], axis=1)
+        out = run_ort(sess, {"hidden": chunk.astype(np.float32)}, name=f"{name}_{start}")
+        outputs.append(out[:, : end - start, :])
+    return np.concatenate(outputs, axis=1)
+
+
+def run_vae_encode_chunked(
+    audio: np.ndarray,
+    sess: ort.InferenceSession,
+    chunk_size: int,
+    block_latent_len: int = 64,
+    overlap_latent: int = 4,
+    name: str = "vae_encode",
+):
+    """Encode audio with fixed-length blocks (latent steps) and overlap to avoid dynamic shapes."""
+    block_samples = block_latent_len * chunk_size
+    overlap_samples = overlap_latent * chunk_size
+    stride_samples = block_samples - overlap_samples
+    outputs = []
+    t = audio.shape[-1]
+    for start in range(0, t, stride_samples):
+        end = min(start + block_samples, t)
+        chunk = audio[..., start:end]
+        chunk_len = chunk.shape[-1]
+        if chunk_len < block_samples:
+            chunk = np.pad(chunk, ((0, 0), (0, 0), (0, block_samples - chunk_len)), mode="constant")
+        out = run_ort(sess, {"audio_wave": chunk.astype(np.float32)}, name=f"{name}_{start}")
+        valid_latent = math.ceil(chunk_len / chunk_size)
+        start_keep = overlap_latent if start > 0 else 0
+        if valid_latent <= start_keep:
+            start_keep = max(valid_latent - 1, 0)
+        end_keep = min(valid_latent, block_latent_len)
+        if end_keep > start_keep:
+            outputs.append(out[:, :, start_keep:end_keep])
+    if not outputs:
+        return np.zeros((audio.shape[0], 0, 0), dtype=audio.dtype)
+    return np.concatenate(outputs, axis=2)
+
+
+def run_vae_decode_chunked(
+    latent: np.ndarray,
+    sess: ort.InferenceSession,
+    chunk_size: int,
+    block_latent_len: int = 64,
+    overlap_latent: int = 4,
+    name: str = "vae_decode",
+):
+    """Decode latent in fixed-length blocks with overlap-add style stitching."""
+    block_len = block_latent_len
+    overlap = overlap_latent
+    stride = block_len - overlap
+    outputs = []
+    total = latent.shape[2]
+    for start in range(0, total, stride):
+        end = min(start + block_len, total)
+        chunk = latent[:, :, start:end]
+        valid_latent = end - start
+        if chunk.shape[2] < block_len:
+            pad = np.zeros(
+                (latent.shape[0], latent.shape[1], block_len - chunk.shape[2]),
+                dtype=latent.dtype,
+            )
+            chunk = np.concatenate([chunk, pad], axis=2)
+        out = run_ort(sess, {"latent": chunk.astype(np.float32)}, name=f"{name}_{start}")
+        start_keep = overlap * chunk_size if start > 0 else 0
+        valid_audio = valid_latent * chunk_size
+        if valid_audio <= start_keep:
+            start_keep = max(valid_audio - chunk_size, 0)
+        end_keep = min(out.shape[-1], start_keep + valid_audio)
+        if end_keep > start_keep:
+            outputs.append(out[..., start_keep:end_keep])
+    return np.concatenate(outputs, axis=-1)
+
+
+def cfm_euler_with_onnx_step(
+    dit_sess: ort.InferenceSession,
+    x: np.ndarray,
+    mu: np.ndarray,
+    cond: np.ndarray,
+    n_timesteps: int,
+    cfg_value: float,
+    use_cfg_zero_star: bool = True,
+    mean_mode: bool = False,
+):
+    """
+    Re-implementation of UnifiedCFM.solve_euler using ONNX DiT single step.
+    Shapes:
+        x: [B, C, P], mu: [B, H_dit], cond: [B, C, P]
+    """
+    dtype = x.dtype
+    t_span = np.linspace(1.0, 0.0, n_timesteps + 1, dtype=dtype)
+    t_span = t_span + 1.0 * (np.cos(np.pi / 2 * t_span) - 1 + t_span)  # sway sampling
+
+    t = t_span[0]
+    dt = t_span[0] - t_span[1]
+    zero_init_steps = max(1, int(len(t_span) * 0.04))
+
+    for step in tqdm(range(1, len(t_span))):
+        if use_cfg_zero_star and step <= zero_init_steps:
+            dphi_dt = np.zeros_like(x, dtype=dtype)
+        else:
+            b = x.shape[0]
+            t_in = np.full((b,), t, dtype=dtype)
+            dt_in = np.full((b,), dt if mean_mode else 0.0, dtype=dtype)
+
+            # run conditional branch (pos)
+            dphi_dt_pos = np.asarray(
+                run_ort(
+                    dit_sess,
+                    {
+                        "x": x,
+                        "mu": mu,
+                        "t": t_in,
+                        "cond": cond,
+                        "dt": dt_in,
+                    },
+                ),
+                dtype=dtype,
+            )
+
+            # run "negative" branch (unconditional: mu=0, cond=0)
+            dphi_dt_neg = np.asarray(
+                run_ort(
+                    dit_sess,
+                    {
+                        "x": x,
+                        "mu": np.zeros_like(mu),
+                        "t": t_in,
+                        "cond": np.zeros_like(cond),
+                        "dt": np.zeros_like(dt_in),
+                    },
+                ),
+                dtype=dtype,
+            )
+
+            if use_cfg_zero_star:
+                positive_flat = dphi_dt_pos.reshape(b, -1)
+                negative_flat = dphi_dt_neg.reshape(b, -1)
+                st_star = np.sum(positive_flat * negative_flat, axis=1, keepdims=True) / (
+                    np.sum(negative_flat ** 2, axis=1, keepdims=True) + 1e-8
+                )
+                st_star = st_star.reshape(b, *([1] * (dphi_dt_pos.ndim - 1)))
+            else:
+                st_star = 1.0
+
+            dphi_dt = dphi_dt_neg * st_star + cfg_value * (dphi_dt_pos - dphi_dt_neg * st_star)
+
+        x = x - dt * dphi_dt
+        t = t - dt
+        if step < len(t_span) - 1:
+            dt = t - t_span[step + 1]
+
+    return x
+
+
+def prepare_audio_features(
+    audio_path: str,
+    sample_rate: int,
+    patch_size: int,
+    chunk_size: int,
+    vae_encode_sess: ort.InferenceSession,
+):
+    audio, sr = sf.read(audio_path, always_2d=False)
+    if audio.ndim > 1:
+        audio = audio.mean(axis=1)
+    audio = audio.astype(np.float32)
+    if sr != sample_rate:
+        gcd = math.gcd(int(sr), int(sample_rate))
+        up = sample_rate // gcd
+        down = sr // gcd
+        audio = signal.resample_poly(audio, up, down)
+
+    # Expect shape [B, 1, T] for ONNX VAE encoder
+    if audio.ndim == 1:
+        audio = audio[np.newaxis, np.newaxis, :]
+    else:
+        audio = audio[np.newaxis, ...]
+
+    patch_len = patch_size * chunk_size
+    t = audio.shape[-1]
+    pad_right = (patch_len - t % patch_len) % patch_len
+    if pad_right > 0:
+        audio = np.pad(audio, ((0, 0), (0, 0), (0, pad_right)), mode="constant")
+
+    latent = run_vae_encode_chunked(
+        audio,
+        vae_encode_sess,
+        chunk_size=chunk_size,
+        block_latent_len=64,
+        overlap_latent=4,
+        name="vae_encode",
+    )
+    latent_dim = latent.shape[1]
+    t_latent = latent.shape[2]
+    if t_latent % patch_size != 0:
+        raise ValueError(f"Encoded latent length {t_latent} not divisible by patch_size={patch_size}")
+
+    audio_feat = latent.reshape(latent_dim, -1, patch_size).transpose(1, 2, 0)
+    audio_feat = audio_feat[:-1, ...]  # remove last padding token
+    return audio_feat
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Hybrid ONNX inference for VoxCPM (non-streaming).")
+    parser.add_argument("--tokenizer-dir", required=True, help="Path to tokenizer (e.g., VoxCPM-0.5B).")
+    parser.add_argument("--base-hf-dir", required=True, help="Path to transformers-formatted base MiniCPM.")
+    parser.add_argument("--residual-hf-dir", required=True, help="Path to transformers-formatted residual MiniCPM.")
+    parser.add_argument("--onnx-dir", required=True, help="Directory containing exported ONNX files.")
+    parser.add_argument("--text", required=True, help="Target text to synthesize.")
+    parser.add_argument("--prompt-audio", default=None, help="Optional prompt audio path.")
+    parser.add_argument("--prompt-text", default=None, help="Text transcript of prompt audio (required if prompt-audio).")
+    parser.add_argument("--output", default="onnx_output.wav", help="Output wav path.")
+    parser.add_argument("--cfg-value", type=float, default=2.0, help="CFG value for diffusion.")
+    parser.add_argument("--inference-timesteps", type=int, default=10, help="Diffusion steps.")
+    parser.add_argument("--min-len", type=int, default=2, help="Minimum generated patch count before stop allowed.")
+    parser.add_argument("--max-len", type=int, default=2000, help="Maximum generated patch count.")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
+    parser.add_argument(
+        "--providers",
+        nargs="+",
+        default=None,
+        help="ONNX Runtime providers (e.g., CUDAExecutionProvider CPUExecutionProvider).",
+    )
+    args = parser.parse_args()
+
+    providers = args.providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+    # Seed
+    if args.seed is not None:
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+
+    # Load tokenizer and HF MiniCPM models
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)
+    tokenizer = mask_multichar_chinese_tokens(tokenizer)
+    base_model = RKLLMRuntime()
+    base_model_params = base_model.create_default_param()
+    base_model_params.model_path = (args.base_hf_dir + "/base_lm.rkllm").encode("utf-8")
+    base_model_params.max_context_len = 1024
+    base_model.init(base_model_params, lambda: print("base_model init"))
+    residual_model = RKLLMRuntime()
+    residual_model_params = residual_model.create_default_param()
+    residual_model_params.model_path = (args.residual_hf_dir + "/residual_lm.rkllm").encode("utf-8")
+    residual_model_params.max_context_len = 1024
+    residual_model.init(residual_model_params, lambda: print("residual_model init"))
+
+    dtype = np.float32
+
+    # constants tied to exported ONNX
+    patch_size = 2
+    latent_dim = 64
+    chunk_size = 640
+    sample_rate = 16000
+    audio_start_token = 101
+    fixed_seq_len = 64  # target platform prefers fixed seq len for locenc/fsq
+
+    # Load ONNX sessions
+    vae_encode_sess = load_rknn(os.path.join(args.onnx_dir, "audio_vae_encode.rknn"), providers)
+    vae_decode_sess = load_rknn(os.path.join(args.onnx_dir, "audio_vae_decode.rknn"), providers)
+    locenc_64_sess = load_rknn(os.path.join(args.onnx_dir, "locenc_64.rknn"), providers)
+    locenc_1_sess = load_rknn(os.path.join(args.onnx_dir, "locenc_1.rknn"), providers)
+    fsq_sess = load_rknn(os.path.join(args.onnx_dir, "fsq_layer.rknn"), providers)
+    stop_sess = load_rknn(os.path.join(args.onnx_dir, "stop_head.rknn"), providers)
+    dit_step_sess = load_rknn(os.path.join(args.onnx_dir, "dit_step.rknn"), providers)
+    lm_to_dit_sess = load_rknn(os.path.join(args.onnx_dir, "lm_to_dit_proj.rknn"), providers)
+    res_to_dit_sess = load_rknn(os.path.join(args.onnx_dir, "res_to_dit_proj.rknn"), providers)
+
+    # Build text/audio features
+    if args.prompt_audio:
+        if not args.prompt_text:
+            raise ValueError("prompt-text is required when prompt-audio is provided.")
+        text = args.prompt_text + args.text
+    else:
+        text = args.text
+
+    tokenized = tokenizer(text, add_special_tokens=False)
+    text_token = tokenized["input_ids"] if isinstance(tokenized, dict) else tokenized
+    text_token = np.asarray(text_token, dtype=np.int64)
+    text_token = np.concatenate([text_token, np.array([audio_start_token], dtype=np.int64)], axis=0)
+    text_length = text_token.shape[0]
+
+    if args.prompt_audio:
+        audio_feat = prepare_audio_features(
+            args.prompt_audio, sample_rate, patch_size, chunk_size, vae_encode_sess
+        )
+        audio_length = audio_feat.shape[0]
+
+        text_pad_token = np.zeros(audio_length, dtype=np.int64)
+        text_token = np.concatenate([text_token, text_pad_token])
+
+        audio_pad_feat = np.zeros(
+            (text_length, patch_size, latent_dim),
+            dtype=np.float32,
+        )
+        audio_feat = np.concatenate([audio_pad_feat, audio_feat], axis=0)
+
+        text_mask = np.concatenate([np.ones(text_length, dtype=np.int32), np.zeros(audio_length, dtype=np.int32)])
+        audio_mask = np.concatenate([np.zeros(text_length, dtype=np.int32), np.ones(audio_length, dtype=np.int32)])
+    else:
+        audio_feat = np.zeros(
+            (text_length, patch_size, latent_dim),
+            dtype=np.float32,
+        )
+        text_mask = np.ones(text_length, dtype=np.int32)
+        audio_mask = np.zeros(text_length, dtype=np.int32)
+
+    text_token = text_token[np.newaxis, ...]
+    text_mask = text_mask[np.newaxis, ...]
+    audio_feat = audio_feat[np.newaxis, ...]
+    audio_mask = audio_mask[np.newaxis, ...]
+
+    # LocEnc (ONNX) with enc_to_lm projection fused
+    feat_embed = run_locenc_fixed_len(
+        audio_feat.astype(np.float32), locenc_64_sess, fixed_seq_len, name="locenc"
+    )
+    feat_embed = np.asarray(feat_embed, dtype=dtype)
+
+    # Text embed using exported embedding npy
+    scale_emb = 1.0
+    embedding_path = os.path.join(args.onnx_dir, "embed_tokens.npy")
+    if not os.path.exists(embedding_path):
+        raise FileNotFoundError(f"Embedding npy file not found: {embedding_path}")
+    embedding_weight = np.load(embedding_path).astype(dtype)
+    text_embed = embedding_weight[text_token] * scale_emb
+    
+    combined_embed = text_mask[..., None] * text_embed + audio_mask[..., None] * feat_embed
+
+    # Base LM forward
+    start = time.perf_counter()
+    enc_outputs = np.asarray(base_model.forward_embed(combined_embed.astype(np.float32), keep_history=True), dtype=dtype)
+    enc_outputs = enc_outputs * 4
+    print(f"[time] base_lm initial: {(time.perf_counter() - start)*1000:.2f} ms")
+
+    # FSQ on audio positions
+    enc_outputs_fsq = run_fsq_fixed_len(enc_outputs.astype(np.float32), fsq_sess, fixed_seq_len, name="fsq_init")
+
+    enc_outputs = enc_outputs_fsq * audio_mask[..., None] + enc_outputs * text_mask[..., None]
+
+    # Residual LM forward
+    residual_inputs = enc_outputs + audio_mask[..., None] * feat_embed
+    start = time.perf_counter()
+    residual_outputs = np.asarray(
+        residual_model.forward_embed(residual_inputs.astype(np.float32), keep_history=True),
+        dtype=dtype,
+    )
+    residual_outputs = residual_outputs * 4
+    print(f"[time] residual_lm initial: {(time.perf_counter() - start)*1000:.2f} ms")
+
+    lm_hidden = enc_outputs[:, -1, :].astype(dtype)
+    res_hidden = residual_outputs[:, -1, :].astype(dtype)
+
+    prefix_feat_cond = audio_feat[:, -1, :, :].astype(dtype)  # [B, P, D]
+    pred_feat_seq = []
+
+    # Generation loop
+    for step_idx in tqdm(range(args.max_len), desc="gen_loop"):
+        dit_hidden_lm = run_ort(lm_to_dit_sess, {"input": lm_hidden}, name="lm_to_dit")
+        dit_hidden_res = run_ort(res_to_dit_sess, {"input": res_hidden}, name="res_to_dit")
+        dit_hidden = np.asarray(dit_hidden_lm + dit_hidden_res, dtype=dtype)
+        cond = np.transpose(prefix_feat_cond, (0, 2, 1))  # [B, D, P]
+
+        # Sample next patch via ONNX DiT
+        x0 = np.random.randn(*cond.shape).astype(dtype)  # [B, D, P]
+        pred_feat = cfm_euler_with_onnx_step(
+            dit_step_sess,
+            x0,
+            dit_hidden,
+            cond,
+            n_timesteps=args.inference_timesteps,
+            cfg_value=args.cfg_value,
+            use_cfg_zero_star=True,
+        ).transpose(0, 2, 1)  # -> [B, P, D]
+
+        pred_feat_seq.append(pred_feat[:, np.newaxis, :, :])  # keep time dimension
+        prefix_feat_cond = pred_feat
+
+        # Encode new patch for next step (ONNX locenc)
+        curr_embed = run_ort(locenc_1_sess, {"x": pred_feat[:, np.newaxis, :, :]}, name="locenc_step")
+        curr_embed = np.asarray(curr_embed, dtype=dtype)
+
+        # Stop check (use lm_hidden BEFORE update, consistent with original)
+        stop_logits_np = run_ort(stop_sess, {"hidden": lm_hidden})
+        stop_flag = int(np.argmax(stop_logits_np, axis=-1)[0])
+        if step_idx > args.min_len and stop_flag == 1:
+            break
+
+        lm_hidden_step = np.asarray(
+            base_model.forward_embed(curr_embed.astype(np.float32), keep_history=True),
+            dtype=dtype,
+        )
+        lm_hidden_step = lm_hidden_step * 4
+        lm_hidden = lm_hidden_step[:, -1, :].astype(dtype)
+
+        # FSQ expects [B, T, H]; expand step dimension then squeeze back
+        lm_hidden_step_for_fsq = lm_hidden[:, np.newaxis, :]
+        lm_hidden_fsq_np = run_ort(fsq_sess, {"hidden": lm_hidden_step_for_fsq.astype(np.float32)})
+        lm_hidden_fsq = np.asarray(lm_hidden_fsq_np, dtype=dtype)[:, 0, :]
+
+        res_step_inputs = (lm_hidden_fsq + curr_embed[:, 0, :]).astype(dtype)
+        res_step_inputs = res_step_inputs[:, np.newaxis, :]
+        res_step = np.asarray(
+            residual_model.forward_embed(res_step_inputs.astype(np.float32), keep_history=True),
+            dtype=dtype,
+        )
+        res_step = res_step * 4
+        res_hidden = res_step[:, -1, :].astype(dtype)
+        lm_hidden = lm_hidden_fsq
+
+    if len(pred_feat_seq) == 0:
+        raise RuntimeError("Generation produced zero patches.")
+
+    pred_feat_seq = np.concatenate(pred_feat_seq, axis=1)  # [B, T_gen, P, D]
+    feat_pred = pred_feat_seq.transpose(0, 3, 1, 2).reshape(
+        pred_feat_seq.shape[0], pred_feat_seq.shape[3], -1
+    )  # [B, D, T_gen*P]
+
+    # Decode audio via ONNX
+    audio = run_vae_decode_chunked(
+        feat_pred,
+        vae_decode_sess,
+        chunk_size=chunk_size,
+        block_latent_len=64,
+        overlap_latent=4,
+        name="vae_decode",
+    )
+    audio = audio[..., 640:-640]  # trim start/end
+
+    wav = audio.squeeze()
+    sf.write(args.output, wav, sample_rate)
+    print(f"Saved: {args.output}")
+
+
+if __name__ == "__main__":
+    main()

onnx_infer.py

@@ -0,0 +1,507 @@
+import argparse
+import os
+import random
+import numpy as np
+import torch
+import torchaudio
+import sys
+import pathlib
+import onnxruntime as ort
+import time
+from tqdm import tqdm
+
+from transformers import AutoTokenizer
+from transformers.cache_utils import DynamicCache
+
+from modeling_minicpm import MiniCPMModel  # noqa: E402
+
+def mask_multichar_chinese_tokens(tokenizer):
+    # Pre-compute multi-character tokens (length >= 2, pure Chinese characters)
+    multichar_tokens = {
+        token for token in tokenizer.vocab.keys() 
+        if len(token) >= 2 and all("\u4e00" <= c <= "\u9fff" for c in token)
+    }
+
+    class CharTokenizerWrapper:
+        """Wrapper class for tokenizers that handles multi-character Chinese tokens.
+        
+        This wrapper automatically splits multi-character Chinese tokens into
+        individual characters while preserving the original tokenizer's interface.
+        """
+        
+        def __init__(self, base_tokenizer) -> None:
+            """Initialize the wrapper with a base tokenizer.
+            
+            Args:
+                base_tokenizer: The tokenizer to wrap
+            """
+            self.tokenizer = base_tokenizer
+            self.multichar_tokens = multichar_tokens
+
+        def tokenize(self, text: str, **kwargs):
+            """Tokenize text and split multi-character Chinese tokens into single characters.
+            
+            Args:
+                text: Input text to tokenize
+                **kwargs: Additional arguments passed to the base tokenizer
+                
+            Returns:
+                List of processed tokens with multi-character Chinese tokens split
+                
+            Example:
+                >>> wrapper = CharTokenizerWrapper(tokenizer)
+                >>> tokens = wrapper.tokenize("你好世界")
+                >>> # Returns ["你", "好", "世", "界"] instead of ["你好", "世界"]
+            """
+            if not isinstance(text, str):
+                raise TypeError(f"Expected string input, got {type(text)}")
+                
+            tokens = self.tokenizer.tokenize(text, **kwargs)
+            processed = []
+            
+            for token in tokens:
+                # Remove possible subword prefix
+                clean_token = token.replace("▁", "")
+
+                if clean_token in self.multichar_tokens:
+                    # Split multi-character token into single characters
+                    chars = list(clean_token)
+                    processed.extend(chars)
+                else:
+                    processed.append(token)
+                    
+            return processed
+
+        def __call__(self, text: str, **kwargs):
+            """Call the tokenizer and return token IDs.
+            
+            This method provides the same interface as the original tokenizer
+            but with multi-character Chinese token handling.
+            
+            Args:
+                text: Input text to tokenize
+                **kwargs: Additional arguments passed to the base tokenizer
+                
+            Returns:
+                List of token IDs
+                
+            Raises:
+                TypeError: If input is not a string
+                ValueError: If tokenization fails
+            """
+            try:
+                tokens = self.tokenize(text, **kwargs)
+                result = self.tokenizer.convert_tokens_to_ids(tokens)
+                return result
+            except Exception as e:
+                raise ValueError(f"Tokenization failed: {str(e)}") from e
+
+    return CharTokenizerWrapper(tokenizer)
+
+
+def load_onnx(path: str, providers):
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"ONNX file not found: {path}")
+    return ort.InferenceSession(path, providers=providers)
+
+
+def to_numpy(t: torch.Tensor):
+    return t.detach().cpu().numpy()
+
+
+def run_ort(session: ort.InferenceSession, inputs: dict, name: str = None):
+    start = time.perf_counter()
+    ort_inputs = {k: (v if isinstance(v, np.ndarray) else to_numpy(v)) for k, v in inputs.items()}
+    outputs = session.run(None, ort_inputs)
+    if name:
+        elapsed_ms = (time.perf_counter() - start) * 1000
+        print(f"[time] {name}: {elapsed_ms:.2f} ms")
+    return outputs[0]
+
+
+def cfm_euler_with_onnx_step(
+    dit_sess: ort.InferenceSession,
+    x: torch.Tensor,
+    mu: torch.Tensor,
+    cond: torch.Tensor,
+    n_timesteps: int,
+    cfg_value: float,
+    use_cfg_zero_star: bool = True,
+    mean_mode: bool = False,
+):
+    """
+    Re-implementation of UnifiedCFM.solve_euler using ONNX DiT single step.
+    Shapes:
+        x: [B, C, P], mu: [B, H_dit], cond: [B, C, P]
+    """
+    device = x.device
+    dtype = x.dtype
+    t_span = torch.linspace(1.0, 0.0, n_timesteps + 1, device=device, dtype=dtype)
+    t_span = t_span + 1.0 * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span)  # sway sampling
+
+    t = t_span[0]
+    dt = t_span[0] - t_span[1]
+    zero_init_steps = max(1, int(len(t_span) * 0.04))
+
+    for step in range(1, len(t_span)):
+        if use_cfg_zero_star and step <= zero_init_steps:
+            dphi_dt = torch.zeros_like(x)
+        else:
+            b = x.size(0)
+            t_in = t.expand(b)
+            dt_in = dt.expand(b)
+            if not mean_mode:
+                dt_in = torch.zeros_like(dt_in)
+
+            # run conditional branch (pos)
+            dphi_dt_pos = run_ort(
+                dit_sess,
+                {
+                    "x": x,
+                    "mu": mu,
+                    "t": t_in,
+                    "cond": cond,
+                    "dt": dt_in,
+                },
+                name=f"dit_step_pos_{step}",
+            )
+            dphi_dt_pos = torch.from_numpy(dphi_dt_pos).to(device=device, dtype=dtype)
+
+            # run "negative" branch (unconditional: mu=0, cond=0)
+            dphi_dt_neg = run_ort(
+                dit_sess,
+                {
+                    "x": x,
+                    "mu": torch.zeros_like(mu),
+                    "t": t_in,
+                    "cond": torch.zeros_like(cond),
+                    "dt": torch.zeros_like(dt_in),
+                },
+                name=f"dit_step_neg_{step}",
+            )
+            dphi_dt_neg = torch.from_numpy(dphi_dt_neg).to(device=device, dtype=dtype)
+
+            if use_cfg_zero_star:
+                positive_flat = dphi_dt_pos.view(b, -1)
+                negative_flat = dphi_dt_neg.view(b, -1)
+                st_star = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) / (
+                    torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
+                )
+                st_star = st_star.view(b, *([1] * (dphi_dt_pos.ndim - 1)))
+            else:
+                st_star = 1.0
+
+            dphi_dt = dphi_dt_neg * st_star + cfg_value * (dphi_dt_pos - dphi_dt_neg * st_star)
+
+        x = x - dt * dphi_dt
+        t = t - dt
+        if step < len(t_span) - 1:
+            dt = t - t_span[step + 1]
+
+    return x
+
+
+def prepare_audio_features(
+    audio_path: str,
+    sample_rate: int,
+    patch_size: int,
+    chunk_size: int,
+    vae_encode_sess: ort.InferenceSession,
+):
+    audio, sr = torchaudio.load(audio_path)
+    if audio.size(0) > 1:
+        audio = audio.mean(dim=0, keepdim=True)
+    if sr != sample_rate:
+        audio = torchaudio.functional.resample(audio, sr, sample_rate)
+
+    # Expect shape [B, 1, T] for ONNX VAE encoder
+    if audio.ndim == 2:
+        audio = audio.unsqueeze(0)
+
+    patch_len = patch_size * chunk_size
+    t = audio.size(-1)
+    pad_right = (patch_len - t % patch_len) % patch_len
+    if pad_right > 0:
+        audio = torch.nn.functional.pad(audio, (0, pad_right))
+
+    latent = run_ort(vae_encode_sess, {"audio_wave": audio}, name="vae_encode")
+    latent = torch.from_numpy(latent)
+    latent_dim = latent.shape[1]
+    t_latent = latent.shape[2]
+    if t_latent % patch_size != 0:
+        raise ValueError(f"Encoded latent length {t_latent} not divisible by patch_size={patch_size}")
+
+    audio_feat = latent.view(latent_dim, -1, patch_size).permute(1, 2, 0)
+    audio_feat = audio_feat[:-1, ...]  # remove last padding token
+    return audio_feat
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Hybrid ONNX/PyTorch inference for VoxCPM (non-streaming).")
+    parser.add_argument("--tokenizer-dir", required=True, help="Path to tokenizer (e.g., VoxCPM-0.5B).")
+    parser.add_argument("--base-hf-dir", required=True, help="Path to transformers-formatted base MiniCPM.")
+    parser.add_argument("--residual-hf-dir", required=True, help="Path to transformers-formatted residual MiniCPM.")
+    parser.add_argument("--onnx-dir", required=True, help="Directory containing exported ONNX files.")
+    parser.add_argument("--text", required=True, help="Target text to synthesize.")
+    parser.add_argument("--prompt-audio", default=None, help="Optional prompt audio path.")
+    parser.add_argument("--prompt-text", default=None, help="Text transcript of prompt audio (required if prompt-audio).")
+    parser.add_argument("--output", default="onnx_output.wav", help="Output wav path.")
+    parser.add_argument("--device", default=None, help="torch device; default auto (cuda if available else cpu).")
+    parser.add_argument("--cfg-value", type=float, default=2.0, help="CFG value for diffusion.")
+    parser.add_argument("--inference-timesteps", type=int, default=10, help="Diffusion steps.")
+    parser.add_argument("--min-len", type=int, default=2, help="Minimum generated patch count before stop allowed.")
+    parser.add_argument("--max-len", type=int, default=2000, help="Maximum generated patch count.")
+    parser.add_argument("--force-fp32", action="store_true", help="Force model dtype to float32 for consistency with ONNX.")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
+    parser.add_argument(
+        "--providers",
+        nargs="+",
+        default=None,
+        help="ONNX Runtime providers (e.g., CUDAExecutionProvider CPUExecutionProvider).",
+    )
+    args = parser.parse_args()
+
+    device = torch.device(args.device) if args.device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    providers = args.providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+    # Seed
+    if args.seed is not None:
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(args.seed)
+
+    # Inference with no grad to avoid graph retention / memory growth
+    with torch.inference_mode():
+        # Load tokenizer and HF MiniCPM models
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)
+        tokenizer = mask_multichar_chinese_tokens(tokenizer)
+        base_model = MiniCPMModel.from_pretrained(args.base_hf_dir).to(device).eval()
+        residual_model = MiniCPMModel.from_pretrained(args.residual_hf_dir).to(device).eval()
+        if args.force_fp32:
+            dtype = torch.float32
+        else:
+            dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+        base_model = base_model.to(dtype)
+        residual_model = residual_model.to(dtype)
+
+        # constants tied to exported ONNX
+        patch_size = 2
+        feat_dim = 64
+        latent_dim = 64
+        chunk_size = 640
+        sample_rate = 16000
+        audio_start_token = 101
+
+        # Load ONNX sessions
+        vae_encode_sess = load_onnx(os.path.join(args.onnx_dir, "audio_vae_encode.onnx"), providers)
+        vae_decode_sess = load_onnx(os.path.join(args.onnx_dir, "audio_vae_decode.onnx"), providers)
+        locenc_sess = load_onnx(os.path.join(args.onnx_dir, "locenc.onnx"), providers)
+        fsq_sess = load_onnx(os.path.join(args.onnx_dir, "fsq_layer.onnx"), providers)
+        stop_sess = load_onnx(os.path.join(args.onnx_dir, "stop_head.onnx"), providers)
+        dit_step_sess = load_onnx(os.path.join(args.onnx_dir, "dit_step.onnx"), providers)
+        enc_to_lm_sess = load_onnx(os.path.join(args.onnx_dir, "enc_to_lm_proj.onnx"), providers)
+        lm_to_dit_sess = load_onnx(os.path.join(args.onnx_dir, "lm_to_dit_proj.onnx"), providers)
+        res_to_dit_sess = load_onnx(os.path.join(args.onnx_dir, "res_to_dit_proj.onnx"), providers)
+
+        # Build text/audio features
+        if args.prompt_audio:
+            if not args.prompt_text:
+                raise ValueError("prompt-text is required when prompt-audio is provided.")
+            text = args.prompt_text + args.text
+        else:
+            text = args.text
+
+        tokenized = tokenizer(text, add_special_tokens=False)
+        if isinstance(tokenized, dict):
+            text_token = tokenized["input_ids"]
+        else:
+            text_token = tokenized
+        text_token = torch.LongTensor(text_token)
+        text_token = torch.cat([text_token, torch.tensor([audio_start_token], dtype=torch.int64)], dim=-1)
+        text_length = text_token.shape[0]
+
+        if args.prompt_audio:
+            audio_feat = prepare_audio_features(
+                args.prompt_audio, sample_rate, patch_size, chunk_size, vae_encode_sess
+            )
+            audio_length = audio_feat.size(0)
+
+            text_pad_token = torch.zeros(audio_length, dtype=torch.int64)
+            text_token = torch.cat([text_token, text_pad_token])
+
+            audio_pad_feat = torch.zeros(
+                (text_length, patch_size, latent_dim),
+                dtype=torch.float32,
+            )
+            audio_feat = torch.cat([audio_pad_feat, audio_feat], dim=0)
+
+            text_mask = torch.cat([torch.ones(text_length), torch.zeros(audio_length)]).type(torch.int32)
+            audio_mask = torch.cat([torch.zeros(text_length), torch.ones(audio_length)]).type(torch.int32)
+        else:
+            audio_feat = torch.zeros(
+                (text_length, patch_size, latent_dim),
+                dtype=torch.float32,
+            )
+            text_mask = torch.ones(text_length).type(torch.int32)
+            audio_mask = torch.zeros(text_length).type(torch.int32)
+
+        text_token = text_token.unsqueeze(0).to(device)
+        text_mask = text_mask.unsqueeze(0).to(device)
+        audio_feat = audio_feat.unsqueeze(0).to(device)
+        audio_mask = audio_mask.unsqueeze(0).to(device)
+
+        # LocEnc (ONNX)
+        feat_embed_np = run_ort(locenc_sess, {"x": audio_feat.float()}, name="locenc")
+        feat_embed = torch.from_numpy(feat_embed_np).to(device=device, dtype=dtype)
+        feat_embed = run_ort(enc_to_lm_sess, {"input": feat_embed.float()}, name="enc_to_lm_init")
+        feat_embed = torch.from_numpy(feat_embed).to(device=device, dtype=dtype)
+
+        # Text embed
+        scale_emb = 1.0
+        text_embed = base_model.embed_tokens(text_token) * scale_emb
+        np.save("text_embed_ref.npy", text_embed.cpu().numpy())
+
+        combined_embed = text_mask.unsqueeze(-1) * text_embed + audio_mask.unsqueeze(-1) * feat_embed
+        attn_mask = torch.ones((combined_embed.size(0), combined_embed.size(1)), device=device, dtype=torch.long)
+        np.save("combined_embed_ref.npy", combined_embed.cpu().numpy())
+
+        # Base LM forward
+        start = time.perf_counter()
+        base_cache = DynamicCache()
+        base_outputs = base_model(
+            inputs_embeds=combined_embed,
+            attention_mask=attn_mask,
+            past_key_values=base_cache,
+            use_cache=True,
+            return_dict=True,
+        )
+        enc_outputs = base_outputs.last_hidden_state
+        base_cache = base_outputs.past_key_values
+        print(f"[time] base_lm initial: {(time.perf_counter() - start)*1000:.2f} ms")
+        np.save("enc_outputs_ref.npy", enc_outputs.cpu().numpy())
+
+        # FSQ on audio positions
+        enc_outputs_fsq_np = run_ort(fsq_sess, {"hidden": enc_outputs.float()}, name="fsq_init")
+        enc_outputs_fsq = torch.from_numpy(enc_outputs_fsq_np).to(device=device, dtype=dtype)
+        np.save("enc_outputs_fsq_ref.npy", enc_outputs_fsq.cpu().numpy())
+
+        enc_outputs = enc_outputs_fsq * audio_mask.unsqueeze(-1) + enc_outputs * text_mask.unsqueeze(-1)
+
+        # Residual LM forward
+        np.save("audio_mask_ref.npy", audio_mask.cpu().numpy())
+        np.save("feat_embed_ref.npy", feat_embed.cpu().numpy())
+        residual_inputs = enc_outputs + audio_mask.unsqueeze(-1) * feat_embed
+        start = time.perf_counter()
+        residual_cache = DynamicCache()
+        res_outputs = residual_model(
+            inputs_embeds=residual_inputs,
+            attention_mask=attn_mask,
+            past_key_values=residual_cache,
+            use_cache=True,
+            return_dict=True,
+        )
+        residual_outputs = res_outputs.last_hidden_state
+        residual_cache = res_outputs.past_key_values
+        print(f"[time] residual_lm initial: {(time.perf_counter() - start)*1000:.2f} ms")
+        np.save("residual_inputs_ref.npy", residual_inputs.cpu().numpy())
+        np.save("residual_outputs_ref.npy", residual_outputs.cpu().numpy())
+        np.save("residual_inputs_ref.npy", residual_inputs.cpu().numpy())
+        np.save("residual_outputs_ref.npy", residual_outputs.cpu().numpy())
+
+        lm_hidden = enc_outputs[:, -1, :].to(dtype)
+        res_hidden = residual_outputs[:, -1, :].to(dtype)
+
+        prefix_feat_cond = audio_feat[:, -1, :, :]  # [B, P, D]
+        pred_feat_seq = []
+
+        # Generation loop
+        for step_idx in tqdm(range(args.max_len), desc="gen_loop"):
+            dit_hidden_lm = run_ort(lm_to_dit_sess, {"input": lm_hidden.float()}, name="lm_to_dit")
+            dit_hidden_res = run_ort(res_to_dit_sess, {"input": res_hidden.float()}, name="res_to_dit")
+            dit_hidden = torch.from_numpy(dit_hidden_lm + dit_hidden_res).to(device=device, dtype=dtype)
+            cond = prefix_feat_cond.transpose(1, 2).contiguous()  # [B, D, P]
+
+            # Sample next patch via ONNX DiT
+            x0 = torch.randn_like(prefix_feat_cond.transpose(1, 2))  # [B, D, P]
+            pred_feat = cfm_euler_with_onnx_step(
+                dit_step_sess,
+                x0,
+                dit_hidden,
+                cond,
+                n_timesteps=args.inference_timesteps,
+                cfg_value=args.cfg_value,
+                use_cfg_zero_star=True,
+            ).transpose(1, 2)  # -> [B, P, D]
+
+            pred_feat_seq.append(pred_feat.unsqueeze(1))  # keep time dimension
+            prefix_feat_cond = pred_feat
+
+            # Encode new patch for next step (ONNX locenc)
+            locenc_step_np = run_ort(locenc_sess, {"x": pred_feat.unsqueeze(1).float()}, name="locenc_step")
+            curr_embed = torch.from_numpy(locenc_step_np).to(device=device, dtype=dtype)
+            curr_embed = run_ort(enc_to_lm_sess, {"input": curr_embed.float()}, name="enc_to_lm_step")
+            curr_embed = torch.from_numpy(curr_embed).to(device=device, dtype=dtype)
+
+            # Stop check (use lm_hidden BEFORE update, consistent with original)
+            stop_logits_np = run_ort(stop_sess, {"hidden": lm_hidden.float()})
+            stop_logits = torch.from_numpy(stop_logits_np)
+            stop_flag = stop_logits.argmax(dim=-1)[0].item()
+            if step_idx > args.min_len and stop_flag == 1:
+                break
+
+            # Update LMs using transformers cache API
+            attn_mask = torch.cat(
+                [attn_mask, torch.ones((attn_mask.size(0), 1), device=device, dtype=torch.long)], dim=1
+            )
+            np.save("curr_embed_ref.npy", curr_embed.cpu().numpy())
+            base_step = base_model(
+                inputs_embeds=curr_embed,
+                attention_mask=attn_mask,
+                past_key_values=base_cache,
+                use_cache=True,
+                return_dict=True,
+            )
+            lm_hidden_step = base_step.last_hidden_state  # [B, 1, H]
+            base_cache = base_step.past_key_values
+            lm_hidden = lm_hidden_step.squeeze(1).to(dtype)
+
+            # FSQ expects [B, T, H]; expand step dimension then squeeze back
+            lm_hidden_step = lm_hidden.unsqueeze(1)
+            lm_hidden_fsq_np = run_ort(fsq_sess, {"hidden": lm_hidden_step.float()})
+            lm_hidden_fsq = torch.from_numpy(lm_hidden_fsq_np).to(device=device, dtype=dtype).squeeze(1)
+
+            res_step_inputs = (lm_hidden_fsq + curr_embed[:, 0, :]).unsqueeze(1)
+            res_step = residual_model(
+                inputs_embeds=res_step_inputs,
+                attention_mask=attn_mask,
+                past_key_values=residual_cache,
+                use_cache=True,
+                return_dict=True,
+            )
+            residual_cache = res_step.past_key_values
+            res_hidden = res_step.last_hidden_state.squeeze(1).to(dtype)
+            lm_hidden = lm_hidden_fsq
+
+        if len(pred_feat_seq) == 0:
+            raise RuntimeError("Generation produced zero patches.")
+
+        pred_feat_seq = torch.cat(pred_feat_seq, dim=1)  # [B, T_gen, P, D]
+        feat_pred = pred_feat_seq.permute(0, 3, 1, 2).reshape(
+            pred_feat_seq.size(0), pred_feat_seq.size(3), -1
+        )  # [B, D, T_gen*P]
+
+        # Decode audio via ONNX
+        audio_np = run_ort(vae_decode_sess, {"latent": feat_pred.float()}, name="vae_decode")
+        audio = torch.from_numpy(audio_np)
+        audio = audio[..., 640:-640]  # trim start/end
+
+        wav = audio.squeeze(0).squeeze(0).cpu().numpy()
+        torchaudio.save(args.output, torch.from_numpy(wav).unsqueeze(0), sample_rate)
+        print(f"Saved: {args.output}")
+
+
+if __name__ == "__main__":
+    main()

res_to_dit_proj.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66e9971b7a59007bc0042430a1fdf0d021f4cf2d4875d77e7241f8a600d3a002
+size 2139756

residual_lm.rkllm

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3770af5abc4159e0f7c88da567b0d7cc56f7571502756492f2d991dd93deb4b6
+size 497735540

rkllm_binding.py

@@ -0,0 +1,1277 @@
+import argparse
+import ctypes
+import enum
+import os
+import threading
+from typing import Optional, Sequence, Tuple
+
+import numpy as np
+
+# Define constants from the header
+CPU0 = (1 << 0)  # 0x01
+CPU1 = (1 << 1)  # 0x02
+CPU2 = (1 << 2)  # 0x04
+CPU3 = (1 << 3)  # 0x08
+CPU4 = (1 << 4)  # 0x10
+CPU5 = (1 << 5)  # 0x20
+CPU6 = (1 << 6)  # 0x40
+CPU7 = (1 << 7)  # 0x80
+
+# --- Enums ---
+class LLMCallState(enum.IntEnum):
+    RKLLM_RUN_NORMAL = 0
+    RKLLM_RUN_WAITING = 1
+    RKLLM_RUN_FINISH = 2
+    RKLLM_RUN_ERROR = 3
+
+class RKLLMInputType(enum.IntEnum):
+    RKLLM_INPUT_PROMPT = 0
+    RKLLM_INPUT_TOKEN = 1
+    RKLLM_INPUT_EMBED = 2
+    RKLLM_INPUT_MULTIMODAL = 3
+
+class RKLLMInferMode(enum.IntEnum):
+    RKLLM_INFER_GENERATE = 0
+    RKLLM_INFER_GET_LAST_HIDDEN_LAYER = 1
+    RKLLM_INFER_GET_LOGITS = 2
+
+# --- Structures ---
+class RKLLMExtendParam(ctypes.Structure):
+    base_domain_id: ctypes.c_int32
+    embed_flash: ctypes.c_int8
+    enabled_cpus_num: ctypes.c_int8
+    enabled_cpus_mask: ctypes.c_uint32
+    n_batch: ctypes.c_uint8
+    use_cross_attn: ctypes.c_int8
+    reserved: ctypes.c_uint8 * 104
+
+    _fields_ = [
+        ("base_domain_id", ctypes.c_int32),     # 基础域ID
+        ("embed_flash", ctypes.c_int8),         # 是否从闪存查询词嵌入向量（1启用，0禁用）
+        ("enabled_cpus_num", ctypes.c_int8),    # 推理启用的CPU数量
+        ("enabled_cpus_mask", ctypes.c_uint32), # 指示启用哪些CPU的位掩码
+        ("n_batch", ctypes.c_uint8),            # 一次前向传播中并发处理的输入样本数，设置>1启用批量推理，默认为1
+        ("use_cross_attn", ctypes.c_int8),      # 是否启用交叉注意力（非零启用，0禁用）
+        ("reserved", ctypes.c_uint8 * 104)     # 保留字段
+    ]
+
+class RKLLMParam(ctypes.Structure):
+    model_path: ctypes.c_char_p
+    max_context_len: ctypes.c_int32
+    max_new_tokens: ctypes.c_int32
+    top_k: ctypes.c_int32
+    n_keep: ctypes.c_int32
+    top_p: ctypes.c_float
+    temperature: ctypes.c_float
+    repeat_penalty: ctypes.c_float
+    frequency_penalty: ctypes.c_float
+    presence_penalty: ctypes.c_float
+    mirostat: ctypes.c_int32
+    mirostat_tau: ctypes.c_float
+    mirostat_eta: ctypes.c_float
+    skip_special_token: ctypes.c_bool
+    is_async: ctypes.c_bool
+    img_start: ctypes.c_char_p
+    img_end: ctypes.c_char_p
+    img_content: ctypes.c_char_p
+    extend_param: RKLLMExtendParam
+
+    _fields_ = [
+        ("model_path", ctypes.c_char_p),         # 模型文件路径
+        ("max_context_len", ctypes.c_int32),     # 上下文窗口最大token数
+        ("max_new_tokens", ctypes.c_int32),      # 最大生成新token数
+        ("top_k", ctypes.c_int32),               # Top-K采样参数
+        ("n_keep", ctypes.c_int32),              # 上下文窗口移动时保留的kv缓存数量
+        ("top_p", ctypes.c_float),               # Top-P（nucleus）采样参数
+        ("temperature", ctypes.c_float),         # 采样温度，影响token选择的随机性
+        ("repeat_penalty", ctypes.c_float),      # 重复token惩罚
+        ("frequency_penalty", ctypes.c_float),   # 频繁token惩罚
+        ("presence_penalty", ctypes.c_float),    # 输入中已存在token的惩罚
+        ("mirostat", ctypes.c_int32),            # Mirostat采样策略标志（0表示禁用）
+        ("mirostat_tau", ctypes.c_float),        # Mirostat采样Tau参数
+        ("mirostat_eta", ctypes.c_float),        # Mirostat采样Eta参数
+        ("skip_special_token", ctypes.c_bool),   # 是否跳过特殊token
+        ("is_async", ctypes.c_bool),             # 是否异步推理
+        ("img_start", ctypes.c_char_p),          # 多模态输入中图像的起始位置
+        ("img_end", ctypes.c_char_p),            # 多模态输入中图像的结束位置
+        ("img_content", ctypes.c_char_p),        # 图像内容指针
+        ("extend_param", RKLLMExtendParam)       # 扩展参数
+    ]
+
+class RKLLMLoraAdapter(ctypes.Structure):
+    lora_adapter_path: ctypes.c_char_p
+    lora_adapter_name: ctypes.c_char_p
+    scale: ctypes.c_float
+
+    _fields_ = [
+        ("lora_adapter_path", ctypes.c_char_p),
+        ("lora_adapter_name", ctypes.c_char_p),
+        ("scale", ctypes.c_float)
+    ]
+
+class RKLLMEmbedInput(ctypes.Structure):
+    embed: ctypes.POINTER(ctypes.c_float)
+    n_tokens: ctypes.c_size_t
+
+    _fields_ = [
+        ("embed", ctypes.POINTER(ctypes.c_float)),
+        ("n_tokens", ctypes.c_size_t)
+    ]
+
+class RKLLMTokenInput(ctypes.Structure):
+    input_ids: ctypes.POINTER(ctypes.c_int32)
+    n_tokens: ctypes.c_size_t
+
+    _fields_ = [
+        ("input_ids", ctypes.POINTER(ctypes.c_int32)),
+        ("n_tokens", ctypes.c_size_t)
+    ]
+
+class RKLLMMultiModelInput(ctypes.Structure):
+    prompt: ctypes.c_char_p
+    image_embed: ctypes.POINTER(ctypes.c_float)
+    n_image_tokens: ctypes.c_size_t
+    n_image: ctypes.c_size_t
+    image_width: ctypes.c_size_t
+    image_height: ctypes.c_size_t
+
+    _fields_ = [
+        ("prompt", ctypes.c_char_p),
+        ("image_embed", ctypes.POINTER(ctypes.c_float)),
+        ("n_image_tokens", ctypes.c_size_t),
+        ("n_image", ctypes.c_size_t),
+        ("image_width", ctypes.c_size_t),
+        ("image_height", ctypes.c_size_t)
+    ]
+
+class RKLLMCrossAttnParam(ctypes.Structure):
+    """
+    交叉注意力参数结构体
+    
+    该结构体用于在解码器中执行交叉注意力时使用。
+    它提供编码器输出（键/值缓存）、位置索引和注意力掩码。
+    
+    - encoder_k_cache必须存储在连续内存中，布局为：
+      [num_layers][num_tokens][num_kv_heads][head_dim]
+    - encoder_v_cache必须存储在连续内存中，布局为：
+      [num_layers][num_kv_heads][head_dim][num_tokens]
+    """
+    encoder_k_cache: ctypes.POINTER(ctypes.c_float)
+    encoder_v_cache: ctypes.POINTER(ctypes.c_float)
+    encoder_mask: ctypes.POINTER(ctypes.c_float)
+    encoder_pos: ctypes.POINTER(ctypes.c_int32)
+    num_tokens: ctypes.c_int
+
+    _fields_ = [
+        ("encoder_k_cache", ctypes.POINTER(ctypes.c_float)),  # 编码器键缓存指针（大小：num_layers * num_tokens * num_kv_heads * head_dim）
+        ("encoder_v_cache", ctypes.POINTER(ctypes.c_float)),  # 编码器值缓存指针（大小：num_layers * num_kv_heads * head_dim * num_tokens）
+        ("encoder_mask", ctypes.POINTER(ctypes.c_float)),     # 编码器注意力掩码指针（大小：num_tokens的数组）
+        ("encoder_pos", ctypes.POINTER(ctypes.c_int32)),      # 编码器token位置指针（大小：num_tokens的数组）
+        ("num_tokens", ctypes.c_int)                          # 编码器序列中的token数量
+    ]
+
+class RKLLMPerfStat(ctypes.Structure):
+    """
+    性能统计结构体
+    
+    用于保存预填充和生成阶段的性能统计信息。
+    """
+    prefill_time_ms: ctypes.c_float
+    prefill_tokens: ctypes.c_int
+    generate_time_ms: ctypes.c_float
+    generate_tokens: ctypes.c_int
+    memory_usage_mb: ctypes.c_float
+
+    _fields_ = [
+        ("prefill_time_ms", ctypes.c_float),   # 预填充阶段总耗时（毫秒）
+        ("prefill_tokens", ctypes.c_int),      # 预填充阶段处理的token数量
+        ("generate_time_ms", ctypes.c_float),  # 生成阶段总耗时（毫秒）
+        ("generate_tokens", ctypes.c_int),     # 生成阶段处理的token数量
+        ("memory_usage_mb", ctypes.c_float)    # 推理期间VmHWM常驻内存使用量（MB）
+    ]
+
+class _RKLLMInputUnion(ctypes.Union):
+    prompt_input: ctypes.c_char_p
+    embed_input: RKLLMEmbedInput
+    token_input: RKLLMTokenInput
+    multimodal_input: RKLLMMultiModelInput
+
+    _fields_ = [
+        ("prompt_input", ctypes.c_char_p),
+        ("embed_input", RKLLMEmbedInput),
+        ("token_input", RKLLMTokenInput),
+        ("multimodal_input", RKLLMMultiModelInput)
+    ]
+
+class RKLLMInput(ctypes.Structure):
+    """
+    LLM输入结构体
+    
+    通过联合体表示不同类型的LLM输入。
+    """
+    role: ctypes.c_char_p
+    enable_thinking: ctypes.c_bool
+    input_type: ctypes.c_int
+    _union_data: _RKLLMInputUnion
+
+    _fields_ = [
+        ("role", ctypes.c_char_p),              # 消息角色："user"（用户输入）、"tool"（函数结果）
+        ("enable_thinking", ctypes.c_bool),     # 控制Qwen3模型是否启用"思考模式"
+        ("input_type", ctypes.c_int),           # 枚举类型，指定输入类型（如prompt、token、embed、multimodal）
+        ("_union_data", _RKLLMInputUnion)       # 联合体数据
+    ]
+    # Properties to make accessing union members easier
+    @property
+    def prompt_input(self) -> bytes: # Assuming c_char_p maps to bytes
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_PROMPT:
+            return self._union_data.prompt_input
+        raise AttributeError("Not a prompt input")
+    @prompt_input.setter
+    def prompt_input(self, value: bytes): # Assuming c_char_p maps to bytes
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_PROMPT:
+            self._union_data.prompt_input = value
+        else:
+            raise AttributeError("Not a prompt input")
+    @property
+    def embed_input(self) -> RKLLMEmbedInput:
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_EMBED:
+            return self._union_data.embed_input
+        raise AttributeError("Not an embed input")
+    @embed_input.setter
+    def embed_input(self, value: RKLLMEmbedInput):
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_EMBED:
+            self._union_data.embed_input = value
+        else:
+            raise AttributeError("Not an embed input")
+
+    @property
+    def token_input(self) -> RKLLMTokenInput:
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_TOKEN:
+            return self._union_data.token_input
+        raise AttributeError("Not a token input")
+    @token_input.setter
+    def token_input(self, value: RKLLMTokenInput):
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_TOKEN:
+            self._union_data.token_input = value
+        else:
+            raise AttributeError("Not a token input")
+
+    @property
+    def multimodal_input(self) -> RKLLMMultiModelInput:
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_MULTIMODAL:
+            return self._union_data.multimodal_input
+        raise AttributeError("Not a multimodal input")
+    @multimodal_input.setter
+    def multimodal_input(self, value: RKLLMMultiModelInput):
+        if self.input_type == RKLLMInputType.RKLLM_INPUT_MULTIMODAL:
+            self._union_data.multimodal_input = value
+        else:
+            raise AttributeError("Not a multimodal input")
+
+class RKLLMLoraParam(ctypes.Structure): # For inference
+    lora_adapter_name: ctypes.c_char_p
+
+    _fields_ = [
+        ("lora_adapter_name", ctypes.c_char_p)
+    ]
+
+class RKLLMPromptCacheParam(ctypes.Structure): # For inference
+    save_prompt_cache: ctypes.c_int # bool-like
+    prompt_cache_path: ctypes.c_char_p
+
+    _fields_ = [
+        ("save_prompt_cache", ctypes.c_int), # bool-like
+        ("prompt_cache_path", ctypes.c_char_p)
+    ]
+
+class RKLLMInferParam(ctypes.Structure):
+    mode: ctypes.c_int
+    lora_params: ctypes.POINTER(RKLLMLoraParam)
+    prompt_cache_params: ctypes.POINTER(RKLLMPromptCacheParam)
+    keep_history: ctypes.c_int # bool-like
+
+    _fields_ = [
+        ("mode", ctypes.c_int), # Enum will be passed as int, changed RKLLMInferMode to ctypes.c_int
+        ("lora_params", ctypes.POINTER(RKLLMLoraParam)),
+        ("prompt_cache_params", ctypes.POINTER(RKLLMPromptCacheParam)),
+        ("keep_history", ctypes.c_int) # bool-like
+    ]
+
+class RKLLMResultLastHiddenLayer(ctypes.Structure):
+    hidden_states: ctypes.POINTER(ctypes.c_float)
+    embd_size: ctypes.c_int
+    num_tokens: ctypes.c_int
+
+    _fields_ = [
+        ("hidden_states", ctypes.POINTER(ctypes.c_float)),
+        ("embd_size", ctypes.c_int),
+        ("num_tokens", ctypes.c_int)
+    ]
+
+class RKLLMResultLogits(ctypes.Structure):
+    logits: ctypes.POINTER(ctypes.c_float)
+    vocab_size: ctypes.c_int
+    num_tokens: ctypes.c_int
+
+    _fields_ = [
+        ("logits", ctypes.POINTER(ctypes.c_float)),
+        ("vocab_size", ctypes.c_int),
+        ("num_tokens", ctypes.c_int)
+    ]
+
+class RKLLMResult(ctypes.Structure):
+    """
+    LLM推理结果结构体
+    
+    表示LLM推理的结果，包含生成的文本、token ID、隐藏层状态、logits和性能统计。
+    """
+    text: ctypes.c_char_p
+    token_id: ctypes.c_int32
+    last_hidden_layer: RKLLMResultLastHiddenLayer
+    logits: RKLLMResultLogits
+    perf: RKLLMPerfStat
+
+    _fields_ = [
+        ("text", ctypes.c_char_p),                                  # 生成的文本结果
+        ("token_id", ctypes.c_int32),                              # 生成的token ID
+        ("last_hidden_layer", RKLLMResultLastHiddenLayer),         # 最后一层的隐藏状态（如果请求的话）
+        ("logits", RKLLMResultLogits),                             # 模型输出的logits
+        ("perf", RKLLMPerfStat)                                    # 性能统计（预填充和生成）
+    ]
+
+# --- Typedefs ---
+LLMHandle = ctypes.c_void_p
+
+# --- Callback Function Type ---
+LLMResultCallback = ctypes.CFUNCTYPE(
+    ctypes.c_int,  # 返回类型：int，表示处理状态
+    ctypes.POINTER(RKLLMResult),  # LLM结果指针
+    ctypes.c_void_p,              # 用户数据指针
+    ctypes.c_int                  # LLM调用状态（LLMCallState枚举值）
+)
+"""
+回调函数类型定义
+
+用于处理LLM结果的回调函数。
+
+参数：
+- result: 指向LLM结果的指针
+- userdata: 回调的用户数据指针  
+- state: LLM调用状态（例如：完成、错误）
+
+返回值：
+- 0: 正常继续推理
+- 1: 暂停推理。如果用户想要修改或干预结果（例如编辑输出、注入新提示），
+     返回1以暂停当前推理。稍后，使用更新的内容调用rkllm_run来恢复推理。
+"""
+
+class RKLLMRuntime:
+    def __init__(self, library_path="./librkllmrt.so"):
+        try:
+            self.lib = ctypes.CDLL(library_path)
+        except OSError as e:
+            raise OSError(f"Failed to load RKLLM library from {library_path}. "
+                          f"Ensure it's in your LD_LIBRARY_PATH or provide the full path. Error: {e}")
+        self._setup_functions()
+        self.llm_handle = LLMHandle()
+        self._c_callback = None # To keep the callback object alive
+        self._user_callback = None
+
+    def _setup_functions(self):
+        # RKLLMParam rkllm_createDefaultParam();
+        self.lib.rkllm_createDefaultParam.restype = RKLLMParam
+        self.lib.rkllm_createDefaultParam.argtypes = []
+
+        # int rkllm_init(LLMHandle* handle, RKLLMParam* param, LLMResultCallback callback);
+        self.lib.rkllm_init.restype = ctypes.c_int
+        self.lib.rkllm_init.argtypes = [
+            ctypes.POINTER(LLMHandle),
+            ctypes.POINTER(RKLLMParam),
+            LLMResultCallback
+        ]
+
+        # int rkllm_load_lora(LLMHandle handle, RKLLMLoraAdapter* lora_adapter);
+        self.lib.rkllm_load_lora.restype = ctypes.c_int
+        self.lib.rkllm_load_lora.argtypes = [LLMHandle, ctypes.POINTER(RKLLMLoraAdapter)]
+
+        # int rkllm_load_prompt_cache(LLMHandle handle, const char* prompt_cache_path);
+        self.lib.rkllm_load_prompt_cache.restype = ctypes.c_int
+        self.lib.rkllm_load_prompt_cache.argtypes = [LLMHandle, ctypes.c_char_p]
+
+        # int rkllm_release_prompt_cache(LLMHandle handle);
+        self.lib.rkllm_release_prompt_cache.restype = ctypes.c_int
+        self.lib.rkllm_release_prompt_cache.argtypes = [LLMHandle]
+
+        # int rkllm_destroy(LLMHandle handle);
+        self.lib.rkllm_destroy.restype = ctypes.c_int
+        self.lib.rkllm_destroy.argtypes = [LLMHandle]
+
+        # int rkllm_run(LLMHandle handle, RKLLMInput* rkllm_input, RKLLMInferParam* rkllm_infer_params, void* userdata);
+        self.lib.rkllm_run.restype = ctypes.c_int
+        self.lib.rkllm_run.argtypes = [
+            LLMHandle,
+            ctypes.POINTER(RKLLMInput),
+            ctypes.POINTER(RKLLMInferParam),
+            ctypes.c_void_p # userdata
+        ]
+
+        # int rkllm_run_async(LLMHandle handle, RKLLMInput* rkllm_input, RKLLMInferParam* rkllm_infer_params, void* userdata);
+        # Assuming async also takes userdata for the callback context
+        self.lib.rkllm_run_async.restype = ctypes.c_int
+        self.lib.rkllm_run_async.argtypes = [
+            LLMHandle,
+            ctypes.POINTER(RKLLMInput),
+            ctypes.POINTER(RKLLMInferParam),
+            ctypes.c_void_p # userdata
+        ]
+
+        # int rkllm_abort(LLMHandle handle);
+        self.lib.rkllm_abort.restype = ctypes.c_int
+        self.lib.rkllm_abort.argtypes = [LLMHandle]
+
+        # int rkllm_is_running(LLMHandle handle);
+        self.lib.rkllm_is_running.restype = ctypes.c_int # 0 if running, non-zero otherwise
+        self.lib.rkllm_is_running.argtypes = [LLMHandle]
+
+        # int rkllm_clear_kv_cache(LLMHandle handle, int keep_system_prompt, int* start_pos, int* end_pos);
+        self.lib.rkllm_clear_kv_cache.restype = ctypes.c_int
+        self.lib.rkllm_clear_kv_cache.argtypes = [
+            LLMHandle, 
+            ctypes.c_int,
+            ctypes.POINTER(ctypes.c_int),  # start_pos
+            ctypes.POINTER(ctypes.c_int)   # end_pos
+        ]
+
+        # int rkllm_get_kv_cache_size(LLMHandle handle, int* cache_sizes);
+        self.lib.rkllm_get_kv_cache_size.restype = ctypes.c_int
+        self.lib.rkllm_get_kv_cache_size.argtypes = [LLMHandle, ctypes.POINTER(ctypes.c_int)]
+
+        # int rkllm_set_chat_template(LLMHandle handle, const char* system_prompt, const char* prompt_prefix, const char* prompt_postfix);
+        self.lib.rkllm_set_chat_template.restype = ctypes.c_int
+        self.lib.rkllm_set_chat_template.argtypes = [
+            LLMHandle,
+            ctypes.c_char_p,
+            ctypes.c_char_p,
+            ctypes.c_char_p
+        ]
+
+        # int rkllm_set_function_tools(LLMHandle handle, const char* system_prompt, const char* tools, const char* tool_response_str);
+        self.lib.rkllm_set_function_tools.restype = ctypes.c_int
+        self.lib.rkllm_set_function_tools.argtypes = [
+            LLMHandle,
+            ctypes.c_char_p,  # system_prompt
+            ctypes.c_char_p,  # tools
+            ctypes.c_char_p   # tool_response_str
+        ]
+
+        # int rkllm_set_cross_attn_params(LLMHandle handle, RKLLMCrossAttnParam* cross_attn_params);
+        self.lib.rkllm_set_cross_attn_params.restype = ctypes.c_int
+        self.lib.rkllm_set_cross_attn_params.argtypes = [LLMHandle, ctypes.POINTER(RKLLMCrossAttnParam)]
+
+    def create_default_param(self) -> RKLLMParam:
+        """Creates a default RKLLMParam structure."""
+        return self.lib.rkllm_createDefaultParam()
+
+    def init(self, param: RKLLMParam, callback_func) -> int:
+        """
+        Initializes the LLM.
+        :param param: RKLLMParam structure.
+        :param callback_func: A Python function that matches the signature:
+                              def my_callback(result_ptr, userdata_ptr, state_enum):
+                                  result = result_ptr.contents # RKLLMResult
+                                  # Process result
+                                  # userdata can be retrieved if passed during run, or ignored
+                                  # state = LLMCallState(state_enum)
+        :return: 0 for success, non-zero for failure.
+        """
+        if not callable(callback_func):
+            raise ValueError("callback_func must be a callable Python function.")
+
+        self._user_callback = callback_func
+
+        # Keep a reference to the ctypes callback object to prevent it from being garbage collected.
+        # Always register a trampoline so we can swap the Python-level handler when needed.
+        self._c_callback = LLMResultCallback(self._callback_trampoline)
+
+        ret = self.lib.rkllm_init(ctypes.byref(self.llm_handle), ctypes.byref(param), self._c_callback)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_init failed with error code {ret}")
+        return ret
+
+    def load_lora(self, lora_adapter: RKLLMLoraAdapter) -> int:
+        """Loads a Lora adapter."""
+        ret = self.lib.rkllm_load_lora(self.llm_handle, ctypes.byref(lora_adapter))
+        if ret != 0:
+            raise RuntimeError(f"rkllm_load_lora failed with error code {ret}")
+        return ret
+
+    def load_prompt_cache(self, prompt_cache_path: str) -> int:
+        """Loads a prompt cache from a file."""
+        c_path = prompt_cache_path.encode('utf-8')
+        ret = self.lib.rkllm_load_prompt_cache(self.llm_handle, c_path)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_load_prompt_cache failed for {prompt_cache_path} with error code {ret}")
+        return ret
+
+    def release_prompt_cache(self) -> int:
+        """Releases the prompt cache from memory."""
+        ret = self.lib.rkllm_release_prompt_cache(self.llm_handle)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_release_prompt_cache failed with error code {ret}")
+        return ret
+
+    def destroy(self) -> int:
+        """Destroys the LLM instance and releases resources."""
+        if self.llm_handle and self.llm_handle.value: # Check if handle is not NULL
+            ret = self.lib.rkllm_destroy(self.llm_handle)
+            self.llm_handle = LLMHandle() # Reset handle
+            if ret != 0:
+                # Don't raise here as it might be called in __del__
+                print(f"Warning: rkllm_destroy failed with error code {ret}") 
+            return ret
+        return 0 # Already destroyed or not initialized
+
+    def run(self, rkllm_input: RKLLMInput, rkllm_infer_params: RKLLMInferParam, userdata=None) -> int:
+        """Runs an LLM inference task synchronously."""
+        # userdata can be a ctypes.py_object if you want to pass Python objects,
+        # then cast to c_void_p. Or simply None.
+        if userdata is not None:
+            # Store the userdata object to keep it alive during the call
+            self._userdata_ref = userdata
+            c_userdata = ctypes.cast(ctypes.pointer(ctypes.py_object(userdata)), ctypes.c_void_p)
+        else:
+            c_userdata = None
+        ret = self.lib.rkllm_run(self.llm_handle, ctypes.byref(rkllm_input), ctypes.byref(rkllm_infer_params), c_userdata)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_run failed with error code {ret}")
+        return ret
+
+    def run_async(self, rkllm_input: RKLLMInput, rkllm_infer_params: RKLLMInferParam, userdata=None) -> int:
+        """Runs an LLM inference task asynchronously."""
+        if userdata is not None:
+            # Store the userdata object to keep it alive during the call
+            self._userdata_ref = userdata
+            c_userdata = ctypes.cast(ctypes.pointer(ctypes.py_object(userdata)), ctypes.c_void_p)
+        else:
+            c_userdata = None
+        ret = self.lib.rkllm_run_async(self.llm_handle, ctypes.byref(rkllm_input), ctypes.byref(rkllm_infer_params), c_userdata)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_run_async failed with error code {ret}")
+        return ret
+
+    def abort(self) -> int:
+        """Aborts an ongoing LLM task."""
+        ret = self.lib.rkllm_abort(self.llm_handle)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_abort failed with error code {ret}")
+        return ret
+
+    def is_running(self) -> bool:
+        """Checks if an LLM task is currently running. Returns True if running."""
+        # The C API returns 0 if running, non-zero otherwise.
+        # This is a bit counter-intuitive for a boolean "is_running".
+        return self.lib.rkllm_is_running(self.llm_handle) == 0
+
+    def clear_kv_cache(self, keep_system_prompt: bool, start_pos: list = None, end_pos: list = None) -> int:
+        """
+        清除键值缓存
+        
+        此函数用于清除部分或全部KV缓存。
+        
+        参数：
+        - keep_system_prompt: 是否在缓存中保留系统提示（True保留，False清除）
+                              如果提供了特定范围[start_pos, end_pos)，此标志将被忽略
+        - start_pos: 要清除的KV缓存范围的起始位置数组（包含），每个批次一个
+        - end_pos: 要清除的KV缓存范围的结束位置数组（不包含），每个批次一个
+                   如果start_pos和end_pos都设置为None，将清除整个缓存，keep_system_prompt将生效
+                   如果start_pos[i] < end_pos[i]，只有指定的范围会被清除，keep_system_prompt将被忽略
+        
+        注意：start_pos或end_pos只有在keep_history == 0且生成已通过在回调中返回1暂停时才有效
+        
+        返回：0表示缓存清除成功，非零表示失败
+        """
+        # 准备C数组参数
+        c_start_pos = None
+        c_end_pos = None
+        
+        if start_pos is not None and end_pos is not None:
+            if len(start_pos) != len(end_pos):
+                raise ValueError("start_pos和end_pos数组长度必须相同")
+            
+            # 创建C数组
+            c_start_pos = (ctypes.c_int * len(start_pos))(*start_pos)
+            c_end_pos = (ctypes.c_int * len(end_pos))(*end_pos)
+        
+        ret = self.lib.rkllm_clear_kv_cache(
+            self.llm_handle, 
+            ctypes.c_int(1 if keep_system_prompt else 0),
+            c_start_pos,
+            c_end_pos
+        )
+        if ret != 0:
+            raise RuntimeError(f"rkllm_clear_kv_cache失败，错误代码：{ret}")
+        return ret
+
+    def set_chat_template(self, system_prompt: str, prompt_prefix: str, prompt_postfix: str) -> int:
+        """Sets the chat template for the LLM."""
+        c_system = system_prompt.encode('utf-8') if system_prompt else b""
+        c_prefix = prompt_prefix.encode('utf-8') if prompt_prefix else b""
+        c_postfix = prompt_postfix.encode('utf-8') if prompt_postfix else b""
+        
+        ret = self.lib.rkllm_set_chat_template(self.llm_handle, c_system, c_prefix, c_postfix)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_set_chat_template failed with error code {ret}")
+        return ret
+
+    def get_kv_cache_size(self, n_batch: int) -> list:
+        """
+        获取给定LLM句柄的键值缓存当前大小
+        
+        此函数返回当前存储在模型KV缓存中的位置总数。
+        
+        参数：
+        - n_batch: 批次数量，用于确定返回数组的大小
+        
+        返回：
+        - list: 每个批次的缓存大小列表
+        """
+        # 预分配数组以存储每个批次的缓存大小
+        cache_sizes = (ctypes.c_int * n_batch)()
+        
+        ret = self.lib.rkllm_get_kv_cache_size(self.llm_handle, cache_sizes)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_get_kv_cache_size失败，错误代码：{ret}")
+        
+        # 转换为Python列表
+        return [cache_sizes[i] for i in range(n_batch)]
+
+    def set_function_tools(self, system_prompt: str, tools: str, tool_response_str: str) -> int:
+        """
+        为LLM设置函数调用配置，包括系统提示、工具定义和工具响应token
+        
+        参数：
+        - system_prompt: 定义语言模型上下文或行为的系统提示
+        - tools: JSON格式的字符串，定义可用的函数，包括它们的名称、描述和参数
+        - tool_response_str: 用于识别对话中函数调用结果的唯一标签。它作为标记标签，
+                            允许分词器将工具输出与正常对话轮次分开识别
+        
+        返回：0表示配置设置成功，非零表示错误
+        """
+        c_system = system_prompt.encode('utf-8') if system_prompt else b""
+        c_tools = tools.encode('utf-8') if tools else b""
+        c_tool_response = tool_response_str.encode('utf-8') if tool_response_str else b""
+        
+        ret = self.lib.rkllm_set_function_tools(self.llm_handle, c_system, c_tools, c_tool_response)
+        if ret != 0:
+            raise RuntimeError(f"rkllm_set_function_tools失败，错误代码：{ret}")
+        return ret
+
+    def set_cross_attn_params(self, cross_attn_params: RKLLMCrossAttnParam) -> int:
+        """
+        为LLM解码器设置交叉注意力参数
+        
+        参数：
+        - cross_attn_params: 包含用于交叉注意力的编码器相关输入数据的结构体
+                            （详见RKLLMCrossAttnParam说明）
+        
+        返回：0表示参数设置成功，非零表示错误
+        """
+        ret = self.lib.rkllm_set_cross_attn_params(self.llm_handle, ctypes.byref(cross_attn_params))
+        if ret != 0:
+            raise RuntimeError(f"rkllm_set_cross_attn_params失败，错误代码：{ret}")
+        return ret
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.destroy()
+
+    def __del__(self):
+        self.destroy() # Ensure resources are freed if object is garbage collected
+
+    def _callback_trampoline(self, result_ptr, userdata_ptr, state_enum):
+        """
+        Bridge callback that forwards to the currently active Python handler.
+        This keeps the C callback pointer stable while allowing per-call overrides.
+        """
+        handler = self._user_callback
+        if handler is None:
+            return 0
+        try:
+            return handler(result_ptr, userdata_ptr, state_enum)
+        except Exception as exc:
+            # Avoid propagating exceptions through the C callback boundary.
+            print(f"[rkllm_binding] Callback raised an exception: {exc}")
+            return 0
+
+    def forward_embed(
+        self,
+        embeds: np.ndarray,
+        *,
+        keep_history: bool = False,
+        timeout: Optional[float] = None,
+        return_last_only: bool = False,
+    ) -> np.ndarray:
+        """
+        Run a single forward pass with embedding input and return the last hidden layer.
+
+        Args:
+            embeds: Float32 embeddings shaped (T, H) or (1, T, H). Batch>1 is not supported.
+            keep_history: When False, KV cache will be cleared after the call. When True,
+                          cache is kept; call clear_kv_cache() manually if needed.
+            timeout: Optional timeout (seconds) for waiting on the callback.
+            return_last_only: If True, return the last token vector shape (H,).
+
+        Returns:
+            np.ndarray containing hidden states (T, H) or the last token (H,).
+        """
+        if embeds is None:
+            raise ValueError("embeds must not be None.")
+
+        np_embeds = np.asarray(embeds, dtype=np.float32)
+        if np_embeds.ndim == 3:
+            if np_embeds.shape[0] != 1:
+                raise ValueError("Only batch size 1 is supported for forward_embed.")
+            num_tokens = np_embeds.shape[1]
+            flat = np_embeds.reshape(-1)
+        elif np_embeds.ndim == 2:
+            num_tokens = np_embeds.shape[0]
+            flat = np_embeds.reshape(-1)
+        else:
+            raise ValueError("embeds must have shape (T, H) or (1, T, H).")
+
+        flat = np.ascontiguousarray(flat, dtype=np.float32)
+        embed_buffer = (ctypes.c_float * flat.size)(*flat)
+
+        rk_input = RKLLMInput()
+        rk_input.input_type = RKLLMInputType.RKLLM_INPUT_EMBED
+        embed_input = RKLLMEmbedInput()
+        embed_input.embed = embed_buffer
+        embed_input.n_tokens = num_tokens
+        rk_input._union_data.embed_input = embed_input
+
+        infer_params = RKLLMInferParam()
+        infer_params.mode = RKLLMInferMode.RKLLM_INFER_GET_LAST_HIDDEN_LAYER
+        infer_params.keep_history = 1 if keep_history else 0
+        infer_params.lora_params = None
+        infer_params.prompt_cache_params = None
+
+        done = threading.Event()
+        result_holder = {"hidden": None, "error": None}
+
+        def _capture_hidden(result_ptr, userdata_ptr, state_enum):
+            state = LLMCallState(state_enum)
+            if state == LLMCallState.RKLLM_RUN_ERROR:
+                result_holder["error"] = "RKLLM reported an error state."
+                done.set()
+                return 0
+
+            if not result_ptr:
+                result_holder["error"] = "Empty result pointer received."
+                done.set()
+                return 0
+
+            result = result_ptr.contents
+            if result.last_hidden_layer.hidden_states and result.last_hidden_layer.embd_size > 0:
+                hidden = np.ctypeslib.as_array(
+                    result.last_hidden_layer.hidden_states,
+                    shape=(1, result.last_hidden_layer.num_tokens, result.last_hidden_layer.embd_size),
+                ).copy()
+                result_holder["hidden"] = hidden[-1].copy() if return_last_only else hidden
+                done.set()
+                return 1  # Pause further work; we already have the hidden states.
+
+            if state == LLMCallState.RKLLM_RUN_FINISH:
+                done.set()
+            return 0
+
+        previous_callback = self._user_callback
+        self._user_callback = _capture_hidden
+        try:
+            self.run(rk_input, infer_params)
+            if not done.wait(timeout):
+                raise TimeoutError("forward_embed timed out waiting for hidden states.")
+        finally:
+            self._user_callback = previous_callback
+
+        if result_holder["error"]:
+            raise RuntimeError(result_holder["error"])
+        if result_holder["hidden"] is None:
+            raise RuntimeError("forward_embed did not receive hidden states.")
+
+        try:
+            if not keep_history:
+                self.clear_kv_cache(True)
+        except Exception:
+            # Cache clearing best-effort; keep the forward result usable even if clearing fails.
+            pass
+
+        return result_holder["hidden"]
+
+# --- Demo CLI ---
+def _cli_parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Demo application showcasing rkllm_binding usage."
+    )
+    parser.add_argument(
+        "model",
+        help="Path to the .rkllm model file used for inference."
+    )
+    parser.add_argument(
+        "--lib",
+        default="./librkllmrt.so",
+        help="Path to librkllmrt.so. Defaults to ./librkllmrt.so."
+    )
+
+    # Core generation parameters
+    parser.add_argument("--max-context-len", type=int, default=512, help="Maximum context length.")
+    parser.add_argument("--max-new-tokens", type=int, default=128, help="Maximum number of new tokens to generate.")
+    parser.add_argument("--top-k", type=int, default=1, help="Top-K sampling parameter.")
+    parser.add_argument("--top-p", type=float, default=0.0, help="Top-P (nucleus) sampling parameter.")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature.")
+    parser.add_argument("--repeat-penalty", type=float, default=1.1, help="Penalty applied to repeated tokens.")
+    parser.add_argument("--n-keep", type=int, default=0, help="Number of tokens to keep when context slides.")
+    parser.add_argument("--mirostat", type=int, default=0, help="Enable Mirostat sampling (0 disables).")
+    parser.add_argument("--mirostat-tau", type=float, default=5.0, help="Mirostat tau parameter.")
+    parser.add_argument("--mirostat-eta", type=float, default=0.1, help="Mirostat eta parameter.")
+    parser.add_argument(
+        "--skip-special-token",
+        action="store_true",
+        help="Skip special tokens when generating output."
+    )
+
+    # Input management
+    parser.add_argument(
+        "--input-type",
+        choices=("prompt", "token", "multimodal"),
+        default="prompt",
+        help="Select prompt, raw token, or multimodal (image + prompt) input."
+    )
+    parser.add_argument("--prompt", help="Prompt text to send to the model.")
+    parser.add_argument("--prompt-file", help="Path to a UTF-8 text file containing the prompt.")
+    parser.add_argument(
+        "--token-ids",
+        type=int,
+        nargs="+",
+        help="Raw token IDs (space separated). Only valid when --input-type token."
+    )
+    parser.add_argument("--role", default="user", help="Role metadata for the input message (e.g., user/system).")
+    parser.add_argument(
+        "--enable-thinking",
+        action="store_true",
+        help="Enable thinking mode for supported models."
+    )
+    parser.add_argument("--image", help="Path to an image file used when --input-type multimodal.")
+    parser.add_argument("--vision-encoder", help="Path to the ONNX vision encoder model.")
+    parser.add_argument(
+        "--encoder-provider",
+        help="Comma separated ONNX Runtime providers (e.g., 'CPUExecutionProvider')."
+    )
+    parser.add_argument(
+        "--encoder-threads",
+        type=int,
+        help="Thread count hint for ONNX Runtime session."
+    )
+    parser.add_argument(
+        "--encoder-input-shape",
+        help="Override encoder input spatial size as HxW or H,W (e.g., 392x392)."
+    )
+    parser.add_argument(
+        "--norm",
+        choices=("imagenet", "divide_255", "divide_128_sub_1"),
+        default="imagenet",
+        help="Image normalization preset."
+    )
+    parser.add_argument(
+        "--norm-mean",
+        type=float,
+        nargs=3,
+        metavar=("R", "G", "B"),
+        help="Override normalization mean (RGB order)."
+    )
+    parser.add_argument(
+        "--norm-std",
+        type=float,
+        nargs=3,
+        metavar=("R", "G", "B"),
+        help="Override normalization std (RGB order)."
+    )
+    parser.add_argument(
+        "--image-background",
+        type=int,
+        nargs=3,
+        metavar=("R", "G", "B"),
+        default=(128, 128, 128),
+        help="Background color used when padding image to target size."
+    )
+    parser.add_argument("--img-start-token", help="Override image start token string passed to the model.")
+    parser.add_argument("--img-end-token", help="Override image end token string passed to the model.")
+    parser.add_argument("--img-content-token", help="Override image content token string passed to the model.")
+
+    # Inference options
+    parser.add_argument(
+        "--mode",
+        choices=("generate", "hidden", "logits"),
+        default="generate",
+        help="Inference mode: generate tokens, return last hidden layer, or logits."
+    )
+    parser.add_argument(
+        "--no-keep-history",
+        action="store_true",
+        help="Do not keep dialogue history on the device."
+    )
+
+    # Output options
+    parser.add_argument(
+        "--stream",
+        action="store_true",
+        default=True,
+        help="Stream tokens to stdout as they arrive from the callback."
+    )
+    parser.add_argument(
+        "--hide-stats",
+        action="store_true",
+        help="Suppress performance statistics after inference."
+    )
+
+    args = parser.parse_args()
+
+    if args.prompt and args.prompt_file:
+        parser.error("Arguments --prompt and --prompt-file cannot be used together.")
+
+    if args.input_type == "prompt":
+        if not args.prompt and not args.prompt_file:
+            parser.error("Provide --prompt or --prompt-file when --input-type is prompt.")
+        if args.token_ids:
+            parser.error("--token-ids is only valid when --input-type token.")
+    elif args.input_type == "token":
+        if not args.token_ids:
+            parser.error("--token-ids is required when --input-type token.")
+        if args.prompt or args.prompt_file:
+            parser.error("--prompt/--prompt-file cannot be combined with --input-type token.")
+    else:  # multimodal
+        if args.token_ids:
+            parser.error("--token-ids cannot be used with --input-type multimodal.")
+        if not args.prompt and not args.prompt_file:
+            parser.error("Provide --prompt or --prompt-file when --input-type is multimodal.")
+        if not args.image:
+            parser.error("--image is required when --input-type multimodal.")
+        if not args.vision_encoder:
+            parser.error("--vision-encoder is required when --input-type multimodal.")
+
+    if args.image_background:
+        for component in args.image_background:
+            if component < 0 or component > 255:
+                parser.error("--image-background values must be in the range [0, 255].")
+
+    return args
+
+
+def _load_prompt_from_args(args: argparse.Namespace) -> str:
+    if args.prompt:
+        return args.prompt
+    if args.prompt_file:
+        try:
+            with open(args.prompt_file, "r", encoding="utf-8") as fp:
+                return fp.read()
+        except OSError as exc:
+            raise RuntimeError(f"Failed to read prompt file '{args.prompt_file}': {exc}") from exc
+    raise RuntimeError("Prompt text is required but not provided.")
+
+
+def _mode_to_enum(mode: str) -> int:
+    mapping = {
+        "generate": RKLLMInferMode.RKLLM_INFER_GENERATE,
+        "hidden": RKLLMInferMode.RKLLM_INFER_GET_LAST_HIDDEN_LAYER,
+        "logits": RKLLMInferMode.RKLLM_INFER_GET_LOGITS,
+    }
+    return mapping[mode]
+
+
+def _parse_hw_string(value: str) -> Tuple[int, int]:
+    separators = ("x", "X", ",", " ")
+    token = value.strip()
+    for sep in separators:
+        if sep in token:
+            parts = [p for p in token.split(sep) if p]
+            break
+    else:
+        parts = [token]
+    if len(parts) != 2:
+        raise ValueError(f"Unable to parse height/width from '{value}'. Expected format like 392x392.")
+    try:
+        height = int(parts[0])
+        width = int(parts[1])
+    except ValueError as exc:
+        raise ValueError(f"Height/width must be integers, got '{value}'.") from exc
+    if height <= 0 or width <= 0:
+        raise ValueError("Height and width must be positive integers.")
+    return height, width
+
+
+def _infer_hw_from_onnx_shape(shape: Sequence) -> Tuple[Optional[int], Optional[int]]:
+    if shape is None or len(shape) < 4:
+        return None, None
+    height = shape[-2]
+    width = shape[-1]
+    if isinstance(height, str) or height is None:
+        height = None
+    if isinstance(width, str) or width is None:
+        width = None
+    return height, width
+
+
+def _parse_providers(provider_str: Optional[str]) -> Optional[list]:
+    if not provider_str:
+        return None
+    providers = [item.strip() for item in provider_str.split(",") if item.strip()]
+    return providers or None
+
+
+def _load_vision_encoder_session(encoder_path: str, providers: Optional[list], threads: Optional[int]):
+    try:
+        import onnxruntime as ort
+    except ImportError as exc:
+        raise RuntimeError("onnxruntime is required for multimodal inference. Please install onnxruntime.") from exc
+
+    sess_options = ort.SessionOptions()
+    if threads and threads > 0:
+        sess_options.intra_op_num_threads = threads
+    try:
+        if providers:
+            session = ort.InferenceSession(encoder_path, sess_options=sess_options, providers=providers)
+        else:
+            session = ort.InferenceSession(encoder_path, sess_options=sess_options)
+    except Exception as exc:
+        raise RuntimeError(f"Failed to load vision encoder '{encoder_path}': {exc}") from exc
+    return session
+
+
+def _letterbox_resize(image, target_hw: Tuple[int, int], background_color: Sequence[int]):
+    try:
+        import cv2
+        import numpy as np
+    except ImportError as exc:
+        raise RuntimeError("OpenCV (cv2) and numpy are required for multimodal preprocessing.") from exc
+
+    target_h, target_w = target_hw
+    if image.ndim != 3 or image.shape[2] != 3:
+        raise RuntimeError("Expected RGB image with 3 channels.")
+
+    src_h, src_w = image.shape[:2]
+    if src_h == 0 or src_w == 0:
+        raise RuntimeError("Loaded image has invalid dimensions.")
+
+    scale = min(target_w / src_w, target_h / src_h)
+    resized_w = max(1, int(round(src_w * scale)))
+    resized_h = max(1, int(round(src_h * scale)))
+    resized = cv2.resize(image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR)
+
+    canvas = np.full((target_h, target_w, 3), background_color, dtype=resized.dtype)
+    top = (target_h - resized_h) // 2
+    left = (target_w - resized_w) // 2
+    canvas[top:top + resized_h, left:left + resized_w] = resized
+    return canvas, resized_h, resized_w
+
+
+def _normalize_image(image, method: str, mean: Optional[Sequence[float]], std: Optional[Sequence[float]]):
+    import numpy as np
+
+    img = image.astype(np.float32)
+    mean_arr = np.array(mean, dtype=np.float32) if mean else None
+    std_arr = np.array(std, dtype=np.float32) if std else None
+
+    if method == "imagenet":
+        img = img / 255.0
+        if mean_arr is None:
+            mean_arr = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
+        if std_arr is None:
+            std_arr = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
+        img = (img - mean_arr) / std_arr
+    elif method == "divide_255":
+        img = img / 255.0
+        if mean_arr is not None:
+            img = img - mean_arr
+        if std_arr is not None:
+            img = img / std_arr
+    elif method == "divide_128_sub_1":
+        img = img / 128.0 - 1.0
+        if mean_arr is not None:
+            img = img - mean_arr
+        if std_arr is not None:
+            img = img / std_arr
+    else:
+        raise RuntimeError(f"Unsupported normalization method '{method}'.")
+
+    return img
+
+
+def _encode_image_to_embedding(
+    session,
+    image_path: str,
+    input_name: str,
+    output_name: str,
+    target_hw: Tuple[int, int],
+    background_color: Sequence[int],
+    norm_method: str,
+    norm_mean: Optional[Sequence[float]],
+    norm_std: Optional[Sequence[float]]
+):
+    try:
+        import cv2
+        import numpy as np
+    except ImportError as exc:
+        raise RuntimeError("OpenCV (cv2) and numpy are required for multimodal preprocessing.") from exc
+
+    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
+    if image is None:
+        raise RuntimeError(f"Failed to read image from '{image_path}'.")
+
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    padded, resized_h, resized_w = _letterbox_resize(image, target_hw, background_color)
+
+    normalized = _normalize_image(padded, norm_method, norm_mean, norm_std)
+    tensor = np.transpose(normalized, (2, 0, 1))  # HWC -> CHW
+    tensor = np.expand_dims(tensor, axis=0)  # Add batch dimension
+    tensor = np.ascontiguousarray(tensor, dtype=np.float32)
+
+    try:
+        output_list = session.run([output_name], {input_name: tensor})
+    except Exception as exc:
+        raise RuntimeError(f"Vision encoder inference failed: {exc}") from exc
+
+    if not output_list:
+        raise RuntimeError("Vision encoder returned no outputs.")
+
+    embedding = output_list[0]
+    if embedding.ndim == 3:
+        if embedding.shape[0] != 1:
+            raise RuntimeError("Vision encoder output batch dimension must be 1 for a single image.")
+        n_tokens = embedding.shape[1]
+    elif embedding.ndim == 2:
+        n_tokens = embedding.shape[0]
+    else:
+        raise RuntimeError(f"Unsupported vision encoder output shape {embedding.shape}.")
+
+    flat_embedding = embedding.reshape(-1).astype(np.float32, copy=False)
+    flat_embedding = np.ascontiguousarray(flat_embedding)
+
+    return flat_embedding, n_tokens, target_hw
+
+if __name__ == "__main__":
+    import os
+    os.environ["RKLLM_LOG_LEVEL"] = "1"
+    args = _cli_parse_arguments()
+
+    prompt_text = None
+    if args.input_type == "prompt":
+        prompt_text = _load_prompt_from_args(args)
+
+    token_id_array = None
+    token_input_struct = None
+
+    generated_chunks = []
+    perf_snapshot = {
+        "prefill_tokens": 0,
+        "prefill_time_ms": 0.0,
+        "generate_tokens": 0,
+        "generate_time_ms": 0.0,
+        "memory_usage_mb": 0.0,
+    }
+
+    def demo_callback(result_ptr, userdata_ptr, state_enum):
+        state = LLMCallState(state_enum)
+        result = result_ptr.contents
+
+        current_text = ""
+        if result.text:
+            current_text = result.text.decode("utf-8", errors="ignore")
+            generated_chunks.append(current_text)
+            if args.stream and current_text:
+                print(current_text, end="", flush=True)
+
+        perf_snapshot.update(
+            prefill_tokens=result.perf.prefill_tokens,
+            prefill_time_ms=result.perf.prefill_time_ms,
+            generate_tokens=result.perf.generate_tokens,
+            generate_time_ms=result.perf.generate_time_ms,
+            memory_usage_mb=result.perf.memory_usage_mb,
+        )
+
+        if state == LLMCallState.RKLLM_RUN_ERROR:
+            print("\n[Callback] 推理过程中出现错误。")
+
+        return 0
+
+    try:
+        with RKLLMRuntime(library_path=args.lib) as rk_llm:
+            params = rk_llm.create_default_param()
+            params.model_path = os.path.abspath(args.model).encode("utf-8")
+            params.max_context_len = args.max_context_len
+            params.max_new_tokens = args.max_new_tokens
+            params.top_k = args.top_k
+            params.top_p = float(args.top_p)
+            params.temperature = float(args.temperature)
+            params.repeat_penalty = float(args.repeat_penalty)
+            params.n_keep = args.n_keep
+            params.mirostat = args.mirostat
+            params.mirostat_tau = float(args.mirostat_tau)
+            params.mirostat_eta = float(args.mirostat_eta)
+            params.skip_special_token = bool(args.skip_special_token)
+            params.is_async = False
+
+            rk_llm.init(params, demo_callback)
+
+            rk_input = RKLLMInput()
+            rk_input.role = args.role.encode("utf-8")
+            rk_input.enable_thinking = bool(args.enable_thinking)
+
+            if args.input_type == "prompt":
+                rk_input.input_type = RKLLMInputType.RKLLM_INPUT_PROMPT
+                rk_input._union_data.prompt_input = prompt_text.encode("utf-8")
+            else:
+                rk_input.input_type = RKLLMInputType.RKLLM_INPUT_TOKEN
+                token_id_array = (ctypes.c_int32 * len(args.token_ids))(*args.token_ids)
+                token_input_struct = RKLLMTokenInput()
+                token_input_struct.input_ids = token_id_array
+                token_input_struct.n_tokens = len(args.token_ids)
+                rk_input._union_data.token_input = token_input_struct
+
+            infer_params = RKLLMInferParam()
+            infer_params.mode = _mode_to_enum(args.mode)
+            infer_params.keep_history = 0 if args.no_keep_history else 1
+            infer_params.lora_params = None
+            infer_params.prompt_cache_params = None
+
+            if args.stream:
+                print("=== Streaming Output ===")
+
+            rk_llm.run(rk_input, infer_params)
+
+    except OSError as exc:
+        print(f"无法加载 RKLLM 运行时库：{exc}")
+    except RuntimeError as exc:
+        print(f"推理失败：{exc}")
+    except Exception as exc:
+        print(f"发生未预期的错误：{exc}")
+    else:
+        if args.stream:
+            print()  # Ensure newline after streaming output
+
+        final_text = "".join(generated_chunks)
+        if final_text:
+            print("=== 生成结果 ===")
+            print(final_text)
+        else:
+            print("未收到生成文本。")
+
+        if not args.hide_stats:
+            print("=== 性能统计 ===")
+            print(
+                f"预填充: {perf_snapshot['prefill_tokens']} tokens / {perf_snapshot['prefill_time_ms']:.2f} ms"
+            )
+            print(
+                f"生成: {perf_snapshot['generate_tokens']} tokens / {perf_snapshot['generate_time_ms']:.2f} ms"
+            )
+            print(f"最大常驻内存: {perf_snapshot['memory_usage_mb']:.2f} MB")

rknn_output.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d58c10f7fa051c0181f4752c14433610c9af9577ebe3306e750865c7a76a2dfa
+size 320044

stop_head.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72d604bcb8531a27d8a2d678d6a51839eff0217cf0ccfd722e61f18f3be9f782
+size 2156714

tokenizer_config.json

@@ -0,0 +1,212 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "<|audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "<|audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "<|audio_prompt_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "104": {
+      "content": "<|audio_prompt_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "105": {
+      "content": "<|background|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "106": {
+      "content": "<|/background|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "107": {
+      "content": "<|characters|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "108": {
+      "content": "<|/characters|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "109": {
+      "content": "<|speaker_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },   
+    "110": {
+      "content": "<|/speaker_id|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "111": {
+      "content": "<|span|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "112": {
+      "content": "<|/span|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73440": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73441": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73442": {
+      "content": "<|tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73443": {
+      "content": "<|execute_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73444": {
+      "content": "<|execute_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73445": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73446": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "73447": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_end|>",
+    "<|im_start|>",
+    "<|tool_call|>",
+    "<|execute_start|>",
+    "<|execute_end|>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>"
+  ],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": null,
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false,
+  "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+}

ztu_somemodelruntime_rknnlite2.py

@@ -0,0 +1,1334 @@
+# 模块级常量和函数
+from rknnlite.api import RKNNLite
+import numpy as np
+import os
+import warnings
+import logging
+from typing import List, Dict, Union, Optional, Tuple
+import ctypes
+
+try:
+    import onnxruntime as ort
+    HAS_ORT = True
+except ImportError:
+    HAS_ORT = False
+    warnings.warn("onnxruntime未安装,只能使用RKNN后端", ImportWarning)
+
+# 常量定义
+RKNN_MAX_DIMS = 16
+RKNN_MAX_NAME_LEN = 256
+RKNN_MAX_DYNAMIC_SHAPE_NUM = 512
+RKNN_QUERY_CUSTOM_STRING = 7
+RKNN_QUERY_INPUT_DYNAMIC_RANGE = 13
+
+class RKNNInputRange(ctypes.Structure):
+    _fields_ = [
+        ('index', ctypes.c_uint32),
+        ('shape_number', ctypes.c_uint32),
+        ('fmt', ctypes.c_int),
+        ('name', ctypes.c_char * RKNN_MAX_NAME_LEN),
+        ('dyn_range', (ctypes.c_uint32 * RKNN_MAX_DIMS) * RKNN_MAX_DYNAMIC_SHAPE_NUM),
+        ('n_dims', ctypes.c_uint32)
+    ]
+
+class RKNNCustomString(ctypes.Structure):
+    _fields_ = [
+        ('string', ctypes.c_char * 1024)
+    ]
+
+# 支持的平台列表
+SUPPORTED_PLATFORMS = [
+    "rv1103", "rv1103b", "rv1106", "rv1106b",
+    "rk2118", "rk3562", "rk3566", "rk3568", "rk3576", "rk3588"
+]
+
+def get_current_platform() -> Optional[str]:
+    """
+    获取当前运行平台
+
+    Returns:
+        Optional[str]: 平台名称,如果文件不存在返回None
+
+    Raises:
+        RuntimeError: 如果平台不在支持列表中
+    """
+    platform_file = "/proc/device-tree/compatible"
+
+    if not os.path.exists(platform_file):
+        logger.debug(f"平台信息文件不存在: {platform_file}")
+        return None
+
+    try:
+        with open(platform_file, 'r') as f:
+            content = f.read()
+
+        # 使用正则匹配r[kv]\d{4}[b]?格式
+        match = re.search(r'r[kv]\d{4}[b]?', content.lower())
+        if not match:
+            raise RuntimeError(f"无法从{platform_file}解析平台信息")
+
+        platform = match.group()
+        if platform not in SUPPORTED_PLATFORMS:
+            raise RuntimeError(f"不支持的平台: {platform}")
+
+        logger.debug(f"当前平台: {platform}")
+        return platform
+
+    except Exception as e:
+        if not isinstance(e, RuntimeError):
+            logger.error(f"读取平台信息时发生错误: {str(e)}")
+            raise RuntimeError(f"读取平台信息失败: {str(e)}")
+        raise
+    
+# 配置日志
+logger = logging.getLogger("somemodelruntime_rknnlite2")
+logger.setLevel(logging.ERROR)  # 默认只输出错误信息
+if not logger.handlers:
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+    logger.addHandler(handler)
+
+_TRUE_VALUES = {"1", "true", "yes", "y", "on", "t"}
+_FALSE_VALUES = {"0", "false", "no", "n", "off", "f"}
+
+def _get_env_bool(name: str, default: bool = False) -> bool:
+    """
+    解析环境变量布尔值,支持1/0/true/false/yes/no/y/n/on/off.
+    """
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    value = raw.strip().lower()
+    if value in _TRUE_VALUES:
+        return True
+    if value in _FALSE_VALUES:
+        return False
+    logger.warning(f"环境变量{name}的值无效: {raw}, 应该是1/0或true/false/yes/no/y/n/on/off")
+    return default
+
+def _parse_model_list(list_str: str) -> Tuple[bool, List[str]]:
+    """
+    解析模型列表,支持以^开头表示反选.
+    返回(是否反选, 过滤后的模型列表).
+    """
+    if not list_str:
+        return False, []
+    items = [item.strip() for item in list_str.split(',') if item.strip()]
+    if not items:
+        return False, []
+    negated = any(item.startswith('^') for item in items)
+    cleaned = []
+    for item in items:
+        if item.startswith('^'):
+            item = item[1:].strip()
+        if item:
+            cleaned.append(item)
+    return negated, cleaned
+
+# ONNX Runtime日志级别到Python logging级别的映射
+_LOGGING_LEVEL_MAP = {
+    0: logging.DEBUG,    # Verbose
+    1: logging.INFO,     # Info
+    2: logging.WARNING,  # Warning
+    3: logging.ERROR,    # Error
+    4: logging.CRITICAL  # Fatal
+}
+
+# 检查环境变量中的日志级别设置
+try:
+    env_log_level = os.getenv('ZTU_MODELRT_RKNNL2_LOG_LEVEL')
+    if env_log_level is not None:
+        log_level = int(env_log_level)
+        if log_level in _LOGGING_LEVEL_MAP:
+            logger.setLevel(_LOGGING_LEVEL_MAP[log_level])
+            logger.info(f"从环境变量设置日志级别: {log_level}")
+        else:
+            logger.warning(f"环境变量ZTU_MODELRT_RKNNL2_LOG_LEVEL的值无效: {log_level}, 应该是0-4之间的整数")
+except ValueError:
+    logger.warning(f"环境变量ZTU_MODELRT_RKNNL2_LOG_LEVEL的值无效: {env_log_level}, 应该是0-4之间的整数")
+
+
+def set_default_logger_severity(level: int) -> None:
+    """
+    Sets the default logging severity. 0:Verbose, 1:Info, 2:Warning, 3:Error, 4:Fatal
+    
+    Args:
+        level: 日志级别(0-4)
+    """
+    if level not in _LOGGING_LEVEL_MAP:
+        raise ValueError(f"无效的日志级别: {level}, 应该是0-4之间的整数")
+    logger.setLevel(_LOGGING_LEVEL_MAP[level])
+
+def set_default_logger_verbosity(level: int) -> None:
+    """
+    Sets the default logging verbosity level. To activate the verbose log, 
+    you need to set the default logging severity to 0:Verbose level.
+    
+    Args:
+        level: 日志级别(0-4)
+    """
+    set_default_logger_severity(level)
+
+# RKNN tensor type到numpy dtype的映射
+RKNN_DTYPE_MAP = {
+    0: np.float32,  # RKNN_TENSOR_FLOAT32
+    1: np.float16,  # RKNN_TENSOR_FLOAT16
+    2: np.int8,     # RKNN_TENSOR_INT8
+    3: np.uint8,    # RKNN_TENSOR_UINT8
+    4: np.int16,    # RKNN_TENSOR_INT16
+    5: np.uint16,   # RKNN_TENSOR_UINT16
+    6: np.int32,    # RKNN_TENSOR_INT32
+    7: np.uint32,   # RKNN_TENSOR_UINT32
+    8: np.int64,    # RKNN_TENSOR_INT64
+    9: bool,        # RKNN_TENSOR_BOOL
+    10: np.int8,    # RKNN_TENSOR_INT4 (用int8表示)
+}
+
+def get_available_providers() -> List[str]:
+    """
+    获取可用的设备提供者列表(为保持接口兼容性的占位函数)
+    
+    Returns:
+        list: 可用的设备提供者列表,总是返回["CPUExecutionProvider", "somemodelruntime_rknnlite2_ExecutionProvider"]
+    """
+    return ["CPUExecutionProvider", "somemodelruntime_rknnlite2_ExecutionProvider"]
+
+
+def get_device() -> str:
+    """
+    获取当前设备
+
+    Returns:
+        str: 当前设备
+    """
+    return "RKNN2"
+
+def get_version_info() -> Dict[str, str]:
+    """
+    获取版本信息
+    
+    Returns:
+        dict: 包含API和驱动版本信息的字典
+    """
+    runtime = RKNNLite()
+    version = runtime.get_sdk_version()
+    return {
+        "api_version": version.split('\n')[2].split(': ')[1].split(' ')[0],
+        "driver_version": version.split('\n')[3].split(': ')[1]
+    }
+
+class IOTensor:
+    """输入/输出张量的信息封装类"""
+    def __init__(self, name, shape, type=None):
+        self.name = name.decode() if isinstance(name, bytes) else name
+        self.shape = shape
+        self.type = type
+
+    def __str__(self):
+        return f"IOTensor(name='{self.name}', shape={self.shape}, type={self.type})"
+
+class SessionOptions:
+    """会话选项类"""
+    def __init__(self):
+        self.enable_profiling = False  # 是否使用性能分析
+        self.intra_op_num_threads = 1  # 设置RKNN的线程数, 对应rknn的core_mask
+        self.log_severity_level = -1 # 另一个设置日志级别的参数
+        self.log_verbosity_level = -1 # 另一个设置日志级别的参数
+
+
+class InferenceSession:
+    """
+    RKNNLite运行时封装类,API风格类似ONNX Runtime
+    """
+
+    def __new__(cls, model_path: str, sess_options: Optional[SessionOptions] = None, **kwargs):
+        processed_path = InferenceSession._process_model_path(model_path, sess_options)
+        if isinstance(processed_path, str) and processed_path.lower().endswith('.onnx'):
+            logger.info("使用ONNX Runtime加载模型")
+            if not HAS_ORT:
+                raise RuntimeError("未安装onnxruntime,无法加载ONNX模型")
+            return ort.InferenceSession(processed_path, sess_options=sess_options, **kwargs)
+        else:
+            # 如果不是 ONNX 模型，则调用父类的 __new__ 创建 InferenceSession 实例
+            instance = super().__new__(cls)
+            # 保存处理后的路径
+            instance._processed_path = processed_path
+            return instance
+
+    def __init__(self, model_path: str, sess_options: Optional[SessionOptions] = None, **kwargs):
+        """
+        初始化运行时并加载模型
+        
+        Args:
+            model_path: 模型文件路径(.rknn或.onnx)
+            sess_options: 会话选项
+            **kwargs: 其他初始化参数
+        """
+        options = sess_options or SessionOptions()
+
+        # 只在未设置环境变量时使用SessionOptions中的日志级别
+        if os.getenv('ZTU_MODELRT_RKNNL2_LOG_LEVEL') is None:
+            if options.log_severity_level != -1:
+                set_default_logger_severity(options.log_severity_level)
+            if options.log_verbosity_level != -1:
+                set_default_logger_verbosity(options.log_verbosity_level)
+            
+        # 使用__new__中处理好的路径
+        model_path = getattr(self, '_processed_path', model_path)
+        if isinstance(model_path, str) and model_path.lower().endswith('.onnx'):
+            # 避免重复加载 ONNX 模型
+            return
+
+        # ... 现有的 RKNN 模型加载和初始化代码 ...
+        self.model_path = model_path
+        if not os.path.exists(self.model_path):
+            logger.error(f"模型文件不存在: {self.model_path}")
+            raise FileNotFoundError(f"模型文件不存在: {self.model_path}")
+
+        self.runtime = RKNNLite(verbose=options.enable_profiling)
+
+        logger.debug(f"正在加载模型: {self.model_path}")
+        ret = self.runtime.load_rknn(self.model_path)
+        if ret != 0:
+            logger.error(f"加载RKNN模型失败: {self.model_path}")
+            raise RuntimeError(f'加载RKNN模型失败: {self.model_path}')
+        logger.debug("模型加载成功")
+
+
+        if options.intra_op_num_threads == 1:
+            core_mask = RKNNLite.NPU_CORE_AUTO
+        elif options.intra_op_num_threads == 2:
+            core_mask = RKNNLite.NPU_CORE_0_1
+        elif options.intra_op_num_threads == 3:
+            core_mask = RKNNLite.NPU_CORE_0_1_2
+        else:
+            raise ValueError(f"intra_op_num_threads的值无效: {options.intra_op_num_threads}, 只能是1,2或3")
+
+        logger.debug("正在初始化运行时环境")
+        ret = self.runtime.init_runtime(core_mask=core_mask)
+        if ret != 0:
+            logger.error("初始化运行时环境失败")
+            raise RuntimeError('初始化运行时环境失败')
+
+        logger.debug("运行时环境初始化成功")
+
+        # 在 runtime 初始化后，按环境变量自动注册自定义算子插件库
+        try:
+            # 注册用户指定路径插件（逗号/分号分隔）
+            env_custom = os.getenv('ZTU_MODELRT_RKNN2_REG_CUSTOM_OP_LIB', '').strip()
+            if env_custom:
+                paths = [seg.strip() for seg in re.split(r"[,;:]", env_custom) if seg.strip()]
+                ok = 0
+                for p in paths:
+                    if self.register_custom_op_lib(p):
+                        ok += 1
+                if ok > 0:
+                    logger.info(f"已注册 {ok}/{len(paths)} 个自定义算子插件")
+                        # 注册系统目录下插件
+            if _get_env_bool('ZTU_MODELRT_RKNN2_REG_SYSTEM_CUSTOM_OP_LIB', True):
+                cnt = self.register_system_custom_op_lib()
+                if cnt > 0:
+                    logger.info(f"已从系统目录注册 {cnt} 个自定义算子插件")
+        except Exception as e:
+            logger.warning(f"自动注册自定义算子插件失败: {e}")
+
+        # 可选：按环境变量注册内置(基于Python)捆绑算子
+        if _get_env_bool('ZTU_MODELRT_RKNN2_REG_BUNDLED_OPS', False):
+            logger.info("根据环境变量注册捆绑算子")
+            self.register_bundled_ops()
+
+        self._init_io_info()
+        self.options = options
+
+    def get_performance_info(self) -> Dict[str, float]:
+        """
+        获取性能信息
+        
+        Returns:
+            dict: 包含性能信息的字典
+        """
+        if not self.options.perf_debug:
+            raise RuntimeError("性能分析未启用,请在SessionOptions中设置perf_debug=True")
+            
+        perf = self.runtime.rknn_runtime.get_run_perf()
+        return {
+            "run_duration": perf.run_duration / 1000.0  # 转换为毫秒
+        }
+
+    def set_core_mask(self, core_mask: int) -> None:
+        """
+        设置NPU核心使用模式
+        
+        Args:
+            core_mask: NPU核心掩码,使用NPU_CORE_*常量
+        """
+        ret = self.runtime.rknn_runtime.set_core_mask(core_mask)
+        if ret != 0:
+            raise RuntimeError("设置NPU核心模式失败")
+
+    @staticmethod
+    def _process_model_path(model_path, sess_options):
+        """
+        处理模型路径,支持.onnx和.rknn文件
+        
+        Args:
+            model_path: 模型文件路径
+        """
+        # 如果是ONNX文件,检查是否需要自动加载RKNN
+        if model_path.lower().endswith('.onnx'):
+            logger.info("检测到ONNX模型文件")
+            
+            # 获取需要跳过自动加载的模型列表
+            skip_models = os.getenv('ZTU_MODELRT_RKNNL2_SKIP', '').strip()
+            if skip_models:
+                # 获取模型文件名(不含路径)用于匹配
+                model_name = os.path.basename(model_path)
+                model_name_lower = model_name.lower()
+                negated, skip_list = _parse_model_list(skip_models)
+                if skip_list:
+                    skip_set = {m.lower() for m in skip_list}
+                    should_skip = model_name_lower not in skip_set if negated else model_name_lower in skip_set
+                else:
+                    should_skip = False
+                if should_skip:
+                    logger.info(f"模型{model_name}在跳过列表中,将使用ONNX Runtime")
+                    return model_path
+            
+            # 构造RKNN文件路径
+            rknn_path = os.path.splitext(model_path)[0] + '.rknn'
+            if os.path.exists(rknn_path):
+                logger.info(f"找到对应的RKNN模型,将使用RKNN: {rknn_path}")
+                return rknn_path
+            else:
+                logger.info("未找到对应的RKNN模型,将使用ONNX Runtime")
+                return model_path
+            
+        return model_path
+        
+    def _convert_nhwc_to_nchw(self, shape):
+        """将NHWC格式的shape转换为NCHW格式"""
+        if len(shape) == 4:
+            # NHWC -> NCHW
+            n, h, w, c = shape
+            return [n, c, h, w]
+        return shape
+        
+    def _init_io_info(self):
+        """初始化模型的输入输出信息"""
+        runtime = self.runtime.rknn_runtime
+        
+        # 获取输入输出数量
+        n_input, n_output = runtime.get_in_out_num()
+        
+        # 获取输入信息
+        self.input_tensors = []
+        for i in range(n_input):
+            attr = runtime.get_tensor_attr(i)
+            shape = [attr.dims[j] for j in range(attr.n_dims)]
+            # 对四维输入进行NHWC到NCHW的转换
+            shape = self._convert_nhwc_to_nchw(shape)
+            # 获取dtype
+            dtype = RKNN_DTYPE_MAP.get(attr.type, None)
+            tensor = IOTensor(attr.name, shape, dtype)
+            self.input_tensors.append(tensor)
+            
+        # 获取输出信息
+        self.output_tensors = []
+        for i in range(n_output):
+            attr = runtime.get_tensor_attr(i, is_output=True)
+            shape = runtime.get_output_shape(i)
+            # 获取dtype
+            dtype = RKNN_DTYPE_MAP.get(attr.type, None)
+            tensor = IOTensor(attr.name, shape, dtype)
+            self.output_tensors.append(tensor)
+        
+    def get_inputs(self):
+        """
+        获取模型输入信息
+        
+        Returns:
+            list: 包含输入信息的列表
+        """
+        return self.input_tensors
+        
+    def get_outputs(self):
+        """
+        获取模型输出信息
+        
+        Returns:
+            list: 包含输出信息的列表
+        """
+        return self.output_tensors
+        
+    def run(self, output_names=None, input_feed=None, data_format="nchw", **kwargs):
+        """
+        执行模型推理
+        
+        Args:
+            output_names: 输出节点名称列表,指定需要返回哪些输出
+            input_feed: 输入数据字典或列表
+            data_format: 输入数据格式,"nchw"或"nhwc"
+            **kwargs: 其他运行时参数
+            
+        Returns:
+            list: 模型输出结果列表,如果指定了output_names则只返回指定的输出
+        """
+        if input_feed is None:
+            logger.error("input_feed不能为None")
+            raise ValueError("input_feed不能为None")
+            
+        # 准备输入数据
+        if isinstance(input_feed, dict):
+            # 如果是字典,按照模型输入顺序排列
+            inputs = []
+            input_map = {tensor.name: i for i, tensor in enumerate(self.input_tensors)}
+            for tensor in self.input_tensors:
+                if tensor.name not in input_feed:
+                    raise ValueError(f"缺少输入: {tensor.name}")
+                inputs.append(input_feed[tensor.name])
+        elif isinstance(input_feed, (list, tuple)):
+            # 如果是列表,确保长度匹配
+            if len(input_feed) != len(self.input_tensors):
+                raise ValueError(f"输入数量不匹配: 期望{len(self.input_tensors)}, 实际{len(input_feed)}")
+            inputs = list(input_feed)
+        else:
+            logger.error("input_feed必须是字典或列表类型")
+            raise ValueError("input_feed必须是字典或列表类型")
+            
+        # 执行推理
+        try:
+            logger.debug("开始执行推理")
+            all_outputs = self.runtime.inference(inputs=inputs, data_format=data_format)
+            
+            # 如果没有指定output_names,返回所有输出
+            if output_names is None:
+                return all_outputs
+                
+            # 获取指定的输出
+            output_map = {tensor.name: i for i, tensor in enumerate(self.output_tensors)}
+            selected_outputs = []
+            for name in output_names:
+                if name not in output_map:
+                    raise ValueError(f"未找到输出节点: {name}")
+                selected_outputs.append(all_outputs[output_map[name]])
+                    
+            return selected_outputs
+            
+        except Exception as e:
+            logger.error(f"推理执行失败: {str(e)}")
+            raise RuntimeError(f"推理执行失败: {str(e)}")
+        
+    def close(self):
+        """
+        关闭会话,释放资源
+        """
+        if self.runtime is not None:
+            logger.info("正在释放运行时资源")
+            self.runtime.release()
+            self.runtime = None
+            
+    def __enter__(self):
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def end_profiling(self) -> Optional[str]:
+        """
+        结束性能分析的存根方法
+        
+        Returns:
+            Optional[str]: None
+        """
+        warnings.warn("end_profiling()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return None
+        
+    def get_profiling_start_time_ns(self) -> int:
+        """
+        获取性能分析开始时间的存根方法
+        
+        Returns:
+            int: 0
+        """
+        warnings.warn("get_profiling_start_time_ns()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return 0
+        
+    def get_modelmeta(self) -> Dict[str, str]:
+        """
+        获取模型元数据的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_modelmeta()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+        
+    def get_session_options(self) -> SessionOptions:
+        """
+        获取会话选项
+        
+        Returns:
+            SessionOptions: 当前会话选项
+        """
+        return self.options
+        
+    def get_providers(self) -> List[str]:
+        """
+        获取当前使用的providers的存根方法
+        
+        Returns:
+            List[str]: ["CPUExecutionProvider"]
+        """
+        warnings.warn("get_providers()是存根方法,始终返回CPUExecutionProvider", RuntimeWarning, stacklevel=2)
+        return ["CPUExecutionProvider"]
+        
+    def get_provider_options(self) -> Dict[str, Dict[str, str]]:
+        """
+        获取provider选项的存根方法
+        
+        Returns:
+            Dict[str, Dict[str, str]]: 空字典
+        """
+        warnings.warn("get_provider_options()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {} 
+
+    def get_session_config(self) -> Dict[str, str]:
+        """
+        获取会话配置的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_session_config()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def get_session_state(self) -> Dict[str, str]:
+        """
+        获取会话状态的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_session_state()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def set_session_config(self, config: Dict[str, str]) -> None:
+        """
+        设置会话配置的存根方法
+        
+        Args:
+            config: 会话配置字典
+        """
+        warnings.warn("set_session_config()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_memory_info(self) -> Dict[str, int]:
+        """
+        获取内存使用信息的存根方法
+        
+        Returns:
+            Dict[str, int]: 空字典
+        """
+        warnings.warn("get_memory_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def set_memory_pattern(self, enable: bool) -> None:
+        """
+        设置内存模式的存根方法
+        
+        Args:
+            enable: 是否启用内存模式
+        """
+        warnings.warn("set_memory_pattern()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def disable_memory_pattern(self) -> None:
+        """
+        禁用内存模式的存根方法
+        """
+        warnings.warn("disable_memory_pattern()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_optimization_level(self) -> int:
+        """
+        获取优化级别的存根方法
+        
+        Returns:
+            int: 0
+        """
+        warnings.warn("get_optimization_level()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return 0
+
+    def set_optimization_level(self, level: int) -> None:
+        """
+        设置优化级别的存根方法
+        
+        Args:
+            level: 优化级别
+        """
+        warnings.warn("set_optimization_level()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_model_metadata(self) -> Dict[str, str]:
+        """
+        获取模型元数据的存根方法(与get_modelmeta不同的接口)
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_model_metadata()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def get_model_path(self) -> str:
+        """
+        获取模型路径
+        
+        Returns:
+            str: 模型文件路径
+        """
+        return self.model_path
+
+    def get_input_type_info(self) -> List[Dict[str, str]]:
+        """
+        获取输入类型信息的存根方法
+        
+        Returns:
+            List[Dict[str, str]]: 空列表
+        """
+        warnings.warn("get_input_type_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return []
+
+    def get_output_type_info(self) -> List[Dict[str, str]]:
+        """
+        获取输出类型信息的存根方法
+        
+        Returns:
+            List[Dict[str, str]]: 空列表
+        """
+        warnings.warn("get_output_type_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return [] 
+
+    ################### 自定义算子 ###################
+
+    def _init_custom_op_types(self):
+        """初始化自定义算子的类型定义"""
+        # 常量
+        self._RKNN_TENSOR_FLOAT32 = 0
+        self._RKNN_TENSOR_UINT8 = 3
+        self._RKNN_TENSOR_INT64 = 8
+        self._RKNN_TARGET_TYPE_CPU = 1
+
+        # 结构体定义
+        class RKNN_TensorAttr(ctypes.Structure):
+            _fields_ = [
+                ("index", ctypes.c_uint32),
+                ("n_dims", ctypes.c_uint32),
+                ("dims", ctypes.c_uint32 * RKNN_MAX_DIMS),
+                ("name", ctypes.c_char * RKNN_MAX_NAME_LEN),
+                ("n_elems", ctypes.c_uint32),
+                ("size", ctypes.c_uint32),
+                ("fmt", ctypes.c_int),
+                ("type", ctypes.c_int),
+                ("qnt_type", ctypes.c_int),
+                ("fl", ctypes.c_int8),
+                ("zp", ctypes.c_int32),
+                ("scale", ctypes.c_float),
+                ("w_stride", ctypes.c_uint32),
+                ("size_with_stride", ctypes.c_uint32),
+                ("pass_through", ctypes.c_uint8),
+                ("h_stride", ctypes.c_uint32),
+            ]
+
+        class RKNN_TensorMem(ctypes.Structure):
+            _fields_ = [
+                ("virt_addr", ctypes.c_void_p),
+                ("phys_addr", ctypes.c_uint64),
+                ("fd", ctypes.c_int32),
+                ("offset", ctypes.c_int32),
+                ("size", ctypes.c_uint32),
+                ("flags", ctypes.c_uint32),
+                ("priv_data", ctypes.c_void_p),
+            ]
+
+        class RKNN_CustomOpTensor(ctypes.Structure):
+            _fields_ = [
+                ("attr", RKNN_TensorAttr),
+                ("mem", RKNN_TensorMem),
+            ]
+
+        class RKNN_GPUOpContext(ctypes.Structure):
+            _fields_ = [
+                ("cl_context", ctypes.c_void_p),
+                ("cl_command_queue", ctypes.c_void_p),
+                ("cl_kernel", ctypes.c_void_p),
+            ]
+
+        InternalCtxType = (
+            ctypes.c_uint64 if ctypes.sizeof(ctypes.c_void_p) == 8 else ctypes.c_uint32
+        )
+
+        class RKNN_CustomOpContext(ctypes.Structure):
+            _fields_ = [
+                ("target", ctypes.c_int),
+                ("internal_ctx", InternalCtxType),
+                ("gpu_ctx", RKNN_GPUOpContext),
+                ("priv_data", ctypes.c_void_p),
+            ]
+
+        class RKNN_CustomOpAttr(ctypes.Structure):
+            _fields_ = [
+                ("name", ctypes.c_char * RKNN_MAX_NAME_LEN),
+                ("dtype", ctypes.c_int),
+                ("n_elems", ctypes.c_uint32),
+                ("data", ctypes.c_void_p),
+            ]
+
+        CB_SIG = ctypes.CFUNCTYPE(
+            ctypes.c_int,
+            ctypes.POINTER(RKNN_CustomOpContext),
+            ctypes.POINTER(RKNN_CustomOpTensor),
+            ctypes.c_uint32,
+            ctypes.POINTER(RKNN_CustomOpTensor),
+            ctypes.c_uint32,
+        )
+
+        DESTROY_SIG = ctypes.CFUNCTYPE(
+            ctypes.c_int, ctypes.POINTER(RKNN_CustomOpContext)
+        )
+
+        class RKNN_CustomOp(ctypes.Structure):
+            _fields_ = [
+                ("version", ctypes.c_uint32),
+                ("target", ctypes.c_int),
+                ("op_type", ctypes.c_char * RKNN_MAX_NAME_LEN),
+                ("cl_kernel_name", ctypes.c_char * RKNN_MAX_NAME_LEN),
+                ("cl_kernel_source", ctypes.c_char_p),
+                ("cl_source_size", ctypes.c_uint64),
+                ("cl_build_options", ctypes.c_char * RKNN_MAX_NAME_LEN),
+                ("init", CB_SIG),
+                ("prepare", CB_SIG),
+                ("compute", CB_SIG),
+                ("compute_native", CB_SIG),
+                ("destroy", DESTROY_SIG),
+            ]
+
+        # 保存类型定义
+        self._RKNN_TensorAttr = RKNN_TensorAttr
+        self._RKNN_TensorMem = RKNN_TensorMem
+        self._RKNN_CustomOpTensor = RKNN_CustomOpTensor
+        self._RKNN_CustomOpContext = RKNN_CustomOpContext
+        self._RKNN_CustomOpAttr = RKNN_CustomOpAttr
+        self._RKNN_CustomOp = RKNN_CustomOp
+        self._CB_SIG = CB_SIG
+        self._DESTROY_SIG = DESTROY_SIG
+
+    def _create_attr_readers(self, get_op_attr):
+        """创建属性读取函数"""
+        def read_attr_int64(op_ctx_ptr, key: str, default: int = 0) -> int:
+            attr = self._RKNN_CustomOpAttr()
+            get_op_attr(op_ctx_ptr, key.encode("utf-8"), ctypes.byref(attr))
+            if attr.n_elems == 1 and attr.dtype == self._RKNN_TENSOR_INT64 and attr.data:
+                return ctypes.c_int64.from_address(attr.data).value
+            return default
+
+        def read_attr_float32(op_ctx_ptr, key: str, default: float = 0) -> float:
+            attr = self._RKNN_CustomOpAttr()
+            get_op_attr(op_ctx_ptr, key.encode("utf-8"), ctypes.byref(attr))
+            if attr.n_elems == 1 and attr.dtype == self._RKNN_TENSOR_FLOAT32 and attr.data:
+                return ctypes.c_float.from_address(attr.data).value
+            return default
+
+        def read_attr_str(op_ctx_ptr, key: str, default: str = "") -> str:
+            attr = self._RKNN_CustomOpAttr()
+            get_op_attr(op_ctx_ptr, key.encode("utf-8"), ctypes.byref(attr))
+            if attr.n_elems > 0 and attr.dtype == self._RKNN_TENSOR_UINT8 and attr.data:
+                buf = (ctypes.c_ubyte * attr.n_elems).from_address(attr.data)
+                try:
+                    return bytes(buf).decode("utf-8", errors="ignore").strip('"')
+                except Exception:
+                    return default
+            return default
+        
+ 
+        return read_attr_int64, read_attr_str, read_attr_float32
+  
+    def _build_py_custom_op(self,
+                            op_type: str,
+                            n_inputs: int,
+                            n_outputs: int,
+                            on_init,
+                            on_compute):
+        """通用的Python自定义算子构造器
+
+        Args:
+            op_type: 算子类型名(字符串)
+            n_inputs: 输入个数
+            n_outputs: 输出个数
+            on_init: 回调,签名 on_init(op_ctx_p, read_attr_int64, read_attr_str) -> state
+            on_compute: 回调,签名 on_compute(op_ctx_p, inputs_p, outputs_p, state) -> int(0成功)
+        Returns:
+            (RKNN_CustomOp对象, 回调tuple)
+        """
+        @self._CB_SIG
+        def _py_init(op_ctx_p, inputs_p, n_inputs_p, outputs_p, n_outputs_p):
+            try:
+                # 允许无需提前读取属性
+                runtime = self.runtime.rknn_runtime
+                read_attr_int64, read_attr_str, read_attr_float32 = self._create_attr_readers(runtime.lib.rknn_custom_op_get_op_attr)
+                user_state = on_init(op_ctx_p, read_attr_int64, read_attr_str, read_attr_float32)
+                # 为该实例分配唯一ID, 并写入priv_data
+                if not hasattr(self, "_custom_op_states"):
+                    self._custom_op_states = {}
+                if not hasattr(self, "_next_custom_op_id"):
+                    self._next_custom_op_id = 1
+                inst_id = int(self._next_custom_op_id)
+                self._next_custom_op_id += 1
+                # 保存Python侧状态
+                self._custom_op_states[inst_id] = user_state
+                # 将实例ID写入priv_data
+                try:
+                    op_ctx_p.contents.priv_data = ctypes.c_void_p(inst_id)
+                except Exception:
+                    # 回退: 直接写入整数
+                    op_ctx_p.contents.priv_data = inst_id
+                return 0
+            except Exception as e:
+                logger.error(f"{op_type} init失败: {e}")
+                return -1
+
+        @self._CB_SIG
+        def _py_prepare(op_ctx_p, inputs_p, n_inputs_p, outputs_p, n_outputs_p):
+            return 0
+
+        @self._CB_SIG
+        def _py_compute(op_ctx_p, inputs_p, n_inputs_p, outputs_p, n_outputs_p):
+            try:
+                if n_inputs_p != n_inputs or n_outputs_p != n_outputs:
+                    return -1
+                # 通过priv_data取回该实例的状态
+                try:
+                    inst_id = int(op_ctx_p.contents.priv_data) if op_ctx_p.contents.priv_data else 0
+                except Exception:
+                    inst_id = 0
+                user_state = None
+                if hasattr(self, "_custom_op_states") and inst_id in self._custom_op_states:
+                    user_state = self._custom_op_states.get(inst_id)
+                else:
+                    logger.error(f"{op_type} compute失败: 找不到实例状态, inst_id={inst_id}")
+                    return -1
+                return on_compute(op_ctx_p, inputs_p, outputs_p, user_state)
+            except Exception as e:
+                logger.error(f"{op_type} compute失败: {e}")
+                import traceback
+                logger.error(f"{op_type} compute失败: {traceback.format_exc()}")
+                return -1
+
+        @self._DESTROY_SIG
+        def _py_destroy(op_ctx_p):
+            try:
+                # 清理该实例的状态
+                try:
+                    inst_id = int(op_ctx_p.contents.priv_data) if op_ctx_p.contents.priv_data else 0
+                except Exception:
+                    inst_id = 0
+                if hasattr(self, "_custom_op_states") and inst_id in self._custom_op_states:
+                    del self._custom_op_states[inst_id]
+                # 将priv_data清空
+                try:
+                    op_ctx_p.contents.priv_data = ctypes.c_void_p(0)
+                except Exception:
+                    op_ctx_p.contents.priv_data = 0
+                return 0
+            except Exception:
+                return -1
+
+        op = self._RKNN_CustomOp()
+        op.version = 1
+        op.target = self._RKNN_TARGET_TYPE_CPU
+        op.op_type = op_type.encode("utf-8")
+        op.cl_kernel_name = b""
+        op.cl_kernel_source = None
+        op.cl_source_size = 0
+        op.cl_build_options = b""
+        op.init = _py_init
+        op.prepare = _py_prepare
+        op.compute = _py_compute
+        op.compute_native = self._CB_SIG()  # NULL
+        op.destroy = _py_destroy
+
+        return op, (_py_init, _py_prepare, _py_compute, _py_destroy)
+
+
+    def _tensor_to_numpy(self, rknn_tensor):
+        """将 RKNN_CustomOpTensor 转换为 Numpy 数组视图"""
+        # 确定Numpy数据类型
+        # 您可以扩展这个映射
+        dtype_map = {
+            self._RKNN_TENSOR_FLOAT32: (ctypes.c_float, np.float32),
+            self._RKNN_TENSOR_UINT8: (ctypes.c_uint8, np.uint8),
+            self._RKNN_TENSOR_INT64: (ctypes.c_int64, np.int64),
+        }
+        c_type, np_dtype = dtype_map.get(rknn_tensor.attr.type, (None, None))
+        if c_type is None:
+            raise TypeError(f"不支持的RKNN张量类型: {rknn_tensor.attr.type}")
+
+        # 获取内存地址和形状
+        addr = (rknn_tensor.mem.virt_addr or 0) + int(rknn_tensor.mem.offset)
+        ptr = ctypes.cast(addr, ctypes.POINTER(c_type))
+        shape = tuple(rknn_tensor.attr.dims[i] for i in range(rknn_tensor.attr.n_dims))
+        
+        # 创建Numpy数组视图
+        return np.ctypeslib.as_array(ptr, shape=shape)
+
+
+    def _create_onnxscript_op_creator(self,
+                                      op_type: str,
+                                      # 现在接收一个"函数模板构造器"
+                                      onnxscript_func_builder,
+                                      n_inputs: int,
+                                      n_outputs: int,
+                                      attributes: dict = {},
+                                      constants: dict = {}):
+        """
+        一个高阶工厂函数，用于创建基于ONNXScript的自定义算子构造器。
+        它在 on_init 阶段动态生成最终的 onnxscript 计算函数。
+
+        Args:
+            op_type (str): 算子类型名。
+            onnxscript_func_builder: 一个函数，它接收所有属性和常量作为关键字参数，
+                                     并返回一个编译好的 onnxscript 函数。
+                                     例如: def builder(mean, scale):
+                                               @onnxscript.script()
+                                               def compute(like):
+                                                  return opset.RandomNormalLike(like, mean=mean, scale=scale)
+                                               return compute
+            attributes (dict): 从模型中读取的属性字典。
+            constants (dict): 编译时常量字典。
+            n_inputs (int): 输入个数。
+            n_outputs (int): 输出个数。
+        """
+
+        def creator_func():
+            def on_init(op_ctx_p, read_i64, read_s, read_f32):
+                # 1. 读取所有动态属性
+                attr_values = {}
+                for name, (attr_type, default) in attributes.items():
+                    if attr_type == 'int64':
+                        attr_values[name] = read_i64(op_ctx_p, name, default)
+                    elif attr_type == 'str':
+                        attr_values[name] = read_s(op_ctx_p, name, default)
+                    elif attr_type == 'float32':
+                        attr_values[name] = read_f32(op_ctx_p, name, default)
+                    else:
+                        raise ValueError(f"不支持的属性类型: {attr_type}")
+
+                # 2. 合并常量和属性
+                final_kwargs = {**constants, **attr_values}
+
+                # 3. 动态构建 onnxscript 函数！ <<<<< 核心修改
+                #    这确保了所有属性值都作为常量被闭包捕获
+                compute_func = onnxscript_func_builder(**final_kwargs)
+
+                # 4. 将最终生成的、已编译的函数存入 state
+                return {"compute_func": compute_func}
+
+            def on_compute(op_ctx_p, inputs_p, outputs_p, state):
+                compute_func = state["compute_func"]
+
+                input_nps = [self._tensor_to_numpy(inputs_p[i]) for i in range(n_inputs)]
+                output_nps = [self._tensor_to_numpy(outputs_p[i]) for i in range(n_outputs)]
+                
+                results = compute_func(*input_nps)
+
+                if n_outputs == 1:
+                    result_val = results[0] if isinstance(results, tuple) else results
+                    output_nps[0][...] = result_val
+                else:
+                    for i in range(n_outputs):
+                        output_nps[i][...] = results[i]
+                
+                return 0
+
+            return self._build_py_custom_op(
+                op_type=op_type,
+                n_inputs=n_inputs,
+                n_outputs=n_outputs,
+                on_init=on_init,
+                on_compute=on_compute
+            )
+        
+        return creator_func
+
+    def _create_gridsample_op(self):
+        import onnxscript
+        from onnxscript import opset17 as opset
+
+        def grid_sample_builder(align_corners, mode, padding_mode):
+            @onnxscript.script()
+            def grid_sample_compute(X, G):
+                return opset.GridSample(X, G, align_corners=align_corners, mode=mode, padding_mode=padding_mode)
+            return grid_sample_compute
+
+        grid_sample_creator = self._create_onnxscript_op_creator(
+            op_type="GridSample",
+            onnxscript_func_builder=grid_sample_builder, # << 传入 builder
+            attributes={
+                "align_corners": ("int64", 0),
+                "mode": ("str", "bilinear"),
+                "padding_mode": ("str", "zeros"),
+            },
+            n_inputs = 2,
+            n_outputs = 1
+        )
+        return grid_sample_creator
+
+    def _create_scatterelements_op(self):
+        import onnxscript
+        from onnxscript import opset17 as opset
+        
+        @onnxscript.script()
+        def scatter_elements_compute(data, indices, updates):
+            indices_i64 = opset.Cast(indices, to=onnxscript.INT64.dtype)
+            return opset.ScatterElements(data, indices_i64, updates)
+        
+        scatter_elements_creator = self._create_onnxscript_op_creator(
+            op_type="ScatterElements",
+            onnxscript_func_builder=lambda: scatter_elements_compute,
+            n_inputs = 3,
+            n_outputs = 1
+        )
+        return scatter_elements_creator
+
+    def _create_randomnormallike_op(self):
+        import onnxscript
+        from onnxscript import opset17 as opset
+
+        def random_normal_like_builder(mean, scale):
+            @onnxscript.script()
+            def random_normal_like_compute(like):
+                return opset.RandomNormalLike(like, mean=mean, scale=scale)
+            
+            return random_normal_like_compute
+        
+        # 3. 使用新的工厂函数
+        random_normal_like_creator = self._create_onnxscript_op_creator(
+            op_type="RandomNormalLike",
+            onnxscript_func_builder=random_normal_like_builder, # << 传入 builder
+            attributes={
+                "mean": ("float32", 0.0),
+                "scale": ("float32", 1.0),
+            },
+            n_inputs = 1,
+            n_outputs = 1
+        )
+        return random_normal_like_creator
+
+    def _create_einsum_op(self): 
+        import onnxscript
+        from onnxscript import opset17 as opset
+
+        def einsum_builder(equation):
+
+            @onnxscript.script()
+            def einsum_compute(in1, in2):
+                return opset.Einsum(in1, in2, equation=equation)
+            
+            return einsum_compute
+        
+        # 3. 使用新的工厂函数
+        einsum_creator = self._create_onnxscript_op_creator(
+            op_type="Einsum",
+            onnxscript_func_builder=einsum_builder, # << 传入 builder
+            attributes={
+                "equation": ("str", ""),
+            },
+            n_inputs = 2,
+            n_outputs = 1
+        )
+        return einsum_creator
+
+    def _create_reducel1_op(self):
+        import onnxscript
+        from onnxscript import opset17 as opset
+
+        def reduce_l1_builder(keepdims, noop_with_empty_axes):
+            @onnxscript.script()
+            def reduce_l1_compute(data, axes):
+                data_abs = opset.Abs(data)
+                return opset.ReduceSum(
+                    data_abs,
+                    axes,
+                    keepdims=keepdims,
+                    noop_with_empty_axes=noop_with_empty_axes,
+                )
+            return reduce_l1_compute
+
+        reduce_l1_creator = self._create_onnxscript_op_creator(
+            op_type="ReduceL1",
+            onnxscript_func_builder=reduce_l1_builder,
+            attributes={
+                "keepdims": ("int64", 1),
+                "noop_with_empty_axes": ("int64", 0),
+            },
+            n_inputs=2,
+            n_outputs=1,
+        )
+        return reduce_l1_creator
+
+    def register_bundled_ops(self) -> None:
+        """注册自定义操作"""
+        if getattr(self, "_custom_ops_registered", False):
+            return
+
+        runtime = self.runtime.rknn_runtime
+        lib = runtime.lib
+        ctx = runtime.context
+
+        try:
+            _ = lib.rknn_register_custom_ops
+            _ = lib.rknn_custom_op_get_op_attr
+        except AttributeError as e:
+            logger.debug(f"SDK不支持自定义算子注册: {e}")
+            return
+
+        self._init_custom_op_types()
+
+        # 注意：插件库注册已在模型加载后由环境变量控制，不在此处重复触发
+
+        # 算子创建函数的列表现在更加清晰
+        op_creator_factories = [
+            self._create_gridsample_op,
+            self._create_scatterelements_op,
+            self._create_randomnormallike_op,
+            self._create_einsum_op,
+            self._create_reducel1_op,
+            # self._create_my_custom_add_op, # 添加新算子非常简单
+        ]
+
+        ops_to_register = []
+        all_callbacks = []
+
+        for factory in op_creator_factories:
+            try:
+                # 调用工厂获得真正的构造器
+                creator_func = factory()
+                # 调用构造器生成算子实例
+                op, callbacks = creator_func()
+                ops_to_register.append(op)
+                all_callbacks.extend(callbacks)
+                logger.debug(f"成功创建自定义算子: {op.op_type.decode()}")
+            except Exception as e:
+                logger.warning(f"创建自定义算子失败: {e}", exc_info=True)
+
+        if not ops_to_register:
+            logger.debug("没有可注册的自定义算子")
+            return
+
+        # 创建一个ctypes数组以包含所有要注册的算子, 然后一次性注册
+        num_ops = len(ops_to_register)
+        op_array = (self._RKNN_CustomOp * num_ops)(*ops_to_register)
+        ret = lib.rknn_register_custom_ops(ctx, op_array, num_ops)
+        if ret != 0:
+            logger.error(f"注册自定义算子失败, ret={ret} (可能是误报, 继续执行...)")
+            # raise RuntimeError(f"rknn_register_custom_ops 失败, ret={ret}")
+
+        logger.info(f"成功注册 {len(ops_to_register)} 个自定义算子")
+
+        self._custom_ops_registered = True
+        self._registered_ops = ops_to_register
+        self._op_callbacks = all_callbacks
+
+    def _load_and_register_plugin_op(self, so_path: str) -> bool:
+        """加载单个插件库并注册其中的自定义算子。
+
+        要求插件实现 get_rknn_custom_op()，返回 rknn_custom_op*。
+        我们将该 C 指针直接传递给 rknn_register_custom_ops，避免复制。
+        """
+        if not os.path.isfile(so_path):
+            logger.warning(f"插件库不存在: {so_path}")
+            return False
+
+        runtime = self.runtime.rknn_runtime
+        lib = runtime.lib
+        ctx = runtime.context
+
+        # 根据平台位宽设置 rknn_context 的 ctypes 类型
+        ContextCType = ctypes.c_uint64 if ctypes.sizeof(ctypes.c_void_p) == 8 else ctypes.c_uint32
+        # 设置 rknn_register_custom_ops(ctx, op_ptr, num) 签名。第二参数按 void* 传递，避免结构体布局不一致
+        try:
+            lib.rknn_register_custom_ops.argtypes = [ContextCType, ctypes.c_void_p, ctypes.c_uint32]
+            lib.rknn_register_custom_ops.restype = ctypes.c_int
+        except Exception:
+            pass
+
+        # 加载插件
+        try:
+            handle = ctypes.CDLL(so_path)
+        except Exception as e:
+            logger.error(f"dlopen 失败: {so_path}, err={e}")
+            return False
+
+        # 获取 get_rknn_custom_op 符号
+        try:
+            get_sym = getattr(handle, "get_rknn_custom_op")
+        except AttributeError:
+            logger.error(f"插件缺少符号 get_rknn_custom_op: {so_path}")
+            return False
+
+        # 返回类型直接使用 void*，避免 Python 解析第三方结构体
+        try:
+            get_sym.argtypes = []
+        except Exception:
+            pass
+        get_sym.restype = ctypes.c_void_p
+
+        op_void_ptr = get_sym()
+        if not op_void_ptr:
+            logger.error(f"get_rknn_custom_op 返回空指针: {so_path}")
+            return False
+
+        # 直接使用原生指针注册（零拷贝）
+        ctx_val = ContextCType(runtime.context)
+        ret = lib.rknn_register_custom_ops(ctx_val, ctypes.c_void_p(op_void_ptr), 1)
+        if ret != 0:
+            logger.error(f"rknn_register_custom_ops 失败, ret={ret}, so={so_path} (可能是误报, 继续执行...)")
+            # return False
+
+        # 保留句柄，避免被垃圾回收卸载
+        if not hasattr(self, "_plugin_handles"):
+            self._plugin_handles = []
+        self._plugin_handles.append(handle)
+        logger.info(f"成功注册插件自定义算子: {so_path}")
+        return True
+
+    def register_plugin_ops(self, plugin_paths: List[str]) -> int:
+        """按给定路径列表注册插件库中的自定义算子。返回成功数量。"""
+        if not plugin_paths:
+            return 0
+        success = 0
+        for path in plugin_paths:
+            try:
+                if self._load_and_register_plugin_op(path):
+                    success += 1
+            except Exception as e:
+                logger.error(f"注册插件失败: {path}, err={e}")
+        return success
+
+    # 对外API：注册单个自定义算子插件库
+    def register_custom_op_lib(self, path: str) -> bool:
+        return self._load_and_register_plugin_op(path)
+
+    # 对外API：扫描并注册 Linux 系统目录下所有插件库（Android 不处理）
+    def register_system_custom_op_lib(self) -> int:
+        if os.name != 'posix':
+            return 0
+        # 仅 Linux：RKNN 官方默认目录
+        system_dir = "/usr/lib/rknpu/op_plugins/"
+        if not os.path.isdir(system_dir):
+            return 0
+        try:
+            entries = os.listdir(system_dir)
+        except Exception:
+            return 0
+        so_list = []
+        for name in entries:
+            # 官方要求文件名以 librkcst_ 开头
+            if name.startswith("librkcst_") and name.endswith('.so'):
+                so_list.append(os.path.join(system_dir, name))
+        return self.register_plugin_ops(so_list)