From d2524d6a19634c069fe91f00cb5eff79c3e1fe97 Mon Sep 17 00:00:00 2001 From: jhqxxx <18280426169@163.com> Date: Sat, 7 Mar 2026 18:08:06 +0800 Subject: [PATCH] merge glm-ocr model --- Cargo.lock | 2 +- Cargo.toml | 4 +- README.md | 3 + README.zh-CN.md | 3 + docs/changelog.md | 3 + docs/changelog.zh-CN.md | 3 + docs/supported-models.md | 1 + docs/supported-models.zh-CN.md | 1 + src/exec/glm_ocr.rs | 4 +- src/models/glm_ocr/config.rs | 322 +++----------------------------- src/models/glm_ocr/generate.rs | 196 +++---------------- src/models/glm_ocr/model.rs | 194 +++++++++---------- src/models/glm_ocr/processor.rs | 182 +++++------------- src/models/mod.rs | 7 +- src/models/qwen3vl/processor.rs | 63 +------ src/position_embed/rope.rs | 2 +- src/utils/mod.rs | 22 ++- src/utils/video_utils.rs | 59 +++++- tests/test_glm_ocr.rs | 99 ++++++++++ 19 files changed, 395 insertions(+), 775 deletions(-) create mode 100644 tests/test_glm_ocr.rs diff --git a/Cargo.lock b/Cargo.lock index f7ea864..1af6034 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -30,7 +30,7 @@ dependencies = [ [[package]] name = "aha" -version = "0.2.1" +version = "0.2.2" dependencies = [ "aha_openai_dive", "anyhow", diff --git a/Cargo.toml b/Cargo.toml index d2fc8d2..134e639 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" } diff --git a/README.md b/README.md index 3284aba..9164227 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/README.zh-CN.md b/README.zh-CN.md index 9a6221a..0c1f177 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -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 模型 diff --git a/docs/changelog.md b/docs/changelog.md index 6044da0..640c49e 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -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 diff --git a/docs/changelog.zh-CN.md b/docs/changelog.zh-CN.md index a8d5500..0655ca1 100644 --- a/docs/changelog.zh-CN.md +++ b/docs/changelog.zh-CN.md @@ -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 模型 diff --git a/docs/supported-models.md b/docs/supported-models.md index 00cc460..6f65d90 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -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) diff --git a/docs/supported-models.zh-CN.md b/docs/supported-models.zh-CN.md index 94bca43..259a6a9 100644 --- a/docs/supported-models.zh-CN.md +++ b/docs/supported-models.zh-CN.md @@ -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) diff --git a/src/exec/glm_ocr.rs b/src/exec/glm_ocr.rs index 273d397..4505071 100644 --- a/src/exec/glm_ocr.rs +++ b/src/exec/glm_ocr.rs @@ -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(()) } -} \ No newline at end of file +} diff --git a/src/models/glm_ocr/config.rs b/src/models/glm_ocr/config.rs index 769af59..f1b246a 100644 --- a/src/models/glm_ocr/config.rs +++ b/src/models/glm_ocr/config.rs @@ -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, - /// 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, + 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, - /// 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, - /// Full RoPE configuration parameters. - #[serde(default)] - pub rope_parameters: Option, - /// End-of-sequence token ID. - #[serde(default)] - pub eos_token_id: Option, -} - -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 { - Some(128) + pub rope_parameters: GlmOcrRopeParameters, + pub eos_token_id: Vec, + 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, - #[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, - /// 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, } -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, - /// Std dev values for image normalization (per channel). - #[serde(default)] pub image_std: Vec, - /// Shortest edge for dynamic image resizing. Default: 448 - #[serde(default)] - pub size: Option, - /// 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, - /// Merge size for spatial merge. - #[serde(default = "default_merge_size")] - pub merge_size: Option, -} - -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 { - Some(14) -} - -fn default_merge_size() -> Option { - Some(2) + pub patch_size: usize, + pub merge_size: usize, + pub temporal_patch_size: usize, } diff --git a/src/models/glm_ocr/generate.rs b/src/models/glm_ocr/generate.rs index 49d79c5..c6c8d41 100644 --- a/src/models/glm_ocr/generate.rs +++ b/src/models/glm_ocr/generate.rs @@ -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) -> Result { - 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 = 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 { - // 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 { - 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 -} + + diff --git a/src/models/glm_ocr/model.rs b/src/models/glm_ocr/model.rs index d4fe3d7..ef09dd1 100644 --- a/src/models/glm_ocr/model.rs +++ b/src/models/glm_ocr/model.rs @@ -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 = Vec::with_capacity((q_len + CHUNK_SIZE - 1) / CHUNK_SIZE); + let mut chunks: Vec = + 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 { 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, - proj: Linear, - norm: LayerNorm, - #[allow(dead_code)] num_queries: usize, -} +// pub struct GlmOcrProjector { +// #[allow(dead_code)] +// query_embed: Option, +// proj: Linear, +// norm: LayerNorm, +// #[allow(dead_code)] +// num_queries: usize, +// } -impl GlmOcrProjector { - pub fn new( - vb: VarBuilder, - vision_config: &GlmOcrVisionConfig, - config: &GlmOcrProjectorConfig, - ) -> Result { - 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 { +// 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 { - let projected = self.proj.forward(image_features)?; - Ok(self.norm.forward(&projected)?) - } -} +// pub fn forward(&self, image_features: &Tensor) -> Result { +// 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 { - 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 = (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 { + pub fn new( + vb: VarBuilder, + config: GlmOcrTextConfig, + spatial_merge_size: usize, + ) -> Result { 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::()?; + let mask_vec = image_mask + .squeeze(0)? + .to_dtype(DType::U8)? + .to_vec1::()?; let mut t_ids: Vec = Vec::with_capacity(seq_len); let mut h_ids: Vec = 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 = 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) { diff --git a/src/models/glm_ocr/processor.rs b/src/models/glm_ocr/processor.rs index ceabe56..950aa64 100644 --- a/src/models/glm_ocr/processor.rs +++ b/src/models/glm_ocr/processor.rs @@ -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, image_std: Vec, - 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 { + 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 { 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 = 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] <|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] <|user|> \n - input_ids_vec.push(59248); // [gMASK] - input_ids_vec.push(59250); // - 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, diff --git a/src/models/mod.rs b/src/models/mod.rs index 7435dbc..5ea49f4 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -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), GlmASRNano(GlmAsrNanoGenerateModel<'a>), FunASRNano(FunAsrNanoGenerateModel), - GlmOCR(GlmOcrGenerateModel<'a>), + GlmOCR(GlmOcrGenerateModel), } impl<'a> GenerateModel for ModelInstance<'a> { diff --git a/src/models/qwen3vl/processor.rs b/src/models/qwen3vl/processor.rs index f69f75c..613f409 100644 --- a/src/models/qwen3vl/processor.rs +++ b/src/models/qwen3vl/processor.rs @@ -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, -) -> 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, diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index 8026fc7..8e55d1f 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -180,7 +180,7 @@ fn rotate_half_llm(x: &Tensor) -> Result { 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) diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 33a5c93..e0c5a3a 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -719,16 +719,20 @@ where pub fn extract_user_text(mes: &ChatCompletionParameters) -> Result { 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" + } + } } - } + }, + _ => {} } } } diff --git a/src/utils/video_utils.rs b/src/utils/video_utils.rs index 0545cb0..aaecf44 100644 --- a/src/utils/video_utils.rs +++ b/src/utils/video_utils.rs @@ -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, +) -> 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, diff --git a/tests/test_glm_ocr.rs b/tests/test_glm_ocr.rs new file mode 100644 index 0000000..2b94192 --- /dev/null +++ b/tests/test_glm_ocr.rs @@ -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(()) +}