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,