add deepseek-ocr2 support

Cargo.lock

@@ -21,7 +21,7 @@ dependencies = [
 
 [[package]]
 name = "aha"
-version = "0.2.2"
+version = "0.2.3"
 dependencies = [
  "aha_openai_dive",
  "ahash",

Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "aha"
-version = "0.2.2"
+version = "0.2.3"
 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, Hunyuan-OCR, PaddleOCR-VL, VoxCPM1.5, RMBG2.0, GLM-ASR-Nano-2512, Fun-ASR-Nano-2512, Qwen3, Qwen3-ASR, Qwen3.5, GLM-OCR"
+description = "aha model inference library, now supports Qwen(2.5VL/3/3VL/3.5/ASR), MiniCPM4, VoxCPM, DeepSeek-OCR/2, Hunyuan-OCR, PaddleOCR-VL/1.5, VoxCPM1.5, RMBG2.0, GLM-ASR-Nano-2512, Fun-ASR-Nano-2512, GLM-OCR"
 
 [dependencies]
 candle-core = { version = "0.9.2" }

README.md

@@ -25,6 +25,9 @@
 aha is a high-performance, cross-platform AI inference engine built with Rust and the Candle framework. It brings state-of-the-art AI models to your local machine—no API keys, no cloud dependencies, just pure, fast AI running directly on your hardware.
 
 ## Changelog
+### 2026-03-18
+- add DeepSeek-OCR-2
+
 ### 2026-03-17
 - add PaddleOCR-VL1.5 model
 - fix qwen3.5 position_ids create bug

README.zh-CN.md

@@ -25,6 +25,9 @@
 aha 是一款基于 Rust 和 Candle 框架构建的高性能跨平台 AI 推理引擎。将最先进的 AI 模型带到您的本地机器——无需 API 密钥，无需云依赖，纯粹、快速的 AI，直接在您的硬件上运行。
 
 ## 更新日志
+### 2026-03-18
+- 新增 DeepSeek-OCR-2
+
 ### 2026-03-17
 - 新增 PaddleOCR-VL1.5 模型
 - 修复 qwen3.5 position_ids 创建错误

docs/changelog.md

@@ -5,6 +5,9 @@ All notable changes to aha will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+### 2026-03-18
+- add DeepSeek-OCR-2
+
 ### 2026-03-17
 - add PaddleOCR-VL1.5 model
 - fix qwen3.5 position_ids create bug

docs/changelog.zh-CN.md

@@ -5,6 +5,9 @@
 格式基于 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.0.0/)，
 本项目遵循 [语义化版本](https://semver.org/lang/zh-CN/spec/v2.0.0.html)。
 
+### 2026-03-18
+- 新增 DeepSeek-OCR-2
+
 ### 2026-03-17
 - 新增 PaddleOCR-VL1.5 模型
 - 修复 qwen3.5 position_ids 创建错误

docs/model-card.md

@@ -7,6 +7,10 @@
 "OCR:" | "Table Recognition:" | "Formula Recognition:" | "Chart Recognition:" | "Spotting:" | "Seal Recognition:"
 
 ## DeepSeekOCR
+#### Main Prompts
+1. document: <image>\n<|grounding|>Convert the document to markdown.
+2. without layouts: <image>\nFree OCR.
+
 #### Metadata
 * base_size: 512 | 640 | 1024 | 1280
 * image_size: 512 | 640 | 1024 | 1280
@@ -25,6 +29,15 @@
 | Large | 1280 | 1280 | false | 
 | Gundam | 1024 | 640 | true | 
 
+## DeepSeekOCR2
+#### Main Prompts
+1. document: <image>\n<|grounding|>Convert the document to markdown.
+2. without layouts: <image>\nFree OCR.
+
+#### Metadata
+* crop_mode: false | true
+
+
 ## GLM-OCR
 #### Prompt Limited
 ##### Document Parsing prompt： 

docs/model-card.zh-CN.md

@@ -7,6 +7,10 @@
 "OCR:" | "Table Recognition:" | "Formula Recognition:" | "Chart Recognition:" | "Spotting:" | "Seal Recognition:"
 
 ## DeepSeekOCR
+#### 主要提示词
+1. 文档: <image>\n<|grounding|>Convert the document to markdown.
+2. 无布局信息: <image>\nFree OCR.
+
 #### Metadata
 * base_size: 512 | 640 | 1024 | 1280
 * image_size: 512 | 640 | 1024 | 1280
@@ -25,6 +29,14 @@
 | Large | 1280 | 1280 | false | 
 | Gundam | 1024 | 640 | true | 
 
+## DeepSeekOCR2
+#### 主要提示词
+1. 文档: <image>\n<|grounding|>Convert the document to markdown.
+2. 无布局信息: <image>\nFree OCR.
+
+#### Metadata
+* crop_mode: false | true
+
 ## GLM-OCR
 #### 限定提示词
 ##### 文档解析： 

src/api/mod.rs

@@ -262,6 +262,7 @@ fn which_model_to_id(which_model: WhichModel) -> &'static str {
         WhichModel::Qwen3vl8B => "qwen3vl-8b",
         WhichModel::Qwen3vl32B => "qwen3vl-32b",
         WhichModel::DeepSeekOCR => "deepseek-ocr",
+        WhichModel::DeepSeekOCR2 => "deepseek-ocr2",
         WhichModel::HunyuanOCR => "hunyuan-ocr",
         WhichModel::PaddleOCRVL => "paddleocr-vl",
         WhichModel::PaddleOCRVL1_5 => "paddleocr-vl1.5",
@@ -289,7 +290,7 @@ fn which_model_to_owner(which_model: WhichModel) -> &'static str {
         | WhichModel::Qwen3_5_2B
         | WhichModel::Qwen3_5_4B
         | WhichModel::Qwen3_5_9B => "Qwen",
-        WhichModel::DeepSeekOCR => "deepseek-ai",
+        WhichModel::DeepSeekOCR | WhichModel::DeepSeekOCR2 => "deepseek-ai",
         WhichModel::HunyuanOCR => "Tencent-Hunyuan",
         WhichModel::PaddleOCRVL | WhichModel::PaddleOCRVL1_5 => "PaddlePaddle",
         WhichModel::RMBG2_0 => "AI-ModelScope",

src/main.rs

@@ -258,6 +258,7 @@ fn run_list(args: ListArgs) -> anyhow::Result<()> {
         WhichModel::Qwen3vl8B,
         WhichModel::Qwen3vl32B,
         WhichModel::DeepSeekOCR,
+        WhichModel::DeepSeekOCR2,
         WhichModel::HunyuanOCR,
         WhichModel::PaddleOCRVL,
         WhichModel::PaddleOCRVL1_5,
@@ -520,6 +521,10 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> {
             use aha::exec::deepseek_ocr::DeepSeekORExec;
             DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?;
         }
+        WhichModel::DeepSeekOCR2 => {
+            use aha::exec::deepseek_ocr::DeepSeekORExec;
+            DeepSeekORExec::run(&input, output.as_deref(), &weight_path)?;
+        }
         WhichModel::HunyuanOCR => {
             use aha::exec::hunyuan_ocr::HunyuanORExec;
             HunyuanORExec::run(&input, output.as_deref(), &weight_path)?;

src/models/deepseek_ocr/config.rs

@@ -90,10 +90,17 @@ pub struct SamVitB {
     pub width: usize,
 }
 
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Qwen2_0_5B {
+    dim: usize,
+}
+
 #[derive(Debug, Clone, PartialEq, serde::Deserialize)]
 pub struct Width {
     #[serde(rename = "clip-l-14-224")]
-    pub clip_l_14_224: ClipL14_224,
+    pub clip_l_14_224: Option<ClipL14_224>,
+    #[serde(rename = "qwen2-0-5b")]
+    pub qwen2_0_5b: Option<Qwen2_0_5B>,
     pub sam_vit_b: SamVitB,
 }
 

src/models/deepseek_ocr/generate.rs

@@ -30,6 +30,7 @@ pub struct DeepseekOCRGenerateModel {
     device: Device,
     size: Vec<u32>,
     model_name: String,
+    version: usize,
 }
 
 impl DeepseekOCRGenerateModel {
@@ -40,13 +41,23 @@ impl DeepseekOCRGenerateModel {
         let cfg_dtype = cfg.language_config.torch_dtype.clone();
         let device = &get_device(device);
         let dtype = get_dtype(dtype, &cfg_dtype);
-        let processor = DeepseekOCRProcessor::new(device, dtype)?;
+        let model_name = std::path::Path::new(path)
+            .file_stem() // 获取文件名主干（不含扩展名）
+            .and_then(|s| s.to_str())
+            .unwrap_or("deepseek-ocr");
+        let version = if model_name.contains("2") || cfg.vision_config.width.qwen2_0_5b.is_some() {
+            2usize
+        } else {
+            1usize
+        };
+        let processor = DeepseekOCRProcessor::new(device, dtype, version)?;
         let eos_token_id = cfg.eos_token_id;
         let bos_token_id = cfg.bos_token_id;
         let model_list = find_type_files(path, "safetensors")?;
         let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
-        let deepseekocr_model = DeepseekOCRModel::new(vb, cfg)?;
+        let deepseekocr_model = DeepseekOCRModel::new(vb, cfg, version)?;
         let size = vec![512u32, 640, 1024, 1280];
+
         Ok(Self {
             tokenizer,
             processor,
@@ -55,7 +66,8 @@ impl DeepseekOCRGenerateModel {
             eos_token_id,
             device: device.clone(),
             size,
-            model_name: "deepseek-ocr".to_string(),
+            model_name: model_name.to_string(),
+            version,
         })
     }
 }
@@ -77,6 +89,8 @@ impl GenerateModel for DeepseekOCRGenerateModel {
         let crop_mode = extract_metadata_value::<bool>(&mes.metadata, "crop_mode").unwrap_or(false);
         let seed = mes.seed.unwrap_or(34562) as u64;
         let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
+        let base_size = if self.version == 2 { 1024 } else { base_size };
+        let image_size = if self.version == 2 { 768 } else { image_size };
         let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self
             .processor
             .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?;
@@ -152,6 +166,8 @@ impl GenerateModel for DeepseekOCRGenerateModel {
         let crop_mode = extract_metadata_value::<bool>(&mes.metadata, "crop_mode").unwrap_or(false);
         let seed = mes.seed.unwrap_or(34562) as u64;
         let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
+        let base_size = if self.version == 2 { 1024 } else { base_size };
+        let image_size = if self.version == 2 { 768 } else { image_size };
         let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self
             .processor
             .process_info(&mes, &self.tokenizer, base_size, image_size, crop_mode)?;

src/models/deepseek_ocr/model.rs

@@ -1,4 +1,6 @@
-use anyhow::Result;
+use core::f32;
+
+use anyhow::{Result, anyhow};
 use candle_core::{D, IndexOp, Tensor};
 use candle_nn::{
     Activation, Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding,
@@ -15,12 +17,15 @@ use crate::{
             get_layer_norm,
         },
         deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config},
+        qwen2::{Qwen2Config, Qwen2Decoder},
     },
     position_embed::rope::RoPE,
-    utils::interpolate::{interpolate_bicubic, interpolate_linear_1d},
+    utils::{
-    utils::tensor_utils::{
+        interpolate::{interpolate_bicubic, interpolate_linear_1d},
-        index_select_2d, masked_scatter_dim0, nonzero, onehot, prepare_causal_attention_mask,
+        tensor_utils::{
-        quick_gelu, topk,
+            attn_masked_fill, index_select_2d, masked_scatter_dim0, nonzero, onehot,
+            prepare_causal_attention_mask, quick_gelu, topk,
+        },
     },
 };
 
@@ -417,6 +422,7 @@ impl ImageEncoderViT {
         // rel_pos_zero_init: bool,
         window_size: usize,
         global_attn_indexes: Vec<usize>,
+        version: usize,
     ) -> Result<Self> {
         let patch_embed = PatchEmbed::new(
             vb.pp("patch_embed"),
@@ -463,7 +469,8 @@ impl ImageEncoderViT {
         let neck = Neck::new(vb.pp("neck"), embed_dim, out_chans)?;
 
         let net_2 = get_conv2d(vb.pp("net_2"), 256, 512, 3, 1, 2, 1, 1, false)?;
-        let net_3 = get_conv2d(vb.pp("net_3"), 512, 1024, 3, 1, 2, 1, 1, false)?;
+        let net_3_out_c = if version == 2 { 896 } else { 1024 };
+        let net_3 = get_conv2d(vb.pp("net_3"), 512, net_3_out_c, 3, 1, 2, 1, 1, false)?;
         Ok(Self {
             // img_size,
             patch_embed,
@@ -559,7 +566,6 @@ impl CLIPVisionEmbeddings {
     }
 
     fn get_abs_pos(&self, tgt_size: usize) -> Result<Tensor> {
-        // println!("self.pos_embeds: {:?}", self.pos_embeds);
         let abs_pos_new = self.pos_embeds.clone();
         let (len, dim) = abs_pos_new.dims2()?;
         let src_size = ((len - 1) as f32).sqrt() as usize;
@@ -1107,19 +1113,101 @@ impl DeepseekV2Model {
     }
 }
 
+pub struct Qwen2Decoder2Encoder {
+    model: Qwen2Decoder,
+    query_768: Embedding,
+    query_1024: Embedding,
+}
+
+impl Qwen2Decoder2Encoder {
+    pub fn new(vb: VarBuilder) -> Result<Self> {
+        let qwen2_config = Qwen2Config {
+            vocab_size: 151936,
+            hidden_size: 896,
+            intermediate_size: 4864,
+            num_hidden_layers: 24,
+            num_attention_heads: 14,
+            num_key_value_heads: 2,
+            max_position_embeddings: 131072,
+            sliding_window: 32768,
+            max_window_layers: 21,
+            tie_word_embeddings: true,
+            rope_theta: 1000000.0,
+            rms_norm_eps: 1e-06,
+            use_sliding_window: false,
+            hidden_act: Activation::Silu,
+        };
+        let model = Qwen2Decoder::new(vb.pp("model.model"), &qwen2_config)?;
+        let query_768 = embedding(144, qwen2_config.hidden_size, vb.pp("query_768"))?;
+        let query_1024 = embedding(256, qwen2_config.hidden_size, vb.pp("query_1024"))?;
+
+        Ok(Self {
+            model,
+            query_768,
+            query_1024,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = xs.flatten_from(2)?.transpose(1, 2)?;
+        let (bs, n_query, _) = xs.dims3()?;
+        let param_img = if n_query == 144 {
+            self.query_768.embeddings()
+        } else if n_query == 256 {
+            self.query_1024.embeddings()
+        } else {
+            return Err(anyhow!(
+                "Only support 144/256 seq_len, data {n_query} illigal"
+            ));
+        };
+        let brach_query_imgs = param_img.unsqueeze(0)?.repeat((bs, 1, 1))?;
+        let x_combined = Tensor::cat(&[&xs, &brach_query_imgs], 1)?;
+        let token_type_ids = Tensor::cat(
+            &[
+                Tensor::ones(n_query, candle_core::DType::U32, xs.device())?,
+                Tensor::zeros(n_query, candle_core::DType::U32, xs.device())?,
+            ],
+            0,
+        )?
+        .unsqueeze(0)?; // (n_query*2, 1)
+        let mask_up = token_type_ids.repeat((n_query, 1))?; //(n_query, n_query*2)
+        let mask_down_1 = Tensor::ones((n_query, n_query), candle_core::DType::U32, xs.device())?;
+        let mask_down_2 = Tensor::tril2(n_query, candle_core::DType::U32, xs.device())?;
+        let mask_down = Tensor::cat(&[mask_down_1, mask_down_2], 1)?;
+        let mask = Tensor::cat(&[mask_up, mask_down], 0)?;
+        let on_true = mask
+            .zeros_like()?
+            .unsqueeze(0)?
+            .unsqueeze(0)?
+            .to_dtype(candle_core::DType::F32)?; // (1, 1, n_query*2, n_query*2)
+        let attn_mask = attn_masked_fill(&on_true, &mask, f32::NEG_INFINITY)?;
+        let xs = self
+            .model
+            .forward_no_cache(&x_combined, Some(&attn_mask), 0)?;
+        let xs = xs.narrow(1, n_query, n_query)?;
+        Ok(xs)
+    }
+}
+
+pub enum VisionModel {
+    Vit(VitModel),               // vb_name: vision_model
+    Qwen2(Qwen2Decoder2Encoder), // vb_name: qwen2_model.model.
+}
+
 pub struct DeepseekOCRModel {
     // config: DeepseekOCRConfig,
     sam_model: ImageEncoderViT,
-    vision_model: VitModel,
+    // vision_model: VitModel,
+    vision_model: VisionModel,
     projector: Linear,
     language_model: DeepseekV2Model,
-    image_newline: Tensor,
+    image_newline: Option<Tensor>,
     view_seperator: Tensor,
     lm_head: Linear,
 }
 
 impl DeepseekOCRModel {
-    pub fn new(vb: VarBuilder, config: DeepseekOCRConfig) -> Result<Self> {
+    pub fn new(vb: VarBuilder, config: DeepseekOCRConfig, version: usize) -> Result<Self> {
         let vb_m = vb.pp("model");
         let sam_model = ImageEncoderViT::new(
             vb_m.pp("sam_model"),
@@ -1143,24 +1231,34 @@ impl DeepseekOCRModel {
                 .sam_vit_b
                 .global_attn_indexes
                 .clone(),
+            version,
         )?;
-        let vision_model = VitModel::new(
+        let (vision_model, image_newline) = if version == 2 {
-            vb_m.pp("vision_model"),
+            // v2
-            224,
+            let qwen2 = Qwen2Decoder2Encoder::new(vb_m.pp("qwen2_model"))?;
-            14,
+            (VisionModel::Qwen2(qwen2), None)
-            3,
+        } else {
-            24,
+            let vision_model = VitModel::new(
-            1024,
+                vb_m.pp("vision_model"),
-            16,
+                224,
-            4096,
+                14,
-            1e-5,
+                3,
-        )?;
+                24,
+                1024,
+                16,
+                4096,
+                1e-5,
+            )?;
+            let image_newline = vb_m.get_with_hints(1280, "image_newline", Init::Const(0.))?;
+            (VisionModel::Vit(vision_model), Some(image_newline))
+        };
+
         let projector = linear(
             config.projector_config.input_dim,
             config.projector_config.n_embed,
             vb_m.pp("projector.layers"),
         )?;
-        let image_newline = vb_m.get_with_hints(1280, "image_newline", Init::Const(0.))?;
+
         let view_seperator = vb_m.get_with_hints(1280, "view_seperator", Init::Const(0.))?;
         let language_model = DeepseekV2Model::new(vb_m, config.language_config.clone())?;
         let lm_head = linear_no_bias(config.hidden_size, config.vocab_size, vb.pp("lm_head"))?;
@@ -1210,41 +1308,88 @@ impl DeepseekOCRModel {
                     let image_crop_i = image_crop.i(last_crop_num..last_crop_num + crop_num)?;
                     last_crop_num += crop_num;
                     let local_feature_1 = self.sam_model.forward(&image_crop_i)?;
-                    let local_feature_2 = self
+                    let local_features = match &self.vision_model {
-                        .vision_model
+                        VisionModel::Vit(vit) => {
-                        .forward(&image_crop_i, Some(&local_feature_1))?;
+                            let local_feature_2 =
-                    let local_feature_1 = local_feature_1.flatten(2, 3)?.permute((0, 2, 1))?;
+                                vit.forward(&image_crop_i, Some(&local_feature_1))?;
-                    let local_feature_2 = local_feature_2.i((.., 1..))?;
+                            let local_feature_1 =
-                    let local_features =
+                                local_feature_1.flatten(2, 3)?.permute((0, 2, 1))?;
-                        Tensor::cat(&[local_feature_2, local_feature_1], D::Minus1)?
+                            let local_feature_2 = local_feature_2.i((.., 1..))?;
-                            .contiguous()?;
+                            Tensor::cat(&[local_feature_2, local_feature_1], D::Minus1)?
+                                .contiguous()?
+                        }
+                        VisionModel::Qwen2(qwen2) => qwen2.forward(&local_feature_1)?,
+                    };
+                    // let local_feature_2 = self
+                    //     .vision_model
+                    //     .forward(&image_crop_i, Some(&local_feature_1))?;
+                    // let local_feature_1 = local_feature_1.flatten(2, 3)?.permute((0, 2, 1))?;
+                    // let local_feature_2 = local_feature_2.i((.., 1..))?;
+                    // let local_features =
+                    //     Tensor::cat(&[local_feature_2, local_feature_1], D::Minus1)?
+                    //         .contiguous()?;
                     let local_features = self.projector.forward(&local_features)?;
                     let global_features_1 = self.sam_model.forward(&image_ori_i)?;
-                    let global_features_2 = self
+                    let global_features = match &self.vision_model {
-                        .vision_model
+                        VisionModel::Vit(vit) => {
-                        .forward(&image_ori_i, Some(&global_features_1))?;
+                            let global_features_2 =
-                    let global_features_1 = global_features_1.flatten(2, 3)?.permute((0, 2, 1))?;
+                                vit.forward(&image_ori_i, Some(&global_features_1))?;
-                    let global_features_2 = global_features_2.i((.., 1..))?;
+                            let global_features_1 =
-                    let global_features =
+                                global_features_1.flatten(2, 3)?.permute((0, 2, 1))?;
-                        Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?;
+                            let global_features_2 = global_features_2.i((.., 1..))?;
+                            Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?
+                        }
+                        VisionModel::Qwen2(qwen2) => qwen2.forward(&global_features_1)?,
+                    };
+                    // let global_features_2 = self
+                    //     .vision_model
+                    //     .forward(&image_ori_i, Some(&global_features_1))?;
+                    // let global_features_1 = global_features_1.flatten(2, 3)?.permute((0, 2, 1))?;
+                    // let global_features_2 = global_features_2.i((.., 1..))?;
+                    // let global_features =
+                    //     Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?;
                     let global_features = self.projector.forward(&global_features)?;
                     let (_, hw, n_dim) = global_features.dims3()?;
-                    let h = (hw as f32).sqrt() as usize;
-                    let w = h;
                     let (_, hw2, n_dim2) = local_features.dims3()?;
-                    let h2 = (hw2 as f32).sqrt() as usize;
+                    let (global_features, local_features) = if let Some(image_newline) =
-                    let w2 = h2;
+                        &self.image_newline
-                    let global_features = global_features.reshape((h, w, n_dim))?;
+                    {
-                    let image_newline = self.image_newline.unsqueeze(0)?.unsqueeze(0)?;
+                        let h = (hw as f32).sqrt() as usize;
-                    let global_cat = image_newline.expand((h, 1, n_dim))?;
+                        let w = h;
-                    let global_features = Tensor::cat(&[&global_features, &global_cat], 1)?;
+                        let h2 = (hw2 as f32).sqrt() as usize;
+                        let w2 = h2;
+                        let global_features = global_features.reshape((h, w, n_dim))?;
+                        let image_newline = image_newline.unsqueeze(0)?.unsqueeze(0)?;
+                        let global_cat = image_newline.expand((h, 1, n_dim))?;
+                        let global_features = Tensor::cat(&[&global_features, &global_cat], 1)?;
+                        let local_features = local_features
+                            .reshape((height_crop_num, width_crop_num, h2, w2, n_dim2))?
+                            .permute((0, 2, 1, 3, 4))?
+                            .reshape((height_crop_num * h2, width_crop_num * w2, n_dim2))?;
+                        let local_cat = image_newline.expand((height_crop_num * h2, 1, n_dim2))?;
+                        let local_features = Tensor::cat(&[&local_features, &local_cat], 1)?;
+                        (global_features, local_features)
+                    } else {
+                        (global_features, local_features)
+                    };
+
+                    // let h = (hw as f32).sqrt() as usize;
+                    // let w = h;
+                    // let (_, hw2, n_dim2) = local_features.dims3()?;
+                    // let h2 = (hw2 as f32).sqrt() as usize;
+                    // let w2 = h2;
+                    // let global_features = global_features.reshape((h, w, n_dim))?;
+                    // let image_newline = self.image_newline.unsqueeze(0)?.unsqueeze(0)?;
+                    // let global_cat = image_newline.expand((h, 1, n_dim))?;
+                    // let global_features = Tensor::cat(&[&global_features, &global_cat], 1)?;
+                    // let global_features = global_features.reshape(((), n_dim))?;
+                    // let local_features = local_features
+                    //     .reshape((height_crop_num, width_crop_num, h2, w2, n_dim2))?
+                    //     .permute((0, 2, 1, 3, 4))?
+                    //     .reshape((height_crop_num * h2, width_crop_num * w2, n_dim2))?;
+                    // let local_cat = image_newline.expand((height_crop_num * h2, 1, n_dim2))?;
+                    // let local_features = Tensor::cat(&[&local_features, &local_cat], 1)?;
                     let global_features = global_features.reshape(((), n_dim))?;
-                    let local_features = local_features
-                        .reshape((height_crop_num, width_crop_num, h2, w2, n_dim2))?
-                        .permute((0, 2, 1, 3, 4))?
-                        .reshape((height_crop_num * h2, width_crop_num * w2, n_dim2))?;
-                    let local_cat = image_newline.expand((height_crop_num * h2, 1, n_dim2))?;
-                    let local_features = Tensor::cat(&[&local_features, &local_cat], 1)?;
                     let local_features = local_features.reshape(((), n_dim2))?;
                     Tensor::cat(
                         &[
@@ -1256,21 +1401,44 @@ impl DeepseekOCRModel {
                     )?
                 } else {
                     let global_features_1 = self.sam_model.forward(&image_ori_i)?;
-                    let global_features_2 = self
+                    let global_features = match &self.vision_model {
-                        .vision_model
+                        VisionModel::Vit(vit) => {
-                        .forward(&image_ori_i, Some(&global_features_1))?;
+                            let global_features_2 =
-                    let global_features_1 = global_features_1.flatten(2, 3)?.permute((0, 2, 1))?;
+                                vit.forward(&image_ori_i, Some(&global_features_1))?;
-                    let global_features_2 = global_features_2.i((.., 1..))?;
+                            let global_features_1 =
-                    let global_features =
+                                global_features_1.flatten(2, 3)?.permute((0, 2, 1))?;
-                        Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?;
+                            let global_features_2 = global_features_2.i((.., 1..))?;
+                            Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?
+                        }
+                        VisionModel::Qwen2(qwen2) => qwen2.forward(&global_features_1)?,
+                    };
+                    // let global_features_2 = self
+                    //     .vision_model
+                    //     .forward(&image_ori_i, Some(&global_features_1))?;
+                    // let global_features_1 = global_features_1.flatten(2, 3)?.permute((0, 2, 1))?;
+                    // let global_features_2 = global_features_2.i((.., 1..))?;
+                    // let global_features =
+                    //     Tensor::cat(&[global_features_2, global_features_1], D::Minus1)?;
                     let global_features = self.projector.forward(&global_features)?;
                     let (_, hw, n_dim) = global_features.dims3()?;
-                    let h = (hw as f32).sqrt() as usize;
+                    let global_features = if let Some(image_newline) = &self.image_newline {
-                    let w = h;
+                        let h = (hw as f32).sqrt() as usize;
-                    let global_features = global_features.reshape((h, w, n_dim))?;
+                        let w = h;
-                    let image_newline = self.image_newline.unsqueeze(0)?.unsqueeze(0)?;
+                        let global_features = global_features.reshape((h, w, n_dim))?;
-                    let global_cat = image_newline.expand((h, 1, n_dim))?;
+                        let image_newline = image_newline.unsqueeze(0)?.unsqueeze(0)?;
-                    let global_features = Tensor::cat(&[&global_features, &global_cat], 1)?;
+                        let global_cat = image_newline.expand((h, 1, n_dim))?;
+                        Tensor::cat(&[&global_features, &global_cat], 1)?
+                    } else {
+                        global_features
+                    };
+                    // let (_, hw, n_dim) = global_features.dims3()?;
+                    // let h = (hw as f32).sqrt() as usize;
+                    // let w = h;
+                    // let global_features = global_features.reshape((h, w, n_dim))?;
+                    // let image_newline = self.image_newline.unsqueeze(0)?.unsqueeze(0)?;
+                    // let global_cat = image_newline.expand((h, 1, n_dim))?;
+                    // let global_features = Tensor::cat(&[&global_features, &global_cat], 1)?;
+
                     let global_features = global_features.reshape(((), n_dim))?;
                     Tensor::cat(&[global_features, self.view_seperator.unsqueeze(0)?], 0)?
                 };

src/models/deepseek_ocr/processor.rs

@@ -18,10 +18,11 @@ pub struct DeepseekOCRProcessor {
     image_token_id: u32,
     patch_size: u32,
     downsample_ratio: u32,
+    version: usize,
 }
 
 impl DeepseekOCRProcessor {
-    pub fn new(device: &Device, dtype: DType) -> Result<Self> {
+    pub fn new(device: &Device, dtype: DType, version: usize) -> Result<Self> {
         Ok(Self {
             device: device.clone(),
             dtype,
@@ -29,6 +30,7 @@ impl DeepseekOCRProcessor {
             image_token_id: 128815,
             patch_size: 16,
             downsample_ratio: 4,
+            version,
         })
     }
 
@@ -70,6 +72,8 @@ impl DeepseekOCRProcessor {
         let mut images_seq_mask = vec![0u32];
         let mut tokenized_id = vec![0u32];
         let mut images_spatial_crop = Vec::new();
+        let min_img_size = if self.version == 2 { 768 } else { 640 };
+        let max_num = if self.version == 2 { 6 } else { 9 };
         for (text_seq, image) in text_splits.iter().zip(imgs) {
             if !text_seq.is_empty() {
                 let token_ids = tokenizer.text_encode_vec(text_seq.to_string(), false)?;
@@ -79,10 +83,12 @@ impl DeepseekOCRProcessor {
             }
             if crop_mode {
                 let mut images_crop_raw = Vec::new();
-                let crop_ratio = if image.height() <= 640 && image.width() <= 640 {
+                let crop_ratio = if image.height() <= min_img_size && image.width() <= min_img_size
+                {
                     (1u32, 1u32)
                 } else {
-                    let (img_crop, ratio) = dynamic_preprocess(&image, image_size, false)?;
+                    let (img_crop, ratio) =
+                        dynamic_preprocess(&image, 2, max_num, min_img_size, false)?;
                     images_crop_raw = img_crop.clone();
                     ratio
                 };
@@ -106,16 +112,25 @@ impl DeepseekOCRProcessor {
 
                 let num_queries = image_size / self.patch_size / self.downsample_ratio;
                 let num_queries_base = base_size / self.patch_size / self.downsample_ratio;
-                let mut token_repeat = num_queries_base.pow(2) + num_queries_base + 1;
+                let mut token_repeat = if self.version == 1 {
+                    num_queries_base.pow(2) + num_queries_base + 1
+                } else {
+                    num_queries_base.pow(2) + 1
+                };
                 if crop_ratio.0 > 1 || crop_ratio.1 > 1 {
-                    token_repeat += (num_queries * crop_ratio.0 + 1) * (num_queries * crop_ratio.1);
+                    let add_num = if self.version == 1 {
+                        (num_queries * crop_ratio.0 + 1) * (num_queries * crop_ratio.1)
+                    } else {
+                        (num_queries * crop_ratio.0) * (num_queries * crop_ratio.1)
+                    };
+                    token_repeat += add_num;
                 }
                 let tokenized_image = vec![self.image_token_id; token_repeat as usize];
                 tokenized_id.extend_from_slice(&tokenized_image);
                 let seq_mask = vec![1u32; tokenized_image.len()];
                 images_seq_mask.extend_from_slice(&seq_mask);
             } else {
-                let global_view = if image_size <= 640 {
+                let global_view = if image_size <= min_img_size {
                     image.resize_exact(
                         image_size,
                         image_size,
@@ -130,7 +145,11 @@ impl DeepseekOCRProcessor {
 
                 images_spatial_crop.push(vec![1, 1]);
                 let num_queries = image_size / self.patch_size / self.downsample_ratio;
-                let token_repeat = num_queries.pow(2) + num_queries + 1;
+                let token_repeat = if self.version == 1 {
+                    num_queries.pow(2) + num_queries + 1
+                } else {
+                    num_queries.pow(2) + 1
+                };
                 let tokenized_image = vec![self.image_token_id; token_repeat as usize];
                 tokenized_id.extend_from_slice(&tokenized_image);
                 let seq_mask = vec![1u32; tokenized_image.len()];

src/models/mod.rs

@@ -10,6 +10,7 @@ pub mod hunyuan_ocr;
 pub mod mask_gct;
 pub mod minicpm4;
 pub mod paddleocr_vl;
+pub mod qwen2;
 pub mod qwen2_5vl;
 pub mod qwen3;
 pub mod qwen3_5;
@@ -70,6 +71,8 @@ pub enum WhichModel {
     Qwen3vl32B,
     #[value(name = "deepseek-ocr", hide = true)]
     DeepSeekOCR,
+    #[value(name = "deepseek-ocr2", hide = true)]
+    DeepSeekOCR2,
     #[value(name = "hunyuan-ocr", hide = true)]
     HunyuanOCR,
     #[value(name = "paddleocr-vl", hide = true)]
@@ -110,6 +113,7 @@ impl WhichModel {
             WhichModel::Qwen3vl8B => "Qwen/Qwen3-VL-8B-Instruct",
             WhichModel::Qwen3vl32B => "Qwen/Qwen3-VL-32B-Instruct",
             WhichModel::DeepSeekOCR => "deepseek-ai/DeepSeek-OCR",
+            WhichModel::DeepSeekOCR2 => "deepseek-ai/DeepSeek-OCR-2",
             WhichModel::HunyuanOCR => "Tencent-Hunyuan/HunyuanOCR",
             WhichModel::PaddleOCRVL => "PaddlePaddle/PaddleOCR-VL",
             WhichModel::PaddleOCRVL1_5 => "PaddlePaddle/PaddleOCR-VL-1.5",
@@ -140,6 +144,7 @@ impl WhichModel {
             | WhichModel::Qwen3_5Gguf => "vlm",
             // OCR models
             WhichModel::DeepSeekOCR
+            | WhichModel::DeepSeekOCR2
             | WhichModel::HunyuanOCR
             | WhichModel::GlmOCR
             | WhichModel::PaddleOCRVL
@@ -312,6 +317,10 @@ pub fn load_model<'a>(
             let model = DeepseekOCRGenerateModel::init(path, None, None)?;
             ModelInstance::DeepSeekOCR(model)
         }
+        WhichModel::DeepSeekOCR2 => {
+            let model = DeepseekOCRGenerateModel::init(path, None, None)?;
+            ModelInstance::DeepSeekOCR(model)
+        }
         WhichModel::HunyuanOCR => {
             let model = HunyuanOCRGenerateModel::init(path, None, None)?;
             ModelInstance::HunyuanOCR(model)

src/models/qwen2/mod.rs

@@ -0,0 +1,275 @@
+use anyhow::Result;
+use candle_core::Tensor;
+use candle_nn::{
+    Activation, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm,
+};
+
+use crate::{
+    models::common::{GateUpDownMLP, eager_attention_forward},
+    position_embed::rope::{RoPE, apply_rotary_pos_emb},
+};
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Qwen2Config {
+    pub vocab_size: usize,
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub num_hidden_layers: usize,
+    pub num_attention_heads: usize,
+    pub num_key_value_heads: usize,
+    pub max_position_embeddings: usize,
+    pub sliding_window: usize,
+    pub max_window_layers: usize,
+    pub tie_word_embeddings: bool,
+    pub rope_theta: f32,
+    pub rms_norm_eps: f64,
+    pub use_sliding_window: bool,
+    pub hidden_act: Activation,
+}
+
+#[derive(Debug, Clone)]
+pub struct Qwen2Attention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+    kv_cache: Option<(Tensor, Tensor)>,
+}
+
+impl Qwen2Attention {
+    pub fn new(cfg: &Qwen2Config, vb: VarBuilder) -> Result<Self> {
+        let hidden_size = cfg.hidden_size;
+        let num_heads = cfg.num_attention_heads;
+        let num_kv_heads = cfg.num_key_value_heads;
+        let num_kv_groups = num_heads / num_kv_heads;
+        let head_dim = hidden_size / num_heads;
+        let q_proj = linear(hidden_size, num_heads * head_dim, vb.pp("q_proj"))?;
+        let k_proj = linear(hidden_size, num_kv_heads * head_dim, vb.pp("k_proj"))?;
+        let v_proj = linear(hidden_size, num_kv_heads * head_dim, vb.pp("v_proj"))?;
+        let o_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("o_proj"))?;
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            num_heads,
+            num_kv_heads,
+            num_kv_groups,
+            head_dim,
+            hidden_size,
+            kv_cache: None,
+        })
+    }
+
+    pub fn forward(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (b_sz, q_len, _) = xs.dims3()?;
+        let query_states = self.q_proj.forward(xs)?;
+        let key_states = self.k_proj.forward(xs)?;
+        let value_states = self.v_proj.forward(xs)?;
+        let query_states = query_states
+            .reshape((b_sz, q_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let key_states = key_states
+            .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        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, false)?;
+        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 scale = 1f64 / f64::sqrt(self.head_dim as f64);
+        let attn_output = eager_attention_forward(
+            &query_states,
+            &key_states,
+            &value_states,
+            Some(self.num_kv_groups),
+            attention_mask,
+            scale,
+        )?;
+        let attn_output = attn_output.reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn forward_no_cache(
+        &self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (b_sz, q_len, _) = xs.dims3()?;
+        let query_states = self.q_proj.forward(xs)?;
+        let key_states = self.k_proj.forward(xs)?;
+        let value_states = self.v_proj.forward(xs)?;
+        let query_states = query_states
+            .reshape((b_sz, q_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let key_states = key_states
+            .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        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, false)?;
+
+        let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+        let attn_output = eager_attention_forward(
+            &query_states,
+            &key_states,
+            &value_states,
+            Some(self.num_kv_groups),
+            attention_mask,
+            scale,
+        )?;
+        let attn_output = attn_output.reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.kv_cache = None
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Qwen2DecoderLayer {
+    self_attn: Qwen2Attention,
+    mlp: GateUpDownMLP,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+}
+
+impl Qwen2DecoderLayer {
+    pub fn new(cfg: &Qwen2Config, vb: VarBuilder) -> Result<Self> {
+        let self_attn = Qwen2Attention::new(cfg, vb.pp("self_attn"))?;
+        let mlp = GateUpDownMLP::new(
+            vb.pp("mlp"),
+            cfg.hidden_size,
+            cfg.intermediate_size,
+            cfg.hidden_act,
+            false,
+            None,
+            None,
+            None,
+        )?;
+        let input_layernorm =
+            rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
+        let post_attention_layernorm = rms_norm(
+            cfg.hidden_size,
+            cfg.rms_norm_eps,
+            vb.pp("post_attention_layernorm"),
+        )?;
+        Ok(Self {
+            self_attn,
+            mlp,
+            input_layernorm,
+            post_attention_layernorm,
+        })
+    }
+
+    pub fn forward(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs;
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?;
+        let xs = (xs + residual)?;
+        let residual = &xs;
+        let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
+        let xs = (residual + xs)?;
+        Ok(xs)
+    }
+
+    pub fn forward_no_cache(
+        &self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs;
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self
+            .self_attn
+            .forward_no_cache(&xs, cos, sin, attention_mask)?;
+        let xs = (xs + residual)?;
+        let residual = &xs;
+        let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
+        let xs = (residual + xs)?;
+        Ok(xs)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.self_attn.clear_kv_cache()
+    }
+}
+
+pub struct Qwen2Decoder {
+    layers: Vec<Qwen2DecoderLayer>,
+    norm: RmsNorm,
+    rotary_emb: RoPE,
+}
+
+impl Qwen2Decoder {
+    pub fn new(vb: VarBuilder, cfg: &Qwen2Config) -> Result<Self> {
+        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
+        let vb_l = vb.pp("layers");
+        for layer_idx in 0..cfg.num_hidden_layers {
+            let layer = Qwen2DecoderLayer::new(cfg, vb_l.pp(layer_idx))?;
+            layers.push(layer)
+        }
+        let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?;
+        let head_dim = cfg.hidden_size / cfg.num_attention_heads;
+        let rotary_emb = RoPE::new(head_dim, cfg.rope_theta, vb.device())?;
+        Ok(Self {
+            layers,
+            norm,
+            rotary_emb,
+        })
+    }
+
+    pub fn forward_no_cache(
+        &self,
+        xs: &Tensor,
+        attention_mask: Option<&Tensor>,
+        seqlen_offset: usize,
+    ) -> Result<Tensor> {
+        let seq_len = xs.dim(1)?;
+        let (cos, sin) = self
+            .rotary_emb
+            .forward(seqlen_offset, seq_len, xs.device())?;
+        let mut xs = xs.clone();
+        for layer in self.layers.iter() {
+            xs = layer.forward_no_cache(&xs, &cos, &sin, attention_mask)?;
+        }
+        let xs = xs.apply(&self.norm)?;
+        Ok(xs)
+    }
+}

src/position_embed/rope.rs

@@ -566,11 +566,13 @@ impl RoPE {
         let positions = Tensor::arange(
             seqlen_offset as f32,
             (seqlen_offset + seq_len) as f32,
-            device,
+            self.inv_freq.device(),
         )?
         .reshape((seq_len, 1))?; // (seq_len, 1)
         let freqs = positions.matmul(&self.inv_freq)?; // (seq_len, dim / 2)
-        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?.contiguous()?; // (seq_len, dim)
+        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?
+            .contiguous()?
+            .to_device(device)?; // (seq_len, dim)
         let cos = emb.cos()?;
         let sin = emb.sin()?;
         Ok((cos, sin))

src/utils/img_utils.rs

@@ -162,6 +162,8 @@ pub fn find_closest_aspect_ratio(
 
 pub fn dynamic_preprocess(
     image: &DynamicImage,
+    min_num: u32,
+    max_num: u32,
     image_size: u32,
     use_thumbnail: bool,
 ) -> Result<(Vec<DynamicImage>, (u32, u32))> {
@@ -169,7 +171,7 @@ pub fn dynamic_preprocess(
     let orig_height = image.height();
     let aspect_ratio = orig_width as f64 / orig_height as f64;
     // 控制分块数量在2-9之间
-    let target_ratios = generate_target_ratios_sorted(2, 9);
+    let target_ratios = generate_target_ratios_sorted(min_num, max_num);
     let target_aspect_ratio = find_closest_aspect_ratio(
         aspect_ratio,
         &target_ratios,

tests/messy_test.rs

@@ -9,7 +9,7 @@
 // use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
 use anyhow::Result;
 // use byteorder::{LittleEndian, ReadBytesExt};
-use candle_core::Tensor;
+// use candle_core::Tensor;
 use modelscope::{DownloadOptions, ModelScope};
 // use sentencepiece::SentencePieceProcessor;
 // use zip::ZipArchive;
@@ -35,12 +35,19 @@ async fn download_test() -> Result<()> {
 #[test]
 fn messy_test() -> Result<()> {
     // RUST_BACKTRACE=1 cargo test -F cuda --test messy_test messy_test -r -- --nocapture
-
+    let save_dir =
-    let device = &candle_core::Device::Cpu;
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
-    let t1 = Tensor::randn(0.0, 1.0, (16, 9, 64, 128), device)?;
+    let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir);
-    let t2 = Tensor::randn(0.0, 1.0, (16, 9, 128, 64), device)?;
+    let stem = std::path::Path::new(&model_path)
-    let out = t1.matmul(&t2)?;
+        .file_stem() // 获取文件名主干（不含扩展名）
-    println!("out shape: {:?}", out);
+        .and_then(|s| s.to_str())
+        .unwrap_or("qwen3.5");
+    println!("stem: {:?}", stem);
+    // let device = &candle_core::Device::Cpu;
+    // let t1 = Tensor::randn(0.0, 1.0, (16, 9, 64, 128), device)?;
+    // let t2 = Tensor::randn(0.0, 1.0, (16, 9, 128, 64), device)?;
+    // let out = t1.matmul(&t2)?;
+    // println!("out shape: {:?}", out);
 
     // let input = Tensor::arange(0.0f32, 25.0f32, device)?.reshape((5, 5))?;
     // println!("input: {}", input);

tests/test_deepseek_ocr.rs

@@ -5,6 +5,55 @@ use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
 use anyhow::Result;
 use rocket::futures::StreamExt;
 
+#[test]
+fn deepseek_ocr2_generate() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda --test test_deepseek_ocr deepseek_ocr2_generate -r -- --nocapture
+    let message = r#"
+    {
+        "model": "deepseek-ocr2",
+        "messages": [
+            {
+                "role": "user",
+                "content": [ 
+                    {
+                        "type": "image",
+                        "image_url": 
+                        {
+                            "url": "file://./assets/img/ocr_test1.png"
+                        }
+                    },              
+                    {
+                        "type": "text", 
+                        "text": "<image>\nConvert the document to markdown. "
+                    }
+                ]
+            }
+        ],
+        "metadata": {"crop_mode": "false"}
+    }
+    "#;
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir);
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut model = DeepseekOCRGenerateModel::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!("generate: \n {:?}", res);
+    if let Some(usage) = &res.usage {
+        let num_token = usage.total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
+    println!("Time elapsed in generate is: {:?}", i_duration);
+    Ok(())
+}
+
 #[test]
 fn deepseek_ocr_generate() -> Result<()> {
     // RUST_BACKTRACE=1 cargo test -F cuda --test test_deepseek_ocr deepseek_ocr_generate -r -- --nocapture

tests/weight_test.rs

@@ -243,3 +243,25 @@ fn index_tts2_weight() -> Result<()> {
     // }
     Ok(())
 }
+
+#[test]
+fn deepseekocrv2_weight() -> Result<()> {
+    // cargo test -F cuda --test weight_test deepseekocrv2_weight -r -- --nocapture
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let model_path = format!("{}/deepseek-ai/DeepSeek-OCR-2/", save_dir);
+    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("qwen2_model") {
+                println!("=== {} === {:?}", key, tensor.shape());
+            }
+            // println!("=== {} === {:?}", key, tensor.shape());
+        }
+    }
+    println!("model_list: {:?}", model_list);
+    Ok(())
+}