tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

replace some function

Browse Source
This commit is contained in:
jhqxxx
2025-12-10 00:05:46 +08:00
parent dc4b593011
commit 515d6da852
12 changed files with 346 additions and 403 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+6 -6
View File
@@ -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
+162 -3
View File
@@ -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<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
+33 -72
View File
@@ -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<Self> {
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<Self> {
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<Self> {
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<Self> {
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<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 {
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,
src/models/hunyuan_ocr/model.rs
+23 -65
View File
@@ -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<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
+42 -134
View File
@@ -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
+18 -73
View File
@@ -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<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
+1 -1
View File
@@ -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
+16 -46
View File
@@ -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<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 {
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,
tests/test_deepseek_ocr.rs
+1 -1
View File
@@ -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
+1 -1
View File
@@ -24,7 +24,7 @@ fn hunyuan_ocr_generate() -> Result<()> {
},
{
"type": "text",
"text": "检测并识别图片中的文字，将文本坐标格式化输出。"
"text": "识别图片中的文字"
}
]
}
tests/test_qwen3vl.rs
+1 -1
View File
@@ -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
+42
View File
@@ -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(())
}