replace some function

README.md

@@ -38,26 +38,26 @@
 项目提供了几个可选的功能特性，您可以根据需要启用它们：
 * flash-attn: 启用 Flash Attention 支持以提升模型推理性能：
 ```bash
-cargo build --features flash-attn
+cargo build -r --features flash-attn
 ```
 
 * cuda: 为 candle 核心组件启用 CUDA 支持，实现 GPU 加速计算：
 ```bash
-cargo build --features cuda
+cargo build -r --features cuda
 ```
 
 * ffmpeg: 启用 FFmpeg 支持，提供多媒体处理功能：
 ```bash
-cargo build --features ffmpeg
+cargo build -r --features ffmpeg
 ```
 * 组合使用功能特性
 
 ```bash
 # 同时启用 CUDA 和 Flash Attention 以获得最佳性能
-cargo build --features "cuda,flash-attn"
+cargo build -r --features "cuda,flash-attn"
 
 # 启用所有功能特性
-cargo build --features "cuda,flash-attn,ffmpeg"
+cargo build -r --features "cuda,flash-attn,ffmpeg"
 ```
 
 ## 安装及使用
@@ -71,7 +71,7 @@ cd aha
 #### cargo run 运行参数说明
 ##### 基本用法
 ```bash
-cargo run -F cuda -- [参数]
+cargo run -F cuda -r -- [参数]
 ```
 ##### 参数详解
 1. 端口设置

src/models/common/mod.rs

@@ -1,8 +1,8 @@
 use anyhow::Result;
 use candle_core::{D, Tensor};
 use candle_nn::{
-    Activation, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder,
+    Activation, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module, RmsNorm,
-    conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias,
+    VarBuilder, conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias, rms_norm,
 };
 
 use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv};
@@ -110,7 +110,6 @@ pub struct NaiveAttention {
     kv_cache: Option<(Tensor, Tensor)>,
 }
 
