tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add deepseek-ocr2 support

Browse Source
This commit is contained in:
jhqxxx
2026-03-18 19:15:21 +08:00
parent ca57a7245c
commit 6890f86030
22 changed files with 708 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/mod.rs
+2 -1
View File
@@ -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
+5
View File
@@ -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
+8 -1
View File
@@ -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
+19 -3
View File
@@ -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
+232 -64
View File
@@ -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::tensor_utils::{
index_select_2d, masked_scatter_dim0, nonzero, onehot, prepare_causal_attention_mask,
quick_gelu, topk,
utils::{
interpolate::{interpolate_bicubic, interpolate_linear_1d},
tensor_utils::{
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(
vb_m.pp("vision_model"),
224,
14,
3,
24,
1024,
16,
4096,
1e-5,
)?;
let (vision_model, image_newline) = if version == 2 {
// v2
let qwen2 = Qwen2Decoder2Encoder::new(vb_m.pp("qwen2_model"))?;
(VisionModel::Qwen2(qwen2), None)
} else {
let vision_model = VitModel::new(
vb_m.pp("vision_model"),
224,
14,
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
.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 = match &self.vision_model {
VisionModel::Vit(vit) => {
let local_feature_2 =
vit.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..))?;
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
.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 = match &self.vision_model {
VisionModel::Vit(vit) => {
let global_features_2 =
vit.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..))?;
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 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, local_features) = if let Some(image_newline) =
&self.image_newline
{
let h = (hw as f32).sqrt() as usize;
let w = h;
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
.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 = match &self.vision_model {
VisionModel::Vit(vit) => {
let global_features_2 =
vit.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..))?;
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 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 = if let Some(image_newline) = &self.image_newline {
let h = (hw as f32).sqrt() as usize;
let w = h;
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))?;
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
+26 -7
View File
@@ -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
+9
View File
@@ -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
+275
View File
@@ -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
+4 -2
View File
@@ -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
+3 -1
View File
@@ -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,