From 515d6da852c54d763923a4f328737a27c440130e Mon Sep 17 00:00:00 2001 From: jhqxxx Date: Wed, 10 Dec 2025 00:05:46 +0800 Subject: [PATCH] replace some function --- README.md | 12 +-- src/models/common/mod.rs | 165 ++++++++++++++++++++++++++++- src/models/deepseek_ocr/model.rs | 105 ++++++------------ src/models/hunyuan_ocr/model.rs | 88 ++++------------ src/models/paddleocr_vl/model.rs | 176 ++++++++----------------------- src/models/qwen2_5vl/model.rs | 91 ++++------------ src/models/qwen3vl/config.rs | 2 +- src/models/qwen3vl/model.rs | 62 +++-------- tests/test_deepseek_ocr.rs | 2 +- tests/test_hunyuan_ocr.rs | 2 +- tests/test_qwen3vl.rs | 2 +- tests/test_robo_brain.rs | 42 ++++++++ 12 files changed, 346 insertions(+), 403 deletions(-) create mode 100644 tests/test_robo_brain.rs diff --git a/README.md b/README.md index 739067b..aee2d83 100644 --- a/README.md +++ b/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. 端口设置 diff --git a/src/models/common/mod.rs b/src/models/common/mod.rs index 68a03e8..315cb7f 100644 --- a/src/models/common/mod.rs +++ b/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, - conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias, + Activation, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module, RmsNorm, + 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, + head_dim: Option, + 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 { + 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 { + 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, + head_dim: Option, + 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 { + 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 { + 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, diff --git a/src/models/deepseek_ocr/model.rs b/src/models/deepseek_ocr/model.rs index 0164b9e..2157498 100644 --- a/src/models/deepseek_ocr/model.rs +++ b/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, - RmsNorm, VarBuilder, conv2d, conv2d_no_bias, embedding, layer_norm, linear, linear_no_bias, + Activation, Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, + 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 { - let cfg = Conv2dConfig { + let proj = get_conv2d( + vb.pp("proj"), + in_chans, + embed_dim, + kernel_size, padding, stride, - dilation: 1, - groups: 1, - cudnn_fwd_algo: None, - }; - let proj = conv2d(in_chans, embed_dim, kernel_size, cfg, vb.pp("proj"))?; + 1, + 1, + true, + )?; Ok(Self { proj }) } @@ -249,12 +255,7 @@ impl Block { window_size: usize, input_size: Option<(usize, usize)>, ) -> Result { - let ln_config = LayerNormConfig { - 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 norm1 = get_layer_norm(vb.pp("norm1"), eps, dim)?; 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 { - let cfg = Conv2dConfig { - 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 conv2d_0 = get_conv2d(vb.pp("0"), embed_dim, out_chans, 1, 0, 1, 1, 1, false)?; let layernorm_1 = LayerNorm2d::new(out_chans, 0.000001, vb.pp("1"))?; - let cfg = Conv2dConfig { - 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 conv2d_2 = get_conv2d(vb.pp("2"), out_chans, out_chans, 3, 1, 1, 1, 1, false)?; 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, - stride: 2, - 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"))?; + + let net_2 = get_conv2d(vb.pp("net_2"), 256, 512, 3, 1, 2, 1, 1, false)?; + let net_3 = get_conv2d(vb.pp("net_3"), 512, 1024, 3, 1, 2, 1, 1, false)?; Ok(Self { // img_size, patch_embed, @@ -548,19 +529,17 @@ impl CLIPVisionEmbeddings { ) -> Result { let class_embedding = vb.get_with_hints(hidden_size, "class_embedding", Init::Const(0.0))?; - let cfg = Conv2dConfig { - padding: 0, - stride: patch_size, - dilation: 1, - groups: 1, - cudnn_fwd_algo: None, - }; - let patch_embedding = conv2d_no_bias( + + let patch_embedding = get_conv2d( + vb.pp("patch_embedding"), num_channels, hidden_size, patch_size, - cfg, - vb.pp("patch_embedding"), + 0, + patch_size, + 1, + 1, + false, )?; let num_patches = (image_size / patch_size).pow(2); @@ -698,13 +677,8 @@ impl NoTPTransformerBlock { ) -> Result { 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 { - eps, - 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"))?; + let layer_norm1 = get_layer_norm(vb.pp("layer_norm1"), eps, hidden_size)?; + let layer_norm2 = get_layer_norm(vb.pp("layer_norm2"), eps, hidden_size)?; Ok(Self { self_attn, mlp, @@ -793,12 +767,7 @@ impl VitModel { ffn_hidden_size, eps, )?; - let ln_config = LayerNormConfig { - 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"))?; + let pre_layrnorm = get_layer_norm(vb.pp("pre_layrnorm"), eps, hidden_size)?; 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, diff --git a/src/models/hunyuan_ocr/model.rs b/src/models/hunyuan_ocr/model.rs index 4e9f5bc..a35c615 100644 --- a/src/models/hunyuan_ocr/model.rs +++ b/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::{ - GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward, get_conv2d, - get_layer_norm, - }, + common::{GateUpDownMLP, NaiveAttnTwoLinearMLPBlock, eager_attention_forward, get_conv2d}, 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 { - 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 { - 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, + layers: Vec, 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 { 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)? { diff --git a/src/models/paddleocr_vl/model.rs b/src/models/paddleocr_vl/model.rs index dc613bf..ebfb86b 100644 --- a/src/models/paddleocr_vl/model.rs +++ b/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 { - 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 { - 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, + layers: Vec, 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 { - 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 { - 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, + layers: Vec, 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"))?; diff --git a/src/models/qwen2_5vl/model.rs b/src/models/qwen2_5vl/model.rs index fefbe3c..457f100 100644 --- a/src/models/qwen2_5vl/model.rs +++ b/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::{ - Activation, Init, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm, -}; +use candle_nn::{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 { - 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 { - 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 { 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 { - 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 { - 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 { 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( diff --git a/src/models/qwen3vl/config.rs b/src/models/qwen3vl/config.rs index ee1a1e4..7de3280 100644 --- a/src/models/qwen3vl/config.rs +++ b/src/models/qwen3vl/config.rs @@ -50,7 +50,7 @@ pub struct Qwen3VLTextConfig { pub struct Qwen3VLVisionConfig { pub deepstack_visual_indexes: Vec, pub depth: usize, - pub hidden_act: String, + pub hidden_act: Activation, pub hidden_size: usize, pub in_channels: usize, pub initializer_range: f32, diff --git a/src/models/qwen3vl/model.rs b/src/models/qwen3vl/model.rs index 6cac190..ca65598 100644 --- a/src/models/qwen3vl/model.rs +++ b/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, - embedding, layer_norm, linear, linear_no_bias, rms_norm, + Activation, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear, + 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 { - 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 { - 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 { 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 { - 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 norm1 = layer_norm(config.hidden_size, ln_config, vb.pp("norm1"))?; - let norm2 = layer_norm(config.hidden_size, ln_config, vb.pp("norm2"))?; + let norm1 = get_layer_norm(vb.pp("norm1"), 1e-6, config.hidden_size)?; + let norm2 = get_layer_norm(vb.pp("norm2"), 1e-6, config.hidden_size)?; 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, diff --git a/tests/test_deepseek_ocr.rs b/tests/test_deepseek_ocr.rs index 329e0cf..115fc07 100644 --- a/tests/test_deepseek_ocr.rs +++ b/tests/test_deepseek_ocr.rs @@ -24,7 +24,7 @@ fn deepseek_ocr_generate() -> Result<()> { }, { "type": "text", - "text": "\n<|grounding|>Convert the document to markdown. " + "text": "\nConvert the document to markdown. " } ] } diff --git a/tests/test_hunyuan_ocr.rs b/tests/test_hunyuan_ocr.rs index ec3f071..5e8c9e1 100644 --- a/tests/test_hunyuan_ocr.rs +++ b/tests/test_hunyuan_ocr.rs @@ -24,7 +24,7 @@ fn hunyuan_ocr_generate() -> Result<()> { }, { "type": "text", - "text": "检测并识别图片中的文字，将文本坐标格式化输出。" + "text": "识别图片中的文字" } ] } diff --git a/tests/test_qwen3vl.rs b/tests/test_qwen3vl.rs index 02ed4bc..609ee59 100644 --- a/tests/test_qwen3vl.rs +++ b/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/"; diff --git a/tests/test_robo_brain.rs b/tests/test_robo_brain.rs new file mode 100644 index 0000000..dd27672 --- /dev/null +++ b/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(()) +}