-// impl AttentionNobias {
 impl NaiveAttention {
     pub fn new(
         vb: VarBuilder,
@@ -260,6 +259,166 @@ impl NaiveAttention {
     }
 }
 
+pub struct NaiveAttnTwoLinearMLPBlock {
+    self_attn: NaiveAttention,
+    mlp: TwoLinearMLP,
+    input_layernorm: LayerNorm,
+    post_attention_layernorm: LayerNorm,
+}
+
+impl NaiveAttnTwoLinearMLPBlock {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        num_attention_heads: usize,
+        num_key_value_heads: Option<usize>,
+        head_dim: Option<usize>,
+        attn_bias: bool,
+        attn_pp_name: &str,
+        o_proj_pp_name: Option<&str>,
+        intermediate_size: usize,
+        hidden_act: Activation,
+        mlp_bias: bool,
+        mlp_pp_name: &str,
+        linear1_pp_name: &str,
+        linear2_pp_name: &str,
+        norm_eps: f64,
+        input_norm_pp_name: &str,
+        post_norm_pp_name: &str,
+    ) -> Result<Self> {
+        let num_key_value_heads = match num_key_value_heads {
+            Some(heads) => heads,
+            None => num_attention_heads,
+        };
+        let self_attn = NaiveAttention::new(
+            vb.pp(attn_pp_name),
+            hidden_size,
+            num_attention_heads,
+            num_key_value_heads,
+            head_dim,
+            attn_bias,
+            o_proj_pp_name,
+        )?;
+        let mlp = TwoLinearMLP::new(
+            vb.pp(mlp_pp_name),
+            hidden_size,
+            intermediate_size,
+            hidden_act,
+            mlp_bias,
+            linear1_pp_name,
+            linear2_pp_name,
+        )?;
+
+        let input_layernorm = get_layer_norm(vb.pp(input_norm_pp_name), norm_eps, hidden_size)?;
+        let post_attention_layernorm =
+            get_layer_norm(vb.pp(post_norm_pp_name), norm_eps, hidden_size)?;
+        Ok(Self {
+            self_attn,
+            mlp,
+            input_layernorm,
+            post_attention_layernorm,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: Option<&Tensor>,
+        sin: Option<&Tensor>,
+        attention_mask: Option<&Tensor>,
+        tof32: bool,
+    ) -> Result<Tensor> {
+        let residual = xs.clone();
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self
+            .self_attn
+            .forward(&xs, cos, sin, attention_mask, tof32)?;
+        let residual = residual.add(&xs)?;
+        let xs = self.post_attention_layernorm.forward(&residual)?;
+        let xs = self.mlp.forward(&xs)?;
+        let xs = residual.add(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct NaiveAttnGateUpDownMLPBlock {
+    self_attn: NaiveAttention,
+    mlp: GateUpDownMLP,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+}
+
+impl NaiveAttnGateUpDownMLPBlock {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        num_attention_heads: usize,
+        num_key_value_heads: Option<usize>,
+        head_dim: Option<usize>,
+        attn_bias: bool,
+        attn_pp_name: &str,
+        o_proj_pp_name: Option<&str>,
+        intermediate_size: usize,
+        hidden_act: Activation,
+        mlp_bias: bool,
+        mlp_pp_name: &str,
+        norm_eps: f64,
+        input_norm_pp_name: &str,
+        post_norm_pp_name: &str,
+    ) -> Result<Self> {
+        let num_key_value_heads = match num_key_value_heads {
+            Some(heads) => heads,
+            None => num_attention_heads,
+        };
+        let self_attn = NaiveAttention::new(
+            vb.pp(attn_pp_name),
+            hidden_size,
+            num_attention_heads,
+            num_key_value_heads,
+            head_dim,
+            attn_bias,
+            o_proj_pp_name,
+        )?;
+        let mlp = GateUpDownMLP::new(
+            vb.pp(mlp_pp_name),
+            hidden_size,
+            intermediate_size,
+            hidden_act,
+            mlp_bias,
+        )?;
+        let input_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(input_norm_pp_name))?;
+        let post_attention_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(post_norm_pp_name))?;
+        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.clone();
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self
+            .self_attn
+            .forward_with_cache(&xs, cos, sin, attention_mask, false)?;
+        let residual = residual.add(&xs)?;
+        let xs = self.post_attention_layernorm.forward(&residual)?;
+        let xs = self.mlp.forward(&xs)?;
+        let xs = residual.add(&xs)?;
+        Ok(xs)
+    }
+    pub fn clear_kv_cache(&mut self) {
+        self.self_attn.clear_kv_cache()
+    }
+}
+
 pub fn eager_attention_forward(
     query_states: &Tensor,
     key_states: &Tensor,

src/models/deepseek_ocr/model.rs

@@ -1,8 +1,8 @@
 use anyhow::Result;
 use candle_core::{D, IndexOp, Tensor};
 use candle_nn::{
-    Activation, Conv2d, Conv2dConfig, Embedding, Init, LayerNorm, LayerNormConfig, Linear, Module,
+    Activation, Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding,
-    RmsNorm, VarBuilder, conv2d, conv2d_no_bias, embedding, layer_norm, linear, linear_no_bias,
+    linear, linear_no_bias,
     ops::{sigmoid, softmax},
     rms_norm,
 };
@@ -10,7 +10,10 @@ use candle_transformers::models::segment_anything::LayerNorm2d;
 
 use crate::{
     models::{
-        common::{GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward},
+        common::{
+            GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv2d,
+            get_layer_norm,
+        },
         deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config},
     },
     position_embed::rope::RoPE,
@@ -33,14 +36,17 @@ impl PatchEmbed {
         stride: usize,
         padding: usize,
     ) -> Result<Self> {
-        let cfg = Conv2dConfig {
+        let proj = get_conv2d(
+            vb.pp("proj"),
+            in_chans,
+            embed_dim,
+            kernel_size,
             padding,
             stride,
-            dilation: 1,
+            1,
-            groups: 1,
+            1,
-            cudnn_fwd_algo: None,
+            true,
-        };
+        )?;
-        let proj = conv2d(in_chans, embed_dim, kernel_size, cfg, vb.pp("proj"))?;
         Ok(Self { proj })
     }
 
@@ -249,12 +255,7 @@ impl Block {
         window_size: usize,
         input_size: Option<(usize, usize)>,
     ) -> Result<Self> {
-        let ln_config = LayerNormConfig {
+        let norm1 = get_layer_norm(vb.pp("norm1"), eps, dim)?;
-            eps,
-            remove_mean: true, // true for layernorm, false for RMSNorm
-            affine: true,      // true for with bias, false for without bias
-        };
-        let norm1 = layer_norm(dim, ln_config, vb.pp("norm1"))?;
         let input_size = if window_size == 0 {
             input_size
         } else {
@@ -268,7 +269,7 @@ impl Block {
             use_rel_pos,
             input_size,
         )?;
-        let norm2 = layer_norm(dim, ln_config, vb.pp("norm2"))?;
+        let norm2 = get_layer_norm(vb.pp("norm2"), eps, dim)?;
         let mlp_dim = (dim as f32 * mlp_ratio) as usize;
         let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, act, true, "lin1", "lin2")?;
         Ok(Self {
@@ -369,23 +370,9 @@ pub struct Neck {
 
 impl Neck {
     pub fn new(vb: VarBuilder, embed_dim: usize, out_chans: usize) -> Result<Self> {
-        let cfg = Conv2dConfig {
+        let conv2d_0 = get_conv2d(vb.pp("0"), embed_dim, out_chans, 1, 0, 1, 1, 1, false)?;
-            padding: 0,
-            stride: 1,
-            dilation: 1,
-            groups: 1,
-            cudnn_fwd_algo: None,
-        };
-        let conv2d_0 = conv2d_no_bias(embed_dim, out_chans, 1, cfg, vb.pp("0"))?;
         let layernorm_1 = LayerNorm2d::new(out_chans, 0.000001, vb.pp("1"))?;
-        let cfg = Conv2dConfig {
+        let conv2d_2 = get_conv2d(vb.pp("2"), out_chans, out_chans, 3, 1, 1, 1, 1, false)?;
-            padding: 1,
-            stride: 1,
-            dilation: 1,
-            groups: 1,
-            cudnn_fwd_algo: None,
-        };
-        let conv2d_2 = conv2d_no_bias(out_chans, out_chans, 3, cfg, vb.pp("2"))?;
         let layernorm_3 = LayerNorm2d::new(out_chans, 0.000001, vb.pp("3"))?;
         Ok(Self {
             conv2d_0,
@@ -476,15 +463,9 @@ impl ImageEncoderViT {
         }
 
         let neck = Neck::new(vb.pp("neck"), embed_dim, out_chans)?;
-        let cfg = Conv2dConfig {
+
-            padding: 1,
+        let net_2 = get_conv2d(vb.pp("net_2"), 256, 512, 3, 1, 2, 1, 1, false)?;
-            stride: 2,
+        let net_3 = get_conv2d(vb.pp("net_3"), 512, 1024, 3, 1, 2, 1, 1, false)?;
-            dilation: 1,
-            groups: 1,
-            cudnn_fwd_algo: None,
-        };
-        let net_2 = conv2d_no_bias(256, 512, 3, cfg, vb.pp("net_2"))?;
-        let net_3 = conv2d_no_bias(512, 1024, 3, cfg, vb.pp("net_3"))?;
         Ok(Self {
             // img_size,
             patch_embed,
@@ -548,19 +529,17 @@ impl CLIPVisionEmbeddings {
     ) -> Result<Self> {
         let class_embedding =
             vb.get_with_hints(hidden_size, "class_embedding", Init::Const(0.0))?;
-        let cfg = Conv2dConfig {
+
-            padding: 0,
+        let patch_embedding = get_conv2d(
-            stride: patch_size,
+            vb.pp("patch_embedding"),
-            dilation: 1,
-            groups: 1,
-            cudnn_fwd_algo: None,
-        };
-        let patch_embedding = conv2d_no_bias(
             num_channels,
             hidden_size,
             patch_size,
-            cfg,
+            0,
-            vb.pp("patch_embedding"),
+            patch_size,
+            1,
+            1,
+            false,
         )?;
 
         let num_patches = (image_size / patch_size).pow(2);
@@ -698,13 +677,8 @@ impl NoTPTransformerBlock {
     ) -> Result<Self> {
         let self_attn = NoTPAttention::new(vb.pp("self_attn"), hidden_size, num_heads)?;
         let mlp = NoTPFeedForward::new(vb.pp("mlp"), hidden_size, ffn_hidden_size)?;
-        let ln_config = LayerNormConfig {
+        let layer_norm1 = get_layer_norm(vb.pp("layer_norm1"), eps, hidden_size)?;
-            eps,
+        let layer_norm2 = get_layer_norm(vb.pp("layer_norm2"), eps, hidden_size)?;
-            remove_mean: true, // true for layernorm, false for RMSNorm
-            affine: true,      // true for with bias, false for without bias
-        };
-        let layer_norm1 = layer_norm(hidden_size, ln_config, vb.pp("layer_norm1"))?;
-        let layer_norm2 = layer_norm(hidden_size, ln_config, vb.pp("layer_norm2"))?;
         Ok(Self {
             self_attn,
             mlp,
@@ -793,12 +767,7 @@ impl VitModel {
             ffn_hidden_size,
             eps,
         )?;
-        let ln_config = LayerNormConfig {
+        let pre_layrnorm = get_layer_norm(vb.pp("pre_layrnorm"), eps, hidden_size)?;
-            eps,
-            remove_mean: true, // true for layernorm, false for RMSNorm
-            affine: true,      // true for with bias, false for without bias
-        };
-        let pre_layrnorm = layer_norm(hidden_size, ln_config, vb.pp("pre_layrnorm"))?;
         Ok(Self {
             embeddings,
             transformer,
@@ -814,10 +783,6 @@ impl VitModel {
     }
 }
 
-// pub struct DeepseekV2MLP {
-
-// }
-
 pub struct MoEGate {
     top_k: usize,
     // n_routed_experts: usize,
@@ -1129,10 +1094,6 @@ impl DeepseekV2Model {
     }
 }
 
-// pub struct MlpProjector {
-//     layers: Linear,
-// }
-
 pub struct DeepseekOCRModel {
     // config: DeepseekOCRConfig,
     sam_model: ImageEncoderViT,

src/models/hunyuan_ocr/model.rs

@@ -1,16 +1,13 @@
 use anyhow::{Result, anyhow};
 use candle_core::{D, IndexOp, Tensor};
 use candle_nn::{
-    Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear,
+    Conv2d, Embedding, Init, Linear, Module, RmsNorm, VarBuilder, embedding, linear,
     linear_no_bias, rms_norm,
 };
 
 use crate::{
     models::{
-        common::{
+        common::{GateUpDownMLP, NaiveAttnTwoLinearMLPBlock, eager_attention_forward, get_conv2d},
-            GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv2d,
-            get_layer_norm,
-        },
         hunyuan_ocr::config::{HunYuanVLConfig, HunYuanVLVisionConfig},
     },
     position_embed::rope::{RoPE, apply_rotary_pos_emb, get_xd_cos_sin},
@@ -98,63 +95,6 @@ impl HunYuanVisionPatchEmbed {
         Ok(embedding)
     }
 }
-pub struct HunYuanVisionBlock {
-    self_attn: NaiveAttention,
-    mlp: TwoLinearMLP,
-    input_layernorm: LayerNorm,
-    post_attention_layernorm: LayerNorm,
-}
-
-impl HunYuanVisionBlock {
-    pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result<Self> {
-        let self_attn = NaiveAttention::new(
-            vb.pp("self_attn"),
-            config.hidden_size,
-            config.num_attention_heads,
-            config.num_attention_heads,
-            None,
-            true,
-            None,
-        )?;
-        let mlp = TwoLinearMLP::new(
-            vb.pp("mlp"),
-            config.hidden_size,
-            config.intermediate_size,
-            config.hidden_act,
-            true,
-            "dense_h_to_4h",
-            "dense_4h_to_h",
-        )?;
-
-        let input_layernorm = get_layer_norm(
-            vb.pp("input_layernorm"),
-            config.rms_norm_eps,
-            config.hidden_size,
-        )?;
-        let post_attention_layernorm = get_layer_norm(
-            vb.pp("post_attention_layernorm"),
-            config.rms_norm_eps,
-            config.hidden_size,
-        )?;
-        Ok(Self {
-            self_attn,
-            mlp,
-            input_layernorm,
-            post_attention_layernorm,
-        })
-    }
-
-    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        let residual = xs.clone();
-        let xs = self.input_layernorm.forward(xs)?;
-        let xs = self.self_attn.forward(&xs, None, None, None, false)?;
-        let residual = residual.add(&xs)?;
-        let xs = self.post_attention_layernorm.forward(&residual)?;
-        let xs = self.mlp.forward(&xs)?;
-        let xs = residual.add(&xs)?;
-        Ok(xs)
-    }
-}
 
 pub struct HunYuanVisionPatchMerger {
     proj_0: Conv2d,
@@ -250,7 +190,7 @@ impl HunYuanVisionPatchMerger {
 
 pub struct HunYuanVisionTransformer {
     embeddings: HunYuanVisionPatchEmbed,
-    layers: Vec<HunYuanVisionBlock>,
+    layers: Vec<NaiveAttnTwoLinearMLPBlock>,
     perceive: HunYuanVisionPatchMerger,
 }
 
@@ -260,7 +200,25 @@ impl HunYuanVisionTransformer {
         let mut layers = vec![];
         let vb_layers = vb.pp("layers");
         for i in 0..config.num_hidden_layers {
-            let layer_i = HunYuanVisionBlock::new(vb_layers.pp(i), config)?;
+            let layer_i = NaiveAttnTwoLinearMLPBlock::new(
+                vb_layers.pp(i),
+                config.hidden_size,
+                config.num_attention_heads,
+                None,
+                None,
+                true,
+                "self_attn",
+                None,
+                config.intermediate_size,
+                config.hidden_act,
+                true,
+                "mlp",
+                "dense_h_to_4h",
+                "dense_4h_to_h",
+                config.rms_norm_eps,
+                "input_layernorm",
+                "post_attention_layernorm",
+            )?;
             layers.push(layer_i);
         }
         let perceive = HunYuanVisionPatchMerger::new(vb.pp("perceive"), config)?;
@@ -274,7 +232,7 @@ impl HunYuanVisionTransformer {
     pub fn forward(&self, xs: &Tensor, grid_thw: &Tensor) -> Result<Tensor> {
         let mut hidden_states = self.embeddings.forward(xs, grid_thw)?;
         for layer in &self.layers {
-            hidden_states = layer.forward(&hidden_states)?;
+            hidden_states = layer.forward(&hidden_states, None, None, None, false)?;
         }
         let mut cu_seqlens = vec![];
         for i in 0..grid_thw.dim(0)? {

src/models/paddleocr_vl/model.rs

@@ -8,7 +8,9 @@ use num::integer::Roots;
 
 use crate::{
     models::{
-        common::{GateUpDownMLP, NaiveAttention, TwoLinearMLP, get_conv2d, get_layer_norm},
+        common::{
+            NaiveAttnGateUpDownMLPBlock, NaiveAttnTwoLinearMLPBlock, get_conv2d, get_layer_norm,
+        },
         paddleocr_vl::config::{
             PaddleOCRVLConfig, PaddleOCRVLRopeScalingConfig, PaddleOCRVLVisionConfig,
         },
@@ -187,70 +189,8 @@ impl SiglipVisionEmbeddings {
     }
 }
 
-pub struct SiglipEncoderLayer {
-    layer_norm1: LayerNorm,
-    self_attn: NaiveAttention,
-    layer_norm2: LayerNorm,
-    mlp: TwoLinearMLP,
-}
-
-impl SiglipEncoderLayer {
-    pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
-        let layer_norm1 = get_layer_norm(
-            vb.pp("layer_norm1"),
-            config.layer_norm_eps,
-            config.hidden_size,
-        )?;
-        let self_attn = NaiveAttention::new(
-            vb.pp("self_attn"),
-            config.hidden_size,
-            config.num_attention_heads,
-            config.num_attention_heads,
-            None,
-            true,
-            Some("out_proj"),
-        )?;
-        let layer_norm2 = get_layer_norm(
-            vb.pp("layer_norm2"),
-            config.layer_norm_eps,
-            config.hidden_size,
-        )?;
-        let mlp = TwoLinearMLP::new(
-            vb.pp("mlp"),
-            config.hidden_size,
-            config.intermediate_size,
-            config.hidden_act,
-            true,
-            "fc1",
-            "fc2",
-        )?;
-        Ok(Self {
-            layer_norm1,
-            self_attn,
-            layer_norm2,
-            mlp,
-        })
-    }
-
-    pub fn forward(
-        &self,
-        xs: &Tensor,
-        cos: Option<&Tensor>,
-        sin: Option<&Tensor>,
-    ) -> Result<Tensor> {
-        let residual = xs.clone();
-        let xs = self.layer_norm1.forward(xs)?;
-        let xs = self.self_attn.forward(&xs, cos, sin, None, true)?;
-        let residual = residual.add(&xs)?;
-        let xs = self.layer_norm2.forward(&residual)?;
-        let xs = self.mlp.forward(&xs)?;
-        let xs = residual.add(&xs)?;
-        Ok(xs)
-    }
-}
-
 pub struct SiglipEncoder {
-    layers: Vec<SiglipEncoderLayer>,
+    layers: Vec<NaiveAttnTwoLinearMLPBlock>,
     rotary_pos_emb: Qwen2_5VisionRotaryEmbedding,
 }
 
@@ -259,7 +199,25 @@ impl SiglipEncoder {
         let vb_layers = vb.pp("layers");
         let mut layers = vec![];
         for i in 0..config.num_hidden_layers {
-            let layer_i = SiglipEncoderLayer::new(vb_layers.pp(i), config)?;
+            let layer_i = NaiveAttnTwoLinearMLPBlock::new(
+                vb_layers.pp(i),
+                config.hidden_size,
+                config.num_attention_heads,
+                None,
+                None,
+                true,
+                "self_attn",
+                Some("out_proj"),
+                config.intermediate_size,
+                config.hidden_act,
+                true,
+                "mlp",
+                "fc1",
+                "fc2",
+                config.layer_norm_eps,
+                "layer_norm1",
+                "layer_norm2",
+            )?;
             layers.push(layer_i);
         }
         let head_dim = config.hidden_size / config.num_attention_heads;
@@ -300,7 +258,7 @@ impl SiglipEncoder {
         let sin = rope_emb.sin()?;
         let mut xs = xs.clone();
         for layer in &self.layers {
-            xs = layer.forward(&xs, Some(&cos), Some(&sin))?;
+            xs = layer.forward(&xs, Some(&cos), Some(&sin), None, false)?;
         }
         Ok(xs)
     }
@@ -347,75 +305,9 @@ impl SiglipVisionModel {
     }
 }
 
-pub struct Ernie4_5DecoderLayer {
-    self_attn: NaiveAttention,
-    mlp: GateUpDownMLP,
-    input_layernorm: RmsNorm,
-    post_attention_layernorm: RmsNorm,
-}
-
-impl Ernie4_5DecoderLayer {
-    pub fn new(vb: VarBuilder, config: &PaddleOCRVLConfig) -> Result<Self> {
-        let self_attn = NaiveAttention::new(
-            vb.pp("self_attn"),
-            config.hidden_size,
-            config.num_attention_heads,
-            config.num_key_value_heads,
-            Some(config.head_dim),
-            config.use_bias,
-            None,
-        )?;
-        let mlp = GateUpDownMLP::new(
-            vb.pp("mlp"),
-            config.hidden_size,
-            config.intermediate_size,
-            config.hidden_act,
-            config.use_bias,
-        )?;
-        let input_layernorm = rms_norm(
-            config.hidden_size,
-            config.rms_norm_eps,
-            vb.pp("input_layernorm"),
-        )?;
-        let post_attention_layernorm = rms_norm(
-            config.hidden_size,
-            config.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.clone();
-        let xs = self.input_layernorm.forward(xs)?;
-        let xs = self
-            .self_attn
-            .forward_with_cache(&xs, cos, sin, attention_mask, false)?;
-        let residual = residual.add(&xs)?;
-        let xs = self.post_attention_layernorm.forward(&residual)?;
-        let xs = self.mlp.forward(&xs)?;
-        let xs = residual.add(&xs)?;
-        Ok(xs)
-    }
-    fn clear_kv_cache(&mut self) {
-        self.self_attn.clear_kv_cache()
-    }
-}
-
 pub struct Ernie4_5Model {
     embed_tokens: Embedding,
-    layers: Vec<Ernie4_5DecoderLayer>,
+    layers: Vec<NaiveAttnGateUpDownMLPBlock>,
     norm: RmsNorm,
     rotary_emb: Qwen2_5VLTextRotaryEmbedding,
     rope_scaling: PaddleOCRVLRopeScalingConfig,
@@ -427,7 +319,23 @@ impl Ernie4_5Model {
         let vb_layers = vb.pp("layers");
         let mut layers = vec![];
         for i in 0..config.num_hidden_layers {
-            let layer_i = Ernie4_5DecoderLayer::new(vb_layers.pp(i), config)?;
+            let layer_i = NaiveAttnGateUpDownMLPBlock::new(
+                vb_layers.pp(i),
+                config.hidden_size,
+                config.num_attention_heads,
+                Some(config.num_key_value_heads),
+                Some(config.head_dim),
+                config.use_bias,
+                "self_attn",
+                None,
+                config.intermediate_size,
+                config.hidden_act,
+                config.use_bias,
+                "mlp",
+                config.rms_norm_eps,
+                "input_layernorm",
+                "post_attention_layernorm",
+            )?;
             layers.push(layer_i);
         }
         let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?;

src/models/qwen2_5vl/model.rs

@@ -1,12 +1,10 @@
 use anyhow::{Result, anyhow};
 use candle_core::{D, DType, Device, IndexOp, Tensor};
-use candle_nn::{
+use candle_nn::{Init, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm};
-    Activation, Init, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm,
-};
 
 use crate::{
     models::{
-        common::eager_attention_forward,
+        common::{GateUpDownMLP, eager_attention_forward},
         qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling},
     },
     position_embed::rope::{
@@ -94,38 +92,6 @@ impl Module for Qwen2_5VLPatchMerger {
     }
 }
 
-#[derive(Debug, Clone)]
-struct Qwen2_5VLVisionMLP {
-    gate_proj: Linear,
-    up_proj: Linear,
-    down_proj: Linear,
-    act_fn: Activation,
-}
-
-impl Qwen2_5VLVisionMLP {
-    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
-        let hidden_sz = cfg.vision_config.hidden_size;
-        let intermediate_sz = cfg.vision_config.intermediate_size;
-        let gate_proj = linear(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?;
-        let up_proj = linear(hidden_sz, intermediate_sz, vb.pp("up_proj"))?;
-        let down_proj = linear(intermediate_sz, hidden_sz, vb.pp("down_proj"))?;
-        Ok(Self {
-            gate_proj,
-            up_proj,
-            down_proj,
-            act_fn: cfg.hidden_act,
-        })
-    }
-}
-
-impl Module for Qwen2_5VLVisionMLP {
-    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
-        let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?;
-        let rhs = xs.apply(&self.up_proj)?;
-        (lhs * rhs)?.apply(&self.down_proj)
-    }
-}
-
 #[derive(Debug, Clone)]
 struct Qwen2_5VLVisionAttention {
     qkv: Linear,
@@ -200,7 +166,7 @@ impl Qwen2_5VLVisionAttention {
 #[derive(Debug, Clone)]
 struct Qwen2_5VLVisionBlock {
     attn: Qwen2_5VLVisionAttention,
-    mlp: Qwen2_5VLVisionMLP,
+    mlp: GateUpDownMLP,
     norm1: RmsNorm,
     norm2: RmsNorm,
 }
@@ -208,7 +174,13 @@ struct Qwen2_5VLVisionBlock {
 impl Qwen2_5VLVisionBlock {
     fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let attn = Qwen2_5VLVisionAttention::new(cfg, vb.pp("attn"))?;
-        let mlp = Qwen2_5VLVisionMLP::new(cfg, vb.pp("mlp"))?;
+        let mlp = GateUpDownMLP::new(
+            vb.pp("mlp"),
+            cfg.vision_config.hidden_size,
+            cfg.vision_config.intermediate_size,
+            cfg.vision_config.hidden_act,
+            true,
+        )?;
         let norm1 = rms_norm(
             cfg.vision_config.hidden_size,
             cfg.rms_norm_eps,
@@ -537,39 +509,6 @@ impl Qwen2_5VLVisionModel {
     }
 }
 
-#[derive(Debug, Clone)]
-struct Qwen2_5VLTextMLP {
-    gate_proj: Linear,
-    up_proj: Linear,
-    down_proj: Linear,
-    act_fn: Activation,
-}
-
-impl Qwen2_5VLTextMLP {
-    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
-        let hidden_sz = cfg.hidden_size;
-        let intermediate_sz = cfg.intermediate_size;
-        let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?;
-        let up_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("up_proj"))?;
-        let down_proj = linear_no_bias(intermediate_sz, hidden_sz, vb.pp("down_proj"))?;
-
-        Ok(Self {
-            gate_proj,
-            up_proj,
-            down_proj,
-            act_fn: cfg.hidden_act,
-        })
-    }
-}
-
-impl Module for Qwen2_5VLTextMLP {
-    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
-        let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?;
-        let rhs = xs.apply(&self.up_proj)?;
-        (lhs * rhs)?.apply(&self.down_proj)
-    }
-}
-
 #[derive(Debug, Clone)]
 struct Qwen2_5VLTextAttention {
     q_proj: Linear,
@@ -663,7 +602,7 @@ impl Qwen2_5VLTextAttention {
 #[derive(Debug, Clone)]
 struct Qwen2_5VLTextDecoderLayer {
     self_attn: Qwen2_5VLTextAttention,
-    mlp: Qwen2_5VLTextMLP,
+    mlp: GateUpDownMLP,
     input_layernorm: RmsNorm,
     post_attention_layernorm: RmsNorm,
 }
@@ -671,7 +610,13 @@ struct Qwen2_5VLTextDecoderLayer {
 impl Qwen2_5VLTextDecoderLayer {
     fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let self_attn = Qwen2_5VLTextAttention::new(cfg, vb.pp("self_attn"))?;
-        let mlp = Qwen2_5VLTextMLP::new(cfg, vb.pp("mlp"))?;
+        let mlp = GateUpDownMLP::new(
+            vb.pp("mlp"),
+            cfg.hidden_size,
+            cfg.intermediate_size,
+            cfg.hidden_act,
+            false,
+        )?;
         let input_layernorm =
             rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
         let post_attention_layernorm = rms_norm(

src/models/qwen3vl/config.rs

@@ -50,7 +50,7 @@ pub struct Qwen3VLTextConfig {
 pub struct Qwen3VLVisionConfig {
     pub deepstack_visual_indexes: Vec<usize>,
     pub depth: usize,
-    pub hidden_act: String,
+    pub hidden_act: Activation,
     pub hidden_size: usize,
     pub in_channels: usize,
     pub initializer_range: f32,

src/models/qwen3vl/model.rs

@@ -1,13 +1,13 @@
 use anyhow::{Result, anyhow};
 use candle_core::{D, DType, IndexOp, Shape, Tensor};
 use candle_nn::{
-    Activation, Embedding, Init, LayerNorm, LayerNormConfig, Linear, Module, RmsNorm, VarBuilder,
+    Activation, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear,
-    embedding, layer_norm, linear, linear_no_bias, rms_norm,
+    linear_no_bias, rms_norm,
 };
 
 use crate::{
     models::{
-        common::{GateUpDownMLP, eager_attention_forward},
+        common::{GateUpDownMLP, TwoLinearMLP, eager_attention_forward, get_layer_norm},
         qwen3vl::config::{Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig},
     },
     position_embed::rope::{
@@ -21,34 +21,6 @@ use crate::{
     },
 };
 
-pub struct Qwen3VLVisionMLP {
-    linear_fc1: Linear,
-    linear_fc2: Linear,
-    act_fn: Activation,
-}
-
-impl Qwen3VLVisionMLP {
-    pub fn new(config: Qwen3VLVisionConfig, vb: VarBuilder) -> Result<Self> {
-        let hidden_size = config.hidden_size;
-        let intermediate_size = config.intermediate_size;
-        let linear_fc1 = linear(hidden_size, intermediate_size, vb.pp("linear_fc1"))?;
-        let linear_fc2 = linear(intermediate_size, hidden_size, vb.pp("linear_fc2"))?;
-        let act_fn = Activation::GeluPytorchTanh;
-        Ok(Self {
-            linear_fc1,
-            linear_fc2,
-            act_fn,
-        })
-    }
-}
-
-impl Module for Qwen3VLVisionMLP {
-    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
-        let xs = xs.apply(&self.linear_fc1)?.apply(&self.act_fn)?;
-        xs.apply(&self.linear_fc2)
-    }
-}
-
 pub struct Qwen3VLVisionPatchEmbed {
     conv3d_weight: Tensor,
     conv3d_bias: Tensor,
@@ -111,17 +83,12 @@ impl Qwen3VLVisionPatchMerger {
         use_postshuffle_norm: bool,
     ) -> Result<Self> {
         let hidden_size = config.hidden_size * config.spatial_merge_size.pow(2);
-        let ln_config = LayerNormConfig {
-            eps: 1e-6,
-            remove_mean: true, // true for layernorm, false for RMSNorm
-            affine: true,      // true for with bias, false for without bias
-        };
         let norm_size = if use_postshuffle_norm {
             hidden_size
         } else {
             config.hidden_size
         };
-        let norm = layer_norm(norm_size, ln_config, vb.pp("norm"))?;
+        let norm = get_layer_norm(vb.pp("norm"), 1e-6, norm_size)?;
         let linear_fc1 = linear(hidden_size, hidden_size, vb.pp("linear_fc1"))?;
         let act_fn = Activation::Gelu;
         let linear_fc2 = linear(hidden_size, config.out_hidden_size, vb.pp("linear_fc2"))?;
@@ -226,20 +193,23 @@ pub struct Qwen3VLVisionBlock {
     norm1: LayerNorm,
     norm2: LayerNorm,
     attn: Qwen3VLVisionAttention,
-    mlp: Qwen3VLVisionMLP,
+    mlp: TwoLinearMLP,
 }
 
 impl Qwen3VLVisionBlock {
     pub fn new(config: Qwen3VLVisionConfig, vb: VarBuilder) -> Result<Self> {
-        let ln_config = LayerNormConfig {
+        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-6, config.hidden_size)?;
-            eps: 1e-6,
+        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-6, config.hidden_size)?;
-            remove_mean: true, // true for layernorm, false for RMSNorm
-            affine: true,      // true for with bias, false for without bias
-        };
-        let norm1 = layer_norm(config.hidden_size, ln_config, vb.pp("norm1"))?;
-        let norm2 = layer_norm(config.hidden_size, ln_config, vb.pp("norm2"))?;
         let attn = Qwen3VLVisionAttention::new(config.clone(), vb.pp("attn"))?;
-        let mlp = Qwen3VLVisionMLP::new(config, vb.pp("mlp"))?;
+        let mlp = TwoLinearMLP::new(
+            vb.pp("mlp"),
+            config.hidden_size,
+            config.intermediate_size,
+            config.hidden_act,
+            true,
+            "linear_fc1",
+            "linear_fc2",
+        )?;
         Ok(Self {
             norm1,
             norm2,

tests/test_deepseek_ocr.rs

@@ -24,7 +24,7 @@ fn deepseek_ocr_generate() -> Result<()> {
                     },              
                     {
                         "type": "text", 
-                        "text": "<image>\n<|grounding|>Convert the document to markdown. "
+                        "text": "<image>\nConvert the document to markdown. "
                     }
                 ]
             }

tests/test_hunyuan_ocr.rs

@@ -24,7 +24,7 @@ fn hunyuan_ocr_generate() -> Result<()> {
                     },              
                     {
                         "type": "text", 
-                        "text": "检测并识别图片中的文字，将文本坐标格式化输出。"
+                        "text": "识别图片中的文字"
                     }
                 ]
             }

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 -r -- --nocapture
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda,ffmpeg qwen3vl_generate -r -- --nocapture
 
     let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-2B-Instruct/";
 

tests/test_robo_brain.rs

@@ -0,0 +1,42 @@
+use std::time::Instant;
+
+use aha::models::{GenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel};
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+
+#[test]
+fn robo_brain_generate() -> Result<()> {
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda robo_brain_generate -r -- --nocapture
+
+    let model_path = "/home/jhq/huggingface_model/BAAI/RoboBrain2.0-3B/";
+
+    let message = r#"
+    {
+        "model": "qwen2.5vl",
+        "messages": [
+            {
+                "role": "user",
+                "content": [           
+                    {
+                        "type": "text", 
+                        "text": "hello RoboBrain"
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut model = Qwen2_5VLGenerateModel::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 result = model.generate(mes)?;
+    println!("generate: \n {:?}", result);
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+
+    Ok(())
+}