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

add hunyuan_ocr

Browse Source
This commit is contained in:
jhqxxx
2025-12-03 17:21:01 +08:00
parent 7d72cb3baf
commit 697484cf23
29 changed files with 1756 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+1
View File
@@ -37,3 +37,4 @@ cuda=["candle-nn/cuda", "candle-core/cuda", "candle-transformers/cuda"]
[lints.clippy]
needless_range_loop = "allow"
single_range_in_vec_init = "allow"
README.md
+15 -1
View File
@@ -17,10 +17,13 @@
* VoxCPM - 面壁智能语音生成模型
* Qwen3VL - 阿里通义千问 3 多模态大语言模型
* DeepSeek-OCR - 深度求索光学文字识别模型
* Hunyuan-OCR - 腾讯混元光学文字识别模型
## 计划支持
我们持续扩展支持的模型列表，欢迎贡献！
⭐ 如果这个项目对你有帮助，请给我们一个 Star！
## 环境依赖
1. ffmpeg:
* ubuntu/WSL
@@ -77,6 +80,10 @@ fn main() -> Result<()> {
git clone https://github.com/jhqxxx/aha.git
cd aha
# 修改测试用例中模型路径
# 运行 Hunyuan-OCR 示例
cargo test -F cuda hunyuan_ocr_generate -r -- --nocapture
# 运行 DeepSeek-OCR 示例
cargo test -F cuda deepseek_ocr_generate -r -- --nocapture
@@ -122,6 +129,7 @@ cargo run -F cuda -- [参数]
* qwen3vl-8b：Qwen/Qwen3-VL-8B-Instruct 模型
* qwen3vl-32b：Qwen/Qwen3-VL-32B-Instruct 模型
* deepseek-ocr: deepseek-ai/DeepSeek-OCR 模型
* hunyuan-ocr: Tencent-Hunyuan/HunyuanOCR 模型
* 示例：--model deepseek-ocr 或 -m qwen3vl-2b
3. 权重路径
@@ -162,6 +170,7 @@ cargo run -F cuda -- [参数]
│ ├── models
│ │ ├── common
│ │ ├── deepseek_ocr
│ │ ├── hunyuan_ocr
│ │ ├── minicpm4
│ │ ├── qwen2_5vl
│ │ ├── qwen3vl
@@ -170,8 +179,10 @@ cargo run -F cuda -- [参数]
│ ├── position_embed
│ ├── tokenizer
│ ├── utils
│ ├── api.rs
│ └── lib.rs
└── tests
├── test_hunyuan_ocr.rs
├── test_deepseek_ocr.rs
├── test_minicpm4.rs
├── test_qwen2_5vl.rs
@@ -196,6 +207,10 @@ cargo run -F cuda -- [参数]
2. 提交新的 Issue，包含详细描述和复现步骤
## 更新日志
### v0.1.3
* 添加 Hunyuan-OCR 模型
### v0.1.2
* 添加 DeepSeek-OCR 模型
@@ -207,4 +222,3 @@ cargo run -F cuda -- [参数]
* 支持 Qwen2.5VL, MiniCPM4, VoxCPM 模型
⭐ 如果这个项目对你有帮助，请给我们一个 Star！
src/main.rs
+1
View File
@@ -84,6 +84,7 @@ async fn main() -> anyhow::Result<()> {
WhichModel::Qwen3vl8B => "Qwen/Qwen3-VL-8B-Instruct",
WhichModel::Qwen3vl32B => "Qwen/Qwen3-VL-32B-Instruct",
WhichModel::DeepSeekOCR => "deepseek-ai/DeepSeek-OCR",
WhichModel::HunyuanOCR => "Tencent-Hunyuan/HunyuanOCR",
};
let model_path = match args.weight_path {
Some(path) => path,
src/models/common/mod.rs
+119 -41
View File
@@ -1,27 +1,41 @@
use anyhow::Result;
use candle_core::{D, Tensor};
use candle_nn::{Activation, Linear, Module, VarBuilder, linear, linear_no_bias};
use candle_nn::{
Activation, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder,
conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias,
};
use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv};
#[derive(Debug, Clone)]
pub struct MLPWithBias {
pub struct GateUpDownMLP {
gate_proj: Linear,
up_proj: Linear,
down_proj: Linear,
act_fn: Activation,
}
impl MLPWithBias {
impl GateUpDownMLP {
pub fn new(
vb: VarBuilder,
hidden_size: usize,
intermediate_size: usize,
act_fn: Activation,
bias: bool,
) -> Result<Self> {
let gate_proj = linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
let up_proj = linear(hidden_size, intermediate_size, vb.pp("up_proj"))?;
let down_proj = linear(intermediate_size, hidden_size, vb.pp("down_proj"))?;
let (gate_proj, up_proj, down_proj) = if bias {
(
linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?,
linear(hidden_size, intermediate_size, vb.pp("up_proj"))?,
linear(intermediate_size, hidden_size, vb.pp("down_proj"))?,
)
} else {
(
linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?,
linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?,
linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?,
)
};
Ok(Self {
gate_proj,
up_proj,
@@ -31,7 +45,7 @@ impl MLPWithBias {
}
}
impl Module for MLPWithBias {
impl Module for GateUpDownMLP {
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)?;
@@ -39,43 +53,51 @@ impl Module for MLPWithBias {
}
}
#[derive(Debug, Clone)]
pub struct MLPNoBias {
gate_proj: Linear,
up_proj: Linear,
down_proj: Linear,
act_fn: Activation,
pub struct TwoLinearMLP {
linear1: Linear,
linear2: Linear,
act: Activation,
}
impl MLPNoBias {
impl TwoLinearMLP {
pub fn new(
vb: VarBuilder,
hidden_size: usize,
intermediate_size: usize,
act_fn: Activation,
embedding_dim: usize,
mlp_dim: usize,
act: Activation,
bias: bool,
linear1_pp_name: &str,
linear2_pp_name: &str,
) -> Result<Self> {
let gate_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
let up_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?;
let down_proj = linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?;
let (linear1, linear2) = if bias {
(
linear(embedding_dim, mlp_dim, vb.pp(linear1_pp_name))?,
linear(mlp_dim, embedding_dim, vb.pp(linear2_pp_name))?,
)
} else {
(
linear_no_bias(embedding_dim, mlp_dim, vb.pp(linear1_pp_name))?,
linear_no_bias(mlp_dim, embedding_dim, vb.pp(linear2_pp_name))?,
)
};
Ok(Self {
gate_proj,
up_proj,
down_proj,
act_fn,
linear1,
linear2,
act,
})
}
}
impl Module for MLPNoBias {
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)
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
let xs = xs
.apply(&self.linear1)?
.apply(&self.act)?
.apply(&self.linear2)?;
Ok(xs)
}
}
#[derive(Debug, Clone)]
pub struct AttentionNobias {
// pub struct AttentionNobias {
pub struct NaiveAttention {
q_proj: Linear,
k_proj: Linear,
v_proj: Linear,
@@ -88,19 +110,33 @@ pub struct AttentionNobias {
kv_cache: Option<(Tensor, Tensor)>,
}
impl AttentionNobias {
// impl AttentionNobias {
impl NaiveAttention {
pub fn new(
vb: VarBuilder,
hidden_size: usize,
num_attention_heads: usize,
num_key_value_heads: usize,
bias: bool,
) -> Result<Self> {
let num_kv_groups = num_attention_heads / num_key_value_heads;
let head_dim = hidden_size / num_attention_heads;
let q_proj = linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
let k_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?;
let v_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?;
let o_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("o_proj"))?;
let (q_proj, k_proj, v_proj, o_proj) = if bias {
(
linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?,
linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?,
linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?,
linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?,
)
} else {
(
linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?,
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?,
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?,
linear_no_bias(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?,
)
};
Ok(Self {
q_proj,
k_proj,
@@ -118,8 +154,8 @@ impl AttentionNobias {
pub fn forward(
&self,
xs: &Tensor,
cos: &Tensor,
sin: &Tensor,
cos: Option<&Tensor>,
sin: Option<&Tensor>,
attention_mask: Option<&Tensor>,
tof32: bool,
) -> Result<Tensor> {
@@ -136,8 +172,14 @@ impl AttentionNobias {
let value_states = value_states
.reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
.transpose(1, 2)?;
let (query_states, key_states) =
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?;
let (query_states, key_states) = if let Some(cos) = cos
&& let Some(sin) = sin
{
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?
} else {
(query_states, key_states)
};
let scale = 1f64 / f64::sqrt(self.head_dim as f64);
let attn_output = eager_attention_forward(
&query_states,
@@ -260,3 +302,39 @@ pub fn eager_attention_forward(
Ok(attn_output)
}
pub fn get_conv2d(
vb: VarBuilder,
in_c: usize,
out_c: usize,
kernel_size: usize,
padding: usize,
stride: usize,
dilation: usize,
groups: usize,
bias: bool,
) -> Result<Conv2d> {
let cfg = Conv2dConfig {
padding,
stride,
dilation,
groups,
cudnn_fwd_algo: None,
};
let conv2d = if bias {
conv2d(in_c, out_c, kernel_size, cfg, vb)?
} else {
conv2d_no_bias(in_c, out_c, kernel_size, cfg, vb)?
};
Ok(conv2d)
}
pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result<LayerNorm> {
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 norm = layer_norm(dim, ln_config, vb)?;
Ok(norm)
}
src/models/deepseek_ocr/generate.rs
+4 -2
View File
@@ -29,6 +29,7 @@ pub struct DeepseekOCRGenerateModel {
eos_token_id: u32,
device: Device,
size: Vec<u32>,
model_name: String,
}
impl DeepseekOCRGenerateModel {
@@ -54,6 +55,7 @@ impl DeepseekOCRGenerateModel {
eos_token_id,
device: device.clone(),
size,
model_name: "deepseek-ocr".to_string(),
})
}
}
@@ -127,7 +129,7 @@ impl GenerateModel for DeepseekOCRGenerateModel {
}
let res = self.tokenizer.token_decode(generate)?;
self.deepseekocr_model.clear_kv_cache();
let response = build_completion_response(res, "deepseek_ocr");
let response = build_completion_response(res, &self.model_name);
Ok(response)
}
@@ -218,7 +220,7 @@ impl GenerateModel for DeepseekOCRGenerateModel {
continue;
}
error_tokens.clear();
let chunk = build_completion_chunk_response(decoded_token, "deepseek_ocr", None, None);
let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
yield Ok(chunk);
if next_token == self.bos_token_id || next_token == self.eos_token_id {
break;
src/models/deepseek_ocr/model.rs
+15 -41
View File
@@ -10,7 +10,7 @@ use candle_transformers::models::segment_anything::LayerNorm2d;
use crate::{
models::{
common::{AttentionNobias, MLPNoBias, eager_attention_forward},
common::{GateUpDownMLP, NaiveAttention, TwoLinearMLP, eager_attention_forward},
deepseek_ocr::config::{DeepseekOCRConfig, DeepseekV2Config},
},
position_embed::rope::RoPE,
@@ -227,41 +227,11 @@ impl Attention {
}
}
pub struct MLPBlock {
linear1: Linear,
linear2: Linear,
act: Activation,
}
impl MLPBlock {
pub fn new(
vb: VarBuilder,
embedding_dim: usize,
mlp_dim: usize,
act: Activation,
) -> Result<Self> {
let linear1 = linear(embedding_dim, mlp_dim, vb.pp("lin1"))?;
let linear2 = linear(mlp_dim, embedding_dim, vb.pp("lin2"))?;
Ok(Self {
linear1,
linear2,
act,
})
}
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
let xs = xs
.apply(&self.linear1)?
.apply(&self.act)?
.apply(&self.linear2)?;
Ok(xs)
}
}
pub struct Block {
norm1: LayerNorm,
attn: Attention,
norm2: LayerNorm,
mlp: MLPBlock,
mlp: TwoLinearMLP,
window_size: usize,
}
@@ -300,7 +270,7 @@ impl Block {
)?;
let norm2 = layer_norm(dim, ln_config, vb.pp("norm2"))?;
let mlp_dim = (dim as f32 * mlp_ratio) as usize;
let mlp = MLPBlock::new(vb.pp("mlp"), dim, mlp_dim, act)?;
let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, act, true, "lin1", "lin2")?;
Ok(Self {
norm1,
attn,
@@ -924,9 +894,9 @@ pub struct DeepseekV2MoE {
// ep_size: usize,
// experts_per_rank: usize,
// ep_rank: usize,
experts: Vec<MLPNoBias>,
experts: Vec<GateUpDownMLP>,
gate: MoEGate,
shared_experts: MLPNoBias,
shared_experts: GateUpDownMLP,
}
impl DeepseekV2MoE {
@@ -937,20 +907,22 @@ impl DeepseekV2MoE {
let mut experts = Vec::new();
let vb_experts = vb.pp("experts");
for i in 0..config.n_routed_experts {
let mlp = MLPNoBias::new(
let mlp = GateUpDownMLP::new(
vb_experts.pp(i),
config.hidden_size,
config.moe_intermediate_size,
Activation::Silu,
false,
)?;
experts.push(mlp);
}
let gate = MoEGate::new(vb.pp("gate"), config)?;
let shared_experts = MLPNoBias::new(
let shared_experts = GateUpDownMLP::new(
vb.pp("shared_experts"),
config.hidden_size,
config.moe_intermediate_size * config.n_shared_experts,
Activation::Silu,
false,
)?;
Ok(Self {
// num_experts_per_tok: config.num_experts_per_tok,
@@ -1014,7 +986,7 @@ impl Module for DeepseekV2MoE {
pub enum DeepseekV2Proj {
MOE(DeepseekV2MoE),
MLP(MLPNoBias),
MLP(GateUpDownMLP),
}
impl DeepseekV2Proj {
@@ -1033,7 +1005,7 @@ impl DeepseekV2Proj {
}
pub struct DeepseekV2DecoderLayer {
self_attn: AttentionNobias,
self_attn: NaiveAttention,
mlp: DeepseekV2Proj,
input_layernorm: RmsNorm,
post_attention_layernorm: RmsNorm,
@@ -1041,22 +1013,24 @@ pub struct DeepseekV2DecoderLayer {
impl DeepseekV2DecoderLayer {
pub fn new(vb: VarBuilder, config: &DeepseekV2Config, layer_id: usize) -> Result<Self> {
let self_attn = AttentionNobias::new(
let self_attn = NaiveAttention::new(
vb.pp("self_attn"),
config.hidden_size,
config.num_attention_heads,
config.num_key_value_heads,
false,
)?;
let mlp = if layer_id >= config.first_k_dense_replace
&& layer_id.is_multiple_of(config.moe_layer_freq)
{
DeepseekV2Proj::MOE(DeepseekV2MoE::new(vb.pp("mlp"), config)?)
} else {
DeepseekV2Proj::MLP(MLPNoBias::new(
DeepseekV2Proj::MLP(GateUpDownMLP::new(
vb.pp("mlp"),
config.hidden_size,
config.intermediate_size,
Activation::Silu,
false,
)?)
};
let input_layernorm = rms_norm(
src/models/hunyuan_ocr/config.rs
+105
View File
@@ -0,0 +1,105 @@
use candle_nn::Activation;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct HunYuanVLConfig {
pub attention_bias: bool,
pub attention_dropout: f64,
pub attention_head_dim: usize,
pub bos_token_id: u32,
pub eod_token_id: u32,
pub eos_token_id: u32,
pub head_dim: usize,
pub hidden_act: Activation,
pub hidden_size: usize,
pub image_start_token_id: u32,
pub image_end_token_id: u32,
pub image_token_id: u32,
pub image_newline_token_id: u32,
pub initializer_range: f64,
pub intermediate_size: usize,
pub max_position_embeddings: usize,
pub mlp_bias: bool,
pub norm_type: String,
pub num_attention_heads: usize,
pub num_experts: usize,
pub num_hidden_layers: usize,
pub num_key_value_heads: usize,
pub org_vocab_size: usize,
pub pad_id: i32,
pub pad_token_id: i32,
pub pretraining_tp: i32,
pub rms_norm_eps: f64,
pub rope_scaling: HunYuanVLRopeScaling,
pub rope_theta: f64,
pub routed_scaling_factor: f64,
pub sep_token_id: u32,
pub text_end_id: u32,
pub text_start_id: u32,
pub tie_word_embeddings: bool,
pub dtype: String,
pub use_cache: bool,
pub use_qk_norm: bool,
pub use_cla: bool,
pub vision_config: HunYuanVLVisionConfig,
pub vocab_size: usize,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct HunYuanVLRopeScaling {
pub alpha: f64,
pub beta_fast: i32,
pub beta_slow: i32,
pub factor: f64,
pub mscale: f64,
pub mscale_all_dim: f64,
#[serde(rename = "type")]
pub type_field: String,
pub xdrope_section: Vec<usize>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct HunYuanVLVisionConfig {
pub add_patchemb_bias: bool,
pub attention_dropout: f64,
pub cat_extra_token: i32,
pub hidden_act: Activation,
pub hidden_dropout: f64,
pub hidden_size: usize,
pub img_max_token_num: usize,
pub intermediate_size: usize,
pub interpolate_mode: String,
pub max_image_size: usize,
pub max_vit_seq_len: usize,
pub num_attention_heads: usize,
pub num_channels: usize,
pub num_hidden_layers: usize,
pub out_hidden_size: usize,
pub patch_size: usize,
pub rms_norm_eps: f64,
pub spatial_merge_size: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct HunyuanOCRGenerationConfig {
pub bos_token_id: usize,
pub pad_token_id: usize,
pub do_sample: bool,
pub eos_token_id: Vec<usize>,
pub top_p: f32,
pub top_k: usize,
pub temperature: f32,
pub repetition_penalty: f32,
}
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct HunyuanOCRPreprocessorConfig {
pub min_pixels: usize,
pub max_pixels: usize,
pub patch_size: usize,
pub resample: usize,
pub temporal_patch_size: usize,
pub merge_size: usize,
pub image_mean: Vec<f32>,
pub image_std: Vec<f32>,
}
src/models/hunyuan_ocr/generate.rs
+218
View File
@@ -0,0 +1,218 @@
use aha_openai_dive::v1::resources::chat::{
ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
};
use anyhow::{Result, anyhow};
use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use rocket::async_stream::stream;
use rocket::futures::Stream;
use crate::{
chat_template::ChatTemplate,
models::{
GenerateModel,
hunyuan_ocr::{
config::{HunYuanVLConfig, HunyuanOCRGenerationConfig},
model::HunyuanVLModel,
processor::HunyuanVLProcessor,
},
},
tokenizer::TokenizerModel,
utils::{
build_completion_chunk_response, build_completion_response, find_type_files, get_device,
get_dtype, get_logit_processor,
},
};
pub struct HunyuanOCRGenerateModel<'a> {
chat_template: ChatTemplate<'a>,
tokenizer: TokenizerModel,
pre_processor: HunyuanVLProcessor,
hunyuan_vl: HunyuanVLModel,
device: Device,
eos_token_id1: u32,
eos_token_id2: u32,
generation_config: HunyuanOCRGenerationConfig,
model_name: String,
}
impl<'a> HunyuanOCRGenerateModel<'a> {
pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
let chat_template = ChatTemplate::init(path)?;
let tokenizer = TokenizerModel::init(path)?;
let config_path = path.to_string() + "/config.json";
let cfg: HunYuanVLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
let device = get_device(device);
let cfg_dtype = cfg.dtype.as_str();
let dtype = get_dtype(dtype, cfg_dtype);
let pre_processor = HunyuanVLProcessor::new(path, &device, dtype)?;
let model_list = find_type_files(path, "safetensors")?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? };
let hunyuan_vl = HunyuanVLModel::new(vb, cfg.clone())?;
let generation_config_path = path.to_string() + "/generation_config.json";
let generation_config: HunyuanOCRGenerationConfig =
serde_json::from_slice(&std::fs::read(generation_config_path)?)?;
Ok(Self {
chat_template,
tokenizer,
pre_processor,
hunyuan_vl,
device,
eos_token_id1: generation_config.eos_token_id[0] as u32,
eos_token_id2: generation_config.eos_token_id[1] as u32,
generation_config,
model_name: "hunyuan_ocr".to_string(),
})
}
}
impl<'a> GenerateModel for HunyuanOCRGenerateModel<'a> {
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
let temperature = match mes.temperature {
None => self.generation_config.temperature,
Some(tem) => tem,
};
let top_p = match mes.top_p {
None => self.generation_config.top_p,
Some(top_p) => top_p,
};
let top_k = self.generation_config.top_k;
let seed = match mes.seed {
None => 34562u64,
Some(s) => s as u64,
};
let mut logit_processor =
get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed);
let mes_render = self.chat_template.apply_chat_template(&mes)?;
let data = self
.pre_processor
.process_info(&mes, &self.tokenizer, &mes_render)?;
let mut input_ids = data.input_ids;
let mut position_ids = Some(&data.position_ids);
let mut image_mask = Some(&data.image_mask);
let mut pixel_values = data.pixel_values;
let mut image_grid_thw = data.image_grid_thw;
let mut seq_len = input_ids.dim(1)?;
let mut seqlen_offset = 0;
let mut generate: Vec<u32> = Vec::new();
let sample_len = mes.max_tokens.unwrap_or(1024);
for _ in 0..sample_len {
let logits = self.hunyuan_vl.forward(
&input_ids,
pixel_values.as_ref(),
image_grid_thw.as_ref(),
image_mask,
position_ids,
seqlen_offset,
)?;
let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
let next_token = logit_processor.sample(&logits)?;
generate.push(next_token);
if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 {
break;
}
seqlen_offset += seq_len;
seq_len = 1;
input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
position_ids = None;
image_mask = None;
pixel_values = None;
image_grid_thw = None;
}
let res = self.tokenizer.token_decode(generate)?;
self.hunyuan_vl.clear_kv_cache();
let response = build_completion_response(res, &self.model_name);
Ok(response)
}
fn generate_stream(
&mut self,
mes: ChatCompletionParameters,
) -> Result<
Box<
dyn Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>
+ Send
+ Unpin
+ '_,
>,
> {
let temperature = match mes.temperature {
None => self.generation_config.temperature,
Some(tem) => tem,
};
let top_p = match mes.top_p {
None => self.generation_config.top_p,
Some(top_p) => top_p,
};
let top_k = self.generation_config.top_k;
let seed = match mes.seed {
None => 34562u64,
Some(s) => s as u64,
};
let mut logit_processor =
get_logit_processor(Some(temperature), Some(top_p), Some(top_k), seed);
let mes_render = self.chat_template.apply_chat_template(&mes)?;
let data = self
.pre_processor
.process_info(&mes, &self.tokenizer, &mes_render)?;
let mut seqlen_offset = 0;
let sample_len = mes.max_tokens.unwrap_or(1024);
let stream = stream! {
let mut error_tokens = Vec::new();
let mut input_ids = data.input_ids;
let mut position_ids = Some(&data.position_ids);
let mut image_mask = Some(&data.image_mask);
let mut pixel_values = data.pixel_values;
let mut image_grid_thw = data.image_grid_thw;
let mut seq_len = input_ids.dim(1)?;
for _ in 0..sample_len {
let logits = self.hunyuan_vl.forward(
&input_ids,
pixel_values.as_ref(),
image_grid_thw.as_ref(),
image_mask,
position_ids,
seqlen_offset,
)?;
let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
let next_token = logit_processor.sample(&logits)?;
let mut decode_ids = Vec::new();
if !error_tokens.is_empty() {
decode_ids.extend_from_slice(&error_tokens);
}
decode_ids.push(next_token);
let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{}", e)))?;
if decoded_token.contains("�") {
error_tokens.push(next_token);
if error_tokens.len() > 3 {
error_tokens.clear();
}
seqlen_offset += seq_len;
seq_len = 1;
input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
position_ids = None;
image_mask = None;
pixel_values = None;
image_grid_thw = None;
continue;
}
error_tokens.clear();
let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
yield Ok(chunk);
if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 {
break;
}
seqlen_offset += seq_len;
seq_len = 1;
input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
position_ids = None;
image_mask = None;
pixel_values = None;
image_grid_thw = None;
}
self.hunyuan_vl.clear_kv_cache();
};
Ok(Box::new(Box::pin(stream)))
}
}
src/models/hunyuan_ocr/mod.rs
+4
View File
@@ -0,0 +1,4 @@
pub mod config;
pub mod generate;
pub mod model;
pub mod processor;
src/models/hunyuan_ocr/model.rs
+621
View File
@@ -0,0 +1,621 @@
use anyhow::{Result, anyhow};
use candle_core::{D, IndexOp, Tensor};
use candle_nn::{
Conv2d, Embedding, Init, LayerNorm, 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,
},
hunyuan_ocr::config::{HunYuanVLConfig, HunYuanVLVisionConfig},
},
position_embed::rope::{RoPE, apply_rotary_pos_emb, get_xd_cos_sin},
utils::tensor_utils::{
interpolate_bilinear, masked_scatter_dim0, prepare_causal_attention_mask, split_tensor,
},
};
pub struct HunYuanVisionPatchEmbed {
patch_embedding: Conv2d,
// position_embedding: Embedding,
num_channels: usize,
patch_size: usize,
// num_positions: usize,
// position_edge: usize,
embed_dim: usize,
patch_pos_embed: Tensor,
}
impl HunYuanVisionPatchEmbed {
pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result<Self> {
let patch_embedding = get_conv2d(
vb.pp("patch_embedding"),
config.num_channels,
config.hidden_size,
config.patch_size,
0,
config.patch_size,
1,
1,
true,
)?;
let num_channels = config.num_channels;
let patch_size = config.patch_size;
let position_edge = config.max_image_size / patch_size;
let num_positions = (position_edge).pow(2) + 1;
let embed_dim = config.hidden_size;
let position_embedding = embedding(num_positions, embed_dim, vb.pp("position_embedding"))?;
let patch_pos_embed = position_embedding
.embeddings()
.i(1..)?
.reshape((1, position_edge, position_edge, embed_dim))?
.permute((0, 3, 1, 2))?;
Ok(Self {
patch_embedding,
// position_embedding,
num_channels,
patch_size,
// num_positions,
// position_edge,
embed_dim,
patch_pos_embed,
})
}
pub fn forward(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result<Tensor> {
let (num_patches, _) = pixel_values.dims2()?;
let pixel_values = pixel_values.reshape((
num_patches,
self.num_channels,
self.patch_size,
self.patch_size,
))?;
let patch_embeds = self.patch_embedding.forward(&pixel_values)?;
let patch_embeds = patch_embeds
.squeeze(D::Minus1)?
.squeeze(D::Minus1)?
.unsqueeze(0)?;
let mut patch_pos_embed_list = vec![];
let img_num = grid_thw.dim(0)?;
for i in 0..img_num {
let grid_i = grid_thw.i(i)?;
let grid_h = grid_i.i(1)?.to_scalar::<u32>()? as usize;
let grid_w = grid_i.i(2)?.to_scalar::<u32>()? as usize;
let patch_pos_embed_ =
interpolate_bilinear(&self.patch_pos_embed, (grid_h, grid_w), Some(false))?;
let patch_pos_embed_ = patch_pos_embed_
.reshape((self.embed_dim, ()))?
.transpose(0, 1)?
.unsqueeze(0)?;
patch_pos_embed_list.push(patch_pos_embed_);
}
let patch_pos_embed = Tensor::cat(&patch_pos_embed_list, 1)?;
let embedding = patch_embeds.add(&patch_pos_embed)?;
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,
true,
)?;
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,
proj_2: Conv2d,
mlp: Linear,
image_newline: Tensor,
image_begin: Tensor,
image_end: Tensor,
// image_sep: Tensor,
before_rms: RmsNorm,
after_rms: RmsNorm,
}
impl HunYuanVisionPatchMerger {
pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result<Self> {
let proj_0 = get_conv2d(
vb.pp("proj.0"),
config.hidden_size,
config.hidden_size * 2,
config.spatial_merge_size,
0,
config.spatial_merge_size,
1,
1,
true,
)?;
let proj_2 = get_conv2d(
vb.pp("proj.2"),
config.hidden_size * 2,
config.hidden_size * 4,
1,
0,
1,
1,
1,
true,
)?;
let mlp = linear(config.hidden_size * 4, config.out_hidden_size, vb.pp("mlp"))?;
let image_newline =
vb.get_with_hints(config.hidden_size * 4, "image_newline", Init::Const(0.))?;
let image_begin =
vb.get_with_hints(config.out_hidden_size, "image_begin", Init::Const(0.))?;
let image_end = vb.get_with_hints(config.out_hidden_size, "image_end", Init::Const(0.))?;
// let image_sep = vb.get_with_hints(config.out_hidden_size, "image_sep", Init::Const(0.))?;
let before_rms = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("before_rms"))?;
let after_rms = rms_norm(
config.out_hidden_size,
config.rms_norm_eps,
vb.pp("after_rms"),
)?;
Ok(Self {
proj_0,
proj_2,
mlp,
image_newline,
image_begin,
image_end,
// image_sep,
before_rms,
after_rms,
})
}
pub fn forward(&self, xs: &Tensor, size: (usize, usize)) -> Result<Tensor> {
let xs = self.before_rms.forward(xs)?;
let (h, w) = size;
let xs = xs.permute((0, 2, 1))?.reshape((xs.dim(0)?, (), h, w))?;
let xs = self.proj_0.forward(&xs)?.gelu()?;
let xs = self.proj_2.forward(&xs)?;
let (b, c, h, _) = xs.dims4()?;
let image_newline = self
.image_newline
.reshape((1, c, 1, 1))?
.broadcast_as((b, c, h, 1))?
.to_dtype(xs.dtype())?;
let xs = Tensor::cat(&[xs, image_newline], D::Minus1)?;
let xs = xs.reshape((b, c, ()))?.permute((0, 2, 1))?;
let xs = self.mlp.forward(&xs)?;
let begin = self
.image_begin
.reshape((1, 1, ()))?
.broadcast_as((b, 1, xs.dim(D::Minus1)?))?
.to_dtype(xs.dtype())?;
let end = self
.image_end
.reshape((1, 1, ()))?
.broadcast_as((b, 1, xs.dim(D::Minus1)?))?
.to_dtype(xs.dtype())?;
let xs = Tensor::cat(&[begin, xs, end], 1)?;
let xs = self.after_rms.forward(&xs)?;
Ok(xs)
}
}
pub struct HunYuanVisionTransformer {
embeddings: HunYuanVisionPatchEmbed,
layers: Vec<HunYuanVisionBlock>,
perceive: HunYuanVisionPatchMerger,
}
impl HunYuanVisionTransformer {
pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result<Self> {
let embeddings = HunYuanVisionPatchEmbed::new(vb.pp("embeddings"), config)?;
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)?;
layers.push(layer_i);
}
let perceive = HunYuanVisionPatchMerger::new(vb.pp("perceive"), config)?;
Ok(Self {
embeddings,
layers,
perceive,
})
}
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)?;
}
let mut cu_seqlens = vec![];
for i in 0..grid_thw.dim(0)? {
let [_, h, w] = grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
};
cu_seqlens.push((h * w) as usize);
}
let split_items = split_tensor(&hidden_states, &cu_seqlens, 1)?;
let mut processed_item = vec![];
for i in 0..grid_thw.dim(0)? {
let [_, h, w] = grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
};
let processed = self
.perceive
.forward(&split_items[i], (h as usize, w as usize))?;
processed_item.push(processed);
}
let xs = Tensor::cat(&processed_item, 1)?;
Ok(xs)
}
}
pub struct HunYuanVLAttention {
q_proj: Linear,
k_proj: Linear,
v_proj: Linear,
o_proj: Linear,
query_layernorm: RmsNorm,
key_layernorm: RmsNorm,
num_attention_heads: usize,
num_key_value_heads: usize,
num_kv_groups: usize,
head_dim: usize,
scaling: f64,
kv_cache: Option<(Tensor, Tensor)>,
}
impl HunYuanVLAttention {
pub fn new(
vb: VarBuilder,
hidden_size: usize,
head_dim: usize,
num_attention_heads: usize,
num_key_value_heads: usize,
attention_bias: bool,
rms_norm_eps: f64,
) -> Result<Self> {
let num_kv_groups = num_attention_heads / num_key_value_heads;
let scaling = 1f64 / f64::sqrt(head_dim as f64);
let (q_proj, k_proj, v_proj, o_proj) = if attention_bias {
let q_proj = linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
let k_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?;
let v_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?;
let o_proj = linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?;
(q_proj, k_proj, v_proj, o_proj)
} else {
let q_proj =
linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
let k_proj =
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?;
let v_proj =
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?;
let o_proj =
linear_no_bias(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?;
(q_proj, k_proj, v_proj, o_proj)
};
let query_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("query_layernorm"))?;
let key_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("key_layernorm"))?;
Ok(Self {
q_proj,
k_proj,
v_proj,
o_proj,
query_layernorm,
key_layernorm,
num_attention_heads,
num_key_value_heads,
num_kv_groups,
head_dim,
scaling,
kv_cache: None,
})
}
pub fn forward(
&mut self,
xs: &Tensor,
cos: &Tensor,
sin: &Tensor,
attention_mask: Option<&Tensor>,
) -> Result<Tensor> {
let (b_sz, q_len, _) = xs.dims3()?;
let query_states = self
.q_proj
.forward(xs)?
.reshape((b_sz, q_len, self.num_attention_heads, self.head_dim))?
.transpose(1, 2)?;
let key_states = self
.k_proj
.forward(xs)?
.reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))?
.transpose(1, 2)?;
let value_states = self.v_proj.forward(xs)?;
let value_states = value_states
.reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))?
.transpose(1, 2)?;
let (query_states, key_states) =
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?;
let query_states = self.query_layernorm.forward(&query_states)?;
let key_states = self.key_layernorm.forward(&key_states)?;
let (key_states, value_states) = match &self.kv_cache {
None => (key_states, value_states),
Some((prev_k, prev_v)) => {
let key_states = Tensor::cat(&[prev_k, &key_states], 2)?;
let value_states = Tensor::cat(&[prev_v, &value_states], 2)?;
(key_states, value_states)
}
};
self.kv_cache = Some((key_states.clone(), value_states.clone()));
let attn_output = eager_attention_forward(
&query_states,
&key_states,
&value_states,
Some(self.num_kv_groups),
attention_mask,
self.scaling,
)?;
let attn_output =
attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?;
let attn_output = attn_output.apply(&self.o_proj)?;
Ok(attn_output)
}
pub fn clear_kv_cache(&mut self) {
self.kv_cache = None
}
}
pub struct HunYuanVLDecoderLayer {
self_attn: HunYuanVLAttention,
mlp: GateUpDownMLP,
input_layernorm: RmsNorm,
post_attention_layernorm: RmsNorm,
}
impl HunYuanVLDecoderLayer {
pub fn new(config: &HunYuanVLConfig, vb: VarBuilder) -> Result<Self> {
let self_attn = HunYuanVLAttention::new(
vb.pp("self_attn"),
config.hidden_size,
config.head_dim,
config.num_attention_heads,
config.num_key_value_heads,
config.attention_bias,
config.rms_norm_eps,
)?;
let mlp = GateUpDownMLP::new(
vb.pp("mlp"),
config.hidden_size,
config.intermediate_size,
config.hidden_act,
false,
)?;
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(&xs, cos, sin, attention_mask)?;
let xs = residual.add(&xs)?;
let residual = xs.clone();
let xs = self.post_attention_layernorm.forward(&xs)?;
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 struct HunYuanVLTextModel {
embed_tokens: Embedding,
layers: Vec<HunYuanVLDecoderLayer>,
norm: RmsNorm,
rope: RoPE,
xdrope_section: Vec<usize>,
}
impl HunYuanVLTextModel {
pub fn new(vb: VarBuilder, config: &HunYuanVLConfig) -> Result<Self> {
let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?;
let mut layers = vec![];
let vb_layers = vb.pp("layers");
for i in 0..config.num_hidden_layers {
let layer = HunYuanVLDecoderLayer::new(config, vb_layers.pp(i))?;
layers.push(layer);
}
let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?;
let base = config.rope_theta
* config
.rope_scaling
.alpha
.powf(config.head_dim as f64 / (config.head_dim - 2) as f64);
let rope = RoPE::new(config.head_dim, base as f32, vb.device())?;
let xdrope_section = config.rope_scaling.xdrope_section.clone();
Ok(Self {
embed_tokens,
layers,
norm,
rope,
xdrope_section,
})
}
pub fn forward(
&mut self,
inputs_embeds: &Tensor,
position_ids: Option<&Tensor>,
seqlen_offset: usize,
) -> Result<Tensor> {
let (b_size, seq_len, _) = inputs_embeds.dims3()?;
// let position_ids = match position_ids {
// Some(ids) => ids.clone(),
// None => Tensor::arange(
// seqlen_offset as u32,
// (seq_len + seqlen_offset) as u32,
// inputs_embeds.device(),
// )?
// .unsqueeze(0)?,
// };
let attention_mask: Option<&Tensor> = {
if seq_len <= 1 {
None
} else {
Some(&prepare_causal_attention_mask(
b_size,
seq_len,
0,
inputs_embeds.device(),
)?)
}
};
let (cos, sin) = self
.rope
.forward(seqlen_offset, seq_len, inputs_embeds.device())?;
let mut xs = inputs_embeds.clone();
for (i, layer) in self.layers.iter_mut().enumerate() {
if i == 0
&& let Some(position_ids) = position_ids
{
let (cos, sin) =
get_xd_cos_sin(&cos, &sin, position_ids, self.xdrope_section.clone())?;
xs = layer.forward(&xs, &cos, &sin, attention_mask)?;
} else {
xs = layer.forward(&xs, &cos, &sin, attention_mask)?;
}
}
let xs = self.norm.forward(&xs)?;
Ok(xs)
}
pub fn clear_kv_cache(&mut self) {
for layer in self.layers.iter_mut() {
layer.clear_kv_cache()
}
}
}
pub struct HunyuanVLModel {
// config: HunYuanVLConfig,
vit: HunYuanVisionTransformer,
model: HunYuanVLTextModel,
lm_head: Linear,
}
impl HunyuanVLModel {
pub fn new(vb: VarBuilder, config: HunYuanVLConfig) -> Result<Self> {
let vit = HunYuanVisionTransformer::new(vb.pp("vit"), &config.vision_config)?;
let model = HunYuanVLTextModel::new(vb.pp("model"), &config)?;
let lm_head = Linear::new(model.embed_tokens.embeddings().clone(), None);
Ok(Self {
// config,
vit,
model,
lm_head,
})
}
pub fn forward(
&mut self,
input_ids: &Tensor,
pixel_values: Option<&Tensor>,
image_grid_thw: Option<&Tensor>,
image_mask: Option<&Tensor>,
position_ids: Option<&Tensor>,
seqlen_offset: usize,
) -> Result<Tensor> {
let mut inputs_embeds = self.model.embed_tokens.forward(input_ids)?;
if let Some(pixel_values) = pixel_values
&& let Some(grid_thw) = image_grid_thw
&& let Some(image_mask) = image_mask
{
let image_embeds = self.vit.forward(pixel_values, grid_thw)?.squeeze(0)?;
inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embeds, image_mask)?;
}
let outputs = self
.model
.forward(&inputs_embeds, position_ids, seqlen_offset)?;
let seq_len = outputs.dim(1)?;
let hidden_state = outputs.narrow(1, seq_len - 1, 1)?;
let logits = self.lm_head.forward(&hidden_state)?;
Ok(logits)
}
pub fn clear_kv_cache(&mut self) {
self.model.clear_kv_cache();
}
}
src/models/hunyuan_ocr/processor.rs
+236
View File
@@ -0,0 +1,236 @@
use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
use anyhow::Result;
use candle_core::{DType, Device, IndexOp, Shape, Tensor};
use image::DynamicImage;
use crate::{
models::hunyuan_ocr::config::HunyuanOCRPreprocessorConfig,
tokenizer::TokenizerModel,
utils::{
img_utils::{extract_images, img_smart_resize, img_transform},
tensor_utils::{get_eq_indices, get_equal_mask},
},
};
pub struct HunyuanVLProcessor {
image_token_id: u32,
image_token: String,
placeholder_token: String,
process_cfg: HunyuanOCRPreprocessorConfig,
device: Device,
dtype: DType,
}
impl HunyuanVLProcessor {
pub fn new(path: &str, device: &Device, dtype: DType) -> Result<Self> {
let path = path.to_string();
assert!(
std::path::Path::new(&path).exists(),
"model path file not exists"
);
let process_cfg_file = path.clone() + "/preprocessor_config.json";
assert!(
std::path::Path::new(&process_cfg_file).exists(),
"preprocessor_config.json not exists in model path"
);
let process_cfg: HunyuanOCRPreprocessorConfig =
serde_json::from_slice(&std::fs::read(process_cfg_file)?)?;
let image_token_id = 120120u32;
let image_token = "<｜hy_place▁holder▁no▁102｜>".to_string();
let placeholder_token = "<｜hy_place▁holder▁no▁799｜>".to_string();
// let pad_id = 120002u32;
Ok(Self {
image_token_id,
image_token,
placeholder_token,
process_cfg,
device: device.clone(),
dtype,
})
}
pub fn process_img(
&self,
img: &DynamicImage,
img_mean: &Tensor,
img_std: &Tensor,
) -> Result<Tensor> {
let img_h = img.height();
let img_w = img.width();
// h,w resize成 32的倍数
let (resize_h, resize_w) = img_smart_resize(
img_h,
img_w,
(self.process_cfg.patch_size * self.process_cfg.merge_size) as u32,
self.process_cfg.min_pixels as u32,
self.process_cfg.max_pixels as u32,
)?;
let img = img.resize_exact(resize_w, resize_h, image::imageops::FilterType::CatmullRom);
let img_tensor = img_transform(&img, img_mean, img_std, &self.device, self.dtype)?;
// (c, h, w) => (1, c, h, w)
let img_tensor = img_tensor.unsqueeze(0)?;
Ok(img_tensor)
}
pub fn process_vision_tensor(&self, img_tensor: &Tensor) -> Result<(Tensor, Tensor)> {
let channel = img_tensor.dim(1)?;
// img_temsor.dim[0] = 1, temporal_patch_size = 1, grid_t = 1
let grid_t = img_tensor.dim(0)? / self.process_cfg.temporal_patch_size;
let grid_h = img_tensor.dim(2)? / self.process_cfg.patch_size;
let grid_w = img_tensor.dim(3)? / self.process_cfg.patch_size;
let shape = Shape::from(vec![
grid_t,
channel,
grid_h / self.process_cfg.merge_size,
self.process_cfg.merge_size,
self.process_cfg.patch_size,
grid_w / self.process_cfg.merge_size,
self.process_cfg.merge_size,
self.process_cfg.patch_size,
]);
let img_tensor = img_tensor.reshape(shape)?;
// shape to // grid_t,
// grid_h / merge_size,
// merge_size,
// grid_w / merge_size,
// merge_size,
// channel,
// patch_size,
// patch_size,
let img_tensor = img_tensor.permute(vec![0, 2, 3, 5, 6, 1, 4, 7])?;
let img_tensor = img_tensor
.reshape((
grid_t * grid_h * grid_w,
channel * self.process_cfg.patch_size * self.process_cfg.patch_size,
))?
.contiguous()?;
let grid_thw = Tensor::from_vec(
vec![grid_t as u32, grid_h as u32, grid_w as u32],
(1, 3),
&self.device,
)?;
Ok((img_tensor, grid_thw))
}
pub fn process_images(
&self,
imgs: &Vec<DynamicImage>,
img_mean: &Tensor,
img_std: &Tensor,
) -> Result<(Tensor, Tensor)> {
let mut pixel_values_vec = Vec::new();
let mut vision_grid_thws_vec = Vec::new();
for img in imgs {
let img_tensor = self.process_img(img, img_mean, img_std)?;
let (img_tensor, grid_thw) = self.process_vision_tensor(&img_tensor)?;
pixel_values_vec.push(img_tensor);
vision_grid_thws_vec.push(grid_thw);
}
let pixel_values = Tensor::cat(&pixel_values_vec, 0)?;
let vision_grid_thws = Tensor::cat(&vision_grid_thws_vec, 0)?;
Ok((pixel_values, vision_grid_thws))
}
pub fn process_info(
&self,
messages: &ChatCompletionParameters,
tokenizer: &TokenizerModel,
text: &str,
) -> Result<HunyuanData> {
let imgs = extract_images(messages)?;
let img_mean = Tensor::from_slice(&self.process_cfg.image_mean, (3, 1, 1), &self.device)?
.to_dtype(self.dtype)?;
let img_std = Tensor::from_slice(&self.process_cfg.image_std, (3, 1, 1), &self.device)?
.to_dtype(self.dtype)?;
let (pixel_values, image_grid_thw) = if !imgs.is_empty() {
let (pixel_values, image_grid_thw) = self.process_images(&imgs, &img_mean, &img_std)?;
(Some(pixel_values), Some(image_grid_thw))
} else {
(None, None)
};
let mut image_tokens_cumsum = vec![0];
let mut text = text.to_string();
if !imgs.is_empty()
&& let Some(grid_thw) = image_grid_thw.as_ref()
{
let mut index = 0;
while text.contains(&self.image_token) {
let grid_i = grid_thw.i(index)?;
let grid_h = grid_i.i(1)?.to_scalar::<u32>()?;
let grid_w = grid_i.i(2)?.to_scalar::<u32>()?;
let patch_h = grid_h / self.process_cfg.merge_size as u32;
let patch_w = grid_w / self.process_cfg.merge_size as u32;
let num_image_tokens = patch_h * (patch_w + 1) + 2;
let num_id = image_tokens_cumsum[image_tokens_cumsum.len() - 1] + num_image_tokens;
image_tokens_cumsum.push(num_id);
let replace = self.placeholder_token.repeat(num_image_tokens as usize);
text = text.replacen(&self.image_token, &replace, 1);
index += 1;
}
}
text = text.replace(&self.placeholder_token, &self.image_token);
let input_ids = tokenizer.text_encode(text, &self.device)?;
let seq_len = input_ids.dim(1)?;
let position_ids = Tensor::arange(0, seq_len as u32, &self.device)?;
let mut position_ids_w = Tensor::arange(0, seq_len as u32, &self.device)?;
let mut position_ids_h = Tensor::arange(0, seq_len as u32, &self.device)?;
let mut position_ids_t = Tensor::arange(0, seq_len as u32, &self.device)?;
if !imgs.is_empty()
&& let Some(grid_thw) = image_grid_thw.as_ref()
{
let image_token_pos_indices = get_eq_indices(&input_ids.i(0)?, self.image_token_id)?;
for i in 0..grid_thw.dim(0)? {
let grid_i = grid_thw.i(i)?;
let grid_h = grid_i.i(1)?.to_scalar::<u32>()?;
let grid_w = grid_i.i(2)?.to_scalar::<u32>()?;
let patch_h = grid_h / self.process_cfg.merge_size as u32;
let patch_w = grid_w / self.process_cfg.merge_size as u32;
let start_pos = image_token_pos_indices
.i(image_tokens_cumsum[i] as usize)?
.to_scalar::<u32>()? as usize
+ 1;
let replace_num = ((patch_w + 1) * patch_h) as usize;
let pos_w: Vec<u32> = (0..patch_h).flat_map(|_| 0u32..patch_w + 1).collect();
position_ids_w = position_ids_w.slice_assign(
&[start_pos..start_pos + replace_num],
&Tensor::new(pos_w, &self.device)?,
)?;
let pos_h: Vec<u32> = (0..patch_h)
.flat_map(|h| vec![h; (patch_w + 1) as usize])
.collect();
position_ids_h = position_ids_h.slice_assign(
&[start_pos..start_pos + replace_num],
&Tensor::new(pos_h, &self.device)?,
)?;
position_ids_t = position_ids_t.slice_assign(
&[start_pos..start_pos + replace_num],
&Tensor::new(vec![0u32; replace_num], &self.device)?,
)?;
}
}
let position_ids = Tensor::stack(
&[position_ids, position_ids_h, position_ids_w, position_ids_t],
0,
)?
.unsqueeze(0)?;
let image_mask = get_equal_mask(&input_ids, self.image_token_id)?;
let data = HunyuanData {
input_ids,
position_ids,
image_mask,
pixel_values,
image_grid_thw,
};
Ok(data)
}
}
pub struct HunyuanData {
pub input_ids: Tensor,
pub position_ids: Tensor,
pub image_mask: Tensor,
pub pixel_values: Option<Tensor>,
pub image_grid_thw: Option<Tensor>,
}
src/models/minicpm4/generate.rs
+4 -2
View File
@@ -23,6 +23,7 @@ pub struct MiniCPMGenerateModel<'a> {
device: Device,
endoftext_id: u32,
im_end_id: u32,
model_name: String,
}
impl<'a> MiniCPMGenerateModel<'a> {
@@ -47,6 +48,7 @@ impl<'a> MiniCPMGenerateModel<'a> {
device: device.clone(),
endoftext_id,
im_end_id,
model_name: "minicpm4".to_string(),
})
}
}
@@ -78,7 +80,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
}
let res = self.tokenizer.token_decode(generate)?;
self.minicpm.clear_kv_cache();
let response = build_completion_response(res, "minicpm");
let response = build_completion_response(res, &self.model_name);
Ok(response)
}
fn generate_stream(
@@ -128,7 +130,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
continue;
}
error_tokens.clear();
let chunk = build_completion_chunk_response(decoded_token, "minicpm", None, None);
let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
yield Ok(chunk);
if next_token == self.endoftext_id || next_token == self.im_end_id {
break;
src/models/minicpm4/model.rs
+8 -6
View File
@@ -4,7 +4,7 @@ use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_n
use crate::{
models::{
common::{AttentionNobias, MLPNoBias},
common::{GateUpDownMLP, NaiveAttention},
minicpm4::config::MiniCPM4Config,
},
position_embed::rope::compute_default_rope_parameters,
@@ -93,8 +93,8 @@ impl MiniCPMLongRoPE {
}
pub struct MiniCPMDecoderLayer {
self_attn: AttentionNobias,
mlp: MLPNoBias,
self_attn: NaiveAttention,
mlp: GateUpDownMLP,
input_layernorm: RmsNorm,
post_attention_layernorm: RmsNorm,
scale_depth: f32,
@@ -103,17 +103,19 @@ pub struct MiniCPMDecoderLayer {
impl MiniCPMDecoderLayer {
pub fn new(vb: VarBuilder, cfg: &MiniCPM4Config) -> Result<Self> {
let self_attn = AttentionNobias::new(
let self_attn = NaiveAttention::new(
vb.pp("self_attn"),
cfg.hidden_size,
cfg.num_attention_heads,
cfg.num_key_value_heads,
false,
)?;
let mlp = MLPNoBias::new(
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"))?;
@@ -143,7 +145,7 @@ impl MiniCPMDecoderLayer {
let xs = self.input_layernorm.forward(xs)?;
let xs = self
.self_attn
.forward(&xs, cos, sin, attention_mask, true)?;
.forward(&xs, Some(cos), Some(sin), attention_mask, true)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
src/models/mod.rs
+12 -1
View File
@@ -1,5 +1,6 @@
pub mod common;
pub mod deepseek_ocr;
pub mod hunyuan_ocr;
pub mod minicpm4;
pub mod qwen2_5vl;
pub mod qwen3vl;
@@ -12,7 +13,8 @@ use anyhow::Result;
use rocket::futures::Stream;
use crate::models::{
deepseek_ocr::generate::DeepseekOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel,
deepseek_ocr::generate::DeepseekOCRGenerateModel,
hunyuan_ocr::generate::HunyuanOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel,
qwen2_5vl::generate::Qwen2_5VLGenerateModel, qwen3vl::generate::Qwen3VLGenerateModel,
};
@@ -34,6 +36,8 @@ pub enum WhichModel {
Qwen3vl32B,
#[value(name = "deepseek-ocr")]
DeepSeekOCR,
#[value(name = "hunyuan-ocr")]
HunyuanOCR,
}
pub trait GenerateModel {
@@ -56,6 +60,7 @@ pub enum ModelInstance<'a> {
Qwen2_5VL(Qwen2_5VLGenerateModel<'a>),
Qwen3VL(Qwen3VLGenerateModel<'a>),
DeepSeekOCR(DeepseekOCRGenerateModel),
HunyuanOCR(HunyuanOCRGenerateModel<'a>),
}
impl<'a> GenerateModel for ModelInstance<'a> {
@@ -65,6 +70,7 @@ impl<'a> GenerateModel for ModelInstance<'a> {
ModelInstance::Qwen2_5VL(model) => model.generate(mes),
ModelInstance::Qwen3VL(model) => model.generate(mes),
ModelInstance::DeepSeekOCR(model) => model.generate(mes),
ModelInstance::HunyuanOCR(model) => model.generate(mes),
}
}
@@ -84,6 +90,7 @@ impl<'a> GenerateModel for ModelInstance<'a> {
ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes),
ModelInstance::Qwen3VL(model) => model.generate_stream(mes),
ModelInstance::DeepSeekOCR(model) => model.generate_stream(mes),
ModelInstance::HunyuanOCR(model) => model.generate_stream(mes),
}
}
}
@@ -122,6 +129,10 @@ pub fn load_model(model_type: WhichModel, path: &str) -> Result<ModelInstance<'_
let model = DeepseekOCRGenerateModel::init(path, None, None)?;
ModelInstance::DeepSeekOCR(model)
}
WhichModel::HunyuanOCR => {
let model = HunyuanOCRGenerateModel::init(path, None, None)?;
ModelInstance::HunyuanOCR(model)
}
};
Ok(model)
}
src/models/qwen2_5vl/generate.rs
+4 -2
View File
@@ -30,6 +30,7 @@ pub struct Qwen2_5VLGenerateModel<'a> {
device: Device,
endoftext_id: u32,
im_end_id: u32,
model_name: String,
}
impl<'a> Qwen2_5VLGenerateModel<'a> {
@@ -57,6 +58,7 @@ impl<'a> Qwen2_5VLGenerateModel<'a> {
device: device.clone(),
endoftext_id,
im_end_id,
model_name: "qwen2.5vl".to_string(),
})
}
}
@@ -119,7 +121,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
}
let res = self.tokenizer.token_decode(generate)?;
self.qwen2_5_vl.clear_kv_cache();
let response = build_completion_response(res, "qwen2.5vl");
let response = build_completion_response(res, &self.model_name);
Ok(response)
}
@@ -202,7 +204,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
continue;
}
error_tokens.clear();
let chunk = build_completion_chunk_response(decoded_token, "qwen2.5vl", None, None);
let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
yield Ok(chunk);
if next_token == self.endoftext_id || next_token == self.im_end_id {
break;
src/models/qwen3vl/generate.rs
+4 -2
View File
@@ -33,6 +33,7 @@ pub struct Qwen3VLGenerateModel<'a> {
eos_token_id1: u32,
eos_token_id2: u32,
generation_config: Qwen3VLGenerationConfig,
model_name: String,
}
impl<'a> Qwen3VLGenerateModel<'a> {
@@ -60,6 +61,7 @@ impl<'a> Qwen3VLGenerateModel<'a> {
eos_token_id1: generation_config.eos_token_id[0] as u32,
eos_token_id2: generation_config.eos_token_id[1] as u32,
generation_config,
model_name: "qwen3vl".to_string(),
})
}
}
@@ -120,7 +122,7 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
}
let res = self.tokenizer.token_decode(generate)?;
self.qwen3_vl.clear_kv_cache();
let response = build_completion_response(res, "qwen3vl");
let response = build_completion_response(res, &self.model_name);
Ok(response)
}
@@ -201,7 +203,7 @@ impl<'a> GenerateModel for Qwen3VLGenerateModel<'a> {
continue;
}
error_tokens.clear();
let chunk = build_completion_chunk_response(decoded_token, "qwen3vl", None, None);
let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
yield Ok(chunk);
if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 {
break;
src/models/qwen3vl/model.rs
+4 -3
View File
@@ -7,7 +7,7 @@ use candle_nn::{
use crate::{
models::{
common::{MLPNoBias, eager_attention_forward},
common::{GateUpDownMLP, eager_attention_forward},
qwen3vl::config::{Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig},
},
position_embed::rope::{
@@ -664,7 +664,7 @@ impl Qwen3VLTextAttention {
pub struct Qwen3VLTextDecoderLayer {
self_attn: Qwen3VLTextAttention,
mlp: MLPNoBias,
mlp: GateUpDownMLP,
input_layernorm: RmsNorm,
post_attention_layernorm: RmsNorm,
}
@@ -672,11 +672,12 @@ pub struct Qwen3VLTextDecoderLayer {
impl Qwen3VLTextDecoderLayer {
pub fn new(config: Qwen3VLTextConfig, vb: VarBuilder) -> Result<Self> {
let self_attn = Qwen3VLTextAttention::new(config.clone(), vb.pp("self_attn"))?;
let mlp = MLPNoBias::new(
let mlp = GateUpDownMLP::new(
vb.pp("mlp"),
config.hidden_size,
config.intermediate_size,
config.hidden_act,
false,
)?;
let input_layernorm = rms_norm(
config.hidden_size,
src/models/qwen3vl/processor.rs
+6 -49
View File
@@ -11,7 +11,11 @@ use num::integer::lcm;
use crate::{
models::qwen3vl::config::PreprocessorConfig,
utils::{ceil_by_factor, floor_by_factor, img_utils::get_image, round_by_factor},
utils::{
ceil_by_factor, floor_by_factor,
img_utils::{get_image, img_smart_resize, img_transform},
round_by_factor,
},
};
#[derive(Clone)]
@@ -136,22 +140,9 @@ impl Qwen3VLProcessor {
(self.img_process_cfg.patch_size * self.img_process_cfg.merge_size) as u32,
self.img_process_cfg.size.shortest_edge as u32,
self.img_process_cfg.size.longest_edge as u32,
None,
)?;
let img = img.resize_exact(resize_w, resize_h, image::imageops::FilterType::CatmullRom);
let img_vec = img.to_rgb8().into_raw();
// (h, w, c) => (c, h, w)
let img_tensor = Tensor::from_slice(
&img_vec,
(resize_h as usize, resize_w as usize, 3),
&self.device,
)?
.permute((2, 0, 1))?
.to_dtype(self.dtype)?;
// 0-255 rescale to 0-1
let img_tensor = img_tensor.affine(1.0 / 255.0, 0.)?;
// normalize
let img_tensor = img_tensor.broadcast_sub(img_mean)?.broadcast_div(img_std)?;
let img_tensor = img_transform(&img, img_mean, img_std, &self.device, self.dtype)?;
// (c, h, w) => (1, c, h, w)
let img_tensor = img_tensor.unsqueeze(0)?;
Ok(img_tensor)
@@ -426,40 +417,6 @@ impl Qwen3VLProcessor {
}
}
pub fn img_smart_resize(
img_h: u32,
img_w: u32,
factor: u32,
min_pixels: u32,
max_pixels: u32,
video_ratio: Option<u32>,
) -> Result<(u32, u32)> {
if std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w) > 200 {
return Err(anyhow!(format!(
"absolute aspect ratio mush be smaller than {}, got {}",
200,
std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w)
)));
}
let mut image_factor = factor;
if let Some(ratio) = video_ratio {
image_factor = lcm(image_factor, ratio);
}
let mut h_bar = std::cmp::max(image_factor, round_by_factor(img_h, image_factor));
let mut w_bar = std::cmp::max(image_factor, round_by_factor(img_w, image_factor));
if h_bar * w_bar > max_pixels {
let beta = ((img_h * img_w) as f32 / max_pixels as f32).sqrt();
h_bar = floor_by_factor(img_h as f32 / beta, image_factor);
w_bar = floor_by_factor(img_w as f32 / beta, image_factor);
} else if h_bar * w_bar < min_pixels {
let beta = (min_pixels as f32 / (img_h * img_w) as f32).sqrt();
h_bar = ceil_by_factor(img_h as f32 * beta, image_factor);
w_bar = ceil_by_factor(img_w as f32 * beta, image_factor);
}
Ok((h_bar, w_bar))
}
pub fn video_smart_resize(
num_frames: u32,
height: u32,
src/models/voxcpm/minicpm4.rs
+8 -6
View File
@@ -4,7 +4,7 @@ use candle_nn::{Embedding, Module, RmsNorm, VarBuilder, embedding, rms_norm};
use crate::{
models::{
common::{AttentionNobias, MLPNoBias},
common::{GateUpDownMLP, NaiveAttention},
voxcpm::config::VoxMiniCPM4Config,
},
position_embed::rope::compute_default_rope_parameters,
@@ -101,8 +101,8 @@ impl MiniCPMLongRoPE {
}
pub struct MiniCPMDecoderLayer {
self_attn: AttentionNobias,
mlp: MLPNoBias,
self_attn: NaiveAttention,
mlp: GateUpDownMLP,
input_layernorm: RmsNorm,
post_attention_layernorm: RmsNorm,
scale_depth: f32,
@@ -112,17 +112,19 @@ pub struct MiniCPMDecoderLayer {
impl MiniCPMDecoderLayer {
pub fn new(vb: VarBuilder, cfg: &VoxMiniCPM4Config) -> Result<Self> {
let self_attn = AttentionNobias::new(
let self_attn = NaiveAttention::new(
vb.pp("self_attn"),
cfg.hidden_size,
cfg.num_attention_heads,
cfg.num_key_value_heads,
false,
)?;
let mlp = MLPNoBias::new(
let mlp = GateUpDownMLP::new(
vb.pp("mlp"),
cfg.hidden_size,
cfg.intermediate_size,
candle_nn::Activation::Silu,
false,
)?;
let input_layernorm =
rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
@@ -153,7 +155,7 @@ impl MiniCPMDecoderLayer {
let xs = self.input_layernorm.forward(xs)?;
let xs = self
.self_attn
.forward(&xs, cos, sin, attention_mask, true)?;
.forward(&xs, Some(cos), Some(sin), attention_mask, true)?;
let xs = if self.use_mup {
(residual
+ xs.affine(
src/position_embed/rope.rs
+44
View File
@@ -2,6 +2,8 @@ use anyhow::Result;
use candle_core::{D, DType, Device, IndexOp, Tensor};
use candle_transformers::models::deepseek2::SplitOp;
use crate::utils::tensor_utils::{index_select_2d, split_tensor};
pub fn compute_default_rope_parameters(dim: usize, base: f32) -> Vec<f32> {
let inv_freq: Vec<f32> = (0..dim)
.step_by(2)
@@ -303,3 +305,45 @@ impl RoPE {
Ok((cos, sin))
}
}
pub fn get_xd_cos_sin(
cos: &Tensor,
sin: &Tensor,
position_ids: &Tensor,
xdrope_section: Vec<usize>,
) -> Result<(Tensor, Tensor)> {
let x_dim = xdrope_section.len();
// position_ids: (bs, 4, seq_len)
let mut cos_vec = vec![];
let mut sin_vec = vec![];
let bs = position_ids.dim(0)?;
for i in 0..bs {
let pos_i = position_ids.i(i)?;
let cos_i = index_select_2d(cos, &pos_i)?;
let sin_i = index_select_2d(sin, &pos_i)?;
cos_vec.push(cos_i);
sin_vec.push(sin_i);
}
// (bs, 4, seq_len, dim) -> (bs, seq_len, 4, dim)
let cos = Tensor::stack(&cos_vec, 0)?
.permute((0, 2, 1, 3))?
.contiguous()?;
let sin = Tensor::stack(&sin_vec, 0)?
.permute((0, 2, 1, 3))?
.contiguous()?;
let xdrope_section: Vec<usize> = xdrope_section.iter().map(|&i| i * 2).collect();
let cos_select: Vec<Tensor> = split_tensor(&cos, &xdrope_section, D::Minus1)?
.iter()
.enumerate()
.map(|(i, m)| m.i((.., .., i % x_dim)).unwrap())
.collect();
let sin_select: Vec<Tensor> = split_tensor(&sin, &xdrope_section, D::Minus1)?
.iter()
.enumerate()
.map(|(i, m)| m.i((.., .., i % x_dim)).unwrap())
.collect();
let cos = Tensor::cat(&cos_select, D::Minus1)?;
let sin = Tensor::cat(&sin_select, D::Minus1)?;
Ok((cos, sin))
}
src/tokenizer/mod.rs
+1 -1
View File
@@ -3,7 +3,7 @@ use candle_core::{Device, Tensor};
use tokenizers::Tokenizer;
pub struct TokenizerModel {
tokenizer: Tokenizer,
pub tokenizer: Tokenizer,
}
impl TokenizerModel {
src/utils/img_utils.rs
+38
View File
@@ -9,6 +9,8 @@ use base64::{Engine, engine::general_purpose};
use candle_core::{DType, Device, Tensor};
use image::{DynamicImage, ImageBuffer, ImageReader, Rgb, RgbImage, imageops};
use crate::utils::{ceil_by_factor, floor_by_factor, round_by_factor};
pub fn load_image_from_url(url: &str) -> Result<DynamicImage> {
tokio::task::block_in_place(|| {
let response = reqwest::blocking::get(url)
@@ -237,3 +239,39 @@ pub fn img_transform(
.to_dtype(dtype)?;
Ok(img_tensor)
}
pub fn img_smart_resize(
img_h: u32,
img_w: u32,
factor: u32,
min_pixels: u32,
max_pixels: u32,
) -> Result<(u32, u32)> {
if std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w) > 200 {
return Err(anyhow!(format!(
"absolute aspect ratio mush be smaller than {}, got {}",
200,
std::cmp::max(img_h, img_w) / std::cmp::min(img_h, img_w)
)));
}
let image_factor = factor;
let mut h_bar = std::cmp::max(image_factor, round_by_factor(img_h, image_factor));
let mut w_bar = std::cmp::max(image_factor, round_by_factor(img_w, image_factor));
if h_bar * w_bar > max_pixels {
let beta = ((img_h * img_w) as f32 / max_pixels as f32).sqrt();
h_bar = std::cmp::max(
image_factor,
floor_by_factor(img_h as f32 / beta, image_factor),
);
w_bar = std::cmp::max(
image_factor,
floor_by_factor(img_w as f32 / beta, image_factor),
);
} else if h_bar * w_bar < min_pixels {
let beta = (min_pixels as f32 / (img_h * img_w) as f32).sqrt();
h_bar = ceil_by_factor(img_h as f32 * beta, image_factor);
w_bar = ceil_by_factor(img_w as f32 * beta, image_factor);
}
Ok((h_bar, w_bar))
}
src/utils/mod.rs
+40 -1
View File
@@ -3,6 +3,8 @@ pub mod img_utils;
pub mod tensor_utils;
pub mod video_utils;
use std::process::Command;
use aha_openai_dive::v1::resources::{
chat::{
ChatCompletionChoice, ChatCompletionChunkChoice, ChatCompletionChunkResponse,
@@ -31,6 +33,33 @@ pub fn get_device(device: Option<&Device>) -> Device {
}
}
pub fn get_gpu_sm_arch() -> Result<f32> {
let output = Command::new("nvidia-smi")
.arg("--query-gpu=compute_cap")
.arg("--format=csv,noheader")
.output()
.map_err(|e| anyhow::anyhow!(format!("Failed to execute nvidia-smi: {}", e)))?;
if !output.status.success() {
return Err(anyhow::anyhow!(format!(
"nvidia-smi failed with status: {}\nError: {}",
output.status,
String::from_utf8_lossy(&output.stderr)
)));
}
let output_str = String::from_utf8_lossy(&output.stdout);
let output_str = output_str.trim();
let sm_float = match output_str.parse::<f32>() {
Ok(num) => num,
Err(_) => {
return Err(anyhow::anyhow!(format!(
"gpr sm arch: {} parse float32 error",
output_str
)));
}
};
Ok(sm_float)
}
pub fn get_dtype(dtype: Option<DType>, cfg_dtype: &str) -> DType {
match dtype {
Some(d) => d,
@@ -41,7 +70,16 @@ pub fn get_dtype(dtype: Option<DType>, cfg_dtype: &str) -> DType {
"float32" | "float" => DType::F32,
"float64" | "double" => DType::F64,
"float16" => DType::F16,
"bfloat16" => DType::BF16,
"bfloat16" => {
let arch = get_gpu_sm_arch();
match arch {
Err(_) => DType::F16,
Ok(a) => {
// nvidia显卡sm架构>=8.0的才支持BF16
if a >= 8.0 { DType::BF16 } else { DType::F16 }
}
}
}
"uint8" => DType::U8,
"int8" | "int16" | "int32" | "int64" => DType::I64,
_ => DType::F32,
@@ -257,6 +295,7 @@ pub fn get_logit_processor(
top_k: Option<usize>,
seed: u64,
) -> LogitsProcessor {
let temperature = temperature.and_then(|v| if v < 1e-7 { None } else { Some(v) });
match top_k {
None => LogitsProcessor::new(
seed,
src/utils/tensor_utils.rs
+97 -17
View File
@@ -221,6 +221,14 @@ pub fn masked_scatter_dim0(original: &Tensor, replace: &Tensor, mask: &Tensor) -
Ok(original)
}
pub fn get_not_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor> {
let image_token_id_tensor = Tensor::new(vec![token_ids], input_ids.device())?;
let mask = input_ids
.broadcast_ne(&image_token_id_tensor)?
.to_dtype(candle_core::DType::U32)?;
Ok(mask)
}
pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor> {
let image_token_id_tensor = Tensor::new(vec![token_ids], input_ids.device())?;
let mask = input_ids
@@ -229,10 +237,16 @@ pub fn get_equal_mask(input_ids: &Tensor, token_ids: u32) -> Result<Tensor> {
Ok(mask)
}
pub fn get_vision_next_indices(input_ids: &Tensor, token_id: u32) -> Result<Tensor> {
pub fn get_eq_indices(input_ids: &Tensor, token_id: u32) -> Result<Tensor> {
// input_ids -> shape: (seq_len)
let mask = get_equal_mask(input_ids, token_id)?;
let indices = nonzero_index(&mask)?;
Ok(indices)
}
pub fn get_vision_next_indices(input_ids: &Tensor, token_id: u32) -> Result<Tensor> {
// input_ids -> shape: (seq_len)
let indices = get_eq_indices(input_ids, token_id)?;
let indices = indices.broadcast_add(&Tensor::new(vec![1u32], input_ids.device())?)?;
Ok(indices)
}
@@ -384,9 +398,9 @@ pub fn interpolate_linear_1d(
align_corner: Option<bool>,
) -> Result<Tensor> {
// t: [b, channels, features]
if t.rank() < 3 {
if t.rank() != 3 {
return Err(anyhow::anyhow!(
"Input rank must have at least 3 dimensions"
"Input rank must have equal to 3 dimensions"
));
}
let shape = t.dims();
@@ -394,19 +408,13 @@ pub fn interpolate_linear_1d(
if orig_size == target_size {
return Ok(t.clone());
}
let mut reshaped = t.clone();
if shape.len() > 3 {
let bs = shape[0];
let channels = shape[1..shape.len() - 1].iter().product::<usize>();
reshaped = reshaped.reshape((bs, channels, orig_size))?;
}
let (bs, channels, _) = reshaped.dims3()?;
let (bs, channels, _) = t.dims3()?;
let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), t.device())?;
let coords = compute_1d_coords(orig_size, target_size, align_corner)?;
for b in 0..bs {
for c in 0..channels {
let input_slice = reshaped.i((b, c))?;
let input_slice = t.i((b, c))?;
let mut out_i = Vec::new();
// for x_out in 0..target_size {
for &coord in coords.iter().take(target_size) {
@@ -424,16 +432,88 @@ pub fn interpolate_linear_1d(
output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?;
}
}
if shape.len() != 3 {
let mut new_shape = shape.to_vec();
let last_dim = new_shape.len() - 1;
new_shape[last_dim] = target_size;
output = output.reshape(new_shape)?
}
output = output.contiguous()?;
Ok(output)
}
pub fn interpolate_bilinear(
input: &Tensor,
target_size: (usize, usize),
align_corner: Option<bool>,
) -> Result<Tensor> {
// input: [b, channels, height, width]
if input.rank() != 4 {
return Err(anyhow::anyhow!(
"Input rank must have equal to 4 dimensions [b, c, h, w]"
));
}
let (bs, channels, input_height, input_width) = input.dims4()?;
let (target_height, target_width) = target_size;
// If size is the same, return clone
if input_height == target_height && input_width == target_width {
return Ok(input.clone());
}
let align_corners = align_corner.unwrap_or(false);
// Compute scaling factors
let height_scale = if align_corners && target_height > 1 {
(input_height - 1) as f64 / (target_height - 1) as f64
} else {
input_height as f64 / target_height as f64
};
let width_scale = if align_corners && target_width > 1 {
(input_width - 1) as f64 / (target_width - 1) as f64
} else {
input_width as f64 / target_width as f64
};
let dim0 = bs * channels;
let input_3dim = input.reshape((dim0, input_height, input_width))?;
let input_data = input_3dim.to_dtype(DType::F32)?.to_vec3::<f32>()?;
let mut output_data = vec![vec![vec![0.0f32; target_width]; target_height]; dim0];
for c in 0..dim0 {
for out_y in 0..target_height {
let src_y = if align_corners {
out_y as f64 * height_scale
} else {
(out_y as f64 + 0.5) * height_scale - 0.5
};
let src_y = src_y.max(0.0).min((input_height - 1) as f64);
let y0 = src_y.floor() as usize;
let y1 = (y0 + 1).min(input_height - 1);
let dy = (src_y - y0 as f64) as f32;
for out_x in 0..target_width {
let src_x = if align_corners {
out_x as f64 * width_scale
} else {
(out_x as f64 + 0.5) * width_scale - 0.5
};
let src_x = src_x.max(0.0).min((input_width - 1) as f64);
let x0 = src_x.floor() as usize;
let x1 = (x0 + 1).min(input_width - 1);
let q00 = input_data[c][y0][x0];
let q01 = input_data[c][y0][x1];
let q10 = input_data[c][y1][x0];
let q11 = input_data[c][y1][x1];
let dx = (src_x - x0 as f64) as f32;
let interpolated = q00 * (1.0 - dx) * (1.0 - dy)
+ q01 * dx * (1.0 - dy)
+ q10 * (1.0 - dx) * dy
+ q11 * dx * dy;
output_data[c][out_y][out_x] = interpolated;
}
}
}
let output = Tensor::new(output_data, input.device())?
.reshape((bs, channels, target_height, target_width))?
.to_dtype(input.dtype())?;
Ok(output.contiguous()?)
}
fn compute_scale(input_size: usize, output_size: usize, align_corners: bool) -> f64 {
if align_corners && output_size > 1 {
(input_size - 1) as f64 / (output_size - 1) as f64
tests/config_tests.rs
+14 -4
View File
@@ -1,7 +1,7 @@
use aha::models::{
deepseek_ocr::config::DeepseekOCRConfig, minicpm4::config::MiniCPM4Config,
qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig,
voxcpm::config::VoxCPMConfig,
deepseek_ocr::config::DeepseekOCRConfig, hunyuan_ocr::config::HunYuanVLConfig,
minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig,
qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig,
};
use anyhow::Result;
@@ -48,10 +48,20 @@ fn qwen3vl_config() -> Result<()> {
#[test]
fn deepseek_ocr_config() -> Result<()> {
// cargo test -F cuda qwen3vl_config -r -- --nocapture
// cargo test -F cuda deepseek_ocr_config -r -- --nocapture
let model_path = "/home/jhq/huggingface_model/deepseek-ai/DeepSeek-OCR/";
let config_path = model_path.to_string() + "/config.json";
let config: DeepseekOCRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
println!("{:?}", config);
Ok(())
}
#[test]
fn hunyuan_ocr_config() -> Result<()> {
// cargo test -F cuda hunyuan_ocr_config -r -- --nocapture
let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/";
let config_path = model_path.to_string() + "/config.json";
let config: HunYuanVLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
println!("{:?}", config);
Ok(())
}
tests/messy_test.rs
+20 -5
View File
@@ -1,4 +1,6 @@
use aha::utils::tensor_utils::interpolate_bicubic;
use std::time::Instant;
use aha::utils::tensor_utils::interpolate_bilinear;
use anyhow::Result;
use candle_core::Tensor;
@@ -6,11 +8,24 @@ use candle_core::Tensor;
fn messy_test() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture
let device = &candle_core::Device::Cpu;
// let t = Tensor::randn(0.0f32, 1.0, (1, 768, 64, 64), device)?;
let t = Tensor::arange(0.0f32, 10.0, device)?.broadcast_as((1, 1, 10, 10))?;
let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 1, 40, 40))?;
println!("t: {}", t);
let t_resized = interpolate_bicubic(&t, (5, 5), Some(true), Some(false))?;
println!("t_resized: {}", t_resized);
let i_start = Instant::now();
let t_inter = interpolate_bilinear(&t, (20, 20), Some(false))?;
let i_duration = i_start.elapsed();
println!("Time elapsed in interpolate_bilinear is: {:?}", i_duration);
println!("t_inter: {}", t_inter);
// let x: Vec<u32> = (0..5).flat_map(|_| 0u32..10).collect();
// let id: Vec<u32> = (0..5).flat_map(|h| vec![h; 10]).collect();
// println!("x: {:?}", id);
// let t = Tensor::randn(0.0f32, 1.0, (1, 768, 64, 64), device)?;
// let t = Tensor::arange(0u32, 10, device)?.broadcast_as((1, 10))?;
// let eq = t.broadcast_eq(&Tensor::new(5u32, device)?)?;
// println!("eq: {}", eq);
// let t = Tensor::arange(0.0f32, 10.0, device)?.broadcast_as((1, 1, 10, 10))?;
// println!("t: {}", t);
// let t_resized = interpolate_bicubic(&t, (5, 5), Some(true), Some(false))?;
// println!("t_resized: {}", t_resized);
// let t1 = Tensor::rand(0.0, 1.0, (1, 5, 5, 10), device)?;
// let t2 = Tensor::rand(0.0, 1.0, (5, 8, 10), device)?;
// let t2 = t2.t()?;
tests/test_hunyuan_ocr.rs
+88
View File
@@ -0,0 +1,88 @@
use std::{pin::pin, time::Instant};
use aha::models::{GenerateModel, hunyuan_ocr::generate::HunyuanOCRGenerateModel};
use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
use anyhow::Result;
use rocket::futures::StreamExt;
#[test]
fn hunyuan_ocr_generate() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda hunyuan_ocr_generate -r -- --nocapture
let message = r#"
{
"model": "hunyuan-ocr",
"messages": [
{
"role": "user",
"content": [
{
"type": "image",
"image_url":
{
"url": "file://./assets/img/ocr_test1.png"
}
},
{
"type": "text",
"text": "检测并识别图片中的文字，将文本坐标格式化输出。"
}
]
}
]
}
"#;
let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/";
let mes: ChatCompletionParameters = serde_json::from_str(message)?;
let i_start = Instant::now();
let mut model = HunyuanOCRGenerateModel::init(model_path, None, None)?;
let i_duration = i_start.elapsed();
println!("Time elapsed in load model is: {:?}", i_duration);
let i_start = Instant::now();
let res = model.generate(mes)?;
let i_duration = i_start.elapsed();
println!("Time elapsed in generate is: {:?}", i_duration);
println!("generate: \n {:?}", res);
Ok(())
}
#[tokio::test]
async fn hunyuan_ocr_stream() -> Result<()> {
// test with cuda: RUST_BACKTRACE=1 cargo test -F cuda hunyuan_ocr_stream -r -- --nocapture
let message = r#"
{
"model": "hunyuan-ocr",
"messages": [
{
"role": "user",
"content": [
{
"type": "image",
"image_url":
{
"url": "file://./assets/img/ocr_test1.png"
}
},
{
"type": "text",
"text": "检测并识别图片中的文字，将文本坐标格式化输出。"
}
]
}
]
}
"#;
let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/";
let mes: ChatCompletionParameters = serde_json::from_str(message)?;
let i_start = Instant::now();
let mut model = HunyuanOCRGenerateModel::init(model_path, None, None)?;
let i_duration = i_start.elapsed();
println!("Time elapsed in load model is: {:?}", i_duration);
let mut stream = pin!(model.generate_stream(mes)?);
while let Some(item) = stream.next().await {
println!("generate: \n {:?}", item);
}
let i_duration = i_start.elapsed();
println!("Time elapsed in generate is: {:?}", i_duration);
Ok(())
}
tests/test_qwen3vl.rs
+4 -4
View File
@@ -19,15 +19,15 @@ fn qwen3vl_generate() -> Result<()> {
"role": "user",
"content": [
{
"type": "video",
"video_url":
"type": "image",
"image_url":
{
"url": "https://www.w3schools.com/html/movie.mp4"
"url": "file://./assets/img/ocr_test1.png"
}
},
{
"type": "text",
"text": "视频中发生了什么"
"text": "请分析图片并提取所有可见文本内容，按从左到右、从上到下的布局，返回纯文本"
}
]
}
tests/weight_test.rs
+19
View File
@@ -80,3 +80,22 @@ fn deepseekocr_weight() -> Result<()> {
println!("model_list: {:?}", model_list);
Ok(())
}
#[test]
fn hunyuanocr_weight() -> Result<()> {
let model_path = "/home/jhq/huggingface_model/Tencent-Hunyuan/HunyuanOCR/";
let model_list = find_type_files(model_path, "safetensors")?;
let device = Device::Cpu;
for m in &model_list {
let weights = safetensors::load(m, &device)?;
for (key, tensor) in weights.iter() {
if key.contains(".image_") {
println!("=== {} === {:?}", key, tensor.shape());
}
// println!("=== {} === {:?}", key, tensor.shape());
}
}
println!("model_list: {:?}", model_list);
Ok(())
}