add paddleocr-vl

Cargo.lock

@@ -19,7 +19,7 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
 
 [[package]]
 name = "aha"
-version = "0.1.3"
+version = "0.1.4"
 dependencies = [
  "aha_openai_dive",
  "anyhow",

Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "aha"
-version = "0.1.3"
+version = "0.1.4"
 edition = "2024"
 repository = "https://github.com/jhqxxx/aha"
 license = "Apache-2.0"
-description = "aha model inference library, now supports Qwen2.5VL, MiniCPM4, VoxCPM, Qwen3VL, DeepSeek-OCR, and Hunyuan-OCR"
+description = "aha model inference library, now supports Qwen2.5VL, MiniCPM4, VoxCPM, Qwen3VL, DeepSeek-OCR, Hunyuan-OCR, and PaddleOCR-VL"
 
 [dependencies]
 candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.9.1", tag = "0.9.1"}

README.md

@@ -18,6 +18,7 @@
 * Qwen3VL - 阿里通义千问 3 多模态大语言模型
 * DeepSeek-OCR - 深度求索光学文字识别模型
 * Hunyuan-OCR - 腾讯混元光学文字识别模型
+* PaddleOCR-VL - 百度飞桨光学文字识别模型
 
 ## 计划支持
 我们持续扩展支持的模型列表，欢迎贡献！
@@ -33,7 +34,71 @@ sudo apt-get install -y clang pkg-config ffmpeg libavutil-dev libavcodec-dev lib
 ```
 * windows参考： https://github.com/zmwangx/rust-ffmpeg/wiki/Notes-on-building
 
-## 安装
+## 安装及使用
+
+### 从源码构建部署
+```bash
+git clone https://github.com/jhqxxx/aha.git
+cd aha
+```
+
+#### cargo run 运行参数说明
+##### 基本用法
+```bash
+cargo run -F cuda -- [参数]
+```
+##### 参数详解
+1. 端口设置
+-----
+    -p, --port <PORT>
+* 设置HTTP服务监听的端口号
+* 默认值：10100
+* 示例：--port 8080 或 -p 8080
+
+2. 模型选择（必选）
+-----
+    -m, --model <MODEL>
+* 指定要加载的模型类型
+* 可选值：
+    * minicpm4-0.5b：OpenBMB/MiniCPM4-0.5B 模型
+    * qwen2.5vl-3b：Qwen/Qwen2.5-VL-3B-Instruct 模型
+    * qwen2.5vl-7b：Qwen/Qwen2.5-VL-7B-Instruct 模型
+    * qwen3vl-2b：Qwen/Qwen3-VL-2B-Instruct 模型
+    * qwen3vl-4b：Qwen/Qwen3-VL-4B-Instruct 模型
+    * qwen3vl-8b：Qwen/Qwen3-VL-8B-Instruct 模型
+    * qwen3vl-32b：Qwen/Qwen3-VL-32B-Instruct 模型
+    * deepseek-ocr: deepseek-ai/DeepSeek-OCR 模型
+    * hunyuan-ocr: Tencent-Hunyuan/HunyuanOCR 模型
+    * paddleocr-vl: PaddlePaddle/PaddleOCR-VL 模型
+* 示例：--model deepseek-ocr 或 -m qwen3vl-2b
+
+3. 权重路径
+-----
+    --weight-path <WEIGHT_PATH>
+* 指定本地模型权重文件路径
+* 如果指定此参数，则跳过模型下载步骤
+* 示例：--weight-path /path/to/model/dir
+
+4. 保存路径
+-----
+    --save-dir <SAVE_DIR>
+* 指定模型下载保存的目录
+* 默认保存在用户主目录下的 .aha 文件夹中
+* 示例：--save-dir /custom/model/path
+
+5. 下载重试次数
+----- 
+    --download-retries <DOWNLOAD_RETRIES>
+* 设置模型下载失败时的最大重试次数
+* 默认值：3次
+* 示例：--download-retries 5
+
+##### 注意事项
+* 参数前需要使用双横线 -- 分隔 cargo 命令和应用程序参数
+* 模型参数 (--model 或 -m) 是必需的
+* 如果未指定 --weight-path，程序会自动下载指定模型
+* 下载的模型默认保存在 ~/.aha/ 目录下（除非指定了 --save-dir）
+
 ### 作为库使用
 * cargo add aha
 * 或者在Cargo.toml中添加
@@ -80,6 +145,8 @@ fn main() -> Result<()> {
 git clone https://github.com/jhqxxx/aha.git
 cd aha
 # 修改测试用例中模型路径
+# 运行 PaddleOCR-Vl 示例
+cargo test -F cuda paddleocr_vl_generate -r -- --nocapture
 
 # 运行 Hunyuan-OCR 示例
 cargo test -F cuda hunyuan_ocr_generate -r -- --nocapture
@@ -97,68 +164,6 @@ cargo test -F cuda minicpm_generate -r -- --nocapture
 cargo test -F cuda voxcpm_generate -r -- --nocapture
 ```
 
-### 从源码构建部署
-```bash
-git clone https://github.com/jhqxxx/aha.git
-cd aha
-```
-
-#### cargo run 运行参数说明
-##### 基本用法
-```bash
-cargo run -F cuda -- [参数]
-```
-##### 参数详解
-1. 端口设置
------
-    -p, --port <PORT>
-* 设置HTTP服务监听的端口号
-* 默认值：10100
-* 示例：--port 8080 或 -p 8080
-
-2. 模型选择（必选）
------
-    -m, --model <MODEL>
-* 指定要加载的模型类型
-* 可选值：
-    * minicpm4-0.5b：OpenBMB/MiniCPM4-0.5B 模型
-    * qwen2.5vl-3b：Qwen/Qwen2.5-VL-3B-Instruct 模型
-    * qwen2.5vl-7b：Qwen/Qwen2.5-VL-7B-Instruct 模型
-    * qwen3vl-2b：Qwen/Qwen3-VL-2B-Instruct 模型
-    * qwen3vl-4b：Qwen/Qwen3-VL-4B-Instruct 模型
-    * qwen3vl-8b：Qwen/Qwen3-VL-8B-Instruct 模型
-    * qwen3vl-32b：Qwen/Qwen3-VL-32B-Instruct 模型
-    * deepseek-ocr: deepseek-ai/DeepSeek-OCR 模型
-    * hunyuan-ocr: Tencent-Hunyuan/HunyuanOCR 模型
-* 示例：--model deepseek-ocr 或 -m qwen3vl-2b
-
-3. 权重路径
------
-    --weight-path <WEIGHT_PATH>
-* 指定本地模型权重文件路径
-* 如果指定此参数，则跳过模型下载步骤
-* 示例：--weight-path /path/to/model/dir
-
-4. 保存路径
------
-    --save-dir <SAVE_DIR>
-* 指定模型下载保存的目录
-* 默认保存在用户主目录下的 .aha 文件夹中
-* 示例：--save-dir /custom/model/path
-
-5. 下载重试次数
------ 
-    --download-retries <DOWNLOAD_RETRIES>
-* 设置模型下载失败时的最大重试次数
-* 默认值：3次
-* 示例：--download-retries 5
-
-##### 注意事项
-* 参数前需要使用双横线 -- 分隔 cargo 命令和应用程序参数
-* 模型参数 (--model 或 -m) 是必需的
-* 如果未指定 --weight-path，程序会自动下载指定模型
-* 下载的模型默认保存在 ~/.aha/ 目录下（除非指定了 --save-dir）
-
 ## 开发
 ### 项目结构
 ```text
@@ -172,6 +177,7 @@ cargo run -F cuda -- [参数]
 │   │   ├── deepseek_ocr
 │   │   ├── hunyuan_ocr
 │   │   ├── minicpm4
+│   │   ├── paddleocr_vl
 │   │   ├── qwen2_5vl
 │   │   ├── qwen3vl
 │   │   ├── voxcpm
@@ -185,6 +191,7 @@ cargo run -F cuda -- [参数]
     ├── test_hunyuan_ocr.rs
     ├── test_deepseek_ocr.rs
     ├── test_minicpm4.rs
+    ├── test_paddleocr_vl.rs
     ├── test_qwen2_5vl.rs
     └── test_voxcpm.rs
 ```
@@ -207,6 +214,8 @@ cargo run -F cuda -- [参数]
 2. 提交新的 Issue，包含详细描述和复现步骤
 
 ## 更新日志
+### v0.1.4
+* 添加PaddleOCR-VL 模型
 
 ### v0.1.3
 * 添加 Hunyuan-OCR 模型

src/chat_template/mod.rs

@@ -51,12 +51,23 @@ pub struct ChatTemplate<'a> {
 
 impl<'a> ChatTemplate<'a> {
     pub fn init(path: &str) -> Result<Self> {
-        let path = path.to_string();
-        assert!(
-            std::path::Path::new(&path).exists(),
-            "model path file not exists"
-        );
-        let template = get_template(path)?;
+        let path: String = path.to_string();
+        if !std::path::Path::new(&path).exists() {
+            return Err(anyhow!("model path not found"));
+        }
+        let template = match get_template(path.clone()) {
+            Ok(template) => template,
+            Err(e) => {
+                let jinja_path = path + "/chat_template.jinja";
+                if !std::path::Path::new(&jinja_path).exists() {
+                    return Err(anyhow!(
+                        "get_template err {e} and chat_template.jinja not found"
+                    ));
+                }
+                std::fs::read_to_string(&jinja_path)
+                    .map_err(|e| anyhow!("Failed to read chat_template.jinja: {}", e))?
+            }
+        };
         let template = string_to_static_str(template);
         // 加载jinjaenv处理chat_template
         let mut env = Environment::new();

src/main.rs

@@ -85,6 +85,7 @@ async fn main() -> anyhow::Result<()> {
         WhichModel::Qwen3vl32B => "Qwen/Qwen3-VL-32B-Instruct",
         WhichModel::DeepSeekOCR => "deepseek-ai/DeepSeek-OCR",
         WhichModel::HunyuanOCR => "Tencent-Hunyuan/HunyuanOCR",
+        WhichModel::PaddleOCRVL => "PaddlePaddle/PaddleOCR-VL",
     };
     let model_path = match args.weight_path {
         Some(path) => path,

src/models/common/mod.rs

@@ -106,7 +106,7 @@ pub struct NaiveAttention {
     num_kv_heads: usize,
     num_kv_groups: usize,
     head_dim: usize,
-    hidden_size: usize,
+    middle_size: usize,
     kv_cache: Option<(Tensor, Tensor)>,
 }
 
@@ -117,23 +117,37 @@ impl NaiveAttention {
         hidden_size: usize,
         num_attention_heads: usize,
         num_key_value_heads: usize,
+        head_dim: Option<usize>,
         bias: bool,
+        o_proj_pp_name: Option<&str>,
     ) -> Result<Self> {
         let num_kv_groups = num_attention_heads / num_key_value_heads;
-        let head_dim = hidden_size / num_attention_heads;
+        let head_dim = match head_dim {
+            None => hidden_size / num_attention_heads,
+            Some(dim) => dim,
+        };
+        let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj");
         let (q_proj, k_proj, v_proj, o_proj) = if bias {
             (
                 linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?,
                 linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?,
                 linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?,
-                linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?,
+                linear(
+                    num_attention_heads * head_dim,
+                    hidden_size,
+                    vb.pp(o_proj_pp_name),
+                )?,
             )
         } else {
             (
                 linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?,
                 linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?,
                 linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?,
-                linear_no_bias(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?,
+                linear_no_bias(
+                    num_attention_heads * head_dim,
+                    hidden_size,
+                    vb.pp(o_proj_pp_name),
+                )?,
             )
         };
 
@@ -146,7 +160,7 @@ impl NaiveAttention {
             num_kv_heads: num_key_value_heads,
             num_kv_groups,
             head_dim,
-            hidden_size,
+            middle_size: num_attention_heads * head_dim,
             kv_cache: None,
         })
     }
@@ -189,7 +203,7 @@ impl NaiveAttention {
             attention_mask,
             scale,
         )?;
-        let attn_output = attn_output.reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.reshape((b_sz, q_len, self.middle_size))?;
         let attn_output = attn_output.apply(&self.o_proj)?;
         Ok(attn_output)
     }
@@ -236,7 +250,7 @@ impl NaiveAttention {
             attention_mask,
             scale,
         )?;
-        let attn_output = attn_output.reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.reshape((b_sz, q_len, self.middle_size))?;
         let attn_output = attn_output.apply(&self.o_proj)?;
         Ok(attn_output)
     }

src/models/deepseek_ocr/model.rs

@@ -1018,7 +1018,9 @@ impl DeepseekV2DecoderLayer {
             config.hidden_size,
             config.num_attention_heads,
             config.num_key_value_heads,
+            None,
             false,
+            None,
         )?;
         let mlp = if layer_id >= config.first_k_dense_replace
             && layer_id.is_multiple_of(config.moe_layer_freq)

src/models/hunyuan_ocr/model.rs

@@ -112,7 +112,9 @@ impl HunYuanVisionBlock {
             config.hidden_size,
             config.num_attention_heads,
             config.num_attention_heads,
+            None,
             true,
+            None,
         )?;
         let mlp = TwoLinearMLP::new(
             vb.pp("mlp"),

src/models/minicpm4/model.rs

@@ -108,7 +108,9 @@ impl MiniCPMDecoderLayer {
             cfg.hidden_size,
             cfg.num_attention_heads,
             cfg.num_key_value_heads,
+            None,
             false,
+            None,
         )?;
         let mlp = GateUpDownMLP::new(
             vb.pp("mlp"),

src/models/mod.rs

@@ -2,6 +2,7 @@ pub mod common;
 pub mod deepseek_ocr;
 pub mod hunyuan_ocr;
 pub mod minicpm4;
+pub mod paddleocr_vl;
 pub mod qwen2_5vl;
 pub mod qwen3vl;
 pub mod voxcpm;
@@ -15,7 +16,8 @@ use rocket::futures::Stream;
 use crate::models::{
     deepseek_ocr::generate::DeepseekOCRGenerateModel,
     hunyuan_ocr::generate::HunyuanOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel,
-    qwen2_5vl::generate::Qwen2_5VLGenerateModel, qwen3vl::generate::Qwen3VLGenerateModel,
+    paddleocr_vl::generate::PaddleOCRVLGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel,
+    qwen3vl::generate::Qwen3VLGenerateModel,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
@@ -38,6 +40,8 @@ pub enum WhichModel {
     DeepSeekOCR,
     #[value(name = "hunyuan-ocr")]
     HunyuanOCR,
+    #[value(name = "paddleocr-vl")]
+    PaddleOCRVL,
 }
 
 pub trait GenerateModel {
@@ -61,6 +65,7 @@ pub enum ModelInstance<'a> {
     Qwen3VL(Qwen3VLGenerateModel<'a>),
     DeepSeekOCR(DeepseekOCRGenerateModel),
     HunyuanOCR(HunyuanOCRGenerateModel<'a>),
+    PaddleOCRVL(Box<PaddleOCRVLGenerateModel<'a>>),
 }
 
 impl<'a> GenerateModel for ModelInstance<'a> {
@@ -71,6 +76,7 @@ impl<'a> GenerateModel for ModelInstance<'a> {
             ModelInstance::Qwen3VL(model) => model.generate(mes),
             ModelInstance::DeepSeekOCR(model) => model.generate(mes),
             ModelInstance::HunyuanOCR(model) => model.generate(mes),
+            ModelInstance::PaddleOCRVL(model) => model.generate(mes),
         }
     }
 
@@ -91,6 +97,7 @@ impl<'a> GenerateModel for ModelInstance<'a> {
             ModelInstance::Qwen3VL(model) => model.generate_stream(mes),
             ModelInstance::DeepSeekOCR(model) => model.generate_stream(mes),
             ModelInstance::HunyuanOCR(model) => model.generate_stream(mes),
+            ModelInstance::PaddleOCRVL(model) => model.generate_stream(mes),
         }
     }
 }
@@ -133,6 +140,10 @@ pub fn load_model(model_type: WhichModel, path: &str) -> Result<ModelInstance<'_
             let model = HunyuanOCRGenerateModel::init(path, None, None)?;
             ModelInstance::HunyuanOCR(model)
         }
+        WhichModel::PaddleOCRVL => {
+            let model = PaddleOCRVLGenerateModel::init(path, None, None)?;
+            ModelInstance::PaddleOCRVL(Box::new(model))
+        }
     };
     Ok(model)
 }

src/models/paddleocr_vl/config.rs

@@ -0,0 +1,87 @@
+use candle_nn::Activation;
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct PaddleOCRVLConfig {
+    pub compression_ratio: f64,
+    pub head_dim: usize,
+    pub hidden_act: Activation,
+    pub hidden_dropout_prob: f64,
+    pub hidden_size: usize,
+    pub ignored_index: i32,
+    pub image_token_id: u32,
+    pub intermediate_size: usize,
+    pub max_position_embeddings: usize,
+    pub max_sequence_length: Option<usize>,
+    pub num_attention_heads: usize,
+    pub num_hidden_layers: usize,
+    pub num_key_value_heads: usize,
+    pub pad_token_id: u32,
+    pub rms_norm_eps: f64,
+    pub rope_scaling: PaddleOCRVLRopeScalingConfig,
+    pub rope_theta: f64,
+    pub sliding_window: Option<u32>,
+    pub tie_word_embeddings: bool,
+    pub torch_dtype: String,
+    pub use_bias: bool,
+    pub use_cache: bool,
+    pub use_flash_attention: bool,
+    pub video_token_id: u32,
+    pub vision_config: PaddleOCRVLVisionConfig,
+    pub vision_start_token_id: u32,
+    pub vocab_size: usize,
+    pub weight_share_add_bias: bool,
+    pub use_3d_rope: bool,
+    pub rope_is_neox_style: bool,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct PaddleOCRVLRopeScalingConfig {
+    pub mrope_section: Vec<usize>,
+    pub rope_type: String,
+    #[serde(rename = "type")]
+    pub scaling_type: String,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct PaddleOCRVLVisionConfig {
+    pub attention_dropout: f64,
+    pub hidden_act: Activation,
+    pub hidden_size: usize,
+    pub image_size: usize,
+    pub intermediate_size: usize,
+    pub layer_norm_eps: f64,
+    pub num_attention_heads: usize,
+    pub num_channels: usize,
+    pub num_hidden_layers: usize,
+    pub pad_token_id: u32,
+    pub patch_size: usize,
+    pub spatial_merge_size: usize,
+    pub temporal_patch_size: usize,
+    pub tokens_per_second: usize,
+    pub torch_dtype: String,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct PaddleOCRVLPreprocessorConfig {
+    pub do_convert_rgb: bool,
+    pub do_normalize: bool,
+    pub do_rescale: bool,
+    pub do_resize: bool,
+    pub image_mean: Vec<f64>,
+    pub image_std: Vec<f64>,
+    pub max_pixels: u32,
+    pub merge_size: usize,
+    pub min_pixels: u32,
+    pub patch_size: usize,
+    pub resample: u32,
+    pub rescale_factor: f64,
+    pub size: Option<SizeConfig>,
+    pub temporal_patch_size: usize,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct SizeConfig {
+    pub max_pixels: usize,
+    pub min_pixels: usize,
+}

src/models/paddleocr_vl/generate.rs

@@ -0,0 +1,195 @@
+use aha_openai_dive::v1::resources::chat::{
+    ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
+};
+use anyhow::{Result, anyhow};
+use candle_core::{D, DType, Device, IndexOp, Tensor};
+use candle_nn::VarBuilder;
+use rocket::async_stream::stream;
+use rocket::futures::Stream;
+
+use crate::models::paddleocr_vl::config::{PaddleOCRVLConfig, PaddleOCRVLPreprocessorConfig};
+use crate::models::paddleocr_vl::model::PaddleOCRVLModel;
+use crate::models::paddleocr_vl::processor::PaddleOCRVLProcessor;
+use crate::utils::tensor_utils::get_equal_mask;
+use crate::utils::{
+    build_completion_chunk_response, build_completion_response, find_type_files, get_device,
+    get_dtype, get_logit_processor,
+};
+use crate::{chat_template::ChatTemplate, models::GenerateModel, tokenizer::TokenizerModel};
+
+pub struct PaddleOCRVLGenerateModel<'a> {
+    chat_template: ChatTemplate<'a>,
+    tokenizer: TokenizerModel,
+    pre_processor: PaddleOCRVLProcessor,
+    paddleocr_vl: PaddleOCRVLModel,
+    cfg: PaddleOCRVLConfig,
+    device: Device,
+    end_token_id: u32,
+    model_name: String,
+}
+
+impl<'a> PaddleOCRVLGenerateModel<'a> {
+    pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
+        let chat_template = ChatTemplate::init(path)?;
+        let tokenizer = TokenizerModel::init(path)?;
+        let config_path = path.to_string() + "/config.json";
+        let cfg: PaddleOCRVLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
+        let device = &get_device(device);
+        let cfg_dtype = cfg.torch_dtype.as_str();
+        let dtype = get_dtype(dtype, cfg_dtype);
+        let processor_cfg_path = path.to_string() + "/preprocessor_config.json";
+        let processor_cfg: PaddleOCRVLPreprocessorConfig =
+            serde_json::from_slice(&std::fs::read(processor_cfg_path)?)?;
+        let pre_processor = PaddleOCRVLProcessor::new(processor_cfg, device, dtype)?;
+        let end_token_id = 2;
+        let model_list = find_type_files(path, "safetensors")?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
+        let paddleocr_vl = PaddleOCRVLModel::new(cfg.clone(), vb)?;
+
+        Ok(PaddleOCRVLGenerateModel {
+            chat_template,
+            tokenizer,
+            pre_processor,
+            paddleocr_vl,
+            cfg,
+            device: device.clone(),
+            end_token_id,
+            model_name: "paddleocr_vl".to_string(),
+        })
+    }
+}
+
+impl<'a> GenerateModel for PaddleOCRVLGenerateModel<'a> {
+    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
+        let seed = match mes.seed {
+            None => 34562u64,
+            Some(s) => s as u64,
+        };
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
+        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+        let (replace_text, mut pixel_values, mut image_grid_thw) =
+            self.pre_processor.process_info(&mes, &mes_render)?;
+        let mut input_ids = self.tokenizer.text_encode(replace_text, &self.device)?;
+        let mut seq_len = input_ids.dim(1)?;
+        let mut seqlen_offset = 0;
+        let image_mask = get_equal_mask(&input_ids, self.cfg.image_token_id)?;
+
+        let mut cache_position = Tensor::ones_like(&input_ids.i(0)?)?
+            .to_dtype(candle_core::DType::F64)?
+            .cumsum(D::Minus1)?
+            .to_dtype(candle_core::DType::U32)?
+            .broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
+
+        let mut generate = Vec::new();
+        let sample_len = mes.max_tokens.unwrap_or(1024);
+        for _ in 0..sample_len {
+            let logits = self.paddleocr_vl.forward(
+                &input_ids,
+                pixel_values.as_ref(),
+                image_grid_thw.as_ref(),
+                &image_mask,
+                Some(&cache_position),
+                seqlen_offset,
+            )?;
+            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+            let next_token = logit_processor.sample(&logits)?;
+            generate.push(next_token);
+            if next_token == self.end_token_id {
+                break;
+            }
+            seqlen_offset += seq_len;
+            seq_len = 1;
+            input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+            cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
+            pixel_values = None;
+            image_grid_thw = None;
+        }
+        let res = self.tokenizer.token_decode(generate)?;
+        self.paddleocr_vl.clear_kv_cache();
+        let response = build_completion_response(res, &self.model_name);
+        Ok(response)
+    }
+
+    fn generate_stream(
+        &mut self,
+        mes: ChatCompletionParameters,
+    ) -> Result<
+        Box<
+            dyn Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>
+                + Send
+                + Unpin
+                + '_,
+        >,
+    > {
+        let seed = match mes.seed {
+            None => 34562u64,
+            Some(s) => s as u64,
+        };
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
+        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+        let (replace_text, pixel_values, image_grid_thw) =
+            self.pre_processor.process_info(&mes, &mes_render)?;
+        let mut input_ids = self.tokenizer.text_encode(replace_text, &self.device)?;
+        let mut seq_len = input_ids.dim(1)?;
+        let mut seqlen_offset = 0;
+        let image_mask = get_equal_mask(&input_ids, self.cfg.image_token_id)?;
+
+        let mut cache_position = Tensor::ones_like(&input_ids.i(0)?)?
+            .to_dtype(candle_core::DType::F64)?
+            .cumsum(D::Minus1)?
+            .to_dtype(candle_core::DType::U32)?
+            .broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
+
+        let sample_len = mes.max_tokens.unwrap_or(1024);
+        let stream = stream! {
+            let mut error_tokens = Vec::new();
+            let mut pixel_values = pixel_values.as_ref();
+            let mut image_grid_thw = image_grid_thw.as_ref();
+            for _ in 0..sample_len {
+                let logits = self.paddleocr_vl.forward(
+                    &input_ids,
+                    pixel_values,
+                    image_grid_thw,
+                    &image_mask,
+                    Some(&cache_position),
+                    seqlen_offset,
+                )?;
+                let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+                let next_token = logit_processor.sample(&logits)?;
+                let mut decode_ids = Vec::new();
+                if !error_tokens.is_empty() {
+                    decode_ids.extend_from_slice(&error_tokens);
+                }
+                decode_ids.push(next_token);
+                let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{}", e)))?;
+                if decoded_token.contains("�") {
+                    error_tokens.push(next_token);
+                    if error_tokens.len() > 3 {
+                        error_tokens.clear();
+                    }
+                    seqlen_offset += seq_len;
+                    seq_len = 1;
+                    input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                    cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
+                    pixel_values = None;
+                    image_grid_thw = None;
+                    continue;
+                }
+                error_tokens.clear();
+                let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
+                yield Ok(chunk);
+                if next_token == self.end_token_id {
+                    break;
+                }
+                seqlen_offset += seq_len;
+                seq_len = 1;
+                input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                cache_position = Tensor::from_vec(vec![seqlen_offset as u32], 1, &self.device)?;
+                pixel_values = None;
+                image_grid_thw = None;
+            }
+            self.paddleocr_vl.clear_kv_cache();
+        };
+        Ok(Box::new(Box::pin(stream)))
+    }
+}

src/models/paddleocr_vl/mod.rs

@@ -0,0 +1,4 @@
+pub mod config;
+pub mod generate;
+pub mod model;
+pub mod processor;

src/models/paddleocr_vl/model.rs

@@ -0,0 +1,829 @@
+use anyhow::{Result, anyhow};
+use candle_core::{D, IndexOp, Shape, Tensor};
+use candle_nn::{
+    Conv2d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear,
+    linear_no_bias, rms_norm,
+};
+use num::integer::Roots;
+
+use crate::{
+    models::{
+        common::{GateUpDownMLP, NaiveAttention, TwoLinearMLP, get_conv2d, get_layer_norm},
+        paddleocr_vl::config::{
+            PaddleOCRVLConfig, PaddleOCRVLRopeScalingConfig, PaddleOCRVLVisionConfig,
+        },
+    },
+    position_embed::rope::{Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding},
+    utils::tensor_utils::{
+        get_vision_next_indices, interpolate_bilinear, masked_scatter_dim0, nonzero_index,
+        prepare_causal_attention_mask, zero_index,
+    },
+};
+
+pub struct Projector {
+    merge_size: usize,
+    pre_norm: LayerNorm,
+    linear_1: Linear,
+    linear_2: Linear,
+}
+
+impl Projector {
+    pub fn new(vb: VarBuilder, config: &PaddleOCRVLConfig) -> Result<Self> {
+        let merge_size = config.vision_config.spatial_merge_size;
+        let hidden_size = config.vision_config.hidden_size * merge_size * merge_size;
+        let pre_norm = get_layer_norm(
+            vb.pp("pre_norm"),
+            config.rms_norm_eps,
+            config.vision_config.hidden_size,
+        )?;
+        let linear_1 = linear(hidden_size, hidden_size, vb.pp("linear_1"))?;
+        let linear_2 = linear(hidden_size, config.hidden_size, vb.pp("linear_2"))?;
+
+        Ok(Self {
+            merge_size,
+            pre_norm,
+            linear_1,
+            linear_2,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor, image_grid_thw: &Tensor) -> Result<Tensor> {
+        let img_num = image_grid_thw.dim(0)?;
+        let mut processed_features = vec![];
+        let start = 0usize;
+        for i in 0..img_num {
+            let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
+                return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
+            };
+            let end = start + (t * h * w) as usize;
+            let xs_i = xs.i((start..end, ..))?;
+            let xs_i = self.pre_norm.forward(&xs_i)?;
+            let dim = xs_i.dim(1)?;
+            let shape = Shape::from(vec![
+                t as usize,
+                h as usize / self.merge_size,
+                self.merge_size,
+                w as usize / self.merge_size,
+                self.merge_size,
+                dim,
+            ]);
+            let xs_i = xs_i
+                .reshape((t as usize, h as usize, w as usize, dim))?
+                .reshape(shape)?
+                .permute((0, 1, 3, 2, 4, 5))?
+                .reshape((
+                    (t * h * w) as usize / self.merge_size / self.merge_size,
+                    self.merge_size * self.merge_size * dim,
+                ))?;
+            let xs_i = self.linear_1.forward(&xs_i)?.gelu()?;
+            let xs_i = self.linear_2.forward(&xs_i)?;
+            processed_features.push(xs_i);
+        }
+        let xs = Tensor::cat(&processed_features, 0)?;
+        Ok(xs)
+    }
+}
+
+pub struct SiglipVisionEmbeddings {
+    embed_dim: usize,
+    patch_size: usize,
+    patch_embedding: Conv2d,
+    num_positions: usize,
+    position_embedding: Embedding,
+    packing_position_embedding: Embedding,
+}
+
+impl SiglipVisionEmbeddings {
+    pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
+        let embed_dim = config.hidden_size;
+        let image_size = config.image_size;
+        let patch_size = config.patch_size;
+        let patch_embedding = get_conv2d(
+            vb.pp("patch_embedding"),
+            config.num_channels,
+            embed_dim,
+            patch_size,
+            0,
+            patch_size,
+            1,
+            1,
+            true,
+        )?;
+        let num_positions = (image_size / patch_size).pow(2);
+        let position_embedding = embedding(num_positions, embed_dim, vb.pp("position_embedding"))?;
+        let packing_position_embedding =
+            embedding(32768, embed_dim, vb.pp("packing_position_embedding"))?;
+        Ok(Self {
+            embed_dim,
+            patch_size,
+            patch_embedding,
+            num_positions,
+            position_embedding,
+            packing_position_embedding,
+        })
+    }
+    fn interpolate_pos_encoding(
+        &self,
+        h: usize,
+        w: usize,
+        is_after_patchify: bool,
+    ) -> Result<Tensor> {
+        let (new_height, new_width) = if is_after_patchify {
+            (h, w)
+        } else {
+            (h / self.patch_size, w / self.patch_size)
+        };
+        let sqrt_num_positions = self.num_positions.sqrt();
+        let patch_pos_embed = self
+            .position_embedding
+            .embeddings()
+            .reshape((1, sqrt_num_positions, sqrt_num_positions, self.embed_dim))?
+            .permute((0, 3, 1, 2))?;
+        let patch_pos_embed =
+            interpolate_bilinear(&patch_pos_embed, (new_height, new_width), Some(false))?;
+        let patch_pos_embed =
+            patch_pos_embed
+                .permute((0, 2, 3, 1))?
+                .reshape((1, (), self.embed_dim))?;
+        Ok(patch_pos_embed)
+    }
+
+    pub fn forward(
+        &self,
+        pixel_values: &Tensor,
+        position_ids: &Tensor,
+        image_grid_thw: &Tensor,
+        interpolate_pos_encoding: bool,
+    ) -> Result<Tensor> {
+        let (bs, seq_len, c, h, w) = pixel_values.dims5()?;
+        let pixel_values = pixel_values.reshape((bs * seq_len, c, h, w))?;
+        let patch_embeds = self.patch_embedding.forward(&pixel_values)?;
+        // (bs*seq_len, c)
+        let mut embeddings = patch_embeds.squeeze(D::Minus1)?.squeeze(D::Minus1)?;
+        if interpolate_pos_encoding {
+            let mut tmp_embeddings = vec![];
+            let img_num = image_grid_thw.dim(0)?;
+            let mut start = 0usize;
+            for i in 0..img_num {
+                let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
+                    return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
+                };
+                let end = start + (t * h * w) as usize;
+                let image_embeddings = embeddings.i(start..end)?;
+                let position_embedding = self
+                    .interpolate_pos_encoding(h as usize, w as usize, true)?
+                    .squeeze(0)?
+                    .repeat((t as usize, 1usize))?;
+                let image_embeddings = image_embeddings.add(&position_embedding)?;
+                tmp_embeddings.push(image_embeddings);
+                start = end;
+            }
+            embeddings = Tensor::cat(&tmp_embeddings, 0)?.unsqueeze(0)?; // add bs dim
+        } else {
+            let packing_pos_embed = self.packing_position_embedding.forward(position_ids)?;
+            embeddings = embeddings.add(&packing_pos_embed)?.unsqueeze(0)?;
+        }
+        Ok(embeddings)
+    }
+}
+
+pub struct SiglipEncoderLayer {
+    layer_norm1: LayerNorm,
+    self_attn: NaiveAttention,
+    layer_norm2: LayerNorm,
+    mlp: TwoLinearMLP,
+}
+
+impl SiglipEncoderLayer {
+    pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
+        let layer_norm1 = get_layer_norm(
+            vb.pp("layer_norm1"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        let self_attn = NaiveAttention::new(
+            vb.pp("self_attn"),
+            config.hidden_size,
+            config.num_attention_heads,
+            config.num_attention_heads,
+            None,
+            true,
+            Some("out_proj"),
+        )?;
+        let layer_norm2 = get_layer_norm(
+            vb.pp("layer_norm2"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        let mlp = TwoLinearMLP::new(
+            vb.pp("mlp"),
+            config.hidden_size,
+            config.intermediate_size,
+            config.hidden_act,
+            true,
+            "fc1",
+            "fc2",
+        )?;
+        Ok(Self {
+            layer_norm1,
+            self_attn,
+            layer_norm2,
+            mlp,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: Option<&Tensor>,
+        sin: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs.clone();
+        let xs = self.layer_norm1.forward(xs)?;
+        let xs = self.self_attn.forward(&xs, cos, sin, None, true)?;
+        let residual = residual.add(&xs)?;
+        let xs = self.layer_norm2.forward(&residual)?;
+        let xs = self.mlp.forward(&xs)?;
+        let xs = residual.add(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct SiglipEncoder {
+    layers: Vec<SiglipEncoderLayer>,
+    rotary_pos_emb: Qwen2_5VisionRotaryEmbedding,
+}
+
+impl SiglipEncoder {
+    pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
+        let vb_layers = vb.pp("layers");
+        let mut layers = vec![];
+        for i in 0..config.num_hidden_layers {
+            let layer_i = SiglipEncoderLayer::new(vb_layers.pp(i), config)?;
+            layers.push(layer_i);
+        }
+        let head_dim = config.hidden_size / config.num_attention_heads;
+        let rotary_pos_emb = Qwen2_5VisionRotaryEmbedding::new(head_dim / 2, Some(10000.0));
+        Ok(Self {
+            layers,
+            rotary_pos_emb,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor, image_grid_thw: &Tensor) -> Result<Tensor> {
+        let mut split_hids = vec![];
+        let mut split_wids = vec![];
+        for i in 0..image_grid_thw.dim(0)? {
+            let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
+                return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
+            };
+            let pos_w: Vec<u32> = (0..h).flat_map(|_| 0u32..w).collect();
+            let pos_w = pos_w.repeat(t as usize);
+            let pos_w = Tensor::new(pos_w, xs.device())?;
+            let pos_h: Vec<u32> = (0..h).flat_map(|h| vec![h; w as usize]).collect();
+            let pos_h = pos_h.repeat(t as usize);
+            let pos_h = Tensor::new(pos_h, xs.device())?;
+            split_hids.push(pos_h);
+            split_wids.push(pos_w);
+        }
+        let width_position_ids = Tensor::cat(&split_wids, 0)?;
+        let height_position_ids = Tensor::cat(&split_hids, 0)?;
+        let max_grid_size = image_grid_thw.i((.., 1..))?.max_all()?.to_scalar::<u32>()?;
+        let rope_emb_max_grid = self
+            .rotary_pos_emb
+            .forward(max_grid_size as usize, xs.device())?;
+        let rotary_pos_emb_h = rope_emb_max_grid.index_select(&height_position_ids, 0)?;
+        let rotary_pos_emb_w = rope_emb_max_grid.index_select(&width_position_ids, 0)?;
+        let rope_emb = Tensor::cat(&[rotary_pos_emb_h, rotary_pos_emb_w], 1)?.contiguous()?;
+        let rope_emb = rope_emb.repeat((1, 2))?;
+        let cos = rope_emb.cos()?;
+        let sin = rope_emb.sin()?;
+        let mut xs = xs.clone();
+        for layer in &self.layers {
+            xs = layer.forward(&xs, Some(&cos), Some(&sin))?;
+        }
+        Ok(xs)
+    }
+}
+
+pub struct SiglipVisionModel {
+    embeddings: SiglipVisionEmbeddings,
+    encoder: SiglipEncoder,
+    post_layernorm: LayerNorm,
+}
+impl SiglipVisionModel {
+    pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
+        let vb = vb.pp("vision_model");
+        let embeddings = SiglipVisionEmbeddings::new(vb.pp("embeddings"), config)?;
+        let encoder = SiglipEncoder::new(vb.pp("encoder"), config)?;
+        let post_layernorm = get_layer_norm(
+            vb.pp("post_layernorm"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        Ok(Self {
+            embeddings,
+            encoder,
+            post_layernorm,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        pixel_values: &Tensor,
+        image_grid_thw: &Tensor,
+        position_ids: &Tensor,
+        interpolate_pos_encoding: bool,
+    ) -> Result<Tensor> {
+        let xs = self.embeddings.forward(
+            pixel_values,
+            position_ids,
+            image_grid_thw,
+            interpolate_pos_encoding,
+        )?;
+        let xs = self.encoder.forward(&xs, image_grid_thw)?;
+        let xs = self.post_layernorm.forward(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct Ernie4_5DecoderLayer {
+    self_attn: NaiveAttention,
+    mlp: GateUpDownMLP,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+}
+
+impl Ernie4_5DecoderLayer {
+    pub fn new(vb: VarBuilder, config: &PaddleOCRVLConfig) -> Result<Self> {
+        let self_attn = NaiveAttention::new(
+            vb.pp("self_attn"),
+            config.hidden_size,
+            config.num_attention_heads,
+            config.num_key_value_heads,
+            Some(config.head_dim),
+            config.use_bias,
+            None,
+        )?;
+        let mlp = GateUpDownMLP::new(
+            vb.pp("mlp"),
+            config.hidden_size,
+            config.intermediate_size,
+            config.hidden_act,
+            config.use_bias,
+        )?;
+        let input_layernorm = rms_norm(
+            config.hidden_size,
+            config.rms_norm_eps,
+            vb.pp("input_layernorm"),
+        )?;
+        let post_attention_layernorm = rms_norm(
+            config.hidden_size,
+            config.rms_norm_eps,
+            vb.pp("post_attention_layernorm"),
+        )?;
+        Ok(Self {
+            self_attn,
+            mlp,
+            input_layernorm,
+            post_attention_layernorm,
+        })
+    }
+
+    pub fn forward(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs.clone();
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self
+            .self_attn
+            .forward_with_cache(&xs, cos, sin, attention_mask, false)?;
+        let residual = residual.add(&xs)?;
+        let xs = self.post_attention_layernorm.forward(&residual)?;
+        let xs = self.mlp.forward(&xs)?;
+        let xs = residual.add(&xs)?;
+        Ok(xs)
+    }
+    fn clear_kv_cache(&mut self) {
+        self.self_attn.clear_kv_cache()
+    }
+}
+
+pub struct Ernie4_5Model {
+    embed_tokens: Embedding,
+    layers: Vec<Ernie4_5DecoderLayer>,
+    norm: RmsNorm,
+    rotary_emb: Qwen2_5VLTextRotaryEmbedding,
+    rope_scaling: PaddleOCRVLRopeScalingConfig,
+}
+
+impl Ernie4_5Model {
+    pub fn new(vb: VarBuilder, config: &PaddleOCRVLConfig) -> Result<Self> {
+        let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?;
+        let vb_layers = vb.pp("layers");
+        let mut layers = vec![];
+        for i in 0..config.num_hidden_layers {
+            let layer_i = Ernie4_5DecoderLayer::new(vb_layers.pp(i), config)?;
+            layers.push(layer_i);
+        }
+        let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?;
+        let rotary_emb =
+            Qwen2_5VLTextRotaryEmbedding::new(config.head_dim, config.rope_theta as f32);
+
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm,
+            rotary_emb,
+            rope_scaling: config.rope_scaling.clone(),
+        })
+    }
+
+    pub fn forward(
+        &mut self,
+        inputs_embeds: &Tensor,
+        seqlen_offset: usize,
+        position_ids: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (b_size, seq_len, _) = inputs_embeds.dims3()?;
+
+        let position_ids = match position_ids {
+            Some(ids) => ids.clone(),
+            None => Tensor::arange(
+                seqlen_offset as u32,
+                (seq_len + seqlen_offset) as u32,
+                inputs_embeds.device(),
+            )?
+            .unsqueeze(0)?
+            .unsqueeze(0)?
+            .broadcast_as((3, b_size, seq_len))?,
+        };
+        let (cos, sin) = self.rotary_emb.forward(
+            &position_ids,
+            inputs_embeds.dtype(),
+            self.rope_scaling.mrope_section.clone(),
+        )?;
+        let mut xs = inputs_embeds.clone();
+        let attention_mask: Option<&Tensor> = {
+            if seq_len <= 1 {
+                None
+            } else {
+                Some(&prepare_causal_attention_mask(
+                    b_size,
+                    seq_len,
+                    0,
+                    xs.device(),
+                )?)
+            }
+        };
+        for layer in self.layers.iter_mut() {
+            xs = layer.forward(&xs, &cos, &sin, attention_mask)?;
+        }
+        let xs = xs.apply(&self.norm)?;
+        Ok(xs)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        for layer in self.layers.iter_mut() {
+            layer.clear_kv_cache()
+        }
+    }
+}
+
+pub struct PaddleOCRVLModel {
+    mlp_ar: Projector,
+    visual: SiglipVisionModel,
+    model: Ernie4_5Model,
+    pub cfg: PaddleOCRVLConfig,
+    lm_head: Linear,
+    rope_deltas: Option<Tensor>,
+}
+
+impl PaddleOCRVLModel {
+    pub fn new(cfg: PaddleOCRVLConfig, vb: VarBuilder) -> Result<Self> {
+        let mlp_ar = Projector::new(vb.pp("mlp_AR"), &cfg)?;
+        let visual = SiglipVisionModel::new(vb.pp("visual"), &cfg.vision_config)?;
+        let model = Ernie4_5Model::new(vb.pp("model"), &cfg)?;
+        let vocab_size = cfg.vocab_size;
+        let lm_head = if cfg.tie_word_embeddings {
+            Linear::new(model.embed_tokens.embeddings().clone(), None)
+        } else {
+            linear_no_bias(cfg.hidden_size, vocab_size, vb.pp("lm_head"))?
+        };
+
+        Ok(Self {
+            mlp_ar,
+            visual,
+            model,
+            cfg,
+            lm_head,
+            rope_deltas: None,
+        })
+    }
+
+    pub fn get_rope_index(
+        &self,
+        input_ids: &Tensor,
+        image_grid_thw: Option<&Tensor>,
+        video_grid_thw: Option<&Tensor>,
+        mask: Option<&Tensor>,
+        second_per_grid_ts: Option<Vec<f32>>,
+    ) -> Result<(Tensor, Tensor)> {
+        let spatial_merge_size = self.cfg.vision_config.spatial_merge_size;
+        let mut mrope_position_deltas: Vec<i64> = Vec::new();
+        if image_grid_thw.is_some() || video_grid_thw.is_some() {
+            let total_input_ids = input_ids.clone();
+            let mask_ = mask
+                .cloned()
+                .unwrap_or(Tensor::ones_like(&total_input_ids)?);
+            let mut position_ids = Tensor::ones(
+                (3, input_ids.dim(0)?, input_ids.dim(1)?),
+                input_ids.dtype(),
+                input_ids.device(),
+            )?;
+            let mut image_index = 0;
+            let mut video_index = 0;
+            for i in 0..total_input_ids.dim(0)? {
+                let mut input_ids_i = total_input_ids.i(i)?;
+                let mask_i = mask_.i(i)?;
+                // 推理时, attention_mask如果是全1向量,取非0索引的操作没必要
+                if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
+                    let nonzero_idx = nonzero_index(&mask_i)?;
+                    input_ids_i = input_ids_i.gather(&nonzero_idx, 0)?;
+                }
+                let mut text_start = 0;
+                let mut text_end = 0;
+                let mut thw = vec![];
+                let mut second_per_grid_t = 0_f32;
+                let mut llm_pos_ids_list: Vec<Tensor> = Vec::new();
+                // vision start的下一个索引
+                let vision_indices =
+                    get_vision_next_indices(&input_ids_i, self.cfg.vision_start_token_id);
+                match vision_indices {
+                    Ok(indeices) => {
+                        let vision_tokens = input_ids_i.gather(&indeices, 0)?.to_vec1::<u32>()?;
+                        let vision_indices_vec = indeices.to_vec1::<u32>()?;
+                        for (j, &token) in vision_tokens.iter().enumerate() {
+                            if token == self.cfg.image_token_id {
+                                thw = image_grid_thw.unwrap().i(image_index)?.to_vec1::<u32>()?;
+                                image_index += 1;
+                                text_end = vision_indices_vec[j];
+                                second_per_grid_t = 0.0;
+                            }
+                            if token == self.cfg.video_token_id {
+                                thw = video_grid_thw.unwrap().i(video_index)?.to_vec1::<u32>()?;
+                                text_end = vision_indices_vec[j];
+                                second_per_grid_t = match second_per_grid_ts {
+                                    None => 1.0,
+                                    Some(ref vec) => vec[video_index],
+                                };
+                                video_index += 1;
+                            }
+                            let llm_grid_t = thw[0];
+                            let llm_grid_h = thw[1] / spatial_merge_size as u32;
+                            let llm_grid_w = thw[2] / spatial_merge_size as u32;
+                            let text_len = text_end - text_start;
+                            let start_idx = if !llm_pos_ids_list.is_empty() {
+                                llm_pos_ids_list[llm_pos_ids_list.len() - 1]
+                                    .max_all()?
+                                    .to_scalar::<u32>()?
+                                    + 1
+                            } else {
+                                0
+                            };
+                            let pos_ids = Tensor::arange(
+                                start_idx,
+                                start_idx + text_len,
+                                input_ids_i.device(),
+                            )?
+                            .unsqueeze(0)?
+                            .broadcast_as((3usize, text_len as usize))?;
+                            llm_pos_ids_list.push(pos_ids);
+                            let range_tensor = Tensor::arange(0, llm_grid_t, input_ids_i.device())?
+                                .unsqueeze(D::Minus1)?;
+                            let expanded_range = range_tensor.broadcast_as((
+                                llm_grid_t as usize,
+                                (llm_grid_h * llm_grid_w) as usize,
+                            ))?;
+                            let time_tensor = expanded_range
+                                .broadcast_mul(&Tensor::new(
+                                    vec![
+                                        (second_per_grid_t
+                                            * self.cfg.vision_config.tokens_per_second as f32)
+                                            as u32,
+                                    ],
+                                    input_ids_i.device(),
+                                )?)?
+                                .broadcast_add(&Tensor::new(
+                                    vec![start_idx + text_len],
+                                    input_ids_i.device(),
+                                )?)?;
+                            let t_index = time_tensor.flatten_all()?;
+                            let h_index = Tensor::arange(
+                                start_idx + text_len,
+                                start_idx + text_len + llm_grid_h,
+                                input_ids_i.device(),
+                            )?
+                            .unsqueeze(0)?
+                            .unsqueeze(D::Minus1)?
+                            .broadcast_as((
+                                llm_grid_t as usize,
+                                llm_grid_h as usize,
+                                llm_grid_w as usize,
+                            ))?
+                            .flatten_all()?;
+                            let w_index = Tensor::arange(
+                                start_idx + text_len,
+                                start_idx + text_len + llm_grid_w,
+                                input_ids_i.device(),
+                            )?
+                            .unsqueeze(0)?
+                            .unsqueeze(0)?
+                            .broadcast_as((
+                                llm_grid_t as usize,
+                                llm_grid_h as usize,
+                                llm_grid_w as usize,
+                            ))?
+                            .flatten_all()?;
+
+                            let thw_index = Tensor::stack(&[t_index, h_index, w_index], 0)?;
+                            llm_pos_ids_list.push(thw_index);
+                            text_start = text_end + llm_grid_t * llm_grid_h * llm_grid_w;
+                        }
+                    }
+                    Err(e) => {
+                        println!("get vision_indices err: {}", e);
+                    }
+                };
+
+                if text_start < input_ids_i.dim(0)? as u32 {
+                    let start_idx = if !llm_pos_ids_list.is_empty() {
+                        llm_pos_ids_list[llm_pos_ids_list.len() - 1]
+                            .max_all()?
+                            .to_scalar::<u32>()?
+                            + 1
+                    } else {
+                        0
+                    };
+                    let text_len = input_ids_i.dim(0)? as u32 - text_start;
+                    let pos_ids =
+                        Tensor::arange(start_idx, start_idx + text_len, input_ids_i.device())?
+                            .unsqueeze(0)?
+                            .broadcast_as((3usize, text_len as usize))?;
+                    llm_pos_ids_list.push(pos_ids);
+                }
+                let llm_position = Tensor::cat(&llm_pos_ids_list, 1)?.reshape((3, 1, ()))?;
+                position_ids = position_ids
+                    .slice_assign(&[(0..3), (i..i + 1), (0..input_ids.dim(1)?)], &llm_position)?;
+                let position_deltas = llm_position.max_all()?.to_scalar::<u32>()? as i64 + 1
+                    - input_ids_i.dim(0)? as i64;
+                mrope_position_deltas.push(position_deltas);
+            }
+
+            let mut mrope_position_deltas = Tensor::new(mrope_position_deltas, input_ids.device())?;
+            if mrope_position_deltas.rank() == 1 {
+                mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
+            }
+            Ok((position_ids.contiguous()?, mrope_position_deltas))
+        } else if let Some(mask) = mask {
+            let mut position_ids = mask
+                .to_dtype(candle_core::DType::F64)?
+                .cumsum(D::Minus1)?
+                .to_dtype(candle_core::DType::U32)?
+                .broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
+            for i in 0..position_ids.dim(0)? {
+                let mut position_ids_i = position_ids.i(i)?;
+                let mask_i = mask.i(i)?;
+                // 如果有pad, 将填充位置置为1
+                // 当bs>1, 可能存在不同序列长度，需要添加pad使seq_len长度一致
+                if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
+                    let zero_indices = zero_index(&mask_i)?;
+                    let replace_1 = Tensor::ones(
+                        zero_indices.dim(0)?,
+                        candle_core::DType::U32,
+                        input_ids.device(),
+                    )?;
+                    position_ids_i = position_ids_i
+                        .scatter(&zero_indices, &replace_1, 0)?
+                        .unsqueeze(0)?;
+                    position_ids = position_ids
+                        .slice_assign(&[(i..i + 1), (0..position_ids.dim(1)?)], &position_ids_i)?;
+                }
+            }
+            position_ids = position_ids
+                .unsqueeze(0)?
+                .broadcast_as((3, input_ids.dim(0)?, input_ids.dim(1)?))?
+                .contiguous()?;
+            let mut mrope_position_deltas = position_ids
+                .max(0)?
+                .max(D::Minus1)?
+                .broadcast_sub(&Tensor::new(
+                    vec![mask.dim(D::Minus1)? as u32 - 1],
+                    input_ids.device(),
+                )?)?
+                .contiguous()?;
+            if mrope_position_deltas.rank() == 1 {
+                mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
+            }
+            Ok((position_ids, mrope_position_deltas))
+        } else {
+            let position_ids =
+                Tensor::arange(0_u32, input_ids.dim(D::Minus1)? as u32, input_ids.device())?
+                    .unsqueeze(0)?
+                    .unsqueeze(0)?
+                    .broadcast_as((3, input_ids.dim(0)?, input_ids.dim(D::Minus1)?))?
+                    .contiguous()?;
+            let mrope_position_deltas = Tensor::zeros(
+                (input_ids.dim(0)?, 1),
+                input_ids.dtype(),
+                input_ids.device(),
+            )?;
+            Ok((position_ids, mrope_position_deltas))
+        }
+    }
+
+    pub fn forward(
+        &mut self,
+        input_ids: &Tensor,
+        pixel_values: Option<&Tensor>,
+        image_grid_thw: Option<&Tensor>,
+        image_mask: &Tensor,
+        cache_position: Option<&Tensor>,
+        seqlen_offset: usize,
+    ) -> Result<Tensor> {
+        let mut inputs_embeds = self.model.embed_tokens.forward(input_ids)?;
+        if let Some(pixel_values) = pixel_values
+            && let Some(image_grid_thw) = image_grid_thw
+        {
+            let pixel_values = pixel_values.unsqueeze(0)?;
+            let mut siglip_position_ids = vec![];
+            let mut sample_indices = vec![];
+            let mut cu_seqlens = vec![0u32];
+            let img_num = image_grid_thw.dim(0)?;
+            for i in 0..img_num {
+                let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
+                    return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
+                };
+                let numel = h * w;
+                let image_position_ids =
+                    Tensor::arange(0, numel, pixel_values.device())?.repeat(t as usize)?;
+                siglip_position_ids.push(image_position_ids);
+                let indices =
+                    Tensor::new(vec![i as u32; (numel * t) as usize], pixel_values.device())?;
+                sample_indices.push(indices);
+                cu_seqlens.push(cu_seqlens[cu_seqlens.len() - 1] + numel * t);
+            }
+            let siglip_position_ids = Tensor::cat(&siglip_position_ids, 0)?;
+
+            let image_embed =
+                self.visual
+                    .forward(&pixel_values, image_grid_thw, &siglip_position_ids, true)?;
+            let image_embed = image_embed.squeeze(0)?;
+            let image_embed = self.mlp_ar.forward(&image_embed, image_grid_thw)?;
+            inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embed, image_mask)?;
+        }
+        let position_ids;
+        let rope_deltas;
+        if (cache_position.is_some() && cache_position.unwrap().i(0)?.to_scalar::<u32>()? == 0)
+            || self.rope_deltas.is_none()
+        {
+            (position_ids, rope_deltas) =
+                self.get_rope_index(input_ids, image_grid_thw, None, None, None)?;
+            self.rope_deltas = Some(rope_deltas);
+        } else {
+            let (bs, seq_len, _) = inputs_embeds.dims3()?;
+            let delta = if let Some(cache_position) = cache_position {
+                cache_position
+                    .i(0)?
+                    .to_dtype(self.rope_deltas.as_ref().unwrap().dtype())?
+                    .broadcast_add(self.rope_deltas.as_ref().unwrap())?
+                    .contiguous()?
+                    .to_dtype(candle_core::DType::U32)?
+            } else {
+                Tensor::zeros(1, inputs_embeds.dtype(), inputs_embeds.device())?
+            };
+            position_ids = Tensor::arange(0u32, seq_len as u32, input_ids.device())?
+                .unsqueeze(0)?
+                .broadcast_as((bs, seq_len))?
+                .broadcast_add(&delta)?
+                .unsqueeze(0)?
+                .broadcast_as((3, bs, seq_len))?
+                .contiguous()?;
+        }
+        let outputs = self
+            .model
+            .forward(&inputs_embeds, seqlen_offset, Some(&position_ids))?;
+        let seq_len = outputs.dim(1)?;
+        let hidden_state = outputs.narrow(1, seq_len - 1, 1)?;
+        let logits = self.lm_head.forward(&hidden_state)?;
+        Ok(logits)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.model.clear_kv_cache();
+    }
+}

src/models/paddleocr_vl/processor.rs

@@ -0,0 +1,148 @@
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{DType, Device, IndexOp, Shape, Tensor};
+use image::DynamicImage;
+
+use crate::{
+    models::paddleocr_vl::config::PaddleOCRVLPreprocessorConfig,
+    utils::img_utils::{extract_images, img_smart_resize, img_transform},
+};
+
+pub struct PaddleOCRVLProcessor {
+    process_cfg: PaddleOCRVLPreprocessorConfig,
+    device: Device,
+    dtype: DType,
+    image_token: String,
+}
+
+impl PaddleOCRVLProcessor {
+    pub fn new(
+        config: PaddleOCRVLPreprocessorConfig,
+        device: &Device,
+        dtype: DType,
+    ) -> Result<Self> {
+        let image_token = "<|IMAGE_PLACEHOLDER|>".to_string();
+        Ok(Self {
+            process_cfg: config,
+            device: device.clone(),
+            dtype,
+            image_token,
+        })
+    }
+
+    pub fn process_img(
+        &self,
+        img: &DynamicImage,
+        img_mean: &Tensor,
+        img_std: &Tensor,
+    ) -> Result<Tensor> {
+        let img_h = img.height();
+        let img_w = img.width();
+        //  h,w resize成 32的倍数
+        let (resize_h, resize_w) = img_smart_resize(
+            img_h,
+            img_w,
+            (self.process_cfg.patch_size * self.process_cfg.merge_size) as u32,
+            self.process_cfg.min_pixels,
+            self.process_cfg.max_pixels,
+        )?;
+        let img = img.resize_exact(resize_w, resize_h, image::imageops::FilterType::CatmullRom);
+        let img_tensor = img_transform(&img, img_mean, img_std, &self.device, self.dtype)?;
+        // (c, h, w) => (1, c, h, w)
+        let img_tensor = img_tensor.unsqueeze(0)?;
+        Ok(img_tensor)
+    }
+
+    pub fn process_vision_tensor(&self, img_tensor: &Tensor) -> Result<(Tensor, Tensor)> {
+        let channel = img_tensor.dim(1)?;
+        // img_temsor.dim[0] = 1, temporal_patch_size = 1, grid_t = 1
+        let grid_t = img_tensor.dim(0)? / self.process_cfg.temporal_patch_size;
+        let grid_h = img_tensor.dim(2)? / self.process_cfg.patch_size;
+        let grid_w = img_tensor.dim(3)? / self.process_cfg.patch_size;
+        let shape = Shape::from(vec![
+            grid_t,
+            self.process_cfg.temporal_patch_size,
+            channel,
+            grid_h,
+            self.process_cfg.patch_size,
+            grid_w,
+            self.process_cfg.patch_size,
+        ]);
+        let img_tensor = img_tensor.reshape(shape)?;
+        // shape to // grid_t,
+        // grid_h,
+        // grid_w,
+        // channel,
+        // temporal_patch_size
+        // patch_size,
+        // patch_size,
+        let img_tensor = img_tensor.permute(vec![0, 3, 5, 2, 1, 4, 6])?;
+        let img_tensor = img_tensor
+            .reshape((
+                grid_t * grid_h * grid_w,
+                channel,
+                self.process_cfg.patch_size,
+                self.process_cfg.patch_size,
+            ))?
+            .contiguous()?;
+        let grid_thw = Tensor::from_vec(
+            vec![grid_t as u32, grid_h as u32, grid_w as u32],
+            (1, 3),
+            &self.device,
+        )?;
+        Ok((img_tensor, grid_thw))
+    }
+
+    pub fn process_images(
+        &self,
+        imgs: &Vec<DynamicImage>,
+        img_mean: &Tensor,
+        img_std: &Tensor,
+    ) -> Result<(Tensor, Tensor)> {
+        let mut pixel_values_vec = Vec::new();
+        let mut vision_grid_thws_vec = Vec::new();
+        for img in imgs {
+            let img_tensor = self.process_img(img, img_mean, img_std)?;
+            let (img_tensor, grid_thw) = self.process_vision_tensor(&img_tensor)?;
+            pixel_values_vec.push(img_tensor);
+            vision_grid_thws_vec.push(grid_thw);
+        }
+        let pixel_values = Tensor::cat(&pixel_values_vec, 0)?;
+        let vision_grid_thws = Tensor::cat(&vision_grid_thws_vec, 0)?;
+        Ok((pixel_values, vision_grid_thws))
+    }
+
+    pub fn process_info(
+        &self,
+        messages: &ChatCompletionParameters,
+        text: &str,
+    ) -> Result<(String, Option<Tensor>, Option<Tensor>)> {
+        let imgs = extract_images(messages)?;
+        let img_mean = Tensor::from_slice(&self.process_cfg.image_mean, (3, 1, 1), &self.device)?
+            .to_dtype(self.dtype)?;
+        let img_std = Tensor::from_slice(&self.process_cfg.image_std, (3, 1, 1), &self.device)?
+            .to_dtype(self.dtype)?;
+        let (pixel_values, image_grid_thw) = if !imgs.is_empty() {
+            let (pixel_values, image_grid_thw) = self.process_images(&imgs, &img_mean, &img_std)?;
+            (Some(pixel_values), Some(image_grid_thw))
+        } else {
+            (None, None)
+        };
+
+        let merge_length = self.process_cfg.merge_size.pow(2);
+        let mut text = text.to_string();
+        if let Some(ref image_grid_thw) = image_grid_thw {
+            let mut index = 0;
+            while text.contains(&self.image_token) {
+                let grid_i = image_grid_thw.i(index)?;
+                let repeat_num =
+                    grid_i.to_vec1::<u32>()?.iter().product::<u32>() as usize / merge_length;
+                let replace = "<|placeholder|>".repeat(repeat_num);
+                text = text.replacen(&self.image_token, &replace, 1);
+                index += 1;
+            }
+            text = text.replace("<|placeholder|>", &self.image_token);
+        }
+        Ok((text, pixel_values, image_grid_thw))
+    }
+}

src/models/qwen2_5vl/generate.rs

@@ -1,4 +1,3 @@
-// use crate::models::GenerateStream;
 use aha_openai_dive::v1::resources::chat::{
     ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
 };

src/models/voxcpm/minicpm4.rs

@@ -117,7 +117,9 @@ impl MiniCPMDecoderLayer {
             cfg.hidden_size,
             cfg.num_attention_heads,
             cfg.num_key_value_heads,
+            None,
             false,
+            None,
         )?;
         let mlp = GateUpDownMLP::new(
             vb.pp("mlp"),

tests/config_tests.rs

@@ -1,7 +1,8 @@
 use aha::models::{
     deepseek_ocr::config::DeepseekOCRConfig, hunyuan_ocr::config::HunYuanVLConfig,
-    minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig,
-    qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig,
+    minicpm4::config::MiniCPM4Config, paddleocr_vl::config::PaddleOCRVLConfig,
+    qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig,
+    voxcpm::config::VoxCPMConfig,
 };
 use anyhow::Result;
 
@@ -65,3 +66,12 @@ fn hunyuan_ocr_config() -> Result<()> {
     println!("{:?}", config);
     Ok(())
 }
+#[test]
+fn paddleocr_vl_config() -> Result<()> {
+    // cargo test -F cuda paddleocr_vl_config -r -- --nocapture
+    let model_path = "/home/jhq/huggingface_model/PaddlePaddle/PaddleOCR-VL/";
+    let config_path = model_path.to_string() + "/config.json";
+    let config: PaddleOCRVLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
+    println!("{:?}", config);
+    Ok(())
+}

tests/test_paddleocr_vl.rs

@@ -0,0 +1,89 @@
+use std::{pin::pin, time::Instant};
+
+use aha::models::{GenerateModel, paddleocr_vl::generate::PaddleOCRVLGenerateModel};
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use rocket::futures::StreamExt;
+
+#[test]
+fn paddleocr_vl_generate() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda paddleocr_vl_generate -r -- --nocapture
+    let message = r#"
+    {
+        "model": "paddleocr_vl",
+        "messages": [
+            {
+                "role": "user",
+                "content": [ 
+                    {
+                        "type": "image",
+                        "image_url": 
+                        {
+                            "url": "file://./assets/img/ocr_test1.png"
+                        }
+                    },              
+                    {
+                        "type": "text", 
+                        "text": "OCR:"
+                    }
+                ]
+            }
+        ],
+        "stream": false
+    }
+    "#;
+    let model_path = "/home/jhq/huggingface_model/PaddlePaddle/PaddleOCR-VL/";
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut model = PaddleOCRVLGenerateModel::init(model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+    let i_start = Instant::now();
+    let res = model.generate(mes)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+    println!("generate: \n {:?}", res);
+    Ok(())
+}
+
+#[tokio::test]
+async fn paddleocr_vl_stream() -> Result<()> {
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda paddleocr_vl_stream -r -- --nocapture
+
+    let message = r#"
+    {
+        "model": "paddleocr_vl",
+        "messages": [
+            {
+                "role": "user",
+                "content": [ 
+                    {
+                        "type": "image",
+                        "image_url": 
+                        {
+                            "url": "file://./assets/img/ocr_test1.png"
+                        }
+                    },              
+                    {
+                        "type": "text", 
+                        "text": "OCR:"
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    let model_path = "/home/jhq/huggingface_model/PaddlePaddle/PaddleOCR-VL/";
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut model = PaddleOCRVLGenerateModel::init(model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+    let mut stream = pin!(model.generate_stream(mes)?);
+    while let Some(item) = stream.next().await {
+        println!("generate: \n {:?}", item);
+    }
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+    Ok(())
+}