add deepseek_ocr

Cargo.toml

@@ -31,3 +31,6 @@ hound = "3.5.1"
 [features]
 flash-attn=["candle-flash-attn"]
 cuda=["candle-nn/cuda", "candle-core/cuda", "candle-transformers/cuda"]
+
+[lints.clippy]
+needless_range_loop = "allow"

clippy.toml

@@ -6,6 +6,6 @@ disallowed-macros = [
     { path = "lazy_static::lazy_static", reason = "Please use `std::sync::LazyLock` instead." },
 ]
 
-too-many-arguments-threshold = 10
+too-many-arguments-threshold = 20
 upper-case-acronyms-aggressive = false
-enum-variant-size-threshold = 200
+enum-variant-size-threshold = 200

src/models/common/mod.rs

@@ -286,6 +286,9 @@ pub fn eager_attention_forward(
         Some(g) => repeat_kv(value_states.clone(), g)?.contiguous()?,
         None => value_states.clone(),
     };
+    let query_states = query_states.contiguous()?;
+    let key_states = key_states.contiguous()?;
+    let value_states = value_states.contiguous()?;
     let attn_output = {
         #[cfg(not(feature = "flash-attn"))]
         {

src/models/deepseek_ocr/config.rs

@@ -1,15 +1,26 @@
-
 #[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 first_k_dense_replace: usize,
     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,
+    #[serde(default = "default_moe_layer_freq")]
+    pub moe_layer_freq: usize,
+    #[serde(default = "default_routed_scaling_factor")]
+    pub routed_scaling_factor: f64,
+    #[serde(default = "default_scoring_func")]
+    pub scoring_func: String,
+    #[serde(default = "default_aux_loss_alpha")]
+    pub aux_loss_alpha: f32,
+    #[serde(default = "default_true")]
+    pub seq_aux: bool,
+    #[serde(default = "default_false")]
+    pub norm_topk_prob: bool,
     pub n_group: usize,
     pub n_routed_experts: usize,
     pub n_shared_experts: usize,
@@ -27,6 +38,30 @@ pub struct DeepseekV2Config {
     pub use_mla: bool,
     pub v_head_dim: usize,
     pub vocab_size: usize,
+    #[serde(default = "default_rms_norm_eps")]
+    pub rms_norm_eps: f64,
+}
+
+fn default_moe_layer_freq() -> usize {
+    1
+}
+fn default_routed_scaling_factor() -> f64 {
+    1.0
+}
+fn default_scoring_func() -> String {
+    "softmax".to_string()
+}
+fn default_aux_loss_alpha() -> f32 {
+    0.001
+}
+fn default_true() -> bool {
+    true
+}
+fn default_false() -> bool {
+    false
+}
+fn default_rms_norm_eps() -> f64 {
+    1e-6
 }
 
 #[derive(Debug, Clone, PartialEq, serde::Deserialize)]
@@ -43,13 +78,13 @@ pub struct ClipL14_224 {
     pub image_size: usize,
     pub layers: usize,
     pub patch_size: usize,
-    pub width: 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 global_attn_indexes: Vec<usize>,
     pub heads: usize,
     pub layers: usize,
     pub width: usize,
@@ -66,7 +101,7 @@ pub struct Width {
 pub struct DeepseekOCRVisionConfig {
     pub image_size: usize,
     pub mlp_ratio: f32,
-    pub width: Width
+    pub width: Width,
 }
 
 #[derive(Debug, Clone, PartialEq, serde::Deserialize)]
@@ -100,4 +135,4 @@ pub struct DeepseekOCRConfig {
     pub use_mla: bool,
     pub v_head_dim: usize,
     pub vocab_size: usize,
-}
+}

src/models/deepseek_ocr/generate.rs

@@ -1,38 +1,168 @@
-use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
-use anyhow::Result;
-use candle_core::{DType, Device};
+use aha_openai_dive::v1::resources::chat::{
+    ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
+};
+use anyhow::{Result, anyhow};
+use candle_core::{DType, Device, Tensor};
+use candle_nn::VarBuilder;
+use rocket::async_stream::stream;
+use rocket::futures::Stream;
 
 use crate::{
-    models::deepseek_ocr::{config::DeepseekOCRConfig, processor::DeepseekOCRProcessor},
+    models::{
+        GenerateModel,
+        deepseek_ocr::{
+            config::DeepseekOCRConfig, model::DeepseekOCRModel, processor::DeepseekOCRProcessor,
+        },
+    },
     tokenizer::TokenizerModel,
-    utils::{get_device, get_dtype},
+    utils::{
+        build_completion_chunk_response, build_completion_response, find_type_files, get_device,
+        get_dtype, get_logit_processor,
+    },
 };
 
 pub struct DeepseekOCRGenerateModel {
     tokenizer: TokenizerModel,
     processor: DeepseekOCRProcessor,
+    deepseekocr_model: DeepseekOCRModel,
+    bos_token_id: u32,
+    eos_token_id: u32,
+    device: Device,
 }
 
 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)?)?;
-        
+        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 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)?;
         Ok(Self {
             tokenizer,
             processor,
+            deepseekocr_model,
+            bos_token_id,
+            eos_token_id,
+            device: device.clone(),
         })
     }
