diff --git a/Cargo.lock b/Cargo.lock index 73ea66b..58fab09 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19,7 +19,7 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "aha" -version = "0.1.3" +version = "0.1.4" dependencies = [ "aha_openai_dive", "anyhow", diff --git a/Cargo.toml b/Cargo.toml index 443ce5a..b8da2b8 100644 --- a/Cargo.toml +++ b/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"} diff --git a/README.md b/README.md index 4722020..255ba3c 100644 --- a/README.md +++ b/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 +* 设置HTTP服务监听的端口号 +* 默认值：10100 +* 示例：--port 8080 或 -p 8080 + +2. 模型选择（必选） +----- + -m, --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 /path/to/model/dir + +4. 保存路径 +----- + --save-dir +* 指定模型下载保存的目录 +* 默认保存在用户主目录下的 .aha 文件夹中 +* 示例：--save-dir /custom/model/path + +5. 下载重试次数 +----- + --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 -* 设置HTTP服务监听的端口号 -* 默认值：10100 -* 示例：--port 8080 或 -p 8080 - -2. 模型选择（必选） ------ - -m, --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 /path/to/model/dir - -4. 保存路径 ------ - --save-dir -* 指定模型下载保存的目录 -* 默认保存在用户主目录下的 .aha 文件夹中 -* 示例：--save-dir /custom/model/path - -5. 下载重试次数 ------ - --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 模型 diff --git a/src/chat_template/mod.rs b/src/chat_template/mod.rs index f1f31f0..24c4bf7 100644 --- a/src/chat_template/mod.rs +++ b/src/chat_template/mod.rs @@ -51,12 +51,23 @@ pub struct ChatTemplate<'a> { impl<'a> ChatTemplate<'a> { pub fn init(path: &str) -> Result { - 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(); diff --git a/src/main.rs b/src/main.rs index 7658745..23a0140 100644 --- a/src/main.rs +++ b/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, diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 015b4f7..68a03e8 100644 --- a/src/models/common/mod.rs +++ b/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, bias: bool, + o_proj_pp_name: Option<&str>, ) -> Result { 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) } diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 950e374..0164b9e 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/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) diff --git a/src/models/hunyuan_ocr/model.rs b/src/models/hunyuan_ocr/model.rs index edc26c0..4e9f5bc 100644 --- a/src/models/hunyuan_ocr/model.rs +++ b/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"), diff --git a/src/models/minicpm4/model.rs b/src/models/minicpm4/model.rs index dd51e46..cf48c5f 100644 --- a/src/models/minicpm4/model.rs +++ b/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"), diff --git a/src/models/mod.rs b/src/models/mod.rs index bc50824..9e1f04d 100644 --- a/src/models/mod.rs +++ b/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>), } 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 { + let model = PaddleOCRVLGenerateModel::init(path, None, None)?; + ModelInstance::PaddleOCRVL(Box::new(model)) + } }; Ok(model) } diff --git a/src/models/paddleocr_vl/config.rs b/src/models/paddleocr_vl/config.rs new file mode 100644 index 0000000..a6a3b59 --- /dev/null +++ b/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, + 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, + 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, + 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, + pub image_std: Vec, + 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, + pub temporal_patch_size: usize, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct SizeConfig { + pub max_pixels: usize, + pub min_pixels: usize, +} diff --git a/src/models/paddleocr_vl/generate.rs b/src/models/paddleocr_vl/generate.rs new file mode 100644 index 0000000..254f377 --- /dev/null +++ b/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) -> Result { + 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 { + 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> + + 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))) + } +} diff --git a/src/models/paddleocr_vl/mod.rs b/src/models/paddleocr_vl/mod.rs new file mode 100644 index 0000000..8b1baf7 --- /dev/null +++ b/src/models/paddleocr_vl/mod.rs @@ -0,0 +1,4 @@ +pub mod config; +pub mod generate; +pub mod model; +pub mod processor; diff --git a/src/models/paddleocr_vl/model.rs b/src/models/paddleocr_vl/model.rs new file mode 100644 index 0000000..dc613bf --- /dev/null +++ b/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 { + 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 { + 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::()?[..] 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 { + 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 { + 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 { + 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::()?[..] 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 { + 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 { + 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, + rotary_pos_emb: Qwen2_5VisionRotaryEmbedding, +} + +impl SiglipEncoder { + pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result { + 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 { + 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::()?[..] else { + return Err(anyhow!(format!("grid_thw Expected exactly 3 elements"))); + }; + let pos_w: Vec = (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 = (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::()?; + 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 { + 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 { + 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 { + 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 { + 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, + norm: RmsNorm, + rotary_emb: Qwen2_5VLTextRotaryEmbedding, + rope_scaling: PaddleOCRVLRopeScalingConfig, +} + +impl Ernie4_5Model { + pub fn new(vb: VarBuilder, config: &PaddleOCRVLConfig) -> Result { + 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 { + 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, +} + +impl PaddleOCRVLModel { + pub fn new(cfg: PaddleOCRVLConfig, vb: VarBuilder) -> Result { + 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>, + ) -> Result<(Tensor, Tensor)> { + let spatial_merge_size = self.cfg.vision_config.spatial_merge_size; + let mut mrope_position_deltas: Vec = 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::()? != 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 = 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::()?; + let vision_indices_vec = indeices.to_vec1::()?; + 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::()?; + 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::()?; + 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::()? + + 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::()? + + 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::()? 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::()? != 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 { + 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::()?[..] 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::()? == 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(); + } +} diff --git a/src/models/paddleocr_vl/processor.rs b/src/models/paddleocr_vl/processor.rs new file mode 100644 index 0000000..594a95c --- /dev/null +++ b/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 { + 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 { + 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, + 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, Option)> { + 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::()?.iter().product::() 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)) + } +} diff --git a/src/models/qwen2_5vl/generate.rs b/src/models/qwen2_5vl/generate.rs index 42a5713..08b41ba 100644 --- a/src/models/qwen2_5vl/generate.rs +++ b/src/models/qwen2_5vl/generate.rs @@ -1,4 +1,3 @@ -// use crate::models::GenerateStream; use aha_openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; diff --git a/src/models/voxcpm/minicpm4.rs b/src/models/voxcpm/minicpm4.rs index 8bcc490..8ee90a7 100644 --- a/src/models/voxcpm/minicpm4.rs +++ b/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"), diff --git a/tests/config_tests.rs b/tests/config_tests.rs index 2e25a14..238b486 100644 --- a/tests/config_tests.rs +++ b/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(()) +} diff --git a/tests/test_paddleocr_vl.rs b/tests/test_paddleocr_vl.rs new file mode 100644 index 0000000..738f686 --- /dev/null +++ b/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(()) +}