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

merge glm-ocr model

Browse Source
This commit is contained in:
jhqxxx
2026-03-07 18:08:06 +08:00
parent cf308f93eb
commit d2524d6a19
19 changed files with 395 additions and 775 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1 -1
View File
@@ -30,7 +30,7 @@ dependencies = [
[[package]]
name = "aha"
version = "0.2.1"
version = "0.2.2"
dependencies = [
"aha_openai_dive",
"anyhow",
Cargo.toml
+2 -2
View File
@@ -1,10 +1,10 @@
[package]
name = "aha"
version = "0.2.1"
version = "0.2.2"
edition = "2024"
repository = "https://github.com/jhqxxx/aha"
license = "Apache-2.0"
description = "aha model inference library, now supports Qwen2.5VL, MiniCPM4, VoxCPM, Qwen3VL, DeepSeek-OCR, Hunyuan-OCR, PaddleOCR-VL, VoxCPM1.5, RMBG2.0, GLM-ASR-Nano-2512, Fun-ASR-Nano-2512, Qwen3, Qwen3-ASR, Qwen3.5"
description = "aha model inference library, now supports Qwen2.5VL, MiniCPM4, VoxCPM, Qwen3VL, DeepSeek-OCR, Hunyuan-OCR, PaddleOCR-VL, VoxCPM1.5, RMBG2.0, GLM-ASR-Nano-2512, Fun-ASR-Nano-2512, Qwen3, Qwen3-ASR, Qwen3.5, GLM-OCR"
[dependencies]
candle-core = { version = "0.9.2" }
README.md
+3
View File
@@ -25,6 +25,9 @@
aha is a high-performance, cross-platform AI inference engine built with Rust and the Candle framework. It brings state-of-the-art AI models to your local machine—no API keys, no cloud dependencies, just pure, fast AI running directly on your hardware.
## Changelog
### v0.2.2 (2026-03-07)
- Added GLM-OCR model
### v0.2.1 (2026-03-05)
- Added Qwen3.5 model
README.zh-CN.md
+3
View File
@@ -25,6 +25,9 @@
aha 是一款基于 Rust 和 Candle 框架构建的高性能跨平台 AI 推理引擎。将最先进的 AI 模型带到您的本地机器——无需 API 密钥，无需云依赖，纯粹、快速的 AI，直接在您的硬件上运行。
## 更新日志
### v0.2.2 (2026-03-07)
- 新增GLM-OCR 模型
### v0.2.1 (2026-03-05)
- 新增Qwen3.5 模型
docs/changelog.md
+3
View File
@@ -5,6 +5,9 @@ All notable changes to aha will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [0.2.2] (2026-03-07)
- Added GLM-OCR model
## [0.2.1] - (2026-03-05)
- Added Qwen3.5 model
docs/changelog.zh-CN.md
+3
View File
@@ -5,6 +5,9 @@
格式基于 [Keep a Changelog](https://keepachangelog.com/zh-CN/1.0.0/)，
本项目遵循 [语义化版本](https://semver.org/lang/zh-CN/spec/v2.0.0.html)。
## [0.2.2] (2026-03-07)
- 新增 GLM-OCR 模型
## [0.2.1] - (2026-03-05)
- 新增Qwen3.5 模型
docs/supported-models.md
+1
View File
@@ -31,6 +31,7 @@ aha supports a growing collection of state-of-the-art AI models across multiple
| **PaddleOCR-VL** | 80+ | Lightweight | General documents | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) |
| **Hunyuan-OCR** | Chinese | Deep learning | Complex layouts | [Tencent Hunyuan Community License](https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE) |
| **DeepSeek-OCR** | Multi | Scene text | Natural images | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) |
| **GLM-OCR** | 8 | Scene text | complex document | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) |
## Speech Recognition (ASR)
docs/supported-models.zh-CN.md
+1
View File
@@ -31,6 +31,7 @@ aha 支持多个领域的最先进 AI 模型集合。
| **PaddleOCR-VL** | 80+ | 轻量级 | 通用文档 | [Apache 2.0](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/apache-2.0.md) |
| **Hunyuan-OCR** | 中文 | 深度学习 | 复杂布局 | [腾讯混元社区许可协议](https://huggingface.co/tencent/HunyuanOCR/blob/main/LICENSE) |
| **DeepSeek-OCR** | 多语言 | 场景文字 | 自然图像 | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) |
| **GLM-OCR** | 8 | 场景文字 | 复杂文档 | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) |
## 语音识别 (ASR)
src/exec/glm_ocr.rs
+2 -2
View File
@@ -49,7 +49,7 @@ impl ExecModel for GlmOcrExec {
}}"#,
input_url
);
let mes = serde_json::from_str(&message)?;
let i_start = Instant::now();
@@ -66,4 +66,4 @@ impl ExecModel for GlmOcrExec {
Ok(())
}
}
}
src/models/glm_ocr/config.rs
+24 -298
View File
@@ -1,369 +1,95 @@
use candle_nn::Activation;
use serde::Deserialize;
use crate::models::qwen3vl::config::Size;
/// Vision encoder configuration for GLM-OCR.
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
pub struct GlmOcrVisionConfig {
#[serde(default)]
pub model_type: String,
/// Number of transformer layers (depth) in the vision encoder. Default: 24
#[serde(default)]
pub depth: usize,
/// Dimensionality of the encoder layers and the pooler layer. Default: 1024
#[serde(default = "default_hidden_size")]
pub hidden_size: usize,
/// Non-linear activation function in the encoder. Default: "silu"
#[serde(default)]
pub hidden_act: Activation,
/// Whether to add bias to queries, keys and values. Default: true
#[serde(default = "default_true")]
pub attention_bias: bool,
/// Dropout probability for attention weights. Default: 0.0
#[serde(default)]
pub attention_dropout: f64,
/// Number of attention heads per layer. Default: 16
#[serde(default = "default_num_heads")]
pub num_heads: usize,
/// Number of input image channels. Default: 3
#[serde(default = "default_in_channels")]
pub in_channels: usize,
/// Input image resolution. Default: 336
#[serde(default = "default_image_size")]
pub image_size: usize,
/// Size of each image patch. Default: 14
#[serde(default = "default_patch_size")]
pub patch_size: usize,
/// Epsilon for RMS normalization layers. Default: 1e-5
#[serde(default = "default_rms_norm_eps")]
pub rms_norm_eps: f64,
/// Size used for merging spatial dimensions. Default: 2
#[serde(default = "default_spatial_merge_size")]
pub spatial_merge_size: usize,
/// Patch size along the temporal dimension (for video). Default: 2
#[serde(default = "default_temporal_patch_size")]
pub temporal_patch_size: usize,
/// Output hidden size of the vision model. Default: 1536
#[serde(alias = "out_hidden_size", default = "default_out_hidden_size")]
pub out_hidden_size: usize,
/// Dimensionality of the feed-forward layer. Default: 4096
#[serde(default = "default_intermediate_size")]
pub intermediate_size: usize,
/// Std dev of truncated normal initializer for weight matrices. Default: 0.02
#[serde(default = "default_initializer_range")]
pub initializer_range: f64,
/// Base frequency for RoPE in vision encoder. Default: 10000.0
#[serde(default = "default_rope_theta")]
pub rope_theta: f32,
}
fn default_hidden_size() -> usize {
1024
}
fn default_true() -> bool {
true
}
fn default_num_heads() -> usize {
16
}
fn default_in_channels() -> usize {
3
}
fn default_image_size() -> usize {
336
}
fn default_patch_size() -> usize {
14
}
fn default_rms_norm_eps() -> f64 {
1e-5
}
fn default_spatial_merge_size() -> usize {
2
}
fn default_temporal_patch_size() -> usize {
2
}
fn default_out_hidden_size() -> usize {
1536
}
fn default_intermediate_size() -> usize {
4096
}
fn default_initializer_range() -> f64 {
0.02
}
fn default_projector_hidden_size() -> usize {
1536
}
/// Projector configuration for mapping vision features to LLM embedding space.
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
pub struct GlmOcrProjectorConfig {
/// Hidden size for the projector. Default: 1536
#[serde(default = "default_projector_hidden_size")]
pub hidden_size: usize,
/// Activation function for the projector.
#[serde(default)]
pub projector_hidden_act: Activation,
/// Number of query tokens for the projector. Default: 256
#[serde(default = "default_num_queries")]
pub num_queries: usize,
}
fn default_num_queries() -> usize {
256
}
/// RoPE (Rotary Position Embedding) configuration parameters.
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
pub struct GlmOcrRopeParameters {
/// Type of RoPE scaling (e.g., "mrope" for multimodal).
#[serde(default)]
pub rope_type: String,
/// Section sizes for M-RoPE (Multimodal RoPE) dimensions.
#[serde(default)]
pub mrope_section: Vec<usize>,
/// Fraction of head dim to apply rotary embedding to.
#[serde(default)]
pub partial_rotary_factor: f32,
/// Base frequency for RoPE. Default: 10000.0
#[serde(default = "default_rope_theta")]
pub rope_theta: f32,
}
fn default_rope_theta() -> f32 {
10000.0
}
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
pub struct GlmOcrRopeParameters {
pub rope_type: String,
pub mrope_section: Vec<usize>,
pub partial_rotary_factor: f32,
pub rope_theta: f32,
}
/// Text decoder configuration for GLM-OCR.
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
pub struct GlmOcrTextConfig {
/// Vocabulary size. Defines the number of different tokens. Default: 59392
#[serde(default = "default_vocab_size")]
pub vocab_size: usize,
/// Dimension of the hidden representations. Default: 1024
#[serde(default = "default_hidden_size")]
pub hidden_size: usize,
/// Dimension of the MLP representations. Default: 4608
#[serde(default = "default_text_intermediate_size")]
pub intermediate_size: usize,
/// Number of hidden layers in the transformer decoder. Default: 16
#[serde(default = "default_num_hidden_layers")]
pub num_hidden_layers: usize,
/// Number of attention heads per layer. Default: 16
#[serde(default = "default_num_attention_heads")]
pub num_attention_heads: usize,
/// Number of key-value heads for Grouped Query Attention.
/// If equal to num_attention_heads, uses MHA. If 1, uses MQA. Default: 8
#[serde(default = "default_num_key_value_heads")]
pub num_key_value_heads: usize,
/// Dimension of each attention head. Default: 128
#[serde(default = "default_head_dim")]
pub head_dim: Option<usize>,
/// Maximum sequence length the model can handle. Default: 131072
#[serde(default = "default_max_position_embeddings")]
pub max_position_embeddings: usize,
/// Epsilon for RMS normalization layers. Default: 1e-5
#[serde(default = "default_rms_norm_eps")]
pub rms_norm_eps: f64,
/// Base frequency for RoPE positional embeddings.
#[serde(default)]
pub rope_theta: f32,
/// Fraction of head dim to apply rotary embedding to.
#[serde(default)]
pub partial_rotary_factor: f32,
/// Non-linear activation function in the decoder. Default: "silu"
#[serde(default)]
pub hidden_act: Activation,
/// Whether to return key/value caches for faster generation. Default: true
#[serde(default = "default_true")]
pub use_cache: bool,
/// Dropout ratio for attention probabilities. Default: 0.0
#[serde(default)]
pub attention_dropout: f64,
/// Type of RoPE scaling.
#[serde(default)]
pub rope_type: String,
/// Section sizes for M-RoPE dimensions.
#[serde(default)]
pub mrope_section: Vec<usize>,
/// Full RoPE configuration parameters.
#[serde(default)]
pub rope_parameters: Option<GlmOcrRopeParameters>,
/// End-of-sequence token ID.
#[serde(default)]
pub eos_token_id: Option<serde_json::Value>,
}
fn default_vocab_size() -> usize {
59392
}
fn default_num_hidden_layers() -> usize {
16
}
fn default_num_attention_heads() -> usize {
16
}
fn default_num_key_value_heads() -> usize {
8
}
fn default_max_position_embeddings() -> usize {
131072
}
fn default_text_intermediate_size() -> usize {
4608
}
fn default_head_dim() -> Option<usize> {
Some(128)
pub rope_parameters: GlmOcrRopeParameters,
pub eos_token_id: Vec<u32>,
pub dtype: String,
}
/// Top-level configuration for GLM-OCR multimodal model.
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct GlmOcrConfig {
#[serde(default)]
pub architectures: Vec<String>,
#[serde(default)]
pub model_type: String,
/// Vision encoder configuration.
#[serde(default)]
pub vision_config: GlmOcrVisionConfig,
/// Projector configuration for vision-to-text mapping.
#[serde(default)]
pub projector_config: GlmOcrProjectorConfig,
/// Text decoder configuration.
#[serde(default)]
pub text_config: GlmOcrTextConfig,
/// Token index to encode image prompts. Default: 59280
#[serde(default = "default_image_token_id")]
pub image_token_id: u32,
/// Token index to encode video prompts. Default: 59281
#[serde(default = "default_video_token_id")]
pub video_token_id: u32,
/// Token index marking start of image. Default: 59256
#[serde(default = "default_image_start_token_id")]
pub image_start_token_id: u32,
/// Token index marking end of image. Default: 59257
#[serde(default = "default_image_end_token_id")]
pub image_end_token_id: u32,
/// Token index marking start of video. Default: 59258
#[serde(default = "default_video_start_token_id")]
pub video_start_token_id: u32,
/// Token index marking end of video. Default: 59259
#[serde(default = "default_video_end_token_id")]
pub video_end_token_id: u32,
/// Beginning-of-sequence token ID.
#[serde(default)]
pub bos_token_id: u32,
/// End-of-sequence token ID.
#[serde(default)]
pub eos_token_id: u32,
/// Padding token ID.
#[serde(default)]
pub pad_token_id: u32,
#[serde(default)]
pub torch_dtype: String,
}
fn default_image_token_id() -> u32 {
59280
}
fn default_video_token_id() -> u32 {
59281
}
fn default_image_start_token_id() -> u32 {
59256
}
fn default_image_end_token_id() -> u32 {
59257
}
fn default_video_start_token_id() -> u32 {
59258
}
fn default_video_end_token_id() -> u32 {
59259
}
/// Generation configuration for controlling text output.
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct GlmOcrGenerationConfig {
/// Beginning-of-sequence token ID.
#[serde(default)]
pub bos_token_id: usize,
/// Padding token ID.
#[serde(default)]
pub pad_token_id: usize,
/// Whether to use sampling (true) or greedy decoding (false). Default: true
#[serde(default = "default_true")]
pub pad_token_id: u32,
pub do_sample: bool,
/// End-of-sequence token ID(s) that stop generation.
#[serde(default)]
pub eos_token_id: Vec<usize>,
/// Nucleus sampling probability threshold. Default: 0.9
#[serde(default = "default_top_p")]
pub top_p: f32,
/// Top-k tokens to consider for sampling. Default: 50
#[serde(default = "default_top_k")]
pub top_k: usize,
/// Sampling temperature (higher = more random). Default: 0.7
#[serde(default = "default_temperature")]
pub temperature: f32,
/// Penalty for repeating tokens. Default: 1.0
#[serde(default = "default_repetition_penalty")]
pub repetition_penalty: f32,
pub eos_token_id: Vec<u32>,
}
fn default_top_p() -> f32 {
0.9
}
fn default_top_k() -> usize {
50
}
fn default_temperature() -> f32 {
0.7
}
fn default_repetition_penalty() -> f32 {
1.0
}
/// Image preprocessor configuration.
#[derive(Debug, Clone, PartialEq, Deserialize, Default)]
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct GlmOcrPreprocessorConfig {
/// Mean values for image normalization (per channel).
#[serde(default)]
pub size: Size,
pub do_rescale: bool,
pub image_mean: Vec<f32>,
/// Std dev values for image normalization (per channel).
#[serde(default)]
pub image_std: Vec<f32>,
/// Shortest edge for dynamic image resizing. Default: 448
#[serde(default)]
pub size: Option<serde_json::Value>,
/// Shortest edge length for resizing (min_pixels in Python).
#[serde(default = "default_shortest_edge")]
pub shortest_edge: usize,
/// Longest edge for resizing (max_pixels in Python).
#[serde(default = "default_longest_edge")]
pub longest_edge: usize,
/// Patch size for vision encoder.
#[serde(default = "default_patch_size_14")]
pub patch_size: Option<usize>,
/// Merge size for spatial merge.
#[serde(default = "default_merge_size")]
pub merge_size: Option<usize>,
}
fn default_shortest_edge() -> usize {
12544 // Python's default min_pixels
}
fn default_longest_edge() -> usize {
9633792 // Python's default max_pixels
}
fn default_patch_size_14() -> Option<usize> {
Some(14)
}
fn default_merge_size() -> Option<usize> {
Some(2)
pub patch_size: usize,
pub merge_size: usize,
pub temporal_patch_size: usize,
}
src/models/glm_ocr/generate.rs
+31 -165
View File
@@ -1,8 +1,8 @@
//! GLM-OCR Inference and Generation
use anyhow::{Result, anyhow};
use aha_openai_dive::v1::resources::chat::{
ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
};
use anyhow::{Result, anyhow};
use candle_core::{DType, Device, IndexOp, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::utils::apply_repeat_penalty;
@@ -10,7 +10,7 @@ use rocket::async_stream::stream;
use rocket::futures::Stream;
use crate::{
chat_template::ChatTemplate,
// chat_template::ChatTemplate,
models::{
GenerateModel,
glm_ocr::{
@@ -20,15 +20,11 @@ use crate::{
},
},
tokenizer::TokenizerModel,
utils::img_utils::extract_image_url,
utils::{
build_completion_chunk_response, build_completion_response, find_type_files, get_device,
get_dtype, get_logit_processor,
},
utils::{build_completion_chunk_response, build_completion_response, extract_user_text, find_type_files, get_device, get_dtype, get_logit_processor, img_utils::extract_image_url},
};
pub struct GlmOcrGenerateModel<'a> {
chat_template: ChatTemplate<'a>,
pub struct GlmOcrGenerateModel {
// chat_template: ChatTemplate<'a>,
tokenizer: TokenizerModel,
processor: GlmOcrProcessor,
model: GlmOcrModel,
@@ -44,86 +40,15 @@ pub struct GlmOcrGenerateModel<'a> {
spatial_merge_size: usize,
}
impl<'a> GlmOcrGenerateModel<'a> {
impl GlmOcrGenerateModel {
pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
let chat_template = ChatTemplate::init(path)?;
// let chat_template = ChatTemplate::init(path)?;
let tokenizer = TokenizerModel::init(path)?;
let config_path = path.to_string() + "/config.json";
let mut cfg: GlmOcrConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
if cfg.projector_config.hidden_size == 0 {
cfg.projector_config.hidden_size = 1536;
}
if cfg.projector_config.num_queries == 0 {
cfg.projector_config.num_queries = 256;
}
// Apply rope_parameters if present (config.json nests these under rope_parameters)
if let Some(ref rope_params) = cfg.text_config.rope_parameters {
if cfg.text_config.rope_theta == 0.0 {
cfg.text_config.rope_theta = rope_params.rope_theta;
}
if cfg.text_config.mrope_section.is_empty() {
cfg.text_config.mrope_section = rope_params.mrope_section.clone();
}
if cfg.text_config.rope_type.is_empty() {
cfg.text_config.rope_type = rope_params.rope_type.clone();
}
if cfg.text_config.partial_rotary_factor == 0.0 {
cfg.text_config.partial_rotary_factor = rope_params.partial_rotary_factor;
}
}
// Fallback rope_theta if still 0
if cfg.text_config.rope_theta == 0.0 {
cfg.text_config.rope_theta = 10000.0;
}
// Collect all EOS token IDs (config may have one or multiple)
let mut eos_token_ids: Vec<u32> = Vec::new();
if cfg.eos_token_id != 0 {
eos_token_ids.push(cfg.eos_token_id);
}
if let Some(ref eos_val) = cfg.text_config.eos_token_id {
match eos_val {
serde_json::Value::Number(n) => {
if let Some(id) = n.as_u64() {
let id = id as u32;
if !eos_token_ids.contains(&id) {
eos_token_ids.push(id);
}
}
}
serde_json::Value::Array(arr) => {
for v in arr {
if let Some(id) = v.as_u64() {
let id = id as u32;
if !eos_token_ids.contains(&id) {
eos_token_ids.push(id);
}
}
}
}
_ => {}
}
}
if eos_token_ids.is_empty() {
eos_token_ids.push(59246); // GLM-OCR default
}
let cfg: GlmOcrConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
let device = get_device(device);
let cfg_dtype = if cfg.torch_dtype.is_empty() {
"bfloat16"
} else {
&cfg.torch_dtype
};
let cfg_dtype = cfg.text_config.dtype.as_str();
let dtype = get_dtype(dtype, cfg_dtype);
// Vision encoder has ops unsupported in F16 on CPU; use F32 for CPU
let dtype = if matches!(device, Device::Cpu) && !matches!(dtype, DType::F32 | DType::F64) {
DType::F32
} else {
dtype
};
let processor = GlmOcrProcessor::new(path, &device, dtype)?;
let model_list = find_type_files(path, "safetensors")?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? };
@@ -133,12 +58,12 @@ impl<'a> GlmOcrGenerateModel<'a> {
serde_json::from_slice(&std::fs::read(generation_config_path)?)?;
Ok(Self {
chat_template,
// chat_template,
tokenizer,
processor,
model,
device,
eos_token_ids,
eos_token_ids: generation_config.eos_token_id.clone(),
generation_config,
model_name: "glm-ocr".to_string(),
image_token_id: cfg.image_token_id,
@@ -151,29 +76,10 @@ impl<'a> GlmOcrGenerateModel<'a> {
}
}
impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
impl GenerateModel for GlmOcrGenerateModel {
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
// Check if sampling is enabled - if do_sample is false, use greedy decoding (temperature = None)
let do_sample = mes.temperature.is_some() || self.generation_config.do_sample;
let temperature = if !do_sample {
None // Greedy decoding
} else {
match mes.temperature {
None => Some(self.generation_config.temperature),
Some(tem) => Some(tem),
}
};
let top_p = match mes.top_p {
None => Some(self.generation_config.top_p),
Some(top_p) => Some(top_p),
};
let top_k = Some(self.generation_config.top_k);
let seed = match mes.seed {
None => 34562u64,
Some(s) => s as u64,
};
let mut logit_processor =
get_logit_processor(temperature, top_p, top_k, seed);
let seed = mes.seed.unwrap_or(34562) as u64;
let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
// Extract image path and prompt from messages
let image_urls = extract_image_url(&mes);
@@ -182,7 +88,10 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
.ok_or_else(|| anyhow!("No image provided"))?;
// Get prompt text from messages
let prompt = extract_text_from_messages(&mes).unwrap_or_else(|| "Extract all text from this image.".to_string());
let mut prompt = extract_user_text(&mes)?;
if prompt.chars().count() == 0 {
prompt = "Extract all text from this image.".to_string()
}
let processed = self.processor.process_info(
image_path,
@@ -202,6 +111,7 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
let image_mask = Some(processed.image_mask);
let mut seqlen_offset = 0;
let mut seq_len = input_ids.dim(1)?;
let prompt_tokens = seq_len as u32;
let mut generate = Vec::new();
let sample_len = mes.max_tokens.unwrap_or(512);
@@ -226,12 +136,7 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
},
seqlen_offset,
)?;
let logits = logits.i((0, seq_len - 1, ..))?.to_dtype(DType::F32)?;
let logits = if self.generation_config.repetition_penalty != 1.0 {
apply_repeat_penalty(&logits, self.generation_config.repetition_penalty, &generate)?
} else {
logits
};
let logits = logits.i((0, seq_len - 1, ..))?.to_dtype(DType::F32)?;
let next_token = logit_processor.sample(&logits)?;
generate.push(next_token);
@@ -246,7 +151,8 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
self.model.clear_kv_cache();
let num_token = generate.len() as u32;
let res = self.tokenizer.token_decode(generate)?;
let response = build_completion_response(res, &self.model_name, Some(num_token));
let response =
build_completion_response(res, &self.model_name, Some(num_token), Some(prompt_tokens));
Ok(response)
}
@@ -261,27 +167,8 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
+ '_,
>,
> {
// Check if sampling is enabled - if do_sample is false, use greedy decoding (temperature = None)
let do_sample = mes.temperature.is_some() || self.generation_config.do_sample;
let temperature = if !do_sample {
None // Greedy decoding
} else {
match mes.temperature {
None => Some(self.generation_config.temperature),
Some(tem) => Some(tem),
}
};
let top_p = match mes.top_p {
None => Some(self.generation_config.top_p),
Some(top_p) => Some(top_p),
};
let top_k = Some(self.generation_config.top_k);
let seed = match mes.seed {
None => 34562u64,
Some(s) => s as u64,
};
let mut logit_processor =
get_logit_processor(temperature, top_p, top_k, seed);
let seed = mes.seed.unwrap_or(34562) as u64;
let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed);
// Extract image path and prompt from messages
let image_urls = extract_image_url(&mes);
@@ -290,7 +177,10 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
.ok_or_else(|| anyhow!("No image provided"))?;
// Get prompt text from messages
let prompt = extract_text_from_messages(&mes).unwrap_or_else(|| "Extract all text from this image.".to_string());
let mut prompt = extract_user_text(&mes)?;
if prompt.chars().count() == 0 {
prompt = "Extract all text from this image.".to_string()
}
let processed = self.processor.process_info(
image_path,
@@ -325,11 +215,7 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
seqlen_offset,
).map_err(|e| anyhow!(format!("forward error: {e}")))?;
let logits = logits.i((0, seq_len - 1, ..)).map_err(|e| anyhow!(format!("index error: {e}")))?.to_dtype(DType::F32).map_err(|e| anyhow!(format!("dtype error: {e}")))?;
let logits = if self.generation_config.repetition_penalty != 1.0 {
apply_repeat_penalty(&logits, self.generation_config.repetition_penalty, &generated).map_err(|e| anyhow!(format!("repeat penalty error: {e}")))?
} else {
logits
};
let next_token = logit_processor.sample(&logits).map_err(|e| anyhow!(format!("sample error: {e}")))?;
generated.push(next_token);
@@ -369,25 +255,5 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> {
}
}
/// Extract text content from chat messages
fn extract_text_from_messages(mes: &ChatCompletionParameters) -> Option<String> {
use aha_openai_dive::v1::resources::chat::{
ChatMessage, ChatMessageContent, ChatMessageContentPart,
};
for msg in &mes.messages {
if let ChatMessage::User { content, .. } = msg {
match content {
ChatMessageContent::Text(text) => return Some(text.clone()),
ChatMessageContent::ContentPart(parts) => {
for part in parts {
if let ChatMessageContentPart::Text(text_part) = part {
return Some(text_part.text.clone());
}
}
}
_ => {}
}
}
}
None
}
src/models/glm_ocr/model.rs
+91 -103
View File
@@ -11,13 +11,11 @@ use crate::{
models::{
common::GateUpDownMLP,
glm_ocr::config::{
GlmOcrConfig, GlmOcrProjectorConfig, GlmOcrTextConfig, GlmOcrVisionConfig,
GlmOcrConfig, GlmOcrTextConfig, GlmOcrVisionConfig,
},
},
position_embed::rope::{apply_rotary_pos_emb_vision, glm_ocr_apply_rotary_pos_emb},
utils::{
tensor_utils::{prepare_causal_attention_mask, repeat_kv},
},
utils::tensor_utils::{prepare_causal_attention_mask, repeat_kv},
};
pub struct GlmOcrRMSNorm(RmsNorm);
@@ -87,7 +85,7 @@ fn eager_attention_forward(
let k = key_states.transpose(1, 2)?.contiguous()?;
let v = value_states.transpose(1, 2)?.contiguous()?;
candle_flash_attn::flash_attn(&q, &k, &v, scaling as f32, attention_mask.is_some())?
// flash_attn returns [batch, q_len, heads, head_dim] — already in final layout
// flash_attn returns [batch, q_len, heads, head_dim] — already in final layout
}
#[cfg(not(feature = "flash-attn"))]
{
@@ -101,7 +99,8 @@ fn eager_attention_forward(
let k_t = key_states.transpose(D::Minus2, D::Minus1)?.contiguous()?;
let raw = if q_len > CHUNK_SIZE {
let mut chunks: Vec<Tensor> = Vec::with_capacity((q_len + CHUNK_SIZE - 1) / CHUNK_SIZE);
let mut chunks: Vec<Tensor> =
Vec::with_capacity((q_len + CHUNK_SIZE - 1) / CHUNK_SIZE);
let mut start = 0;
while start < q_len {
let len = CHUNK_SIZE.min(q_len - start);
@@ -109,9 +108,8 @@ fn eager_attention_forward(
let attn = (q_chunk.matmul(&k_t)? * scaling)?;
let attn = match attention_mask {
None => attn,
Some(mask) => {
attn.broadcast_add(&mask.narrow(2, start, len)?.to_dtype(attn.dtype())?)?
}
Some(mask) => attn
.broadcast_add(&mask.narrow(2, start, len)?.to_dtype(attn.dtype())?)?,
};
// Softmax computation: Optimize dtype conversions for CPU (which uses F32)
let attn = if query_states.dtype() == DType::F32 {
@@ -282,8 +280,7 @@ impl GlmOcrVisionRotaryEmbedding {
.step_by(2)
.map(|i| 1.0 / theta.powf(i as f32 / dim as f32))
.collect();
let inv_freq =
Tensor::from_vec(inv_freq, (dim / 2,), device)?.to_dtype(dtype)?;
let inv_freq = Tensor::from_vec(inv_freq, (dim / 2,), device)?.to_dtype(dtype)?;
Ok(Self { inv_freq })
}
@@ -316,7 +313,7 @@ impl GlmOcrVisionRotaryEmbedding {
let _si = hi % sms;
let _wb = wi / sms;
let _sj = wi % sms;
// After permute(0,2,1,3): position = (hb, wb, si, sj)
// Flatten: idx = hb * w_blocks * sms * sms + wb * sms * sms + si * sms + sj
// But we just need the h and w positions for rotary embedding
@@ -630,21 +627,9 @@ impl GlmOcrVisionPatchMerger {
)?;
let context_dim = config.out_hidden_size * config.in_channels;
let gate_proj = linear_no_bias(
config.out_hidden_size,
context_dim,
vb.pp("gate_proj"),
)?;
let up_proj = linear_no_bias(
config.out_hidden_size,
context_dim,
vb.pp("up_proj"),
)?;
let down_proj = linear_no_bias(
context_dim,
config.out_hidden_size,
vb.pp("down_proj"),
)?;
let gate_proj = linear_no_bias(config.out_hidden_size, context_dim, vb.pp("gate_proj"))?;
let up_proj = linear_no_bias(config.out_hidden_size, context_dim, vb.pp("up_proj"))?;
let down_proj = linear_no_bias(context_dim, config.out_hidden_size, vb.pp("down_proj"))?;
Ok(Self {
proj,
@@ -674,7 +659,8 @@ pub struct GlmOcrVisionPatchEmbed {
patch_size: usize,
temporal_patch_size: usize,
in_channels: usize,
#[allow(dead_code)] embed_dim: usize,
#[allow(dead_code)]
embed_dim: usize,
proj: Linear,
}
@@ -711,7 +697,7 @@ impl GlmOcrVisionPatchEmbed {
pub fn forward(&self, pixel_values: &Tensor) -> Result<Tensor> {
let rank = pixel_values.rank();
if rank == 2 {
let hidden_states = self.proj.forward(pixel_values)?;
Ok(hidden_states)
@@ -767,8 +753,12 @@ impl GlmOcrVisionModel {
let patch_embed = GlmOcrVisionPatchEmbed::new(vb.pp("patch_embed"), config)?;
let head_dim = config.hidden_size / config.num_heads;
let rotary_pos_emb =
GlmOcrVisionRotaryEmbedding::new(head_dim / 2, config.rope_theta, vb.device(), vb.dtype())?;
let rotary_pos_emb = GlmOcrVisionRotaryEmbedding::new(
head_dim / 2,
config.rope_theta,
vb.device(),
vb.dtype(),
)?;
let mut blocks = Vec::new();
let depth = config.depth;
@@ -865,13 +855,13 @@ impl GlmOcrVisionModel {
let sms = self.config.spatial_merge_size;
let hidden_dim = hidden_states.dim(hidden_states.dims().len() - 1)?;
let total_patches = hidden_states.dim(0)?; // 2816
let merged_patches = total_patches / (sms * sms); // 704
let total_patches = hidden_states.dim(0)?; // 2816
let merged_patches = total_patches / (sms * sms); // 704
let hidden_states = hidden_states.reshape((merged_patches, sms, sms, hidden_dim))?;
let hidden_states = hidden_states.permute((0, 3, 1, 2))?; // [704, 1024, 2, 2]
let hidden_states = self.downsample.forward(&hidden_states)?; // [704, 1536, 1, 1]
let hidden_states = hidden_states.reshape((merged_patches, self.config.out_hidden_size))?; // [704, 1536]
let hidden_states = hidden_states.permute((0, 3, 1, 2))?; // [704, 1024, 2, 2]
let hidden_states = self.downsample.forward(&hidden_states)?; // [704, 1536, 1, 1]
let hidden_states = hidden_states.reshape((merged_patches, self.config.out_hidden_size))?; // [704, 1536]
let merged = self.merger.forward(&hidden_states)?;
@@ -880,46 +870,48 @@ impl GlmOcrVisionModel {
}
}
pub struct GlmOcrProjector {
#[allow(dead_code)] query_embed: Option<Tensor>,
proj: Linear,
norm: LayerNorm,
#[allow(dead_code)] num_queries: usize,
}
// pub struct GlmOcrProjector {
// #[allow(dead_code)]
// query_embed: Option<Tensor>,
// proj: Linear,
// norm: LayerNorm,
// #[allow(dead_code)]
// num_queries: usize,
// }
impl GlmOcrProjector {
pub fn new(
vb: VarBuilder,
vision_config: &GlmOcrVisionConfig,
config: &GlmOcrProjectorConfig,
) -> Result<Self> {
let query_embed = vb
.get(
(1, config.num_queries, vision_config.out_hidden_size),
"query_embed",
)
.ok();
// impl GlmOcrProjector {
// pub fn new(
// vb: VarBuilder,
// vision_config: &GlmOcrVisionConfig,
// config: &GlmOcrProjectorConfig,
// ) -> Result<Self> {
// let query_embed = vb
// .get(
// (1, config.num_queries, vision_config.out_hidden_size),
// "query_embed",
// )
// .ok();
let proj = linear_no_bias(
vision_config.out_hidden_size,
config.hidden_size,
vb.pp("proj"),
)?;
let norm = layer_norm(config.hidden_size, 1e-5, vb.pp("norm"))?;
// let proj = linear_no_bias(
// vision_config.out_hidden_size,
// config.hidden_size,
// vb.pp("proj"),
// )?;
// let norm = layer_norm(config.hidden_size, 1e-5, vb.pp("norm"))?;
Ok(Self {
query_embed,
proj,
norm,
num_queries: config.num_queries,
})
}
// Ok(Self {
// query_embed,
// proj,
// norm,
// num_queries: config.num_queries,
// })
// }
pub fn forward(&self, image_features: &Tensor) -> Result<Tensor> {
let projected = self.proj.forward(image_features)?;
Ok(self.norm.forward(&projected)?)
}
}
// pub fn forward(&self, image_features: &Tensor) -> Result<Tensor> {
// let projected = self.proj.forward(image_features)?;
// Ok(self.norm.forward(&projected)?)
// }
// }
pub struct GlmOcrTextRotaryEmbedding {
inv_freq: Tensor,
@@ -932,17 +924,12 @@ impl GlmOcrTextRotaryEmbedding {
device: &candle_core::Device,
dtype: DType,
) -> Result<Self> {
let rope_theta = config.rope_theta;
let rope_theta = config.rope_parameters.rope_theta;
let head_dim = config.head_dim.unwrap_or_else(|| {
// Integer division, panics if num_attention_heads is 0 (like Python)
config.hidden_size / config.num_attention_heads
});
let partial_rotary_factor = if config.partial_rotary_factor == 0.0 {
1.0
} else {
config.partial_rotary_factor
};
let dim = (head_dim as f32 * partial_rotary_factor) as usize;
let dim = (head_dim as f32 * config.rope_parameters.partial_rotary_factor) as usize;
let inv_freq: Vec<f32> = (0..dim)
.step_by(2)
@@ -951,15 +938,9 @@ impl GlmOcrTextRotaryEmbedding {
let inv_freq =
Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?.to_dtype(dtype)?;
let mrope_section = if config.mrope_section.is_empty() {
vec![8, 12, 12]
} else {
config.mrope_section.clone()
};
Ok(Self {
inv_freq,
mrope_section,
mrope_section: config.rope_parameters.mrope_section.clone(),
})
}
@@ -982,10 +963,10 @@ impl GlmOcrTextRotaryEmbedding {
}
Ok(Tensor::cat(&result_parts, D::Minus1)?)
}
pub fn forward_with_position_ids(&self, position_ids: &Tensor) -> Result<(Tensor, Tensor)> {
pub fn forward_with_position_ids(&self, position_ids: &Tensor) -> Result<(Tensor, Tensor)> {
let (_, bs, _seq_len) = position_ids.dims3()?;
let inv_freq_len = self.inv_freq.dim(1)?;
let inv_freq_len = self.inv_freq.dim(1)?;
// inv_freq: (1, inv_freq_len) -> broadcast to (3, bs, inv_freq_len, 1)
let inv_freq = self.inv_freq.unsqueeze(0)?.unsqueeze(D::Minus1)?; // (1, 1, hd/2, 1)
@@ -1064,7 +1045,7 @@ pub struct GlmOcrTextModel {
norm: GlmOcrRMSNorm,
lm_head: Linear,
rotary_emb: GlmOcrTextRotaryEmbedding,
config: GlmOcrTextConfig,
// config: GlmOcrTextConfig,
spatial_merge_size: usize,
/// max_mrope_position + 1 after prefill (stored for decode-pass position computation)
next_mrope_pos: usize,
@@ -1073,7 +1054,11 @@ pub struct GlmOcrTextModel {
}
impl GlmOcrTextModel {
pub fn new(vb: VarBuilder, config: GlmOcrTextConfig, spatial_merge_size: usize) -> Result<Self> {
pub fn new(
vb: VarBuilder,
config: GlmOcrTextConfig,
spatial_merge_size: usize,
) -> Result<Self> {
let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?;
let mut layers = Vec::new();
@@ -1095,7 +1080,7 @@ impl GlmOcrTextModel {
norm,
lm_head,
rotary_emb,
config,
// config,
spatial_merge_size,
next_mrope_pos: 0,
prefill_seq_len: 0,
@@ -1121,7 +1106,10 @@ impl GlmOcrTextModel {
let num_image_tokens = llm_grid_t * llm_grid_h * llm_grid_w;
// Image mask as bool vec (shape (1, seq_len) -> (seq_len,))
let mask_vec = image_mask.squeeze(0)?.to_dtype(DType::U8)?.to_vec1::<u8>()?;
let mask_vec = image_mask
.squeeze(0)?
.to_dtype(DType::U8)?
.to_vec1::<u8>()?;
let mut t_ids: Vec<i64> = Vec::with_capacity(seq_len);
let mut h_ids: Vec<i64> = Vec::with_capacity(seq_len);
@@ -1215,14 +1203,15 @@ impl GlmOcrTextModel {
let embeds_flat = inputs_embeds.squeeze(0)?; // (seq_len, hidden_size)
let mut embeds_vec: Vec<Tensor> = Vec::new();
let mut feat_idx = 0;
// let mut feat_idx = 0;
let mut pos = 0;
for &img_pos in image_indices.iter().take(num_to_replace) {
// for &img_pos in image_indices.iter().take(num_to_replace) {
for (feat_idx, &img_pos) in image_indices.iter().take(num_to_replace).enumerate() {
if img_pos > pos {
embeds_vec.push(embeds_flat.narrow(0, pos, img_pos - pos)?);
}
embeds_vec.push(img_feats.i((0, feat_idx, ..))?.unsqueeze(0)?);
feat_idx += 1;
// feat_idx += 1;
pos = img_pos + 1;
}
if pos < seq_len {
@@ -1247,7 +1236,8 @@ impl GlmOcrTextModel {
let (cos, sin) = if seqlen_offset == 0 {
if let (Some(mask), Some(thw)) = (image_mask, image_grid_thw) {
// Prefill with image: compute 3D M-RoPE position IDs
let pos_ids = self.compute_mrope_position_ids(mask, thw, seq_len, input_ids.device())?;
let pos_ids =
self.compute_mrope_position_ids(mask, thw, seq_len, input_ids.device())?;
self.prefill_seq_len = seq_len;
self.rotary_emb.forward_with_position_ids(&pos_ids)?
} else {
@@ -1316,7 +1306,7 @@ impl GlmOcrModel {
Tensor::new(
&[
1u32,
(pixels.dim(0)? / 44) as u32, // Approximate
(pixels.dim(0)? / 44) as u32, // Approximate
(pixels.dim(1)? / 44) as u32,
],
input_ids.device(),
@@ -1330,15 +1320,13 @@ impl GlmOcrModel {
None
};
let result = self.language_model.forward(
self.language_model.forward(
input_ids,
image_features.as_ref(),
image_mask,
image_grid_thw,
seqlen_offset,
);
result
)
}
pub fn clear_kv_cache(&mut self) {
src/models/glm_ocr/processor.rs
+50 -132
View File
@@ -4,6 +4,7 @@ use candle_core::{DType, Device, Tensor};
use super::config::GlmOcrPreprocessorConfig;
use crate::tokenizer::TokenizerModel;
use crate::utils::img_utils::get_image;
use crate::utils::video_utils::video_smart_resize;
/// GLM-OCR Processor for image and text preprocessing.
///
@@ -13,8 +14,8 @@ use crate::utils::img_utils::get_image;
pub struct GlmOcrProcessor {
image_mean: Vec<f32>,
image_std: Vec<f32>,
shortest_edge: usize, // min_pixels in Python
longest_edge: usize, // max_pixels in Python
shortest_edge: usize, // min_pixels in Python
longest_edge: usize, // max_pixels in Python
patch_size: usize,
merge_size: usize,
temporal_patch_size: usize,
@@ -23,117 +24,45 @@ pub struct GlmOcrProcessor {
}
pub struct ProcessedImage {
pub pixel_values: Tensor, // Shape: [num_patches, patch_dim]
pub pixel_values: Tensor, // Shape: [num_patches, patch_dim]
pub grid_h: usize,
pub grid_w: usize,
}
pub struct ProcessedInput {
pub input_ids: Tensor,
pub pixel_values: Tensor, // Shape: [num_patches, patch_dim]
pub pixel_values: Tensor, // Shape: [num_patches, patch_dim]
pub image_mask: Tensor,
pub grid_thw: Tensor,
}
impl GlmOcrProcessor {
pub fn new(path: &str, device: &Device, dtype: DType) -> Result<Self> {
assert!(
std::path::Path::new(path).exists(),
"model path file not exists"
);
let config_path = format!("{}/preprocessor_config.json", path);
// Load preprocessor config
let (image_mean, image_std, shortest_edge, longest_edge, patch_size, merge_size) =
if std::path::Path::new(&config_path).exists() {
let config: GlmOcrPreprocessorConfig = serde_json::from_slice(&std::fs::read(&config_path)?)?;
// Parse size object - Python uses shortest_edge: 12544, longest_edge: 9633792
let (shortest, longest) = if let Some(size_val) = &config.size {
if let Some(obj) = size_val.as_object() {
let s = obj.get("shortest_edge")
.and_then(|v| v.as_u64())
.map(|v| v as usize)
.unwrap_or(12544);
let l = obj.get("longest_edge")
.and_then(|v| v.as_u64())
.map(|v| v as usize)
.unwrap_or(9_633_792);
(s, l)
} else {
(12544, 9_633_792)
}
} else {
(12544, 9_633_792)
};
let patch_size = config.patch_size.unwrap_or(14);
let merge_size = config.merge_size.unwrap_or(2);
(config.image_mean, config.image_std, shortest, longest, patch_size, merge_size)
} else {
(
vec![0.48145466, 0.4578275, 0.40821073],
vec![0.26862954, 0.26130258, 0.27577711],
12544, // Python's min_pixels
9_633_792, // Python's max_pixels
14, // patch_size
2, // merge_size
)
};
assert!(
std::path::Path::new(&config_path).exists(),
"preprocessor_config.json not exists in model path"
);
let process_cfg: GlmOcrPreprocessorConfig =
serde_json::from_slice(&std::fs::read(config_path)?)?;
Ok(Self {
image_mean,
image_std,
shortest_edge,
longest_edge,
patch_size,
merge_size,
temporal_patch_size: 2, // Fixed for images
image_mean: process_cfg.image_mean.clone(),
image_std: process_cfg.image_std.clone(),
shortest_edge: process_cfg.size.shortest_edge,
longest_edge: process_cfg.size.longest_edge,
patch_size: process_cfg.patch_size,
merge_size: process_cfg.merge_size,
temporal_patch_size: process_cfg.temporal_patch_size, // Fixed for images
device: device.clone(),
dtype,
})
}
/// Python's smart_resize implementation
/// Returns (resized_height, resized_width)
fn smart_resize(&self, height: usize, width: usize) -> (usize, usize) {
let factor = self.patch_size * self.merge_size; // 28
let temporal_factor = self.temporal_patch_size; // 2
// Ensure minimum size
let mut h = height;
let mut w = width;
if h < factor || w < factor {
let scale = (factor as f32 / h.min(w) as f32).max(1.0);
h = (h as f32 * scale).round() as usize;
w = (w as f32 * scale).round() as usize;
}
// Check aspect ratio constraint
if w.max(h) as f32 / w.min(h) as f32 > 200.0 {
// Would raise error in Python
// For now, just proceed
}
// Round to nearest multiple of factor
let h_bar = ((h + factor / 2) / factor) * factor;
let w_bar = ((w + factor / 2) / factor) * factor;
let t_bar = ((temporal_factor + temporal_factor / 2) / temporal_factor) * temporal_factor;
// Check max_pixels constraint
let max_pixels = self.longest_edge;
let min_pixels = self.shortest_edge;
let mut final_h = h_bar;
let mut final_w = w_bar;
if t_bar * h_bar * w_bar > max_pixels {
let beta = ((h * w) as f32 / max_pixels as f32).sqrt();
final_h = (factor.max((h as f32 / beta / factor as f32).floor() as usize * factor)) as usize;
final_w = (factor.max((w as f32 / beta / factor as f32).floor() as usize * factor)) as usize;
} else if t_bar * h_bar * w_bar < min_pixels {
let beta = (min_pixels as f32 / (h * w) as f32).sqrt();
final_h = ((h as f32 * beta / factor as f32).ceil() as usize * factor) as usize;
final_w = ((w as f32 * beta / factor as f32).ceil() as usize * factor) as usize;
}
(final_h, final_w)
}
/// Process image for vision encoder.
///
/// Matches Python's Glm46VImageProcessor._preprocess():
@@ -145,18 +74,27 @@ impl GlmOcrProcessor {
/// For images: grid_t = 1, so [grid_h * grid_w, 3 * 2 * 14 * 14] = [num_patches, 1176]
pub fn process_image(&self, image_path: &str) -> Result<ProcessedImage> {
let img = get_image(image_path)?;
let (orig_w, orig_h) = (img.width() as usize, img.height() as usize);
// Use smart_resize to compute target dimensions
let (target_h, target_w) = self.smart_resize(orig_h, orig_w);
let (target_h, target_w) = video_smart_resize(
self.temporal_patch_size as u32,
img.height(),
img.width(),
self.temporal_patch_size as u32,
(self.patch_size * self.merge_size) as u32,
self.shortest_edge as u32,
self.longest_edge as u32,
None,
)?;
// Resize image
let img = img.resize_exact(
target_w as u32,
target_h as u32,
target_w,
target_h,
image::imageops::FilterType::Lanczos3,
);
let target_h = target_h as usize;
let target_w = target_w as usize;
// Convert to RGB and normalize
let img = img.to_rgb8();
let pixels: Vec<f32> = img
@@ -171,37 +109,26 @@ impl GlmOcrProcessor {
.collect();
let tensor = Tensor::from_vec(pixels, (target_h, target_w, 3), &self.device)?;
let mean = Tensor::new(self.image_mean.clone(), &self.device)?.reshape((1, 1, 3))?;
let std = Tensor::new(self.image_std.clone(), &self.device)?.reshape((1, 1, 3))?;
let tensor = tensor.broadcast_sub(&mean)?.broadcast_div(&std)?;
// Now reshape into flattened patches like Python
let grid_h = target_h / self.patch_size;
let grid_w = target_w / self.patch_size;
let patch_size = self.patch_size;
let channels = 3;
let temporal_patch_size = 2;
let tensor = tensor.reshape((
grid_h, patch_size,
grid_w, patch_size,
channels,
))?;
let tensor = tensor.reshape((grid_h, self.patch_size, grid_w, self.patch_size, channels))?;
let tensor = tensor.permute((0, 2, 4, 1, 3))?;
let num_patches = grid_h * grid_w;
let tensor = tensor.reshape((num_patches, channels, patch_size, patch_size))?;
let tensor = tensor.reshape((num_patches, channels, self.patch_size, self.patch_size))?;
let tensor = tensor.unsqueeze(2)?;
let tensor = tensor.repeat((1, 1, temporal_patch_size, 1, 1))?;
let patch_dim = channels * temporal_patch_size * patch_size * patch_size;
let tensor = tensor.repeat((1, 1, self.temporal_patch_size, 1, 1))?;
let patch_dim = channels * self.temporal_patch_size * self.patch_size * self.patch_size;
let tensor = tensor.reshape((num_patches, patch_dim))?;
let tensor = tensor.to_dtype(self.dtype)?;
Ok(ProcessedImage {
pixel_values: tensor,
grid_h,
@@ -243,16 +170,10 @@ impl GlmOcrProcessor {
// 59254 = <|assistant|>
// Build input_ids following Python format
let mut input_ids_vec = Vec::new();
// Header: [gMASK] <sop> <|user|> \n Image tokens: <|begin_of_image|> <|image|>*N <|end_of_image|>
// 59248 59250 59253 10
let mut input_ids_vec = vec![59248, 59250, 59253, 10, image_start_token_id];
// Header: [gMASK] <sop> <|user|> \n
input_ids_vec.push(59248); // [gMASK]
input_ids_vec.push(59250); // <sop>
input_ids_vec.push(59253); // <|user|>
input_ids_vec.push(10); // newline
// Image tokens: <|begin_of_image|> <|image|>*N <|end_of_image|>
input_ids_vec.push(image_start_token_id); // <|begin_of_image|>
for _ in 0..num_image_tokens {
input_ids_vec.push(image_token_id); // <|image|>
}
@@ -264,7 +185,7 @@ impl GlmOcrProcessor {
// Generation prompt: <|assistant|> \n
input_ids_vec.push(59254); // <|assistant|>
input_ids_vec.push(10); // newline
input_ids_vec.push(10); // newline
let input_ids = Tensor::from_vec(
input_ids_vec.clone(),
@@ -282,11 +203,8 @@ impl GlmOcrProcessor {
let image_mask = Tensor::from_vec(image_mask_vec, (1, input_ids_vec.len()), &self.device)?;
// Compute grid_thw for RoPE
let grid_thw = Tensor::from_vec(
vec![1u32, grid_h as u32, grid_w as u32],
(3,),
&self.device,
)?;
let grid_thw =
Tensor::from_vec(vec![1u32, grid_h as u32, grid_w as u32], (3,), &self.device)?;
Ok(ProcessedInput {
input_ids,
src/models/mod.rs
+5 -2
View File
@@ -132,7 +132,10 @@ impl WhichModel {
| WhichModel::Qwen3_5_4B
| WhichModel::Qwen3_5_9B => "vlm",
// OCR models
WhichModel::DeepSeekOCR | WhichModel::HunyuanOCR | WhichModel::GlmOCR | WhichModel::PaddleOCRVL => "ocr",
WhichModel::DeepSeekOCR
| WhichModel::HunyuanOCR
| WhichModel::GlmOCR
| WhichModel::PaddleOCRVL => "ocr",
// ASR models
WhichModel::Qwen3ASR0_6B
| WhichModel::Qwen3ASR1_7B
@@ -173,7 +176,7 @@ pub enum ModelInstance<'a> {
VoxCPM(Box<VoxCPMGenerate>),
GlmASRNano(GlmAsrNanoGenerateModel<'a>),
FunASRNano(FunAsrNanoGenerateModel),
GlmOCR(GlmOcrGenerateModel<'a>),
GlmOCR(GlmOcrGenerateModel),
}
impl<'a> GenerateModel for ModelInstance<'a> {
src/models/qwen3vl/processor.rs
+4 -59
View File
@@ -3,20 +3,17 @@ use std::collections::HashMap;
use aha_openai_dive::v1::resources::chat::{
ChatCompletionParameters, ChatMessage, ChatMessageContent, ChatMessageContentPart,
};
use anyhow::{Result, anyhow};
use anyhow::Result;
use candle_core::{DType, Device, IndexOp, Shape, Tensor};
#[cfg(feature = "ffmpeg")]
use ffmpeg_next as ffmpeg;
use image::DynamicImage;
use num::integer::lcm;
#[cfg(feature = "ffmpeg")]
use crate::utils::video_utils::video_smart_resize;
use crate::{
models::qwen3vl::config::PreprocessorConfig,
utils::{
ceil_by_factor, floor_by_factor,
img_utils::{get_image, img_smart_resize, img_transform},
round_by_factor,
},
utils::img_utils::{get_image, img_smart_resize, img_transform},
};
#[derive(Clone)]
@@ -423,58 +420,6 @@ impl Qwen3VLProcessor {
}
}
pub fn video_smart_resize(
num_frames: u32,
height: u32,
width: u32,
temporal_factor: u32,
factor: u32,
min_pixels: u32,
max_pixels: u32,
video_ratio: Option<u32>,
) -> Result<(u32, u32)> {
if num_frames < temporal_factor {
return Err(anyhow!(format!(
"{num_frames} must be larger than temporal_factor {temporal_factor}"
)));
}
if height < factor || width < factor {
return Err(anyhow!(format!(
"height:{height} or width:{width} must be larger than factor:{factor}"
)));
}
if std::cmp::max(height, width) / std::cmp::min(height, width) > 200 {
return Err(anyhow!(format!(
"absolute aspect ratio mush be smaller than {}, got {}",
200,
std::cmp::max(height, width) / std::cmp::min(height, width)
)));
}
let mut image_factor = factor;
if let Some(ratio) = video_ratio {
image_factor = lcm(image_factor, ratio);
}
let mut h_bar = round_by_factor(height, image_factor);
let mut w_bar = round_by_factor(width, image_factor);
let t_bar = round_by_factor(num_frames, temporal_factor);
if t_bar * h_bar * w_bar > max_pixels {
let beta = ((num_frames * height * width) as f32 / max_pixels as f32).sqrt();
h_bar = std::cmp::max(
image_factor,
floor_by_factor(height as f32 / beta, image_factor),
);
w_bar = std::cmp::max(
image_factor,
floor_by_factor(width as f32 / beta, image_factor),
);
} else if t_bar * h_bar * w_bar < min_pixels {
let beta = (min_pixels as f32 / (num_frames * height * width) as f32).sqrt();
h_bar = ceil_by_factor(height as f32 * beta, image_factor);
w_bar = ceil_by_factor(width as f32 * beta, image_factor);
}
Ok((h_bar, w_bar))
}
#[cfg(feature = "ffmpeg")]
pub fn get_video_data(
file: &String,
src/position_embed/rope.rs
+1 -1
View File
@@ -180,7 +180,7 @@ fn rotate_half_llm(x: &Tensor) -> Result<Tensor> {
let x_pairs = x.reshape(pair_shape)?; // (..., half, 2)
// col 0 = even elements [x0, x2, ...], col 1 = odd elements [x1, x3, ...]
let x_even = x_pairs.narrow(D::Minus1, 0, 1)?; // (..., half, 1)
let x_odd = x_pairs.narrow(D::Minus1, 1, 1)?; // (..., half, 1)
let x_odd = x_pairs.narrow(D::Minus1, 1, 1)?; // (..., half, 1)
let neg_x_odd = x_odd.affine(-1.0, 0.0)?;
// Concatenate [-x_odd, x_even] → [[-x1,x0], [-x3,x2], ...]
let result_pairs = Tensor::cat(&[&neg_x_odd, &x_even], D::Minus1)?; // (..., half, 2)
src/utils/mod.rs
+13 -9
View File
@@ -719,16 +719,20 @@ where
pub fn extract_user_text(mes: &ChatCompletionParameters) -> Result<String> {
let mut ret = "".to_string();
for chat_mes in mes.messages.clone() {
if let ChatMessage::User { content, .. } = chat_mes.clone()
&& let ChatMessageContent::ContentPart(part_vec) = content
{
for part in part_vec {
if let ChatMessageContentPart::Text(text_part) = part {
let text = text_part.text;
if text.chars().count() > 0 {
ret = ret + &text + "\n"
if let ChatMessage::User { content, .. } = chat_mes.clone() {
match content {
ChatMessageContent::Text(text) => ret = ret + &text + "\n",
ChatMessageContent::ContentPart(part_vec) => {
for part in part_vec {
if let ChatMessageContentPart::Text(text_part) = part {
let text = text_part.text;
if text.chars().count() > 0 {
ret = ret + &text + "\n"
}
}
}
}
},
_ => {}
}
}
}
src/utils/video_utils.rs
+58 -1
View File
@@ -1,7 +1,64 @@
// use std::{fs::File, io::Write};
use anyhow::{Result, anyhow};
use num::integer::lcm;
use crate::utils::{ceil_by_factor, floor_by_factor, round_by_factor};
// use std::{fs::File, io::Write};
// use ffmpeg_next as ffmpeg;
pub fn video_smart_resize(
num_frames: u32,
height: u32,
width: u32,
temporal_factor: u32,
factor: u32,
min_pixels: u32,
max_pixels: u32,
video_ratio: Option<u32>,
) -> Result<(u32, u32)> {
if num_frames < temporal_factor {
return Err(anyhow!(format!(
"{num_frames} must be larger than temporal_factor {temporal_factor}"
)));
}
if height < factor || width < factor {
return Err(anyhow!(format!(
"height:{height} or width:{width} must be larger than factor:{factor}"
)));
}
if std::cmp::max(height, width) / std::cmp::min(height, width) > 200 {
return Err(anyhow!(format!(
"absolute aspect ratio mush be smaller than {}, got {}",
200,
std::cmp::max(height, width) / std::cmp::min(height, width)
)));
}
let mut image_factor = factor;
if let Some(ratio) = video_ratio {
image_factor = lcm(image_factor, ratio);
}
let mut h_bar = round_by_factor(height, image_factor);
let mut w_bar = round_by_factor(width, image_factor);
let t_bar = round_by_factor(num_frames, temporal_factor);
if t_bar * h_bar * w_bar > max_pixels {
let beta = ((num_frames * height * width) as f32 / max_pixels as f32).sqrt();
h_bar = std::cmp::max(
image_factor,
floor_by_factor(height as f32 / beta, image_factor),
);
w_bar = std::cmp::max(
image_factor,
floor_by_factor(width as f32 / beta, image_factor),
);
} else if t_bar * h_bar * w_bar < min_pixels {
let beta = (min_pixels as f32 / (num_frames * height * width) as f32).sqrt();
h_bar = ceil_by_factor(height as f32 * beta, image_factor);
w_bar = ceil_by_factor(width as f32 * beta, image_factor);
}
Ok((h_bar, w_bar))
}
// #[allow(unused)]
// fn save_file(
// frame: &ffmpeg::frame::Video,
tests/test_glm_ocr.rs
+99
View File
@@ -0,0 +1,99 @@
use std::{pin::pin, time::Instant};
use aha::models::{GenerateModel, glm_ocr::generate::GlmOcrGenerateModel};
use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
use anyhow::Result;
use rocket::futures::StreamExt;
#[test]
fn glm_ocr_generate() -> Result<()> {
// RUST_BACKTRACE=1 cargo test -F cuda --test test_glm_ocr glm_ocr_generate -r -- --nocapture
let message = r#"
{
"model": "glm-ocr",
"messages": [
{
"role": "user",
"content": [
{
"type": "image",
"image_url":
{
"url": "file://./assets/img/ocr_test1.png"
}
},
{
"type": "text",
"text": "Text Recognition:"
}
]
}
]
}
"#;
let save_dir =
aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
let model_path = format!("{}/ZhipuAI/GLM-OCR/", save_dir);
let mes: ChatCompletionParameters = serde_json::from_str(message)?;
let i_start = Instant::now();
let mut model = GlmOcrGenerateModel::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!("generate: \n {:?}", res);
if res.usage.is_some() {
let num_token = res.usage.as_ref().unwrap().total_tokens;
let duration_secs = i_duration.as_secs_f64();
let tps = num_token as f64 / duration_secs;
println!("Tokens per second (TPS): {:.2}", tps);
}
println!("Time elapsed in generate is: {:?}", i_duration);
Ok(())
}
#[tokio::test]
async fn glm_ocr_stream() -> Result<()> {
// test with cuda: RUST_BACKTRACE=1 cargo test -F cuda glm_ocr_stream -r -- --nocapture
let message = r#"
{
"model": "glm-ocr",
"messages": [
{
"role": "user",
"content": [
{
"type": "image",
"image_url":
{
"url": "file://./assets/img/ocr_test1.png"
}
},
{
"type": "text",
"text": "Text Recognition:"
}
]
}
]
}
"#;
let save_dir =
aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
let model_path = format!("{}/ZhipuAI/GLM-OCR/", save_dir);
let mes: ChatCompletionParameters = serde_json::from_str(message)?;
let i_start = Instant::now();
let mut model = GlmOcrGenerateModel::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(())
}