+}
 
-    pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> {
-        let (input_ids, images_ori, image_crop, image_seq_mask, images_spatial_crop_t) = self
+impl GenerateModel for DeepseekOCRGenerateModel {
+    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None);
+        let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self
+            .processor
+            .process_info(&mes, &self.tokenizer, 640, 640, true)?;
+        let mut images_ori = Some(&images_ori);
+        let mut image_crop = Some(&image_crop);
+        let mut images_seq_mask = Some(&images_seq_mask);
+        let mut images_spatial_crop_t = Some(&images_spatial_crop_t);
+        let mut seqlen_offset = 0;
+        let mut seq_len = input_ids.dim(1)?;
+        let mut generate = Vec::new();
+        let sample_len = mes.max_tokens.unwrap_or(1024);
+        for _ in 0..sample_len {
+            let logits = self.deepseekocr_model.forward(
+                &input_ids,
+                images_ori,
+                image_crop,
+                images_seq_mask,
+                images_spatial_crop_t,
+                seqlen_offset,
+            )?;
+            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+            let next_token = logit_processor.sample(&logits)?;
+            generate.push(next_token);
+            if next_token == self.bos_token_id || next_token == self.eos_token_id {
+                break;
+            }
+            seqlen_offset += seq_len;
+            seq_len = 1;
+            input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+            images_ori = None;
+            image_crop = None;
+            images_seq_mask = None;
+            images_spatial_crop_t = None;
+        }
+        let res = self.tokenizer.token_decode(generate)?;
+        self.deepseekocr_model.clear_kv_cache();
+        let response = build_completion_response(res, "deepseek_ocr");
+        Ok(response)
+    }
+
+    fn generate_stream(
+        &mut self,
+        mes: ChatCompletionParameters,
+    ) -> Result<impl Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>> {
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None);
+        let (mut input_ids, images_ori, image_crop, images_seq_mask, images_spatial_crop_t) = self
             .processor
             .process_info(&mes, &self.tokenizer, 640, 640, true)?;
 
-        Ok(())
+        let mut seqlen_offset = 0;
+        let mut seq_len = input_ids.dim(1)?;
+        let sample_len = mes.max_tokens.unwrap_or(1024);
+        let stream = stream! {
+            let mut error_tokens = Vec::new();
+            let mut images_ori = Some(&images_ori);
+            let mut image_crop = Some(&image_crop);
+            let mut images_seq_mask = Some(&images_seq_mask);
+            let mut images_spatial_crop_t = Some(&images_spatial_crop_t);
+            for _ in 0..sample_len {
+                let logits = self.deepseekocr_model.forward(
+                    &input_ids,
+                    images_ori,
+                    image_crop,
+                    images_seq_mask,
+                    images_spatial_crop_t,
+                    seqlen_offset,
+                )?;
+                let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+                let next_token = logit_processor.sample(&logits)?;
+                let mut decode_ids = Vec::new();
+                if !error_tokens.is_empty() {
+                    decode_ids.extend_from_slice(&error_tokens);
+                }
+                decode_ids.push(next_token);
+                let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{}", e)))?;
+                if decoded_token.contains("�") {
+                    error_tokens.push(next_token);
+                    if error_tokens.len() > 3 {
+                        error_tokens.clear();
+                    }
+                    seqlen_offset += seq_len;
+                    seq_len = 1;
+                    input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                    images_ori = None;
+                    image_crop = None;
+                    images_seq_mask = None;
+                    images_spatial_crop_t = None;
+                    continue;
+                }
+                error_tokens.clear();
+                let chunk = build_completion_chunk_response(decoded_token, "deepseek_ocr", None, None);
+                yield Ok(chunk);
+                if next_token == self.bos_token_id || next_token == self.eos_token_id {
+                    break;
+                }
+                seqlen_offset += seq_len;
+                seq_len = 1;
+                input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                images_ori = None;
+                image_crop = None;
+                images_seq_mask = None;
+                images_spatial_crop_t = None;
+            }
+            self.deepseekocr_model.clear_kv_cache();
+        };
+        Ok(stream)
     }
 }

src/models/deepseek_ocr/mod.rs

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

src/models/deepseek_ocr/processor.rs

