From 697484cf23be250db1556fb9d17748af4032ae3d Mon Sep 17 00:00:00 2001 From: jhqxxx Date: Wed, 3 Dec 2025 17:21:01 +0800 Subject: [PATCH] add hunyuan_ocr --- Cargo.toml | 3 +- README.md | 16 +- src/main.rs | 1 + src/models/common/mod.rs | 160 +++++-- src/models/deepseek_ocr/generate.rs | 6 +- src/models/deepseek_ocr/model.rs | 56 +-- src/models/hunyuan_ocr/config.rs | 105 +++++ src/models/hunyuan_ocr/generate.rs | 218 ++++++++++ src/models/hunyuan_ocr/mod.rs | 4 + src/models/hunyuan_ocr/model.rs | 621 ++++++++++++++++++++++++++++ src/models/hunyuan_ocr/processor.rs | 236 +++++++++++ src/models/minicpm4/generate.rs | 6 +- src/models/minicpm4/model.rs | 14 +- src/models/mod.rs | 13 +- src/models/qwen2_5vl/generate.rs | 6 +- src/models/qwen3vl/generate.rs | 6 +- src/models/qwen3vl/model.rs | 7 +- src/models/qwen3vl/processor.rs | 55 +-- src/models/voxcpm/minicpm4.rs | 14 +- src/position_embed/rope.rs | 44 ++ src/tokenizer/mod.rs | 2 +- src/utils/img_utils.rs | 38 ++ src/utils/mod.rs | 41 +- src/utils/tensor_utils.rs | 114 ++++- tests/config_tests.rs | 18 +- tests/messy_test.rs | 25 +- tests/test_hunyuan_ocr.rs | 88 ++++ tests/test_qwen3vl.rs | 10 +- tests/weight_test.rs | 19 + 29 files changed, 1756 insertions(+), 190 deletions(-) create mode 100644 src/models/hunyuan_ocr/config.rs create mode 100644 src/models/hunyuan_ocr/generate.rs create mode 100644 src/models/hunyuan_ocr/mod.rs create mode 100644 src/models/hunyuan_ocr/model.rs create mode 100644 src/models/hunyuan_ocr/processor.rs create mode 100644 tests/test_hunyuan_ocr.rs diff --git a/Cargo.toml b/Cargo.toml index 60949d4..0346f9b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,4 +36,5 @@ flash-attn=["candle-flash-attn"] cuda=["candle-nn/cuda", "candle-core/cuda", "candle-transformers/cuda"] [lints.clippy] -needless_range_loop = "allow" \ No newline at end of file +needless_range_loop = "allow" +single_range_in_vec_init = "allow" \ No newline at end of file diff --git a/README.md b/README.md index 183c62a..4722020 100644 --- a/README.md +++ b/README.md @@ -17,10 +17,13 @@ * VoxCPM - 面壁智能语音生成模型 * Qwen3VL - 阿里通义千问 3 多模态大语言模型 * DeepSeek-OCR - 深度求索光学文字识别模型 +* Hunyuan-OCR - 腾讯混元光学文字识别模型 ## 计划支持 我们持续扩展支持的模型列表，欢迎贡献！ +⭐ 如果这个项目对你有帮助，请给我们一个 Star！ + ## 环境依赖 1. ffmpeg: * ubuntu/WSL @@ -77,6 +80,10 @@ fn main() -> Result<()> { git clone https://github.com/jhqxxx/aha.git cd aha # 修改测试用例中模型路径 + +# 运行 Hunyuan-OCR 示例 +cargo test -F cuda hunyuan_ocr_generate -r -- --nocapture + # 运行 DeepSeek-OCR 示例 cargo test -F cuda deepseek_ocr_generate -r -- --nocapture @@ -122,6 +129,7 @@ cargo run -F cuda -- [参数] * 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. 权重路径 @@ -162,6 +170,7 @@ cargo run -F cuda -- [参数] │ ├── models │ │ ├── common │ │ ├── deepseek_ocr +│ │ ├── hunyuan_ocr │ │ ├── minicpm4 │ │ ├── qwen2_5vl │ │ ├── qwen3vl @@ -170,8 +179,10 @@ cargo run -F cuda -- [参数] │ ├── position_embed │ ├── tokenizer │ ├── utils +│ ├── api.rs │ └── lib.rs └── tests + ├── test_hunyuan_ocr.rs ├── test_deepseek_ocr.rs ├── test_minicpm4.rs ├── test_qwen2_5vl.rs @@ -196,6 +207,10 @@ cargo run -F cuda -- [参数] 2. 提交新的 Issue，包含详细描述和复现步骤 ## 更新日志 + +### v0.1.3 +* 添加 Hunyuan-OCR 模型 + ### v0.1.2 * 添加 DeepSeek-OCR 模型 @@ -207,4 +222,3 @@ cargo run -F cuda -- [参数] * 支持 Qwen2.5VL, MiniCPM4, VoxCPM 模型 -⭐ 如果这个项目对你有帮助，请给我们一个 Star！ \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 7a487ae..7658745 100644 --- a/src/main.rs +++ b/src/main.rs @@ -84,6 +84,7 @@ async fn main() -> anyhow::Result<()> { WhichModel::Qwen3vl8B => "Qwen/Qwen3-VL-8B-Instruct", WhichModel::Qwen3vl32B => "Qwen/Qwen3-VL-32B-Instruct", WhichModel::DeepSeekOCR => "deepseek-ai/DeepSeek-OCR", + WhichModel::HunyuanOCR => "Tencent-Hunyuan/HunyuanOCR", }; 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 81fb395..015b4f7 100644 --- a/src/models/common/mod.rs +++ b/src/models/common/mod.rs @@ -1,27 +1,41 @@ use anyhow::Result; use candle_core::{D, Tensor}; -use candle_nn::{Activation, Linear, Module, VarBuilder, linear, linear_no_bias}; +use candle_nn::{ + Activation, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder, + conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias, +}; use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv}; #[derive(Debug, Clone)] -pub struct MLPWithBias { +pub struct GateUpDownMLP { gate_proj: Linear, up_proj: Linear, down_proj: Linear, act_fn: Activation, } -impl MLPWithBias { +impl GateUpDownMLP { pub fn new( vb: VarBuilder, hidden_size: usize, intermediate_size: usize, act_fn: Activation, + bias: bool, ) -> Result { - let gate_proj = linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?; - let up_proj = linear(hidden_size, intermediate_size, vb.pp("up_proj"))?; - let down_proj = linear(intermediate_size, hidden_size, vb.pp("down_proj"))?; + let (gate_proj, up_proj, down_proj) = if bias { + ( + linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?, + linear(hidden_size, intermediate_size, vb.pp("up_proj"))?, + linear(intermediate_size, hidden_size, vb.pp("down_proj"))?, + ) + } else { + ( + linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?, + linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?, + linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?, + ) + }; Ok(Self { gate_proj, up_proj, @@ -31,7 +45,7 @@ impl MLPWithBias { } } -impl Module for MLPWithBias { +impl Module for GateUpDownMLP { fn forward(&self, xs: &Tensor) -> candle_core::Result { let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; let rhs = xs.apply(&self.up_proj)?; @@ -39,43 +53,51 @@ impl Module for MLPWithBias { } } -#[derive(Debug, Clone)] -pub struct MLPNoBias { - gate_proj: Linear, - up_proj: Linear, - down_proj: Linear, - act_fn: Activation, +pub struct TwoLinearMLP { + linear1: Linear, + linear2: Linear, + act: Activation, } -impl MLPNoBias { +impl TwoLinearMLP { pub fn new( vb: VarBuilder, - hidden_size: usize, - intermediate_size: usize, - act_fn: Activation, + embedding_dim: usize, + mlp_dim: usize, + act: Activation, + bias: bool, + linear1_pp_name: &str, + linear2_pp_name: &str, ) -> Result { - let gate_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?; - let up_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?; - let down_proj = linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?; + let (linear1, linear2) = if bias { + ( + linear(embedding_dim, mlp_dim, vb.pp(linear1_pp_name))?, + linear(mlp_dim, embedding_dim, vb.pp(linear2_pp_name))?, + ) + } else { + ( + linear_no_bias(embedding_dim, mlp_dim, vb.pp(linear1_pp_name))?, + linear_no_bias(mlp_dim, embedding_dim, vb.pp(linear2_pp_name))?, + ) + }; Ok(Self { - gate_proj, - up_proj, - down_proj, - act_fn, + linear1, + linear2, + act, }) } -} - -impl Module for MLPNoBias { - fn forward(&self, xs: &Tensor) -> candle_core::Result { - let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?; - let rhs = xs.apply(&self.up_proj)?; - (lhs * rhs)?.apply(&self.down_proj) + pub fn forward(&self, xs: &Tensor) -> Result { + let xs = xs + .apply(&self.linear1)? + .apply(&self.act)? + .apply(&self.linear2)?; + Ok(xs) } } #[derive(Debug, Clone)] -pub struct AttentionNobias { +// pub struct AttentionNobias { +pub struct NaiveAttention { q_proj: Linear, k_proj: Linear, v_proj: Linear, @@ -88,19 +110,33 @@ pub struct AttentionNobias { kv_cache: Option<(Tensor, Tensor)>, } -impl AttentionNobias { +// impl AttentionNobias { +impl NaiveAttention { pub fn new( vb: VarBuilder, hidden_size: usize, num_attention_heads: usize, num_key_value_heads: usize, + bias: bool, ) -> Result { let num_kv_groups = num_attention_heads / num_key_value_heads; let head_dim = hidden_size / num_attention_heads; - let q_proj = linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; - let k_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?; - let v_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?; - let o_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("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"))?, + ) + } 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"))?, + ) + }; + Ok(Self { q_proj, k_proj, @@ -118,8 +154,8 @@ impl AttentionNobias { pub fn forward( &self, xs: &Tensor, - cos: &Tensor, - sin: &Tensor, + cos: Option<&Tensor>, + sin: Option<&Tensor>, attention_mask: Option<&Tensor>, tof32: bool, ) -> Result { @@ -136,8 +172,14 @@ impl AttentionNobias { let value_states = value_states .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))? .transpose(1, 2)?; - let (query_states, key_states) = - apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?; + let (query_states, key_states) = if let Some(cos) = cos + && let Some(sin) = sin + { + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)? + } else { + (query_states, key_states) + }; + let scale = 1f64 / f64::sqrt(self.head_dim as f64); let attn_output = eager_attention_forward( &query_states, @@ -260,3 +302,39 @@ pub fn eager_attention_forward( Ok(attn_output) } + +pub fn get_conv2d( + vb: VarBuilder, + in_c: usize, + out_c: usize, + kernel_size: usize, + padding: usize, + stride: usize, + dilation: usize, + groups: usize, + bias: bool, +) -> Result { + let cfg = Conv2dConfig { + padding, + stride, + dilation, + groups, + cudnn_fwd_algo: None, + }; + let conv2d = if bias { + conv2d(in_c, out_c, kernel_size, cfg, vb)? + } else { + conv2d_no_bias(in_c, out_c, kernel_size, cfg, vb)? + }; + Ok(conv2d) +} + +pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result { + let ln_config = LayerNormConfig { + eps, + remove_mean: true, // true for layernorm, false for RMSNorm + affine: true, // true for with bias, false for without bias + }; + let norm = layer_norm(dim, ln_config, vb)?; + Ok(norm) +} diff --git a/src/models/deepseek_ocr/generate.rs b/src/models/deepseek_ocr/generate.rs index d5281d2..fc14c4e 100644 --- a/src/models/deepseek_ocr/generate.rs +++ b/src/models/deepseek_ocr/generate.rs @@ -29,6 +29,7 @@ pub struct DeepseekOCRGenerateModel { eos_token_id: u32, device: Device, size: Vec, + model_name: String, } impl DeepseekOCRGenerateModel { @@ -54,6 +55,7 @@ impl DeepseekOCRGenerateModel { eos_token_id, device: device.clone(), size, + model_name: "deepseek-ocr".to_string(), }) } } @@ -127,7 +129,7 @@ impl GenerateModel for DeepseekOCRGenerateModel { } let res = self.tokenizer.token_decode(generate)?; self.deepseekocr_model.clear_kv_cache(); - let response = build_completion_response(res, "deepseek_ocr"); + let response = build_completion_response(res, &self.model_name); Ok(response) } @@ -218,7 +220,7 @@ impl GenerateModel for DeepseekOCRGenerateModel { continue; } error_tokens.clear(); - let chunk = build_completion_chunk_response(decoded_token, "deepseek_ocr", None, None); + let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); yield Ok(chunk); if next_token == self.bos_token_id || next_token == self.eos_token_id { break; diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index d2c5135..950e374 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/src/models/deepseek_ocr/model.rs @@ -10,7 +10,7 @@ use candle_transformers::models::segment_anything::LayerNorm2d; use crate::{ models::{ - common::{AttentionNobias, MLPNoBias, eager_attention_forward}, + common::{GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward}, deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config}, }, position_embed::rope::RoPE, @@ -227,41 +227,11 @@ impl Attention { } } -pub struct MLPBlock { - linear1: Linear, - linear2: Linear, - act: Activation, -} - -impl MLPBlock { - pub fn new( - vb: VarBuilder, - embedding_dim: usize, - mlp_dim: usize, - act: Activation, - ) -> Result { - let linear1 = linear(embedding_dim, mlp_dim, vb.pp("lin1"))?; - let linear2 = linear(mlp_dim, embedding_dim, vb.pp("lin2"))?; - Ok(Self { - linear1, - linear2, - act, - }) - } - pub fn forward(&self, xs: &Tensor) -> Result { - let xs = xs - .apply(&self.linear1)? - .apply(&self.act)? - .apply(&self.linear2)?; - Ok(xs) - } -} - pub struct Block { norm1: LayerNorm, attn: Attention, norm2: LayerNorm, - mlp: MLPBlock, + mlp: TwoLinearMLP, window_size: usize, } @@ -300,7 +270,7 @@ impl Block { )?; let norm2 = layer_norm(dim, ln_config, vb.pp("norm2"))?; let mlp_dim = (dim as f32 * mlp_ratio) as usize; - let mlp = MLPBlock::new(vb.pp("mlp"), dim, mlp_dim, act)?; + let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, act, true, "lin1", "lin2")?; Ok(Self { norm1, attn, @@ -924,9 +894,9 @@ pub struct DeepseekV2MoE { // ep_size: usize, // experts_per_rank: usize, // ep_rank: usize, - experts: Vec, + experts: Vec, gate: MoEGate, - shared_experts: MLPNoBias, + shared_experts: GateUpDownMLP, } impl DeepseekV2MoE { @@ -937,20 +907,22 @@ impl DeepseekV2MoE { let mut experts = Vec::new(); let vb_experts = vb.pp("experts"); for i in 0..config.n_routed_experts { - let mlp = MLPNoBias::new( + let mlp = GateUpDownMLP::new( vb_experts.pp(i), config.hidden_size, config.moe_intermediate_size, Activation::Silu, + false, )?; experts.push(mlp); } let gate = MoEGate::new(vb.pp("gate"), config)?; - let shared_experts = MLPNoBias::new( + let shared_experts = GateUpDownMLP::new( vb.pp("shared_experts"), config.hidden_size, config.moe_intermediate_size * config.n_shared_experts, Activation::Silu, + false, )?; Ok(Self { // num_experts_per_tok: config.num_experts_per_tok, @@ -1014,7 +986,7 @@ impl Module for DeepseekV2MoE { pub enum DeepseekV2Proj { MOE(DeepseekV2MoE), - MLP(MLPNoBias), + MLP(GateUpDownMLP), } impl DeepseekV2Proj { @@ -1033,7 +1005,7 @@ impl DeepseekV2Proj { } pub struct DeepseekV2DecoderLayer { - self_attn: AttentionNobias, + self_attn: NaiveAttention, mlp: DeepseekV2Proj, input_layernorm: RmsNorm, post_attention_layernorm: RmsNorm, @@ -1041,22 +1013,24 @@ pub struct DeepseekV2DecoderLayer { impl DeepseekV2DecoderLayer { pub fn new(vb: VarBuilder, config: &DeepseekV2Config, layer_id: usize) -> Result { - let self_attn = AttentionNobias::new( + let self_attn = NaiveAttention::new( vb.pp("self_attn"), config.hidden_size, config.num_attention_heads, config.num_key_value_heads, + false, )?; let mlp = if layer_id >= config.first_k_dense_replace && layer_id.is_multiple_of(config.moe_layer_freq) { DeepseekV2Proj::MOE(DeepseekV2MoE::new(vb.pp("mlp"), config)?) } else { - DeepseekV2Proj::MLP(MLPNoBias::new( + DeepseekV2Proj::MLP(GateUpDownMLP::new( vb.pp("mlp"), config.hidden_size, config.intermediate_size, Activation::Silu, + false, )?) }; let input_layernorm = rms_norm( diff --git a/src/models/hunyuan_ocr/config.rs b/src/models/hunyuan_ocr/config.rs new file mode 100644 index 0000000..7cc172c --- /dev/null +++ b/src/models/hunyuan_ocr/config.rs @@ -0,0 +1,105 @@ +use candle_nn::Activation; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct HunYuanVLConfig { + pub attention_bias: bool, + pub attention_dropout: f64, + pub attention_head_dim: usize, + pub bos_token_id: u32, + pub eod_token_id: u32, + pub eos_token_id: u32, + pub head_dim: usize, + pub hidden_act: Activation, + pub hidden_size: usize, + pub image_start_token_id: u32, + pub image_end_token_id: u32, + pub image_token_id: u32, + pub image_newline_token_id: u32, + pub initializer_range: f64, + pub intermediate_size: usize, + pub max_position_embeddings: usize, + pub mlp_bias: bool, + pub norm_type: String, + pub num_attention_heads: usize, + pub num_experts: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub org_vocab_size: usize, + pub pad_id: i32, + pub pad_token_id: i32, + pub pretraining_tp: i32, + pub rms_norm_eps: f64, + pub rope_scaling: HunYuanVLRopeScaling, + pub rope_theta: f64, + pub routed_scaling_factor: f64, + pub sep_token_id: u32, + pub text_end_id: u32, + pub text_start_id: u32, + pub tie_word_embeddings: bool, + pub dtype: String, + pub use_cache: bool, + pub use_qk_norm: bool, + pub use_cla: bool, + pub vision_config: HunYuanVLVisionConfig, + pub vocab_size: usize, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct HunYuanVLRopeScaling { + pub alpha: f64, + pub beta_fast: i32, + pub beta_slow: i32, + pub factor: f64, + pub mscale: f64, + pub mscale_all_dim: f64, + #[serde(rename = "type")] + pub type_field: String, + pub xdrope_section: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct HunYuanVLVisionConfig { + pub add_patchemb_bias: bool, + pub attention_dropout: f64, + pub cat_extra_token: i32, + pub hidden_act: Activation, + pub hidden_dropout: f64, + pub hidden_size: usize, + pub img_max_token_num: usize, + pub intermediate_size: usize, + pub interpolate_mode: String, + pub max_image_size: usize, + pub max_vit_seq_len: usize, + pub num_attention_heads: usize, + pub num_channels: usize, + pub num_hidden_layers: usize, + pub out_hidden_size: usize, + pub patch_size: usize, + pub rms_norm_eps: f64, + pub spatial_merge_size: usize, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct HunyuanOCRGenerationConfig { + pub bos_token_id: usize, + pub pad_token_id: usize, + pub do_sample: bool, + pub eos_token_id: Vec, + pub top_p: f32, + pub top_k: usize, + pub temperature: f32, + pub repetition_penalty: f32, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct HunyuanOCRPreprocessorConfig { + pub min_pixels: usize, + pub max_pixels: usize, + pub patch_size: usize, + pub resample: usize, + pub temporal_patch_size: usize, + pub merge_size: usize, + pub image_mean: Vec, + pub image_std: Vec, +} diff --git a/src/models/hunyuan_ocr/generate.rs b/src/models/hunyuan_ocr/generate.rs new file mode 100644 index 0000000..0f35df4 --- /dev/null +++ b/src/models/hunyuan_ocr/generate.rs @@ -0,0 +1,218 @@ +use aha_openai_dive::v1::resources::chat::{ + ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +}; +use anyhow::{Result, anyhow}; +use candle_core::{DType, Device, Tensor}; +use candle_nn::VarBuilder; +use rocket::async_stream::stream; +use rocket::futures::Stream; + +use crate::{ + chat_template::ChatTemplate, + models::{ + GenerateModel, + hunyuan_ocr::{ + config::{HunYuanVLConfig, HunyuanOCRGenerationConfig}, + model::HunyuanVLModel, + processor::HunyuanVLProcessor, + }, + }, + tokenizer::TokenizerModel, + utils::{ + build_completion_chunk_response, build_completion_response, find_type_files, get_device, + get_dtype, get_logit_processor, + }, +}; + +pub struct HunyuanOCRGenerateModel<'a> { + chat_template: ChatTemplate<'a>, + tokenizer: TokenizerModel, + pre_processor: HunyuanVLProcessor, + hunyuan_vl: HunyuanVLModel, + device: Device, + eos_token_id1: u32, + eos_token_id2: u32, + generation_config: HunyuanOCRGenerationConfig, + model_name: String, +} + +impl<'a> HunyuanOCRGenerateModel<'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: HunYuanVLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + let device = get_device(device); + let cfg_dtype = cfg.dtype.as_str(); + let dtype = get_dtype(dtype, cfg_dtype); + let pre_processor = HunyuanVLProcessor::new(path, &device, dtype)?; + let model_list = find_type_files(path, "safetensors")?; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? }; + let hunyuan_vl = HunyuanVLModel::new(vb, cfg.clone())?; + let generation_config_path = path.to_string() + "/generation_config.json"; + let generation_config: HunyuanOCRGenerationConfig = + serde_json::from_slice(&std::fs::read(generation_config_path)?)?; + Ok(Self { + chat_template, + tokenizer, + pre_processor, + hunyuan_vl, + device, + eos_token_id1: generation_config.eos_token_id[0] as u32, + eos_token_id2: generation_config.eos_token_id[1] as u32, + generation_config, + model_name: "hunyuan_ocr".to_string(), + }) + } +} + +impl<'a> GenerateModel for HunyuanOCRGenerateModel<'a> { + fn generate(&mut self, mes: ChatCompletionParameters) -> Result { + let temperature = match mes.temperature { + None => self.generation_config.temperature, + Some(tem) => tem, + }; + let top_p = match mes.top_p { + None => self.generation_config.top_p, + Some(top_p) => top_p, + }; + let top_k = self.generation_config.top_k; + let seed = match mes.seed { + None => 34562u64, + Some(s) => s as u64, + }; + let mut logit_processor = + get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed); + let mes_render = self.chat_template.apply_chat_template(&mes)?; + let data = self + .pre_processor + .process_info(&mes, &self.tokenizer, &mes_render)?; + let mut input_ids = data.input_ids; + let mut position_ids = Some(&data.position_ids); + let mut image_mask = Some(&data.image_mask); + let mut pixel_values = data.pixel_values; + let mut image_grid_thw = data.image_grid_thw; + let mut seq_len = input_ids.dim(1)?; + let mut seqlen_offset = 0; + let mut generate: Vec = Vec::new(); + let sample_len = mes.max_tokens.unwrap_or(1024); + for _ in 0..sample_len { + let logits = self.hunyuan_vl.forward( + &input_ids, + pixel_values.as_ref(), + image_grid_thw.as_ref(), + image_mask, + position_ids, + 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.eos_token_id1 || next_token == self.eos_token_id2 { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + position_ids = None; + image_mask = None; + pixel_values = None; + image_grid_thw = None; + } + let res = self.tokenizer.token_decode(generate)?; + self.hunyuan_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 temperature = match mes.temperature { + None => self.generation_config.temperature, + Some(tem) => tem, + }; + let top_p = match mes.top_p { + None => self.generation_config.top_p, + Some(top_p) => top_p, + }; + let top_k = self.generation_config.top_k; + let seed = match mes.seed { + None => 34562u64, + Some(s) => s as u64, + }; + let mut logit_processor = + get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed); + let mes_render = self.chat_template.apply_chat_template(&mes)?; + let data = self + .pre_processor + .process_info(&mes, &self.tokenizer, &mes_render)?; + + let mut seqlen_offset = 0; + let sample_len = mes.max_tokens.unwrap_or(1024); + let stream = stream! { + let mut error_tokens = Vec::new(); + let mut input_ids = data.input_ids; + let mut position_ids = Some(&data.position_ids); + let mut image_mask = Some(&data.image_mask); + let mut pixel_values = data.pixel_values; + let mut image_grid_thw = data.image_grid_thw; + let mut seq_len = input_ids.dim(1)?; + for _ in 0..sample_len { + let logits = self.hunyuan_vl.forward( + &input_ids, + pixel_values.as_ref(), + image_grid_thw.as_ref(), + image_mask, + position_ids, + 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)?; + position_ids = None; + image_mask = None; + 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.eos_token_id1 || next_token == self.eos_token_id2 { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + position_ids = None; + image_mask = None; + pixel_values = None; + image_grid_thw = None; + } + self.hunyuan_vl.clear_kv_cache(); + }; + Ok(Box::new(Box::pin(stream))) + } +} diff --git a/src/models/hunyuan_ocr/mod.rs b/src/models/hunyuan_ocr/mod.rs new file mode 100644 index 0000000..8b1baf7 --- /dev/null +++ b/src/models/hunyuan_ocr/mod.rs @@ -0,0 +1,4 @@ +pub mod config; +pub mod generate; +pub mod model; +pub mod processor; diff --git a/src/models/hunyuan_ocr/model.rs b/src/models/hunyuan_ocr/model.rs new file mode 100644 index 0000000..edc26c0 --- /dev/null +++ b/src/models/hunyuan_ocr/model.rs @@ -0,0 +1,621 @@ +use anyhow::{Result, anyhow}; +use candle_core::{D, IndexOp, Tensor}; +use candle_nn::{ + Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear, + linear_no_bias, rms_norm, +}; + +use crate::{ + models::{ + common::{ + GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv2d, + get_layer_norm, + }, + hunyuan_ocr::config::{HunYuanVLConfig, HunYuanVLVisionConfig}, + }, + position_embed::rope::{RoPE, apply_rotary_pos_emb, get_xd_cos_sin}, + utils::tensor_utils::{ + interpolate_bilinear, masked_scatter_dim0, prepare_causal_attention_mask, split_tensor, + }, +}; + +pub struct HunYuanVisionPatchEmbed { + patch_embedding: Conv2d, + // position_embedding: Embedding, + num_channels: usize, + patch_size: usize, + // num_positions: usize, + // position_edge: usize, + embed_dim: usize, + patch_pos_embed: Tensor, +} + +impl HunYuanVisionPatchEmbed { + pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result { + let patch_embedding = get_conv2d( + vb.pp("patch_embedding"), + config.num_channels, + config.hidden_size, + config.patch_size, + 0, + config.patch_size, + 1, + 1, + true, + )?; + let num_channels = config.num_channels; + let patch_size = config.patch_size; + let position_edge = config.max_image_size / patch_size; + let num_positions = (position_edge).pow(2) + 1; + let embed_dim = config.hidden_size; + let position_embedding = embedding(num_positions, embed_dim, vb.pp("position_embedding"))?; + let patch_pos_embed = position_embedding + .embeddings() + .i(1..)? + .reshape((1, position_edge, position_edge, embed_dim))? + .permute((0, 3, 1, 2))?; + Ok(Self { + patch_embedding, + // position_embedding, + num_channels, + patch_size, + // num_positions, + // position_edge, + embed_dim, + patch_pos_embed, + }) + } + + pub fn forward(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result { + let (num_patches, _) = pixel_values.dims2()?; + let pixel_values = pixel_values.reshape(( + num_patches, + self.num_channels, + self.patch_size, + self.patch_size, + ))?; + let patch_embeds = self.patch_embedding.forward(&pixel_values)?; + let patch_embeds = patch_embeds + .squeeze(D::Minus1)? + .squeeze(D::Minus1)? + .unsqueeze(0)?; + let mut patch_pos_embed_list = vec![]; + let img_num = grid_thw.dim(0)?; + for i in 0..img_num { + let grid_i = grid_thw.i(i)?; + let grid_h = grid_i.i(1)?.to_scalar::()? as usize; + let grid_w = grid_i.i(2)?.to_scalar::()? as usize; + let patch_pos_embed_ = + interpolate_bilinear(&self.patch_pos_embed, (grid_h, grid_w), Some(false))?; + let patch_pos_embed_ = patch_pos_embed_ + .reshape((self.embed_dim, ()))? + .transpose(0, 1)? + .unsqueeze(0)?; + patch_pos_embed_list.push(patch_pos_embed_); + } + let patch_pos_embed = Tensor::cat(&patch_pos_embed_list, 1)?; + let embedding = patch_embeds.add(&patch_pos_embed)?; + Ok(embedding) + } +} +pub struct HunYuanVisionBlock { + self_attn: NaiveAttention, + mlp: TwoLinearMLP, + input_layernorm: LayerNorm, + post_attention_layernorm: LayerNorm, +} + +impl HunYuanVisionBlock { + pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result { + let self_attn = NaiveAttention::new( + vb.pp("self_attn"), + config.hidden_size, + config.num_attention_heads, + config.num_attention_heads, + true, + )?; + let mlp = TwoLinearMLP::new( + vb.pp("mlp"), + config.hidden_size, + config.intermediate_size, + config.hidden_act, + true, + "dense_h_to_4h", + "dense_4h_to_h", + )?; + + let input_layernorm = get_layer_norm( + vb.pp("input_layernorm"), + config.rms_norm_eps, + config.hidden_size, + )?; + let post_attention_layernorm = get_layer_norm( + vb.pp("post_attention_layernorm"), + config.rms_norm_eps, + config.hidden_size, + )?; + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let residual = xs.clone(); + let xs = self.input_layernorm.forward(xs)?; + let xs = self.self_attn.forward(&xs, None, None, None, 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) + } +} + +pub struct HunYuanVisionPatchMerger { + proj_0: Conv2d, + proj_2: Conv2d, + mlp: Linear, + image_newline: Tensor, + image_begin: Tensor, + image_end: Tensor, + // image_sep: Tensor, + before_rms: RmsNorm, + after_rms: RmsNorm, +} + +impl HunYuanVisionPatchMerger { + pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result { + let proj_0 = get_conv2d( + vb.pp("proj.0"), + config.hidden_size, + config.hidden_size * 2, + config.spatial_merge_size, + 0, + config.spatial_merge_size, + 1, + 1, + true, + )?; + let proj_2 = get_conv2d( + vb.pp("proj.2"), + config.hidden_size * 2, + config.hidden_size * 4, + 1, + 0, + 1, + 1, + 1, + true, + )?; + let mlp = linear(config.hidden_size * 4, config.out_hidden_size, vb.pp("mlp"))?; + let image_newline = + vb.get_with_hints(config.hidden_size * 4, "image_newline", Init::Const(0.))?; + let image_begin = + vb.get_with_hints(config.out_hidden_size, "image_begin", Init::Const(0.))?; + let image_end = vb.get_with_hints(config.out_hidden_size, "image_end", Init::Const(0.))?; + // let image_sep = vb.get_with_hints(config.out_hidden_size, "image_sep", Init::Const(0.))?; + let before_rms = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("before_rms"))?; + let after_rms = rms_norm( + config.out_hidden_size, + config.rms_norm_eps, + vb.pp("after_rms"), + )?; + Ok(Self { + proj_0, + proj_2, + mlp, + image_newline, + image_begin, + image_end, + // image_sep, + before_rms, + after_rms, + }) + } + pub fn forward(&self, xs: &Tensor, size: (usize, usize)) -> Result { + let xs = self.before_rms.forward(xs)?; + let (h, w) = size; + let xs = xs.permute((0, 2, 1))?.reshape((xs.dim(0)?, (), h, w))?; + let xs = self.proj_0.forward(&xs)?.gelu()?; + let xs = self.proj_2.forward(&xs)?; + let (b, c, h, _) = xs.dims4()?; + let image_newline = self + .image_newline + .reshape((1, c, 1, 1))? + .broadcast_as((b, c, h, 1))? + .to_dtype(xs.dtype())?; + let xs = Tensor::cat(&[xs, image_newline], D::Minus1)?; + let xs = xs.reshape((b, c, ()))?.permute((0, 2, 1))?; + let xs = self.mlp.forward(&xs)?; + let begin = self + .image_begin + .reshape((1, 1, ()))? + .broadcast_as((b, 1, xs.dim(D::Minus1)?))? + .to_dtype(xs.dtype())?; + let end = self + .image_end + .reshape((1, 1, ()))? + .broadcast_as((b, 1, xs.dim(D::Minus1)?))? + .to_dtype(xs.dtype())?; + let xs = Tensor::cat(&[begin, xs, end], 1)?; + let xs = self.after_rms.forward(&xs)?; + Ok(xs) + } +} + +pub struct HunYuanVisionTransformer { + embeddings: HunYuanVisionPatchEmbed, + layers: Vec, + perceive: HunYuanVisionPatchMerger, +} + +impl HunYuanVisionTransformer { + pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result { + let embeddings = HunYuanVisionPatchEmbed::new(vb.pp("embeddings"), config)?; + let mut layers = vec![]; + let vb_layers = vb.pp("layers"); + for i in 0..config.num_hidden_layers { + let layer_i = HunYuanVisionBlock::new(vb_layers.pp(i), config)?; + layers.push(layer_i); + } + let perceive = HunYuanVisionPatchMerger::new(vb.pp("perceive"), config)?; + Ok(Self { + embeddings, + layers, + perceive, + }) + } + + pub fn forward(&self, xs: &Tensor, grid_thw: &Tensor) -> Result { + let mut hidden_states = self.embeddings.forward(xs, grid_thw)?; + for layer in &self.layers { + hidden_states = layer.forward(&hidden_states)?; + } + let mut cu_seqlens = vec![]; + for i in 0..grid_thw.dim(0)? { + let [_, h, w] = grid_thw.i(i)?.to_vec1::()?[..] else { + return Err(anyhow!(format!("grid_thw Expected exactly 3 elements"))); + }; + cu_seqlens.push((h * w) as usize); + } + let split_items = split_tensor(&hidden_states, &cu_seqlens, 1)?; + let mut processed_item = vec![]; + for i in 0..grid_thw.dim(0)? { + let [_, h, w] = grid_thw.i(i)?.to_vec1::()?[..] else { + return Err(anyhow!(format!("grid_thw Expected exactly 3 elements"))); + }; + let processed = self + .perceive + .forward(&split_items[i], (h as usize, w as usize))?; + processed_item.push(processed); + } + let xs = Tensor::cat(&processed_item, 1)?; + Ok(xs) + } +} + +pub struct HunYuanVLAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + query_layernorm: RmsNorm, + key_layernorm: RmsNorm, + num_attention_heads: usize, + num_key_value_heads: usize, + num_kv_groups: usize, + head_dim: usize, + scaling: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl HunYuanVLAttention { + pub fn new( + vb: VarBuilder, + hidden_size: usize, + head_dim: usize, + num_attention_heads: usize, + num_key_value_heads: usize, + attention_bias: bool, + rms_norm_eps: f64, + ) -> Result { + let num_kv_groups = num_attention_heads / num_key_value_heads; + let scaling = 1f64 / f64::sqrt(head_dim as f64); + let (q_proj, k_proj, v_proj, o_proj) = if attention_bias { + let q_proj = linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?; + (q_proj, k_proj, v_proj, o_proj) + } else { + let q_proj = + linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; + let k_proj = + linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?; + let v_proj = + linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?; + let o_proj = + linear_no_bias(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?; + (q_proj, k_proj, v_proj, o_proj) + }; + let query_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("query_layernorm"))?; + let key_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("key_layernorm"))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + query_layernorm, + key_layernorm, + num_attention_heads, + num_key_value_heads, + num_kv_groups, + head_dim, + scaling, + kv_cache: None, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let (b_sz, q_len, _) = xs.dims3()?; + let query_states = self + .q_proj + .forward(xs)? + .reshape((b_sz, q_len, self.num_attention_heads, self.head_dim))? + .transpose(1, 2)?; + + let key_states = self + .k_proj + .forward(xs)? + .reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = self.v_proj.forward(xs)?; + let value_states = value_states + .reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))? + .transpose(1, 2)?; + let (query_states, key_states) = + apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?; + let query_states = self.query_layernorm.forward(&query_states)?; + let key_states = self.key_layernorm.forward(&key_states)?; + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + let attn_output = eager_attention_forward( + &query_states, + &key_states, + &value_states, + Some(self.num_kv_groups), + attention_mask, + self.scaling, + )?; + let attn_output = + attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?; + let attn_output = attn_output.apply(&self.o_proj)?; + Ok(attn_output) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache = None + } +} + +pub struct HunYuanVLDecoderLayer { + self_attn: HunYuanVLAttention, + mlp: GateUpDownMLP, + input_layernorm: RmsNorm, + post_attention_layernorm: RmsNorm, +} + +impl HunYuanVLDecoderLayer { + pub fn new(config: &HunYuanVLConfig, vb: VarBuilder) -> Result { + let self_attn = HunYuanVLAttention::new( + vb.pp("self_attn"), + config.hidden_size, + config.head_dim, + config.num_attention_heads, + config.num_key_value_heads, + config.attention_bias, + config.rms_norm_eps, + )?; + let mlp = GateUpDownMLP::new( + vb.pp("mlp"), + config.hidden_size, + config.intermediate_size, + config.hidden_act, + false, + )?; + 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(&xs, cos, sin, attention_mask)?; + let xs = residual.add(&xs)?; + let residual = xs.clone(); + let xs = self.post_attention_layernorm.forward(&xs)?; + let xs = self.mlp.forward(&xs)?; + let xs = residual.add(&xs)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +pub struct HunYuanVLTextModel { + embed_tokens: Embedding, + layers: Vec, + norm: RmsNorm, + rope: RoPE, + xdrope_section: Vec, +} + +impl HunYuanVLTextModel { + pub fn new(vb: VarBuilder, config: &HunYuanVLConfig) -> Result { + let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?; + let mut layers = vec![]; + let vb_layers = vb.pp("layers"); + for i in 0..config.num_hidden_layers { + let layer = HunYuanVLDecoderLayer::new(config, vb_layers.pp(i))?; + layers.push(layer); + } + let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?; + let base = config.rope_theta + * config + .rope_scaling + .alpha + .powf(config.head_dim as f64 / (config.head_dim - 2) as f64); + let rope = RoPE::new(config.head_dim, base as f32, vb.device())?; + let xdrope_section = config.rope_scaling.xdrope_section.clone(); + Ok(Self { + embed_tokens, + layers, + norm, + rope, + xdrope_section, + }) + } + + pub fn forward( + &mut self, + inputs_embeds: &Tensor, + position_ids: Option<&Tensor>, + seqlen_offset: usize, + ) -> 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)?, + // }; + let attention_mask: Option<&Tensor> = { + if seq_len <= 1 { + None + } else { + Some(&prepare_causal_attention_mask( + b_size, + seq_len, + 0, + inputs_embeds.device(), + )?) + } + }; + + let (cos, sin) = self + .rope + .forward(seqlen_offset, seq_len, inputs_embeds.device())?; + let mut xs = inputs_embeds.clone(); + for (i, layer) in self.layers.iter_mut().enumerate() { + if i == 0 + && let Some(position_ids) = position_ids + { + let (cos, sin) = + get_xd_cos_sin(&cos, &sin, position_ids, self.xdrope_section.clone())?; + xs = layer.forward(&xs, &cos, &sin, attention_mask)?; + } else { + xs = layer.forward(&xs, &cos, &sin, attention_mask)?; + } + } + let xs = self.norm.forward(&xs)?; + Ok(xs) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } +} + +pub struct HunyuanVLModel { + // config: HunYuanVLConfig, + vit: HunYuanVisionTransformer, + model: HunYuanVLTextModel, + lm_head: Linear, +} + +impl HunyuanVLModel { + pub fn new(vb: VarBuilder, config: HunYuanVLConfig) -> Result { + let vit = HunYuanVisionTransformer::new(vb.pp("vit"), &config.vision_config)?; + let model = HunYuanVLTextModel::new(vb.pp("model"), &config)?; + let lm_head = Linear::new(model.embed_tokens.embeddings().clone(), None); + Ok(Self { + // config, + vit, + model, + lm_head, + }) + } + pub fn forward( + &mut self, + input_ids: &Tensor, + pixel_values: Option<&Tensor>, + image_grid_thw: Option<&Tensor>, + image_mask: Option<&Tensor>, + position_ids: 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(grid_thw) = image_grid_thw + && let Some(image_mask) = image_mask + { + let image_embeds = self.vit.forward(pixel_values, grid_thw)?.squeeze(0)?; + inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embeds, image_mask)?; + } + let outputs = self + .model + .forward(&inputs_embeds, position_ids, seqlen_offset)?; + 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/hunyuan_ocr/processor.rs b/src/models/hunyuan_ocr/processor.rs new file mode 100644 index 0000000..0865fb5 --- /dev/null +++ b/src/models/hunyuan_ocr/processor.rs @@ -0,0 +1,236 @@ +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::hunyuan_ocr::config::HunyuanOCRPreprocessorConfig, + tokenizer::TokenizerModel, + utils::{ + img_utils::{extract_images, img_smart_resize, img_transform}, + tensor_utils::{get_eq_indices, get_equal_mask}, + }, +}; + +pub struct HunyuanVLProcessor { + image_token_id: u32, + image_token: String, + placeholder_token: String, + process_cfg: HunyuanOCRPreprocessorConfig, + device: Device, + dtype: DType, +} + +impl HunyuanVLProcessor { + pub fn new(path: &str, device: &Device, dtype: DType) -> Result { + let path = path.to_string(); + assert!( + std::path::Path::new(&path).exists(), + "model path file not exists" + ); + let process_cfg_file = path.clone() + "/preprocessor_config.json"; + assert!( + std::path::Path::new(&process_cfg_file).exists(), + "preprocessor_config.json not exists in model path" + ); + let process_cfg: HunyuanOCRPreprocessorConfig = + serde_json::from_slice(&std::fs::read(process_cfg_file)?)?; + let image_token_id = 120120u32; + let image_token = "<｜hy_place▁holder▁no▁102｜>".to_string(); + let placeholder_token = "<｜hy_place▁holder▁no▁799｜>".to_string(); + // let pad_id = 120002u32; + Ok(Self { + image_token_id, + image_token, + placeholder_token, + process_cfg, + device: device.clone(), + dtype, + }) + } + + 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 as u32, + self.process_cfg.max_pixels as u32, + )?; + 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, + channel, + grid_h / self.process_cfg.merge_size, + self.process_cfg.merge_size, + self.process_cfg.patch_size, + grid_w / self.process_cfg.merge_size, + self.process_cfg.merge_size, + self.process_cfg.patch_size, + ]); + let img_tensor = img_tensor.reshape(shape)?; + // shape to // grid_t, + // grid_h / merge_size, + // merge_size, + // grid_w / merge_size, + // merge_size, + // channel, + // patch_size, + // patch_size, + let img_tensor = img_tensor.permute(vec![0, 2, 3, 5, 6, 1, 4, 7])?; + 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, + tokenizer: &TokenizerModel, + text: &str, + ) -> Result { + 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 mut image_tokens_cumsum = vec![0]; + let mut text = text.to_string(); + if !imgs.is_empty() + && let Some(grid_thw) = image_grid_thw.as_ref() + { + let mut index = 0; + while text.contains(&self.image_token) { + let grid_i = grid_thw.i(index)?; + let grid_h = grid_i.i(1)?.to_scalar::()?; + let grid_w = grid_i.i(2)?.to_scalar::()?; + let patch_h = grid_h / self.process_cfg.merge_size as u32; + let patch_w = grid_w / self.process_cfg.merge_size as u32; + let num_image_tokens = patch_h * (patch_w + 1) + 2; + let num_id = image_tokens_cumsum[image_tokens_cumsum.len() - 1] + num_image_tokens; + image_tokens_cumsum.push(num_id); + let replace = self.placeholder_token.repeat(num_image_tokens as usize); + text = text.replacen(&self.image_token, &replace, 1); + index += 1; + } + } + + text = text.replace(&self.placeholder_token, &self.image_token); + let input_ids = tokenizer.text_encode(text, &self.device)?; + let seq_len = input_ids.dim(1)?; + let position_ids = Tensor::arange(0, seq_len as u32, &self.device)?; + let mut position_ids_w = Tensor::arange(0, seq_len as u32, &self.device)?; + let mut position_ids_h = Tensor::arange(0, seq_len as u32, &self.device)?; + let mut position_ids_t = Tensor::arange(0, seq_len as u32, &self.device)?; + if !imgs.is_empty() + && let Some(grid_thw) = image_grid_thw.as_ref() + { + let image_token_pos_indices = get_eq_indices(&input_ids.i(0)?, self.image_token_id)?; + for i in 0..grid_thw.dim(0)? { + let grid_i = grid_thw.i(i)?; + let grid_h = grid_i.i(1)?.to_scalar::()?; + let grid_w = grid_i.i(2)?.to_scalar::()?; + let patch_h = grid_h / self.process_cfg.merge_size as u32; + let patch_w = grid_w / self.process_cfg.merge_size as u32; + let start_pos = image_token_pos_indices + .i(image_tokens_cumsum[i] as usize)? + .to_scalar::()? as usize + + 1; + let replace_num = ((patch_w + 1) * patch_h) as usize; + let pos_w: Vec = (0..patch_h).flat_map(|_| 0u32..patch_w + 1).collect(); + position_ids_w = position_ids_w.slice_assign( + &[start_pos..start_pos + replace_num], + &Tensor::new(pos_w, &self.device)?, + )?; + let pos_h: Vec = (0..patch_h) + .flat_map(|h| vec![h; (patch_w + 1) as usize]) + .collect(); + position_ids_h = position_ids_h.slice_assign( + &[start_pos..start_pos + replace_num], + &Tensor::new(pos_h, &self.device)?, + )?; + position_ids_t = position_ids_t.slice_assign( + &[start_pos..start_pos + replace_num], + &Tensor::new(vec![0u32; replace_num], &self.device)?, + )?; + } + } + let position_ids = Tensor::stack( + &[position_ids, position_ids_h, position_ids_w, position_ids_t], + 0, + )? + .unsqueeze(0)?; + let image_mask = get_equal_mask(&input_ids, self.image_token_id)?; + let data = HunyuanData { + input_ids, + position_ids, + image_mask, + pixel_values, + image_grid_thw, + }; + Ok(data) + } +} + +pub struct HunyuanData { + pub input_ids: Tensor, + pub position_ids: Tensor, + pub image_mask: Tensor, + pub pixel_values: Option, + pub image_grid_thw: Option, +} diff --git a/src/models/minicpm4/generate.rs b/src/models/minicpm4/generate.rs index c5edabd..b8af4c3 100644 --- a/src/models/minicpm4/generate.rs +++ b/src/models/minicpm4/generate.rs @@ -23,6 +23,7 @@ pub struct MiniCPMGenerateModel<'a> { device: Device, endoftext_id: u32, im_end_id: u32, + model_name: String, } impl<'a> MiniCPMGenerateModel<'a> { @@ -47,6 +48,7 @@ impl<'a> MiniCPMGenerateModel<'a> { device: device.clone(), endoftext_id, im_end_id, + model_name: "minicpm4".to_string(), }) } } @@ -78,7 +80,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> { } let res = self.tokenizer.token_decode(generate)?; self.minicpm.clear_kv_cache(); - let response = build_completion_response(res, "minicpm"); + let response = build_completion_response(res, &self.model_name); Ok(response) } fn generate_stream( @@ -128,7 +130,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> { continue; } error_tokens.clear(); - let chunk = build_completion_chunk_response(decoded_token, "minicpm", None, None); + let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); yield Ok(chunk); if next_token == self.endoftext_id || next_token == self.im_end_id { break; diff --git a/src/models/minicpm4/model.rs b/src/models/minicpm4/model.rs index 3696e55..dd51e46 100644 --- a/src/models/minicpm4/model.rs +++ b/src/models/minicpm4/model.rs @@ -4,7 +4,7 @@ use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_n use crate::{ models::{ - common::{AttentionNobias, MLPNoBias}, + common::{GateUpDownMLP, NaiveAttention}, minicpm4::config::MiniCPM4Config, }, position_embed::rope::compute_default_rope_parameters, @@ -93,8 +93,8 @@ impl MiniCPMLongRoPE { } pub struct MiniCPMDecoderLayer { - self_attn: AttentionNobias, - mlp: MLPNoBias, + self_attn: NaiveAttention, + mlp: GateUpDownMLP, input_layernorm: RmsNorm, post_attention_layernorm: RmsNorm, scale_depth: f32, @@ -103,17 +103,19 @@ pub struct MiniCPMDecoderLayer { impl MiniCPMDecoderLayer { pub fn new(vb: VarBuilder, cfg: &MiniCPM4Config) -> Result { - let self_attn = AttentionNobias::new( + let self_attn = NaiveAttention::new( vb.pp("self_attn"), cfg.hidden_size, cfg.num_attention_heads, cfg.num_key_value_heads, + false, )?; - let mlp = MLPNoBias::new( + let mlp = GateUpDownMLP::new( vb.pp("mlp"), cfg.hidden_size, cfg.intermediate_size, cfg.hidden_act, + false, )?; let input_layernorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; @@ -143,7 +145,7 @@ impl MiniCPMDecoderLayer { let xs = self.input_layernorm.forward(xs)?; let xs = self .self_attn - .forward(&xs, cos, sin, attention_mask, true)?; + .forward(&xs, Some(cos), Some(sin), attention_mask, true)?; let xs = (residual + xs.affine( self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(), diff --git a/src/models/mod.rs b/src/models/mod.rs index dabf2f0..bc50824 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,5 +1,6 @@ pub mod common; pub mod deepseek_ocr; +pub mod hunyuan_ocr; pub mod minicpm4; pub mod qwen2_5vl; pub mod qwen3vl; @@ -12,7 +13,8 @@ use anyhow::Result; use rocket::futures::Stream; use crate::models::{ - deepseek_ocr::generate::DeepseekOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel, + deepseek_ocr::generate::DeepseekOCRGenerateModel, + hunyuan_ocr::generate::HunyuanOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, qwen3vl::generate::Qwen3VLGenerateModel, }; @@ -34,6 +36,8 @@ pub enum WhichModel { Qwen3vl32B, #[value(name = "deepseek-ocr")] DeepSeekOCR, + #[value(name = "hunyuan-ocr")] + HunyuanOCR, } pub trait GenerateModel { @@ -56,6 +60,7 @@ pub enum ModelInstance<'a> { Qwen2_5VL(Qwen2_5VLGenerateModel<'a>), Qwen3VL(Qwen3VLGenerateModel<'a>), DeepSeekOCR(DeepseekOCRGenerateModel), + HunyuanOCR(HunyuanOCRGenerateModel<'a>), } impl<'a> GenerateModel for ModelInstance<'a> { @@ -65,6 +70,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::Qwen2_5VL(model) => model.generate(mes), ModelInstance::Qwen3VL(model) => model.generate(mes), ModelInstance::DeepSeekOCR(model) => model.generate(mes), + ModelInstance::HunyuanOCR(model) => model.generate(mes), } } @@ -84,6 +90,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes), ModelInstance::Qwen3VL(model) => model.generate_stream(mes), ModelInstance::DeepSeekOCR(model) => model.generate_stream(mes), + ModelInstance::HunyuanOCR(model) => model.generate_stream(mes), } } } @@ -122,6 +129,10 @@ pub fn load_model(model_type: WhichModel, path: &str) -> Result { + let model = HunyuanOCRGenerateModel::init(path, None, None)?; + ModelInstance::HunyuanOCR(model) + } }; Ok(model) } diff --git a/src/models/qwen2_5vl/generate.rs b/src/models/qwen2_5vl/generate.rs index 7468abe..42a5713 100644 --- a/src/models/qwen2_5vl/generate.rs +++ b/src/models/qwen2_5vl/generate.rs @@ -30,6 +30,7 @@ pub struct Qwen2_5VLGenerateModel<'a> { device: Device, endoftext_id: u32, im_end_id: u32, + model_name: String, } impl<'a> Qwen2_5VLGenerateModel<'a> { @@ -57,6 +58,7 @@ impl<'a> Qwen2_5VLGenerateModel<'a> { device: device.clone(), endoftext_id, im_end_id, + model_name: "qwen2.5vl".to_string(), }) } } @@ -119,7 +121,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> { } let res = self.tokenizer.token_decode(generate)?; self.qwen2_5_vl.clear_kv_cache(); - let response = build_completion_response(res, "qwen2.5vl"); + let response = build_completion_response(res, &self.model_name); Ok(response) } @@ -202,7 +204,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> { continue; } error_tokens.clear(); - let chunk = build_completion_chunk_response(decoded_token, "qwen2.5vl", None, None); + let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); yield Ok(chunk); if next_token == self.endoftext_id || next_token == self.im_end_id { break; diff --git a/src/models/qwen3vl/generate.rs b/src/models/qwen3vl/generate.rs index e8669b3..775b822 100644 --- a/src/models/qwen3vl/generate.rs +++ b/src/models/qwen3vl/generate.rs @@ -33,6 +33,7 @@ pub struct Qwen3VLGenerateModel<'a> { eos_token_id1: u32, eos_token_id2: u32, generation_config: Qwen3VLGenerationConfig, + model_name: String, } impl<'a> Qwen3VLGenerateModel<'a> { @@ -60,6 +61,7 @@ impl<'a> Qwen3VLGenerateModel<'a> { eos_token_id1: generation_config.eos_token_id[0] as u32, eos_token_id2: generation_config.eos_token_id[1] as u32, generation_config, + model_name: "qwen3vl".to_string(), }) } } @@ -120,7 +122,7 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> { } let res = self.tokenizer.token_decode(generate)?; self.qwen3_vl.clear_kv_cache(); - let response = build_completion_response(res, "qwen3vl"); + let response = build_completion_response(res, &self.model_name); Ok(response) } @@ -201,7 +203,7 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> { continue; } error_tokens.clear(); - let chunk = build_completion_chunk_response(decoded_token, "qwen3vl", None, None); + let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); yield Ok(chunk); if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 { break; diff --git a/src/models/qwen3vl/model.rs b/src/models/qwen3vl/model.rs index a9736da..6cac190 100644 --- a/src/models/qwen3vl/model.rs +++ b/src/models/qwen3vl/model.rs @@ -7,7 +7,7 @@ use candle_nn::{ use crate::{ models::{ - common::{MLPNoBias, eager_attention_forward}, + common::{GateUpDownMLP, eager_attention_forward}, qwen3vl::config::{Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig}, }, position_embed::rope::{ @@ -664,7 +664,7 @@ impl Qwen3VLTextAttention { pub struct Qwen3VLTextDecoderLayer { self_attn: Qwen3VLTextAttention, - mlp: MLPNoBias, + mlp: GateUpDownMLP, input_layernorm: RmsNorm, post_attention_layernorm: RmsNorm, } @@ -672,11 +672,12 @@ pub struct Qwen3VLTextDecoderLayer { impl Qwen3VLTextDecoderLayer { pub fn new(config: Qwen3VLTextConfig, vb: VarBuilder) -> Result { let self_attn = Qwen3VLTextAttention::new(config.clone(), vb.pp("self_attn"))?; - let mlp = MLPNoBias::new( + let mlp = GateUpDownMLP::new( vb.pp("mlp"), config.hidden_size, config.intermediate_size, config.hidden_act, + false, )?; let input_layernorm = rms_norm( config.hidden_size, diff --git a/src/models/qwen3vl/processor.rs b/src/models/qwen3vl/processor.rs index 04dfadc..3125c6c 100644 --- a/src/models/qwen3vl/processor.rs +++ b/src/models/qwen3vl/processor.rs @@ -11,7 +11,11 @@ use num::integer::lcm; use crate::{ models::qwen3vl::config::PreprocessorConfig, - utils::{ceil_by_factor, floor_by_factor, img_utils::get_image, round_by_factor}, + utils::{ + ceil_by_factor, floor_by_factor, + img_utils::{get_image, img_smart_resize, img_transform}, + round_by_factor, + }, }; #[derive(Clone)] @@ -136,22 +140,9 @@ impl Qwen3VLProcessor { (self.img_process_cfg.patch_size * self.img_process_cfg.merge_size) as u32, self.img_process_cfg.size.shortest_edge as u32, self.img_process_cfg.size.longest_edge as u32, - None, )?; let img = img.resize_exact(resize_w, resize_h, image::imageops::FilterType::CatmullRom); - let img_vec = img.to_rgb8().into_raw(); - // (h, w, c) => (c, h, w) - let img_tensor = Tensor::from_slice( - &img_vec, - (resize_h as usize, resize_w as usize, 3), - &self.device, - )? - .permute((2, 0, 1))? - .to_dtype(self.dtype)?; - // 0-255 rescale to 0-1 - let img_tensor = img_tensor.affine(1.0 / 255.0, 0.)?; - // normalize - let img_tensor = img_tensor.broadcast_sub(img_mean)?.broadcast_div(img_std)?; + 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) @@ -426,40 +417,6 @@ impl Qwen3VLProcessor { } } -pub fn img_smart_resize( - img_h: u32, - img_w: u32, - factor: u32, - min_pixels: u32, - max_pixels: u32, - video_ratio: Option, -) -> Result<(u32, u32)> { - if std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w) > 200 { - return Err(anyhow!(format!( - "absolute aspect ratio mush be smaller than {}, got {}", - 200, - std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w) - ))); - } - let mut image_factor = factor; - if let Some(ratio) = video_ratio { - image_factor = lcm(image_factor, ratio); - } - let mut h_bar = std::cmp::max(image_factor, round_by_factor(img_h, image_factor)); - let mut w_bar = std::cmp::max(image_factor, round_by_factor(img_w, image_factor)); - - if h_bar * w_bar > max_pixels { - let beta = ((img_h * img_w) as f32 / max_pixels as f32).sqrt(); - h_bar = floor_by_factor(img_h as f32 / beta, image_factor); - w_bar = floor_by_factor(img_w as f32 / beta, image_factor); - } else if h_bar * w_bar < min_pixels { - let beta = (min_pixels as f32 / (img_h * img_w) as f32).sqrt(); - h_bar = ceil_by_factor(img_h as f32 * beta, image_factor); - w_bar = ceil_by_factor(img_w as f32 * beta, image_factor); - } - Ok((h_bar, w_bar)) -} - pub fn video_smart_resize( num_frames: u32, height: u32, diff --git a/src/models/voxcpm/minicpm4.rs b/src/models/voxcpm/minicpm4.rs index e4789fd..8bcc490 100644 --- a/src/models/voxcpm/minicpm4.rs +++ b/src/models/voxcpm/minicpm4.rs @@ -4,7 +4,7 @@ use candle_nn::{Embedding, Module, RmsNorm, VarBuilder, embedding, rms_norm}; use crate::{ models::{ - common::{AttentionNobias, MLPNoBias}, + common::{GateUpDownMLP, NaiveAttention}, voxcpm::config::VoxMiniCPM4Config, }, position_embed::rope::compute_default_rope_parameters, @@ -101,8 +101,8 @@ impl MiniCPMLongRoPE { } pub struct MiniCPMDecoderLayer { - self_attn: AttentionNobias, - mlp: MLPNoBias, + self_attn: NaiveAttention, + mlp: GateUpDownMLP, input_layernorm: RmsNorm, post_attention_layernorm: RmsNorm, scale_depth: f32, @@ -112,17 +112,19 @@ pub struct MiniCPMDecoderLayer { impl MiniCPMDecoderLayer { pub fn new(vb: VarBuilder, cfg: &VoxMiniCPM4Config) -> Result { - let self_attn = AttentionNobias::new( + let self_attn = NaiveAttention::new( vb.pp("self_attn"), cfg.hidden_size, cfg.num_attention_heads, cfg.num_key_value_heads, + false, )?; - let mlp = MLPNoBias::new( + let mlp = GateUpDownMLP::new( vb.pp("mlp"), cfg.hidden_size, cfg.intermediate_size, candle_nn::Activation::Silu, + false, )?; let input_layernorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?; @@ -153,7 +155,7 @@ impl MiniCPMDecoderLayer { let xs = self.input_layernorm.forward(xs)?; let xs = self .self_attn - .forward(&xs, cos, sin, attention_mask, true)?; + .forward(&xs, Some(cos), Some(sin), attention_mask, true)?; let xs = if self.use_mup { (residual + xs.affine( diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index ff8c35f..2189059 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -2,6 +2,8 @@ use anyhow::Result; use candle_core::{D, DType, Device, IndexOp, Tensor}; use candle_transformers::models::deepseek2::SplitOp; +use crate::utils::tensor_utils::{index_select_2d, split_tensor}; + pub fn compute_default_rope_parameters(dim: usize, base: f32) -> Vec { let inv_freq: Vec = (0..dim) .step_by(2) @@ -303,3 +305,45 @@ impl RoPE { Ok((cos, sin)) } } + +pub fn get_xd_cos_sin( + cos: &Tensor, + sin: &Tensor, + position_ids: &Tensor, + xdrope_section: Vec, +) -> Result<(Tensor, Tensor)> { + let x_dim = xdrope_section.len(); + // position_ids: (bs, 4, seq_len) + let mut cos_vec = vec![]; + let mut sin_vec = vec![]; + let bs = position_ids.dim(0)?; + for i in 0..bs { + let pos_i = position_ids.i(i)?; + let cos_i = index_select_2d(cos, &pos_i)?; + let sin_i = index_select_2d(sin, &pos_i)?; + cos_vec.push(cos_i); + sin_vec.push(sin_i); + } + // (bs, 4, seq_len, dim) -> (bs, seq_len, 4, dim) + let cos = Tensor::stack(&cos_vec, 0)? + .permute((0, 2, 1, 3))? + .contiguous()?; + let sin = Tensor::stack(&sin_vec, 0)? + .permute((0, 2, 1, 3))? + .contiguous()?; + let xdrope_section: Vec = xdrope_section.iter().map(|&i| i * 2).collect(); + let cos_select: Vec = split_tensor(&cos, &xdrope_section, D::Minus1)? + .iter() + .enumerate() + .map(|(i, m)| m.i((.., .., i % x_dim)).unwrap()) + .collect(); + let sin_select: Vec = split_tensor(&sin, &xdrope_section, D::Minus1)? + .iter() + .enumerate() + .map(|(i, m)| m.i((.., .., i % x_dim)).unwrap()) + .collect(); + + let cos = Tensor::cat(&cos_select, D::Minus1)?; + let sin = Tensor::cat(&sin_select, D::Minus1)?; + Ok((cos, sin)) +} diff --git a/src/tokenizer/mod.rs b/src/tokenizer/mod.rs index 4e0faf8..4687455 100644 --- a/src/tokenizer/mod.rs +++ b/src/tokenizer/mod.rs @@ -3,7 +3,7 @@ use candle_core::{Device, Tensor}; use tokenizers::Tokenizer; pub struct TokenizerModel { - tokenizer: Tokenizer, + pub tokenizer: Tokenizer, } impl TokenizerModel { diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs index 082faf9..cf75cc1 100644 --- a/src/utils/img_utils.rs +++ b/src/utils/img_utils.rs @@ -9,6 +9,8 @@ use base64::{Engine, engine::general_purpose}; use candle_core::{DType, Device, Tensor}; use image::{DynamicImage, ImageBuffer, ImageReader, Rgb, RgbImage, imageops}; +use crate::utils::{ceil_by_factor, floor_by_factor, round_by_factor}; + pub fn load_image_from_url(url: &str) -> Result { tokio::task::block_in_place(|| { let response = reqwest::blocking::get(url) @@ -237,3 +239,39 @@ pub fn img_transform( .to_dtype(dtype)?; Ok(img_tensor) } + +pub fn img_smart_resize( + img_h: u32, + img_w: u32, + factor: u32, + min_pixels: u32, + max_pixels: u32, +) -> Result<(u32, u32)> { + if std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w) > 200 { + return Err(anyhow!(format!( + "absolute aspect ratio mush be smaller than {}, got {}", + 200, + std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w) + ))); + } + let image_factor = factor; + let mut h_bar = std::cmp::max(image_factor, round_by_factor(img_h, image_factor)); + let mut w_bar = std::cmp::max(image_factor, round_by_factor(img_w, image_factor)); + + if h_bar * w_bar > max_pixels { + let beta = ((img_h * img_w) as f32 / max_pixels as f32).sqrt(); + h_bar = std::cmp::max( + image_factor, + floor_by_factor(img_h as f32 / beta, image_factor), + ); + w_bar = std::cmp::max( + image_factor, + floor_by_factor(img_w as f32 / beta, image_factor), + ); + } else if h_bar * w_bar < min_pixels { + let beta = (min_pixels as f32 / (img_h * img_w) as f32).sqrt(); + h_bar = ceil_by_factor(img_h as f32 * beta, image_factor); + w_bar = ceil_by_factor(img_w as f32 * beta, image_factor); + } + Ok((h_bar, w_bar)) +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index a7c34c2..653689f 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -3,6 +3,8 @@ pub mod img_utils; pub mod tensor_utils; pub mod video_utils; +use std::process::Command; + use aha_openai_dive::v1::resources::{ chat::{ ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, @@ -31,6 +33,33 @@ pub fn get_device(device: Option<&Device>) -> Device { } } +pub fn get_gpu_sm_arch() -> Result { + let output = Command::new("nvidia-smi") + .arg("--query-gpu=compute_cap") + .arg("--format=csv,noheader") + .output() + .map_err(|e| anyhow::anyhow!(format!("Failed to execute nvidia-smi: {}", e)))?; + if !output.status.success() { + return Err(anyhow::anyhow!(format!( + "nvidia-smi failed with status: {}\nError: {}", + output.status, + String::from_utf8_lossy(&output.stderr) + ))); + } + let output_str = String::from_utf8_lossy(&output.stdout); + let output_str = output_str.trim(); + let sm_float = match output_str.parse::() { + Ok(num) => num, + Err(_) => { + return Err(anyhow::anyhow!(format!( + "gpr sm arch: {} parse float32 error", + output_str + ))); + } + }; + Ok(sm_float) +} + pub fn get_dtype(dtype: Option, cfg_dtype: &str) -> DType { match dtype { Some(d) => d, @@ -41,7 +70,16 @@ pub fn get_dtype(dtype: Option, cfg_dtype: &str) -> DType { "float32" | "float" => DType::F32, "float64" | "double" => DType::F64, "float16" => DType::F16, - "bfloat16" => DType::BF16, + "bfloat16" => { + let arch = get_gpu_sm_arch(); + match arch { + Err(_) => DType::F16, + Ok(a) => { + // nvidia显卡sm架构>=8.0的才支持BF16 + if a >= 8.0 { DType::BF16 } else { DType::F16 } + } + } + } "uint8" => DType::U8, "int8" | "int16" | "int32" | "int64" => DType::I64, _ => DType::F32, @@ -257,6 +295,7 @@ pub fn get_logit_processor( top_k: Option, seed: u64, ) -> LogitsProcessor { + let temperature = temperature.and_then(|v| if v < 1e-7 { None } else { Some(v) }); match top_k { None => LogitsProcessor::new( seed, diff --git a/src/utils/tensor_utils.rs b/src/utils/tensor_utils.rs index 341f3e2..1057147 100644 --- a/src/utils/tensor_utils.rs +++ b/src/utils/tensor_utils.rs @@ -221,6 +221,14 @@ pub fn masked_scatter_dim0(original: &Tensor, replace: &Tensor, mask: &Tensor) - Ok(original) } +pub fn get_not_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result { + let image_token_id_tensor = Tensor::new(vec![token_ids], input_ids.device())?; + let mask = input_ids + .broadcast_ne(&image_token_id_tensor)? + .to_dtype(candle_core::DType::U32)?; + Ok(mask) +} + pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result { let image_token_id_tensor = Tensor::new(vec![token_ids], input_ids.device())?; let mask = input_ids @@ -229,10 +237,16 @@ pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result { Ok(mask) } -pub fn get_vision_next_indices(input_ids: &Tensor, token_id: u32) -> Result { +pub fn get_eq_indices(input_ids: &Tensor, token_id: u32) -> Result { // input_ids -> shape: (seq_len) let mask = get_equal_mask(input_ids, token_id)?; let indices = nonzero_index(&mask)?; + Ok(indices) +} + +pub fn get_vision_next_indices(input_ids: &Tensor, token_id: u32) -> Result { + // input_ids -> shape: (seq_len) + let indices = get_eq_indices(input_ids, token_id)?; let indices = indices.broadcast_add(&Tensor::new(vec![1u32], input_ids.device())?)?; Ok(indices) } @@ -384,9 +398,9 @@ pub fn interpolate_linear_1d( align_corner: Option, ) -> Result { // t: [b, channels, features] - if t.rank() < 3 { + if t.rank() != 3 { return Err(anyhow::anyhow!( - "Input rank must have at least 3 dimensions" + "Input rank must have equal to 3 dimensions" )); } let shape = t.dims(); @@ -394,19 +408,13 @@ pub fn interpolate_linear_1d( if orig_size == target_size { return Ok(t.clone()); } - let mut reshaped = t.clone(); - if shape.len() > 3 { - let bs = shape[0]; - let channels = shape[1..shape.len() - 1].iter().product::(); - reshaped = reshaped.reshape((bs, channels, orig_size))?; - } - let (bs, channels, _) = reshaped.dims3()?; + let (bs, channels, _) = t.dims3()?; let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), t.device())?; let coords = compute_1d_coords(orig_size, target_size, align_corner)?; for b in 0..bs { for c in 0..channels { - let input_slice = reshaped.i((b, c))?; + let input_slice = t.i((b, c))?; let mut out_i = Vec::new(); // for x_out in 0..target_size { for &coord in coords.iter().take(target_size) { @@ -424,16 +432,88 @@ pub fn interpolate_linear_1d( output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?; } } - if shape.len() != 3 { - let mut new_shape = shape.to_vec(); - let last_dim = new_shape.len() - 1; - new_shape[last_dim] = target_size; - output = output.reshape(new_shape)? - } output = output.contiguous()?; Ok(output) } +pub fn interpolate_bilinear( + input: &Tensor, + target_size: (usize, usize), + align_corner: Option, +) -> Result { + // input: [b, channels, height, width] + if input.rank() != 4 { + return Err(anyhow::anyhow!( + "Input rank must have equal to 4 dimensions [b, c, h, w]" + )); + } + + let (bs, channels, input_height, input_width) = input.dims4()?; + let (target_height, target_width) = target_size; + + // If size is the same, return clone + if input_height == target_height && input_width == target_width { + return Ok(input.clone()); + } + + let align_corners = align_corner.unwrap_or(false); + + // Compute scaling factors + let height_scale = if align_corners && target_height > 1 { + (input_height - 1) as f64 / (target_height - 1) as f64 + } else { + input_height as f64 / target_height as f64 + }; + + let width_scale = if align_corners && target_width > 1 { + (input_width - 1) as f64 / (target_width - 1) as f64 + } else { + input_width as f64 / target_width as f64 + }; + let dim0 = bs * channels; + let input_3dim = input.reshape((dim0, input_height, input_width))?; + let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::()?; + let mut output_data = vec![vec![vec![0.0f32; target_width]; target_height]; dim0]; + + for c in 0..dim0 { + for out_y in 0..target_height { + let src_y = if align_corners { + out_y as f64 * height_scale + } else { + (out_y as f64 + 0.5) * height_scale - 0.5 + }; + let src_y = src_y.max(0.0).min((input_height - 1) as f64); + let y0 = src_y.floor() as usize; + let y1 = (y0 + 1).min(input_height - 1); + let dy = (src_y - y0 as f64) as f32; + for out_x in 0..target_width { + let src_x = if align_corners { + out_x as f64 * width_scale + } else { + (out_x as f64 + 0.5) * width_scale - 0.5 + }; + let src_x = src_x.max(0.0).min((input_width - 1) as f64); + let x0 = src_x.floor() as usize; + let x1 = (x0 + 1).min(input_width - 1); + let q00 = input_data[c][y0][x0]; + let q01 = input_data[c][y0][x1]; + let q10 = input_data[c][y1][x0]; + let q11 = input_data[c][y1][x1]; + let dx = (src_x - x0 as f64) as f32; + let interpolated = q00 * (1.0 - dx) * (1.0 - dy) + + q01 * dx * (1.0 - dy) + + q10 * (1.0 - dx) * dy + + q11 * dx * dy; + output_data[c][out_y][out_x] = interpolated; + } + } + } + let output = Tensor::new(output_data, input.device())? + .reshape((bs, channels, target_height, target_width))? + .to_dtype(input.dtype())?; + Ok(output.contiguous()?) +} + fn compute_scale(input_size: usize, output_size: usize, align_corners: bool) -> f64 { if align_corners && output_size > 1 { (input_size - 1) as f64 / (output_size - 1) as f64 diff --git a/tests/config_tests.rs b/tests/config_tests.rs index e846485..2e25a14 100644 --- a/tests/config_tests.rs +++ b/tests/config_tests.rs @@ -1,7 +1,7 @@ use aha::models::{ - deepseek_ocr::config::DeepseekOCRConfig, minicpm4::config::MiniCPM4Config, - qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig, - voxcpm::config::VoxCPMConfig, + deepseek_ocr::config::DeepseekOCRConfig, hunyuan_ocr::config::HunYuanVLConfig, + minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig, + qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig, }; use anyhow::Result; @@ -48,10 +48,20 @@ fn qwen3vl_config() -> Result<()> { #[test] fn deepseek_ocr_config() -> Result<()> { - // cargo test -F cuda qwen3vl_config -r -- --nocapture + // cargo test -F cuda deepseek_ocr_config -r -- --nocapture let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/"; let config_path = model_path.to_string() + "/config.json"; let config: DeepseekOCRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; println!("{:?}", config); Ok(()) } + +#[test] +fn hunyuan_ocr_config() -> Result<()> { + // cargo test -F cuda hunyuan_ocr_config -r -- --nocapture + let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/"; + let config_path = model_path.to_string() + "/config.json"; + let config: HunYuanVLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + println!("{:?}", config); + Ok(()) +} diff --git a/tests/messy_test.rs b/tests/messy_test.rs index 0202681..3b74720 100644 --- a/tests/messy_test.rs +++ b/tests/messy_test.rs @@ -1,4 +1,6 @@ -use aha::utils::tensor_utils::interpolate_bicubic; +use std::time::Instant; + +use aha::utils::tensor_utils::interpolate_bilinear; use anyhow::Result; use candle_core::Tensor; @@ -6,11 +8,24 @@ use candle_core::Tensor; fn messy_test() -> Result<()> { // RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture let device = &candle_core::Device::Cpu; - // let t = Tensor::randn(0.0f32, 1.0, (1, 768, 64, 64), device)?; - let t = Tensor::arange(0.0f32, 10.0, device)?.broadcast_as((1, 1, 10, 10))?; + let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 1, 40, 40))?; println!("t: {}", t); - let t_resized = interpolate_bicubic(&t, (5, 5), Some(true), Some(false))?; - println!("t_resized: {}", t_resized); + let i_start = Instant::now(); + let t_inter = interpolate_bilinear(&t, (20, 20), Some(false))?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in interpolate_bilinear is: {:?}", i_duration); + println!("t_inter: {}", t_inter); + // let x: Vec = (0..5).flat_map(|_| 0u32..10).collect(); + // let id: Vec = (0..5).flat_map(|h| vec![h; 10]).collect(); + // println!("x: {:?}", id); + // let t = Tensor::randn(0.0f32, 1.0, (1, 768, 64, 64), device)?; + // let t = Tensor::arange(0u32, 10, device)?.broadcast_as((1, 10))?; + // let eq = t.broadcast_eq(&Tensor::new(5u32, device)?)?; + // println!("eq: {}", eq); + // let t = Tensor::arange(0.0f32, 10.0, device)?.broadcast_as((1, 1, 10, 10))?; + // println!("t: {}", t); + // let t_resized = interpolate_bicubic(&t, (5, 5), Some(true), Some(false))?; + // println!("t_resized: {}", t_resized); // let t1 = Tensor::rand(0.0, 1.0, (1, 5, 5, 10), device)?; // let t2 = Tensor::rand(0.0, 1.0, (5, 8, 10), device)?; // let t2 = t2.t()?; diff --git a/tests/test_hunyuan_ocr.rs b/tests/test_hunyuan_ocr.rs new file mode 100644 index 0000000..ec3f071 --- /dev/null +++ b/tests/test_hunyuan_ocr.rs @@ -0,0 +1,88 @@ +use std::{pin::pin, time::Instant}; + +use aha::models::{GenerateModel, hunyuan_ocr::generate::HunyuanOCRGenerateModel}; +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use rocket::futures::StreamExt; + +#[test] +fn hunyuan_ocr_generate() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda hunyuan_ocr_generate -r -- --nocapture + let message = r#" + { + "model": "hunyuan-ocr", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "image_url": + { + "url": "file://./assets/img/ocr_test1.png" + } + }, + { + "type": "text", + "text": "检测并识别图片中的文字，将文本坐标格式化输出。" + } + ] + } + ] + } + "#; + let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/"; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + let mut model = HunyuanOCRGenerateModel::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 hunyuan_ocr_stream() -> Result<()> { + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda hunyuan_ocr_stream -r -- --nocapture + + let message = r#" + { + "model": "hunyuan-ocr", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "image_url": + { + "url": "file://./assets/img/ocr_test1.png" + } + }, + { + "type": "text", + "text": "检测并识别图片中的文字，将文本坐标格式化输出。" + } + ] + } + ] + } + "#; + let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/"; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + let mut model = HunyuanOCRGenerateModel::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(()) +} diff --git a/tests/test_qwen3vl.rs b/tests/test_qwen3vl.rs index bc18348..31c20ce 100644 --- a/tests/test_qwen3vl.rs +++ b/tests/test_qwen3vl.rs @@ -19,15 +19,15 @@ fn qwen3vl_generate() -> Result<()> { "role": "user", "content": [ { - "type": "video", - "video_url": + "type": "image", + "image_url": { - "url": "https://www.w3schools.com/html/movie.mp4" + "url": "file://./assets/img/ocr_test1.png" } - }, + }, { "type": "text", - "text": "视频中发生了什么" + "text": "请分析图片并提取所有可见文本内容，按从左到右、从上到下的布局，返回纯文本" } ] } diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 57607aa..32e3736 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -80,3 +80,22 @@ fn deepseekocr_weight() -> Result<()> { println!("model_list: {:?}", model_list); Ok(()) } + +#[test] +fn hunyuanocr_weight() -> Result<()> { + let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/"; + let model_list = find_type_files(model_path, "safetensors")?; + + let device = Device::Cpu; + for m in &model_list { + let weights = safetensors::load(m, &device)?; + for (key, tensor) in weights.iter() { + if key.contains(".image_") { + println!("=== {} === {:?}", key, tensor.shape()); + } + // println!("=== {} === {:?}", key, tensor.shape()); + } + } + println!("model_list: {:?}", model_list); + Ok(()) +}