temporary save

src/models/deepseek_ocr/config.rs

@@ -0,0 +1,103 @@
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct DeepseekV2Config {
+    pub bos_token_id: u32,
+    pub eos_token_id: u32,
+    pub first_k_dense_replace: u32,
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub kv_lora_rank: Option<usize>,
+    pub lm_head: bool,
+    pub max_position_embeddings: usize,
+    pub moe_intermediate_size: usize,
+    pub n_group: usize,
+    pub n_routed_experts: usize,
+    pub n_shared_experts: usize,
+    pub num_attention_heads: usize,
+    pub num_experts_per_tok: usize,
+    pub num_hidden_layers: usize,
+    pub num_key_value_heads: usize,
+    pub q_lora_rank: Option<usize>,
+    pub qk_nope_head_dim: usize,
+    pub qk_rope_head_dim: usize,
+    pub rm_head: bool,
+    pub topk_group: usize,
+    pub topk_method: String,
+    pub torch_dtype: String,
+    pub use_mla: bool,
+    pub v_head_dim: usize,
+    pub vocab_size: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct ProjectorConfig {
+    pub input_dim: usize,
+    pub model_type: String,
+    pub n_embed: usize,
+    pub projector_type: String,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct ClipL14_224 {
+    pub heads: usize,
+    pub image_size: usize,
+    pub layers: usize,
+    pub patch_size: usize,
+    pub width: usize
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct SamVitB {
+    pub downsample_channels: Vec<usize>,
+    pub global_attn_indexes: Vec<u32>,
+    pub heads: usize,
+    pub layers: usize,
+    pub width: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Width {
+    #[serde(rename = "clip-l-14-224")]
+    pub clip_l_14_224: ClipL14_224,
+    pub sam_vit_b: SamVitB,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct DeepseekOCRVisionConfig {
+    pub image_size: usize,
+    pub mlp_ratio: f32,
+    pub width: Width
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct DeepseekOCRConfig {
+    pub language_config: DeepseekV2Config,
+    pub projector_config: ProjectorConfig,
+    pub torch_dtype: String,
+    pub vision_config: DeepseekOCRVisionConfig,
+    pub bos_token_id: u32,
+    pub eos_token_id: u32,
+    pub first_k_dense_replace: u32,
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub kv_lora_rank: Option<usize>,
+    pub lm_head: bool,
+    pub max_position_embeddings: usize,
+    pub moe_intermediate_size: usize,
+    pub n_group: usize,
+    pub n_routed_experts: usize,
+    pub n_shared_experts: usize,
+    pub num_attention_heads: usize,
+    pub num_experts_per_tok: usize,
+    pub num_hidden_layers: usize,
+    pub num_key_value_heads: usize,
+    pub q_lora_rank: Option<usize>,
+    pub qk_nope_head_dim: usize,
+    pub qk_rope_head_dim: usize,
+    pub rm_head: bool,
+    pub topk_group: usize,
+    pub topk_method: String,
+    pub use_mla: bool,
+    pub v_head_dim: usize,
+    pub vocab_size: usize,
+}

src/models/deepseek_ocr/generate.rs

@@ -0,0 +1,38 @@
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{DType, Device};
+
+use crate::{
+    models::deepseek_ocr::{config::DeepseekOCRConfig, processor::DeepseekOCRProcessor},
+    tokenizer::TokenizerModel,
+    utils::{get_device, get_dtype},
+};
+
+pub struct DeepseekOCRGenerateModel {
+    tokenizer: TokenizerModel,
+    processor: DeepseekOCRProcessor,
+}
+
+impl DeepseekOCRGenerateModel {
+    pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
+        let tokenizer = TokenizerModel::init(path)?;
+        let device = &get_device(device);
+        let dtype = get_dtype(dtype, "bfloat16");
+        let processor = DeepseekOCRProcessor::new(device, dtype)?;
+        let config_path = path.to_string() + "/config.json";
+        let cfg: DeepseekOCRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
+        
+        Ok(Self {
+            tokenizer,
+            processor,
+        })
+    }
+
+    pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> {
+        let (input_ids, images_ori, image_crop, image_seq_mask, images_spatial_crop_t) = self
+            .processor
+            .process_info(&mes, &self.tokenizer, 640, 640, true)?;
+
+        Ok(())
+    }
+}

src/models/deepseek_ocr/mod.rs

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

src/models/deepseek_ocr/model.rs

@@ -0,0 +1,158 @@
+use anyhow::{Ok, Result};
+use candle_core::{IndexOp, Tensor};
+use candle_nn::{
+    Conv2d, Conv2dConfig, Init, LayerNorm, Linear, Module, VarBuilder, conv2d, linear,
+    linear_no_bias,
+};
+
+use crate::models::deepseek_ocr::config::DeepseekOCRConfig;
+
+pub struct PatchEmbed {
+    proj: Conv2d,
+}
+
+impl PatchEmbed {
+    pub fn new(
+        vb: VarBuilder,
+        in_chans: usize,
+        embed_dim: usize,
+        kernel_size: usize,
+        stride: usize,
+        padding: usize,
+    ) -> Result<Self> {
+        let cfg = Conv2dConfig {
+            padding,
+            stride,
+            dilation: 1,
+            groups: 1,
+            cudnn_fwd_algo: None,
+        };
+        let proj = conv2d(in_chans, embed_dim, kernel_size, cfg, vb.pp("proj"))?;
+        Ok(Self { proj })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.proj.forward(xs)?;
+        let xs = xs.permute((0, 2, 3, 1))?;
+        Ok(xs)
+    }
+}
+
+pub struct Attention {
+    num_heads: usize,
+    head_dim: usize,
+    qkv: Linear,
+    proj: Linear,
+    scaling: f64,
+    use_rel_pos: bool,
+    rel_pos_h: Option<Tensor>,
+    rel_pos_w: Option<Tensor>,
+}
+
+impl Attention {
+    pub fn new(
+        vb: VarBuilder,
+        dim: usize,
+        num_heads: usize,
+        qkv_bias: bool,
+        use_rel_pos: bool,
+        input_size: Option<(usize, usize)>,
+    ) -> Result<Self> {
+        let head_dim = dim / num_heads;
+        let scaling = 1.0 / (head_dim as f64).sqrt();
+        let qkv = if qkv_bias {
+            linear(dim, dim * 3, vb.pp("qkv"))?
+        } else {
+            linear_no_bias(dim, dim * 3, vb.pp("qkv"))?
+        };
+        let proj = linear(dim, dim, vb.pp("proj"))?;
+        let mut rel_pos_h = None;
+        let mut rel_pos_w = None;
+        if use_rel_pos {
+            if input_size.is_none() {
+                return Err(anyhow::anyhow!(
+                    "Input size must be provided if using relative positional encoding."
+                ));
+            }
+            let input_size = input_size.unwrap();
+            let h_len = 2 * input_size.0 - 1;
+            let w_len = 2 * input_size.1 - 1;
+            rel_pos_h = Some(vb.get_with_hints((h_len, head_dim), "rel_pos_h", Init::Const(0.))?);
+            rel_pos_w = Some(vb.get_with_hints((w_len, head_dim), "rel_pos_w", Init::Const(0.))?);
+        }
+
+        Ok(Self {
+            num_heads,
+            head_dim,
+            qkv,
+            proj,
+            scaling,
+            use_rel_pos,
+            rel_pos_h,
+            rel_pos_w,
+        })
+    }
+
+    // fn get_rel_pos(q_size: usize, k_size: usize, rel_pos: &Tensor) -> Result<Tensor> {
+    //     let max_rel_dist = 2 * std::cmp::max(q_size, k_size) - 1;
+    //     let rel_pos_resized = if rel_pos.dim(0)? != max_rel_dist {
+    //         let dtype = rel_pos.dtype();
+    //         let rel_pos = rel_pos.to_dtype(candle_core::DType::F32)?;
+    //         let rel_pos_resized = 
+    //     }
+    // }
+    // fn add_decomposed_rel_pos(&self, q: &Tensor, rel_pos_h: &Tensor, rel_pos_w: &Tensor, q_size: (usize, usize), k_size: (usize, usize)) -> Result<Tensor> {
+    //     let (q_h, q_w) = q_size;
+    //     let (k_h, k_w) = k_size;
+        
+    // }
+
+    // pub fn forward(&mut self, xs: &Tensor) -> Result<Tensor> {
+    //     let (b, h, w, _) = xs.dims4()?;
+    //     // (3, B, n_head, h*w, head_dim)
+    //     let qkv = self
+    //         .qkv
+    //         .forward(xs)?
+    //         .reshape((b, h * w, 3, self.num_heads, ()))?
+    //         .permute((2, 0, 3, 1, 4))?
+    //         .contiguous()?;
+    //     let query_states = qkv.i(0)?.contiguous()?;
+    //     let key_states = qkv.i(1)?.contiguous()?;
+    //     let value_states = qkv.i(2)?.contiguous()?;
+    //     let xs = if self.use_rel_pos {
+    //         let (rel_h, rel_w) = 
+    //     } else {
+
+    //     }
+    // }
+}
+
+pub struct Block {
+    norm1: LayerNorm,
+    attn: Attention,
+}
+
+pub struct ImageEncoderViT {
+    img_size: usize,
+    patch_embed: PatchEmbed,
+    pos_embed: Option<Tensor>,
+    blocks: Vec<Block>,
+}
+
+pub struct VitModel {}
+
+pub struct DeepseekV2Model {}
+
+pub struct MlpProjector {}
+
+pub struct DeepseekOCRModel {
+    config: DeepseekOCRConfig,
+    sam_model: ImageEncoderViT,
+    vision_model: VitModel,
+    language_model: DeepseekV2Model,
+    projector: MlpProjector,
+    embed_std: f64,
+    image_newline: Tensor,
+    view_seperator: Tensor,
+    lm_head: Linear,
+}

src/models/deepseek_ocr/processor.rs

@@ -0,0 +1,183 @@
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{DType, Device, Tensor};
+
+use crate::utils::img_utils::dynamic_preprocess;
+use crate::{
+    tokenizer::TokenizerModel,
+    utils::{
+        extract_mes,
+        img_utils::{extract_images, img_transform, resize_with_edge_padding},
+    },
+};
+
+pub struct DeepseekOCRProcessor {
+    device: Device,
+    dtype: DType,
+    image_token: String,
+    image_token_id: u32,
+    patch_size: u32,
+    downsample_ratio: u32,
+}
+
+impl DeepseekOCRProcessor {
+    pub fn new(device: &Device, dtype: DType) -> Result<Self> {
+        Ok(Self {
+            device: device.clone(),
+            dtype,
+            image_token: "<image>".to_string(),
+            image_token_id: 128815,
+            patch_size: 16,
+            downsample_ratio: 4,
+        })
+    }
+
+    fn get_prompt(&self, mes_vec: Vec<(String, String)>) -> Result<String> {
+        let sep = "\n";
+        let sep2 = "";
+        let mut ret = "".to_string();
+        for (i, (_, message)) in mes_vec.iter().enumerate() {
+            if message.chars().count() > 0 {
+                if i % 2 == 0 {
+                    ret = ret + message + sep;
+                } else {
+                    ret = ret + message + sep2;
+                }
+            }
+        }
+        ret = ret.trim().to_string();
+        Ok(ret)
+    }
+
+    pub fn process_info(
+        &self,
+        mes: &ChatCompletionParameters,
+        tokenizer: &TokenizerModel,
+        base_size: u32,
+        image_size: u32,
+        crop_mode: bool,
+    ) -> Result<(Tensor, Tensor, Tensor, Tensor, Tensor)> {
+        let imgs = extract_images(mes)?;
+        let mes_vec = extract_mes(mes)?;
+        let prompt = self.get_prompt(mes_vec.clone())?;
+        let text_splits: Vec<&str> = prompt.split(&self.image_token).collect();
+        let img_mean =
+            Tensor::from_slice(&[0.5, 0.5, 0.5], (3, 1, 1), &self.device)?.to_dtype(self.dtype)?;
+        let img_std =
+            Tensor::from_slice(&[0.5, 0.5, 0.5], (3, 1, 1), &self.device)?.to_dtype(self.dtype)?;
+        let mut images_list = Vec::new();
+        let mut images_crop_list = Vec::new();
+        let mut images_seq_mask = vec![0u32];
+        let mut tokenized_id = vec![0u32];
+        let mut images_spatial_crop = Vec::new();
+        for (text_seq, image) in text_splits.iter().zip(imgs) {
+            if text_seq.len() > 0 {
+                let token_ids = tokenizer.text_encode_vec(text_seq.to_string(), false)?;
+                tokenized_id.extend_from_slice(&token_ids);
+                let seq_mask = vec![0u32; token_ids.len()];
+                images_seq_mask.extend_from_slice(&seq_mask);
+            }
+            if crop_mode {
+                let mut images_crop_raw = Vec::new();
+                let crop_ratio = if image.height() <= 640 && image.width() <= 640 {
+                    (1u32, 1u32)
+                } else {
+                    let (img_crop, ratio) = dynamic_preprocess(&image, image_size, false)?;
+                    images_crop_raw = img_crop.clone();
+                    ratio
+                };
+
+                let gloabal_view =
+                    resize_with_edge_padding(&image, base_size, base_size, [127u8; 3]);
+
+                let global_img_trans =
+                    img_transform(&gloabal_view, &img_mean, &img_std, &self.device, self.dtype)?;
+                images_list.push(global_img_trans);
+
+                images_spatial_crop.push(vec![crop_ratio.0, crop_ratio.1]);
+
+                if crop_ratio.0 > 1 || crop_ratio.1 > 1 {
+                    for img in images_crop_raw {
+                        let img_t =
+                            img_transform(&img, &img_mean, &img_std, &self.device, self.dtype)?;
+                        images_crop_list.push(img_t);
+                    }
+                }
+
+                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;
+                if crop_ratio.0 > 1 || crop_ratio.1 > 1 {
+                    token_repeat += (num_queries * crop_ratio.0 + 1) * (num_queries * crop_ratio.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()];
+                images_seq_mask.extend_from_slice(&seq_mask);
+            } else {
+                let global_view = if image_size <= 640 {
+                    image.resize_exact(
+                        image_size,
+                        image_size,
+                        image::imageops::FilterType::CatmullRom,
+                    )
+                } else {
+                    resize_with_edge_padding(&image, image_size, image_size, [127u8; 3])
+                };
+                let global_img_trans =
+                    img_transform(&global_view, &img_mean, &img_std, &self.device, self.dtype)?;
+                images_list.push(global_img_trans);
+
+                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 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);
+            }
+        }
+        let token_ids =
+            tokenizer.text_encode_vec(text_splits[text_splits.len() - 1].to_string(), false)?;
+        tokenized_id.extend_from_slice(&token_ids);
+        let seq_mask = vec![0u32; token_ids.len()];
+        images_seq_mask.extend_from_slice(&seq_mask);
+        let input_ids = Tensor::new(tokenized_id, &self.device)?.unsqueeze(0)?;
+        let image_seq_mask = Tensor::new(images_seq_mask, &self.device)?;
+        let (images_ori, images_spatial_crop_t, image_crop) = if images_list.len() == 0 {
+            let images_ori = Tensor::zeros(
+                (1usize, 3usize, image_size as usize, image_size as usize),
+                self.dtype,
+                &self.device,
+            )?;
+            let images_spatial_crop_t = Tensor::zeros((1, 2), DType::F64, &self.device)?;
+            let image_crop = Tensor::zeros(
+                (1usize, 3usize, base_size as usize, base_size as usize),
+                self.dtype,
+                &self.device,
+            )?;
+            (images_ori, images_spatial_crop_t, image_crop)
+        } else {
+            let images_ori = Tensor::stack(&images_list, 0)?;
+            let images_spatial_crop_t = Tensor::new(images_spatial_crop, &self.device)?;
+            let image_crop = if images_crop_list.len() > 0 {
+                Tensor::stack(&images_crop_list, 0)?
+            } else {
+                Tensor::zeros(
+                    (1usize, 3usize, base_size as usize, base_size as usize),
+                    self.dtype,
+                    &self.device,
+                )?
+            };
+            (images_ori, images_spatial_crop_t, image_crop)
+        };
+
+        Ok((
+            input_ids,
+            images_ori,
+            image_crop,
+            image_seq_mask,
+            images_spatial_crop_t,
+        ))
+    }
+}

src/models/mod.rs

@@ -3,6 +3,7 @@ pub mod minicpm4;
 pub mod qwen2_5vl;
 pub mod qwen3vl;
 pub mod voxcpm;
+pub mod deepseek_ocr;
 
 use aha_openai_dive::v1::resources::chat::{
     ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,

src/tokenizer/mod.rs

@@ -1,4 +1,4 @@
-use anyhow::{Result, anyhow};
+use anyhow::{Ok, Result, anyhow};
 use candle_core::{Device, Tensor};
 use tokenizers::Tokenizer;
 
@@ -23,13 +23,23 @@ impl TokenizerModel {
         Ok(Self { tokenizer })
     }
 
-    pub fn text_encode(&self, text: String, device: &Device) -> Result<Tensor> {
+    pub fn text_encode_vec(&self, text: String, add_special_token: bool) -> Result<Vec<u32>> {
         let token_id = self
             .tokenizer
-            .encode(text, true)
+            .encode(text, add_special_token)
             .map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))?
             .get_ids()
             .to_vec();
+        Ok(token_id)
+    }
+    pub fn text_encode(&self, text: String, device: &Device) -> Result<Tensor> {
+        // let token_id = self
+        //     .tokenizer
+        //     .encode(text, true)
+        //     .map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))?
+        //     .get_ids()
+        //     .to_vec();
+        let token_id = self.text_encode_vec(text, true)?;
         let token_tensor = Tensor::from_slice(&token_id, (1, token_id.len()), device)?;
         Ok(token_tensor)
     }

src/utils/img_utils.rs

@@ -1,8 +1,13 @@
+use std::collections::HashSet;
 use std::io::Cursor;
 
-use anyhow::{Result, anyhow};
+use aha_openai_dive::v1::resources::chat::{
+    ChatCompletionParameters, ChatMessage, ChatMessageContent, ChatMessageContentPart,
+};
+use anyhow::{Ok, Result, anyhow};
 use base64::{Engine, engine::general_purpose};
-use image::{DynamicImage, ImageReader};
+use candle_core::{DType, Device, Tensor};
+use image::{DynamicImage, ImageBuffer, ImageReader, Rgb, RgbImage, imageops};
 
 pub fn load_image_from_url(url: &str) -> Result<DynamicImage> {
     let response = reqwest::blocking::get(url)
@@ -58,3 +63,173 @@ pub fn get_image(file: &str) -> Result<DynamicImage> {
     }
     Err(anyhow!("get image from message failed".to_string()))
 }
+
+pub fn extract_image_url(mes: &ChatCompletionParameters) -> Result<Vec<String>> {
+    let mut img_vec = Vec::new();
+    for chat_mes in mes.messages.clone() {
+        if let ChatMessage::User { content, .. } = chat_mes
+            && let ChatMessageContent::ContentPart(part_vec) = content
+        {
+            for part in part_vec {
+                if let ChatMessageContentPart::Image(img_part) = part {
+                    let img_url = img_part.image_url;
+                    img_vec.push(img_url.url);
+                }
+            }
+        }
+    }
+    Ok(img_vec)
+}
+
+pub fn extract_images(mes: &ChatCompletionParameters) -> Result<Vec<DynamicImage>> {
+    let img_url_vec = extract_image_url(mes)?;
+    let mut img_vec = Vec::new();
+    for url in img_url_vec {
+        let img = get_image(&url)?;
+        img_vec.push(img);
+    }
+    Ok(img_vec)
+}
+
+pub fn generate_target_ratios_sorted(min_num: u32, max_num: u32) -> Vec<(u32, u32)> {
+    let mut target_ratios = HashSet::new();
+
+    for n in min_num..=max_num {
+        for i in 1..=n {
+            for j in 1..=n {
+                let product = i * j;
+                if product <= max_num && product >= min_num {
+                    target_ratios.insert((i, j));
+                }
+            }
+        }
+    }
+    // Convert to vector and sort by the product of elements (i*j)
+    let mut sorted_ratios: Vec<(u32, u32)> = target_ratios.into_iter().collect();
+    sorted_ratios.sort_by_key(|&(i, j)| i * j);
+
+    sorted_ratios
+}
+
+pub fn find_closest_aspect_ratio(
+    aspect_ratio: f64,
+    target_ratios: &[(u32, u32)],
+    width: u32,
+    height: u32,
+    image_size: u32,
+) -> (u32, u32) {
+    let mut best_ratio_diff = f64::INFINITY;
+    let mut best_ratio = (1, 1);
+    let area = width * height;
+
+    for &ratio in target_ratios {
+        let target_aspect_ratio = ratio.0 as f64 / ratio.1 as f64;
+        let ratio_diff = (aspect_ratio - target_aspect_ratio).abs();
+
+        if ratio_diff < best_ratio_diff {
+            best_ratio_diff = ratio_diff;
+            best_ratio = ratio;
+        } else if (ratio_diff - best_ratio_diff).abs() < 1e-10 {
+            let target_area = 0.5 * (image_size as f64).powi(2) * (ratio.0 * ratio.1) as f64;
+            if area as f64 > target_area {
+                best_ratio = ratio;
+            }
+        }
+    }
+
+    best_ratio
+}
+
+pub fn dynamic_preprocess(
+    image: &DynamicImage,
+    image_size: u32,
+    use_thumbnail: bool,
+) -> Result<(Vec<DynamicImage>, (u32, u32))> {
+    let orig_width = image.width();
+    let orig_height = image.height();
+    let aspect_ratio = orig_width as f64 / orig_height as f64;
+    let target_ratios = generate_target_ratios_sorted(2, 9);
+    let target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio,
+        &target_ratios,
+        orig_width,
+        orig_height,
+        image_size,
+    );
+    let target_width = image_size * target_aspect_ratio.0;
+    let target_height = image_size * target_aspect_ratio.1;
+    let blocks = target_aspect_ratio.0 * target_aspect_ratio.1;
+    let mut resized_img = image.resize_exact(
+        target_width,
+        target_height,
+        image::imageops::FilterType::CatmullRom,
+    );
+    let mut processed_images = Vec::new();
+    let grid_width = target_width / image_size;
+    for i in 0..blocks {
+        // Calculate box coordinates
+        let x1 = (i % grid_width) * image_size;
+        let y1 = (i / grid_width) * image_size;
+
+        // Crop the image
+        let split_img = resized_img.crop(x1, y1, image_size, image_size);
+        processed_images.push(split_img);
+    }
+    assert_eq!(processed_images.len() as u32, blocks);
+
+    if use_thumbnail && processed_images.len() != 1 {
+        let thumbnail_img = image.resize_exact(
+            image_size,
+            image_size,
+            image::imageops::FilterType::CatmullRom,
+        );
+        processed_images.push(thumbnail_img);
+    }
+    Ok((processed_images, target_aspect_ratio))
+}
+
+pub fn resize_with_edge_padding(
+    img: &DynamicImage,
+    width: u32,
+    height: u32,
+    color: [u8; 3],
+) -> DynamicImage {
+    // 按图像原比例resize,可能不是输入的宽高
+    let mut img = img.resize(width, height, image::imageops::FilterType::CatmullRom);
+    // 使用全0像素填充为输入宽高
+    if img.height() != height || img.width() != width {
+        let (img_h, img_w) = (img.height(), img.width());
+        let img_buffer = img.to_rgb8();
+        let mut canvas: ImageBuffer<Rgb<u8>, Vec<u8>> =
+            RgbImage::from_pixel(width, height, Rgb(color));
+        let x_offset = (width - img_w) / 2;
+        let y_offset = (height - img_h) / 2;
+        imageops::overlay(&mut canvas, &img_buffer, x_offset as i64, y_offset as i64);
+        img = DynamicImage::ImageRgb8(canvas);
+    }
+    img
+}
+
+pub fn img_transform(
+    img: &DynamicImage,
+    mean: &Tensor,
+    std: &Tensor,
+    device: &Device,
+    dtype: DType,
+) -> Result<Tensor> {
+    let img_h = img.height();
+    let img_w = img.width();
+    let img_vec = img.to_rgb8().into_raw();
+    // (h, w, c) => (c, h, w)
+    let img_tensor = Tensor::from_slice(&img_vec, (img_h as usize, img_w as usize, 3), device)?
+        .permute((2, 0, 1))?
+        .to_dtype(DType::F32)?;
+    // 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(&mean.to_dtype(DType::F32)?)?
+        .broadcast_div(&std.to_dtype(DType::F32)?)?
+        .to_dtype(dtype)?;
+    Ok(img_tensor)
+}

src/utils/mod.rs

@@ -5,9 +5,7 @@ pub mod video_utils;
 
 use aha_openai_dive::v1::resources::{
     chat::{
-        ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse,
-        ChatCompletionResponse, ChatMessage, ChatMessageContent, DeltaChatMessage, DeltaFunction,
-        DeltaToolCall, Function, ToolCall,
+        ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent, ChatMessageContentPart, DeltaChatMessage, DeltaFunction, DeltaToolCall, Function, ToolCall
     },
     shared::FinishReason,
 };
@@ -281,3 +279,25 @@ pub fn get_logit_processor(
         }
     }
 }
+
+pub fn extract_mes(mes: &ChatCompletionParameters) -> Result<Vec<(String, String)>> {
+    let mut mes_vec = Vec::new();
+    for chat_mes in mes.messages.clone() {
+        if let ChatMessage::User { content, .. } = chat_mes.clone()
+            && let ChatMessageContent::ContentPart(part_vec) = content
+        {
+            for part in part_vec {
+                if let ChatMessageContentPart::Text(text_part) = part {
+                    let text = text_part.text;
+                    mes_vec.push(("<|User|>".to_string(), text));
+                }
+            }
+        } else if let ChatMessage::Assistant { content, .. } = chat_mes.clone()
+            && let Some(cont) = content
+            && let ChatMessageContent::Text(c) = cont
+        {
+            mes_vec.push(("<|Assistant|>".to_string(), c));
+        }
+    }
+    Ok(mes_vec)
+}

tests/test_deepseek_ocr.rs

@@ -0,0 +1,43 @@
+use aha::models::deepseek_ocr::{generate::DeepseekOCRGenerateModel, processor::DeepseekOCRProcessor};
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{DType, Device, IndexOp, Tensor};
+
+#[test]
+fn deepseek_ocr_test() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_test -- --nocapture
+    let message = r#"
+    {
+        "model": "deepseek-ocr",
+        "messages": [
+            {
+                "role": "user",
+                "content": [ 
+                    {
+                        "type": "image",
+                        "image_url": 
+                        {
+                            "url": "file://./assets/img/ocr_test1.png"
+                        }
+                    },              
+                    {
+                        "type": "text", 
+                        "text": "<image>\n<|grounding|>Convert the document to markdown. "
+                    }
+                ]
+            },
+            {
+                "role": "assistant",
+                "content": ""
+            }
+        ]
+    }
+    "#;
+    let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/";
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let device = Device::cuda_if_available(0)?;
+    let dtype = DType::BF16;
+    let mut model = DeepseekOCRGenerateModel::init(model_path, Some(&device), Some(dtype))?;
+    let res = model.generate(mes)?;
+    Ok(())
+}