@@ -71,7 +71,7 @@ impl DeepseekOCRProcessor {
         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 {
+            if !text_seq.is_empty() {
                 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()];
@@ -143,8 +143,8 @@ impl DeepseekOCRProcessor {
         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 image_seq_mask = Tensor::new(images_seq_mask, &self.device)?.unsqueeze(0)?;
+        let (images_ori, images_spatial_crop_t, image_crop) = if images_list.is_empty() {
             let images_ori = Tensor::zeros(
                 (1usize, 3usize, image_size as usize, image_size as usize),
                 self.dtype,
@@ -160,7 +160,7 @@ impl DeepseekOCRProcessor {
         } 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 {
+            let image_crop = if !images_crop_list.is_empty() {
                 Tensor::stack(&images_crop_list, 0)?
             } else {
                 Tensor::zeros(

src/models/mod.rs

@@ -1,9 +1,9 @@
 pub mod common;
+pub mod deepseek_ocr;
 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/position_embed/rope.rs

@@ -272,3 +272,34 @@ impl Qwen3VLTextRotaryEmbedding {
         Ok((cos.to_dtype(dtype)?, sin.to_dtype(dtype)?))
     }
 }
+
+pub struct RoPE {
+    inv_freq: Tensor, // (1, dim / 2)
+}
+
+impl RoPE {
+    pub fn new(dim: usize, theta_base: f32, device: &Device) -> Result<Self> {
+        let inv_freq = compute_default_rope_parameters(dim, theta_base);
+        let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?;
+
+        Ok(Self { inv_freq })
+    }
+    pub fn forward(
+        &self,
+        seqlen_offset: usize,
+        seq_len: usize,
+        device: &Device,
+    ) -> Result<(Tensor, Tensor)> {
+        let positions = Tensor::arange(
+            seqlen_offset as f32,
+            (seqlen_offset + seq_len) as f32,
+            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 cos = emb.cos()?;
+        let sin = emb.sin()?;
+        Ok((cos, sin))
+    }
+}

src/utils/mod.rs

@@ -5,7 +5,9 @@ pub mod video_utils;
 
 use aha_openai_dive::v1::resources::{
     chat::{
-        ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent, ChatMessageContentPart, DeltaChatMessage, DeltaFunction, DeltaToolCall, Function, ToolCall
+        ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse,
+        ChatCompletionParameters, ChatCompletionResponse, ChatMessage, ChatMessageContent,
+        ChatMessageContentPart, DeltaChatMessage, DeltaFunction, DeltaToolCall, Function, ToolCall,
     },
     shared::FinishReason,
 };

src/utils/tensor_utils.rs

@@ -1,6 +1,6 @@
-use anyhow::{Ok, Result, anyhow};
+use anyhow::{Result, anyhow};
 use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim};
-use rocket::figment::value;
+use candle_nn::ops::sigmoid;
 
 pub fn prepare_causal_attention_mask(
     b_size: usize,
@@ -15,8 +15,11 @@ pub fn prepare_causal_attention_mask(
     // let mask = Tensor::from_vec(mask, (tgt_len, tgt_len), device)?;
     let arange = Tensor::arange(0u32, tgt_len as u32, device)?;
     let arange = arange.unsqueeze(1)?.broadcast_as((tgt_len, tgt_len))?;
-    let upper_triangle = arange.t()?.lt(&arange)?.to_dtype(DType::F32)?;
-    let mask = upper_triangle.where_cond(&Tensor::new(f32::NEG_INFINITY, device)?, &Tensor::new(0f32, device)?)?;
+    let upper_triangle = arange.t()?.gt(&arange)?;
+    let mask = upper_triangle.where_cond(
+        &Tensor::new(f32::NEG_INFINITY, device)?.broadcast_as(arange.shape())?,
+        &Tensor::new(0f32, device)?.broadcast_as(arange.shape())?,
+    )?;
     let mask = if seqlen_offset > 0 {
         let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::F32, device)?;
         Tensor::cat(&[&mask0, &mask], D::Minus1)?
@@ -352,51 +355,62 @@ pub fn mask_index_add(original: &Tensor, mask: &Tensor, add: &Tensor) -> Result<
     Ok(xs)
 }
 
-pub fn interpolate_linear(
+pub fn compute_1d_coords(
+    input_size: usize,
+    output_size: usize,
+    align_corner: Option<bool>,
+) -> Result<Vec<f32>> {
+    if input_size == 1 {
+        Ok(vec![0f32; output_size])
+    } else if let Some(align_) = align_corner
+        && align_
+    {
+        Ok((0..output_size)
+            .map(|i| i as f32 * (input_size - 1) as f32 / (output_size - 1) as f32)
+            .collect())
+    } else {
+        Ok((0..output_size)
+            .map(|i| {
+                (i as f32 + 0.5) * (input_size as f32 / output_size as f32) - 0.5
+                // coord.max(0.0).min((input_size - 1) as f32)
+            })
+            .collect())
+    }
+}
+
+pub fn interpolate_linear_1d(
     t: &Tensor,
     target_size: usize,
     align_corner: Option<bool>,
 ) -> Result<Tensor> {
     // t: [b, channels, features]
+    if t.rank() < 3 {
+        return Err(anyhow::anyhow!(
+            "Input rank must have at least 3 dimensions"
+        ));
+    }
     let shape = t.dims();
     let orig_size = shape[shape.len() - 1];
     if orig_size == target_size {
         return Ok(t.clone());
     }
     let mut reshaped = t.clone();
-    if shape.len() != 3 {
+    if shape.len() > 3 {
         let bs = shape[0];
         let channels = shape[1..shape.len() - 1].iter().product::<usize>();
         reshaped = reshaped.reshape((bs, channels, orig_size))?;
     }
     let (bs, channels, _) = reshaped.dims3()?;
-    let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), &t.device())?;
-    let coords = if orig_size == 1 {
-        vec![0f32; target_size]
-    } else {
-        let coords_vec = if let Some(align_) = align_corner
-            && align_
-        {
-            (0..target_size)
-                .map(|i| i as f32 * (orig_size - 1) as f32 / (target_size - 1) as f32)
-                .collect()
-        } else {            
-            (0..target_size)
-                .map(|i| {
-                    let coord = (i as f32 + 0.5) * (orig_size as f32 / target_size as f32) - 0.5;
-                    coord.max(0.0).min((orig_size-1) as f32)
-                })
-                .collect()
-        };
-        coords_vec
-    };
+    let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), t.device())?;
+    let coords = compute_1d_coords(orig_size, target_size, align_corner)?;
 
     for b in 0..bs {
         for c in 0..channels {
             let input_slice = reshaped.i((b, c))?;
             let mut out_i = Vec::new();
-            for x_out in 0..target_size {
-                let coord = coords[x_out];
+            // for x_out in 0..target_size {
+            for &coord in coords.iter().take(target_size) {
+                let coord = if coord < 0.0 { 0.0 } else { coord };
                 let x0 = coord.floor() as usize;
                 let x1 = std::cmp::min(x0 + 1, orig_size - 1);
                 let weight = (coord - x0 as f32) as f64;
@@ -407,25 +421,268 @@ pub fn interpolate_linear(
                 out_i.push(interpolated);
             }
             let out_i = Tensor::stack(&out_i, 0)?.unsqueeze(0)?.unsqueeze(0)?;
-            output = output.slice_assign(&[(b..b+1), (c..c+1), (0..target_size)], &out_i)?;
+            output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?;
         }
     }
     if shape.len() != 3 {
         let mut new_shape = shape.to_vec();
-        let last_dim = new_shape.len()-1;
+        let last_dim = new_shape.len() - 1;
         new_shape[last_dim] = target_size;
         output = output.reshape(new_shape)?
-
     }
     output = output.contiguous()?;
     Ok(output)
 }
 
+fn compute_scale(input_size: usize, output_size: usize, align_corners: bool) -> f64 {
+    if align_corners && output_size > 1 {
+        (input_size - 1) as f64 / (output_size - 1) as f64
+    } else {
+        input_size as f64 / output_size as f64
+    }
+}
+
+fn bicubic_filter(x: f64) -> f64 {
+    let a = -0.75;
+    let x = x.abs();
+    if x < 1.0 {
+        ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0
+    } else if x < 2.0 {
+        (((x - 5.0) * x + 8.0) * x - 4.0) * a
+    } else {
+        0.0
+    }
+}
+
+pub fn interpolate_bicubic_antialias(
+    input: &Tensor,
+    batch_size: usize,
+    channels: usize,
+    input_height: usize,
+    input_width: usize,
+    output_height: usize,
+    output_width: usize,
+    height_scale: f64,
+    width_scale: f64,
+    align_corners: bool,
+) -> Result<Tensor> {
+    // tensor没有to_vec4, 所以把bs和channels先合在一起
+    let dim0 = batch_size * channels;
+    let input_3dim = input.reshape((dim0, input_height, input_width))?;
+    let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
+    let mut output_data = vec![vec![vec![0.0f32; output_width]; output_height]; dim0];
+    let support = 2.0 * height_scale.max(width_scale);
+    for c in 0..dim0 {
+        for out_y in 0..output_height {
+            let center_y = if align_corners {
+                out_y as f64 * height_scale
+            } else {
+                (out_y as f64 + 0.5) * height_scale - 0.5
+            };
+            let start_y = (center_y - support).ceil() as isize;
+            let end_y = (center_y + support).floor() as isize;
+            for out_x in 0..output_width {
+                let center_x = if align_corners {
+                    out_x as f64 * width_scale
+                } else {
+                    (out_x as f64 + 0.5) * width_scale - 0.5
+                };
+                let mut sum = 0.0;
+                let mut weight_sum = 0.0;
+                let start_x = (center_x - support).ceil() as isize;
+                let end_x = (center_x + support).floor() as isize;
+                for iy in start_y..end_y {
+                    for ix in start_x..end_x {
+                        if iy >= 0
+                            && iy < input_height as isize
+                            && ix >= 0
+                            && ix < input_width as isize
+                        {
+                            let dx = (ix as f64 - center_x).abs();
+                            let dy = (iy as f64 - center_y).abs();
+                            let wx = bicubic_filter(dx / width_scale.max(1.0));
+                            let wy = bicubic_filter(dy / height_scale.max(1.0));
+                            let weight = (wx * wy) as f32;
+                            sum += input_data[c][iy as usize][ix as usize] * weight;
+                            weight_sum += weight;
+                        }
+                    }
+                }
+                if weight_sum > 0.0 {
+                    output_data[c][out_y][out_x] = sum / weight_sum;
+                } else {
+                    output_data[c][out_y][out_x] = 0.0;
+                }
+            }
+        }
+    }
+    let output = Tensor::new(output_data, input.device())?
+        .reshape((batch_size, channels, output_height, output_width))?
+        .to_dtype(input.dtype())?;
+    Ok(output)
+}
+
+fn get_cubic_coefficients(t: f64) -> [f64; 4] {
+    let a = -0.75;
+
+    let x1 = t;
+    let coeff0 = cubic_convolution2(x1 + 1.0, a);
+    let coeff1 = cubic_convolution1(x1, a);
+
+    let x2 = 1.0 - t;
+    let coeff2 = cubic_convolution1(x2, a);
+    let coeff3 = cubic_convolution2(x2 + 1.0, a);
+
+    [coeff0, coeff1, coeff2, coeff3]
+}
+
+// 三次卷积函数1
+fn cubic_convolution1(x: f64, a: f64) -> f64 {
+    ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0
+}
+
+// 三次卷积函数2
+fn cubic_convolution2(x: f64, a: f64) -> f64 {
+    ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a
+}
+
+fn cubic_interp1d(x0: f32, x1: f32, x2: f32, x3: f32, t: f64) -> f32 {
+    let coeffs = get_cubic_coefficients(t);
+    x0 * coeffs[0] as f32 + x1 * coeffs[1] as f32 + x2 * coeffs[2] as f32 + x3 * coeffs[3] as f32
+}
+
+pub fn interpolate_bicubic_standard(
+    input: &Tensor,
+    batch_size: usize,
+    channels: usize,
+    input_height: usize,
+    input_width: usize,
+    output_height: usize,
+    output_width: usize,
+    height_scale: f64,
+    width_scale: f64,
+    align_corners: bool,
+) -> Result<Tensor> {
+    // tensor没有to_vec4, 所以把bs和channels先合在一起
+    let dim0 = batch_size * channels;
+    let input_3dim = input.reshape((dim0, input_height, input_width))?;
+    let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
+    let mut output_data = vec![vec![vec![0.0f32; output_width]; output_height]; dim0];
+    for c in 0..dim0 {
+        for out_y in 0..output_height {
+            let center_y = if align_corners {
+                out_y as f64 * height_scale
+            } else {
+                (out_y as f64 + 0.5) * height_scale - 0.5
+            };
+            let in_y = center_y.floor() as isize;
+            let t_y = center_y - in_y as f64;
+            for out_x in 0..output_width {
+                let center_x = if align_corners {
+                    out_x as f64 * width_scale
+                } else {
+                    (out_x as f64 + 0.5) * width_scale - 0.5
+                };
+                let in_x = center_x.floor() as isize;
+                let t_x = center_x - in_x as f64;
+                let mut coefficients = [0.0; 4];
+                // for k in 0..4 {
+                for (k, coefficients_k) in coefficients.iter_mut().enumerate() {
+                    let row = (in_y - 1 + k as isize)
+                        .max(0)
+                        .min(input_height as isize - 1) as usize;
+                    let x_minus_1 = input_data[c][row]
+                        [(in_x - 1).max(0).min(input_width as isize - 1) as usize];
+                    let x_plus_0 =
+                        input_data[c][row][in_x.max(0).min(input_width as isize - 1) as usize];
+                    let x_plus_1 = input_data[c][row]
+                        [(in_x + 1).max(0).min(input_width as isize - 1) as usize];
+                    let x_plus_2 = input_data[c][row]
+                        [(in_x + 2).max(0).min(input_width as isize - 1) as usize];
+
+                    // coefficients[k] = cubic_interp1d(x_minus_1, x_plus_0, x_plus_1, x_plus_2, t_x);
+                    *coefficients_k = cubic_interp1d(x_minus_1, x_plus_0, x_plus_1, x_plus_2, t_x);
+                }
+                output_data[c][out_y][out_x] = cubic_interp1d(
+                    coefficients[0],
+                    coefficients[1],
+                    coefficients[2],
+                    coefficients[3],
+                    t_y,
+                );
+            }
+        }
+    }
+    let output = Tensor::new(output_data, input.device())?
+        .reshape((batch_size, channels, output_height, output_width))?
+        .to_dtype(input.dtype())?;
+    Ok(output)
+}
+
+pub fn interpolate_bicubic(
+    input: &Tensor,
+    target_size: (usize, usize),
+    antialias: Option<bool>,
+    align_corner: Option<bool>,
+) -> Result<Tensor> {
+    if input.rank() != 4 {
+        return Err(anyhow::anyhow!(
+            "Input rank must have at least 3 dimensions"
+        ));
+    }
+    // if input.dim(0)? != 1 {
+    //     return Err(anyhow::anyhow!("Input batch_size must be 1"));
+    // }
+    let (batch_size, channels, input_height, input_width) = input.dims4()?;
+    let (output_height, output_width) = target_size;
+    if output_height == input_height && output_width == input_width {
+        return Ok(input.clone());
+    }
+    let align_corners = match align_corner {
+        Some(true) => true,
+        Some(false) => false,
+        None => false,
+    };
+    let height_scale = compute_scale(input_height, output_height, align_corners);
+    let width_scale = compute_scale(input_width, output_width, align_corners);
+    // let input_squeeze = input.squeeze(0)?;
+    let output = if let Some(antialias_) = antialias
+        && antialias_
+        && (input_height > output_height || input_width > output_width)
+    {
+        interpolate_bicubic_antialias(
+            input,
+            batch_size,
+            channels,
+            input_height,
+            input_width,
+            output_height,
+            output_width,
+            height_scale,
+            width_scale,
+            align_corners,
+        )?
+    } else {
+        interpolate_bicubic_standard(
+            input,
+            batch_size,
+            channels,
+            input_height,
+            input_width,
+            output_height,
+            output_width,
+            height_scale,
+            width_scale,
+            align_corners,
+        )?
+    };
+    let output = output.to_dtype(input.dtype())?.to_device(input.device())?;
+    Ok(output)
+}
+
 pub fn index_select_2d(t: &Tensor, index: &Tensor) -> Result<Tensor> {
     if t.rank() != 2 && index.rank() != 2 {
-        return Err(anyhow::anyhow!(
-                    "t and index rank must be equal to 2"
-                ));
+        return Err(anyhow::anyhow!("t and index rank must be equal to 2"));
     }
     let mut res_vec = Vec::new();
     let index_dim0 = index.dim(0)?;
@@ -433,7 +690,51 @@ pub fn index_select_2d(t: &Tensor, index: &Tensor) -> Result<Tensor> {
         let index_i = index.i(i)?;
         let rel_i = t.index_select(&index_i, 0)?;
         res_vec.push(rel_i);
-    } 
+    }
     let res = Tensor::stack(&res_vec, 0)?;
     Ok(res)
-}
+}
+
+pub fn quick_gelu(xs: &Tensor) -> Result<Tensor> {
+    let x = xs.affine(1.702, 0.0)?;
+    let x = sigmoid(&x)?;
+    Ok(xs.mul(&x)?)
+}
+
+pub fn topk(weight: &Tensor, topk: usize) -> Result<(Tensor, Tensor)> {
+    let topk_idx = weight
+        .arg_sort_last_dim(false)?
+        .narrow(D::Minus1, 0, topk)?
+        .contiguous()?;
+    let topk_weight = weight.gather(&topk_idx, D::Minus1)?;
+    Ok((topk_weight, topk_idx))
+}
+
+pub fn onehot(input: &Tensor, len: usize) -> Result<Tensor> {
+    let mut shape = input.dims().to_vec();
+    shape.push(len);
+    let expand_input = input.unsqueeze(D::Minus1)?.broadcast_as(shape)?;
+    let range =
+        Tensor::arange(0u32, len as u32, input.device())?.broadcast_as(expand_input.dims())?;
+    let onehot = expand_input.eq(&range)?;
+    Ok(onehot)
+}
+
+pub fn nonzero(input: &Tensor) -> Result<(Vec<u32>, Vec<u32>)> {
+    assert!(input.rank() == 2, "input rank must be 2!");
+    let mut topk_ids = Vec::new();
+    let mut token_ids_all = Vec::new();
+    let topk = input.dim(0)?;
+    let input_vec = input.to_vec2::<u32>()?;
+    for (i, vec) in input_vec.iter().enumerate().take(topk) {
+        let token_ids: Vec<u32> = vec
+            .iter()
+            .enumerate()
+            .filter_map(|(idx, &val)| if val > 0 { Some(idx as u32) } else { None })
+            .collect();
+        let token_len = token_ids.len();
+        topk_ids.extend_from_slice(&vec![i as u32; token_len]);
+        token_ids_all.extend_from_slice(&token_ids);
+    }
+    Ok((topk_ids, token_ids_all))
+}

tests/config_tests.rs

@@ -1,12 +1,13 @@
 use aha::models::{
-    minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig,
-    qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig,
+    deepseek_ocr::config::DeepseekOCRConfig, minicpm4::config::MiniCPM4Config,
+    qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig,
+    voxcpm::config::VoxCPMConfig,
 };
 use anyhow::Result;
 
 #[test]
 fn qwen2_5_vl_config() -> Result<()> {
-    // cargo test -F cuda,flash-attn qwen2_5vl_config -- --nocapture
+    // cargo test -F cuda,flash-attn qwen2_5vl_config -r -- --nocapture
     let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/";
     let config_path = model_path.to_string() + "/config.json";
     let config: Qwen2_5VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
@@ -16,7 +17,7 @@ fn qwen2_5_vl_config() -> Result<()> {
 
 #[test]
 fn minicpm4_config() -> Result<()> {
-    // cargo test -F cuda,flash-attn minicpm4_config -- --nocapture
+    // cargo test -F cuda,flash-attn minicpm4_config -r -- --nocapture
     let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/";
     let config_path = model_path.to_string() + "/config.json";
     let config: MiniCPM4Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
@@ -26,7 +27,7 @@ fn minicpm4_config() -> Result<()> {
 
 #[test]
 fn voxcpm_config() -> Result<()> {
-    // cargo test -F cuda,flash-attn minicpm4_config -- --nocapture
+    // cargo test -F cuda,flash-attn minicpm4_config -r -- --nocapture
     // cargo test -F cuda minicpm4_config -- --nocapture
     let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/";
     let config_path = model_path.to_string() + "/config.json";
@@ -37,10 +38,20 @@ fn voxcpm_config() -> Result<()> {
 
 #[test]
 fn qwen3vl_config() -> Result<()> {
-    // cargo test -F cuda qwen3vl_config -- --nocapture
+    // cargo test -F cuda qwen3vl_config -r -- --nocapture
     let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-4B-Instruct/";
     let config_path = model_path.to_string() + "/config.json";
     let config: Qwen3VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
     println!("{:?}", config);
     Ok(())
 }
+
+#[test]
+fn deepseek_ocr_config() -> Result<()> {
+    // cargo test -F cuda qwen3vl_config -r -- --nocapture
+    let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/";
+    let config_path = model_path.to_string() + "/config.json";
+    let config: DeepseekOCRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
+    println!("{:?}", config);
+    Ok(())
+}

tests/messy_test.rs

@@ -1,17 +1,22 @@
-use aha::utils::tensor_utils::{index_select_2d, interpolate_linear};
+use aha::utils::tensor_utils::interpolate_bicubic;
 use anyhow::Result;
-use candle_core::{IndexOp, Tensor};
+use candle_core::Tensor;
 
 #[test]
 fn messy_test() -> Result<()> {
-    // RUST_BACKTRACE=1 cargo test -F cuda messy_test -- --nocapture
+    // RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture
     let device = &candle_core::Device::Cpu;
-    let t1 = Tensor::rand(0.0, 1.0, (1, 5, 5, 10), device)?;
-    let t2 = Tensor::rand(0.0, 1.0, (5, 8, 10), device)?;
-    let t2 = t2.t()?;
-    println!("t2: {:?}", t2);
-    let re = t1.broadcast_matmul(&t2)?;
-    println!("re: {:?}", re);
+    // let t = Tensor::randn(0.0f32, 1.0, (1, 768, 64, 64), device)?;
+    let t = Tensor::arange(0.0f32, 10.0, device)?.broadcast_as((1, 1, 10, 10))?;
+    println!("t: {}", t);
+    let t_resized = interpolate_bicubic(&t, (5, 5), Some(true), Some(false))?;
+    println!("t_resized: {}", t_resized);
+    // let t1 = Tensor::rand(0.0, 1.0, (1, 5, 5, 10), device)?;
+    // let t2 = Tensor::rand(0.0, 1.0, (5, 8, 10), device)?;
+    // let t2 = t2.t()?;
+    // println!("t2: {:?}", t2);
+    // let re = t1.broadcast_matmul(&t2)?;
+    // println!("re: {:?}", re);
     // let index = Tensor::arange(0u32, 10u32, device)?;
     // let index_2d_vec = vec![index;5];
     // let index_2d = Tensor::stack(&index_2d_vec, 0)?;

tests/test_deepseek_ocr.rs

@@ -1,11 +1,13 @@
-use aha::models::deepseek_ocr::{generate::DeepseekOCRGenerateModel, processor::DeepseekOCRProcessor};
+use std::{pin::pin, time::Instant};
+
+use aha::models::{GenerateModel, deepseek_ocr::generate::DeepseekOCRGenerateModel};
 use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
 use anyhow::Result;
-use candle_core::{DType, Device, IndexOp, Tensor};
+use rocket::futures::StreamExt;
 
 #[test]
-fn deepseek_ocr_test() -> Result<()> {
-    // RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_test -- --nocapture
+fn deepseek_ocr_generate() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_generate -r -- --nocapture
     let message = r#"
     {
         "model": "deepseek-ocr",
@@ -35,9 +37,60 @@ fn deepseek_ocr_test() -> Result<()> {
     "#;
     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 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!("Time elapsed in generate is: {:?}", i_duration);
+    println!("generate: \n {:?}", res);
+    Ok(())
+}
+
+#[tokio::test]
+async fn deepseek_ocr_stream() -> Result<()> {
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda deepseek_ocr_stream -r -- --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 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 mut stream = pin!(model.generate_stream(mes)?);
+    while let Some(item) = stream.next().await {
+        println!("generate: \n {:?}", item);
+    }
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
     Ok(())
 }

tests/test_minicpm4.rs

@@ -7,9 +7,9 @@ use rocket::futures::StreamExt;
 
 #[test]
 fn minicpm_generate() -> Result<()> {
-    // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test minicpm_generate -- --nocapture
-    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda minicpm_generate -- --nocapture
-    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_generate -- --nocapture
+    // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test minicpm_generate -r -- --nocapture
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda minicpm_generate -r -- --nocapture
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_generate -r -- --nocapture
 
     let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/";
     let message = r#"
@@ -42,7 +42,7 @@ fn minicpm_generate() -> Result<()> {
 
 #[tokio::test]
 async fn minicpm_stream() -> Result<()> {
-    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_stream -- --nocapture
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn minicpm_stream -r -- --nocapture
 
     let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/";
 

tests/test_qwen2_5vl.rs

@@ -7,9 +7,9 @@ use rocket::futures::StreamExt;
 
 #[test]
 fn qwen2_5vl_generate() -> Result<()> {
-    // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test qwen2_5vl_generate -- --nocapture
-    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_generate -- --nocapture
-    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture
+    // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test qwen2_5vl_generate -r -- --nocapture
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_generate -r -- --nocapture
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -r -- --nocapture
     // let device = Device::cuda_if_available(0)?;
     // let dtype = DType::BF16;
 
@@ -55,7 +55,7 @@ fn qwen2_5vl_generate() -> Result<()> {
 
 #[tokio::test]
 async fn qwen2_5vl_stream() -> Result<()> {
-    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -r -- --nocapture
     // let device = Device::cuda_if_available(0)?;
     // let dtype = DType::BF16;
 

tests/test_qwen3vl.rs

@@ -7,7 +7,7 @@ use rocket::futures::StreamExt;
 
 #[test]
 fn qwen3vl_generate() -> Result<()> {
-    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_generate -- --nocapture
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_generate -r -- --nocapture
 
     let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-2B-Instruct/";
 
@@ -52,7 +52,7 @@ fn qwen3vl_generate() -> Result<()> {
 
 #[tokio::test]
 async fn qwen3vl_stream() -> Result<()> {
-    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_stream -- --nocapture
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen3vl_stream -r -- --nocapture
 
     let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-2B-Instruct/";
 

tests/test_voxcpm.rs

@@ -8,7 +8,7 @@ use anyhow::{Ok, Result};
 
 #[test]
 fn voxcpm_generate() -> Result<()> {
-    // RUST_BACKTRACE=1 cargo test -F cuda,flash-attn voxcpm_generate -- --nocapture
+    // RUST_BACKTRACE=1 cargo test -F cuda,flash-attn voxcpm_generate -r -- --nocapture
     let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/";
 
     let i_start = Instant::now();

tests/weight_test.rs

@@ -61,3 +61,22 @@ fn qwen3vl_weight() -> Result<()> {
     println!("model_list: {:?}", model_list);
     Ok(())
 }
+
+#[test]
+fn deepseekocr_weight() -> Result<()> {
+    let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/";
+    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("lm_head") {
+                println!("=== {} === {:?}", key, tensor.shape());
+            }
+            // println!("=== {} === {:?}", key, tensor.shape());
+        }
+    }
+    println!("model_list: {:?}", model_list);
+    Ok(())
+}