From 3156818a0c2b0977d2f735794a2e844740d0b516 Mon Sep 17 00:00:00 2001 From: jason Date: Thu, 5 Mar 2026 23:49:49 -0500 Subject: [PATCH 1/2] glm-ocr --- src/api.rs | 3 +- src/exec/glm_ocr.rs | 69 ++ src/exec/mod.rs | 1 + src/main.rs | 5 + src/models/glm_ocr/config.rs | 369 +++++++ src/models/glm_ocr/generate.rs | 450 ++++++++ src/models/glm_ocr/mod.rs | 4 + src/models/glm_ocr/model.rs | 1801 +++++++++++++++++++++++++++++++ src/models/glm_ocr/processor.rs | 344 ++++++ src/models/mod.rs | 15 +- src/position_embed/rope.rs | 99 ++ 11 files changed, 3157 insertions(+), 3 deletions(-) create mode 100644 src/exec/glm_ocr.rs create mode 100644 src/models/glm_ocr/config.rs create mode 100644 src/models/glm_ocr/generate.rs create mode 100644 src/models/glm_ocr/mod.rs create mode 100644 src/models/glm_ocr/model.rs create mode 100644 src/models/glm_ocr/processor.rs diff --git a/src/api.rs b/src/api.rs index 060324a..0fe666e 100644 --- a/src/api.rs +++ b/src/api.rs @@ -253,6 +253,7 @@ fn which_model_to_id(which_model: WhichModel) -> &'static str { WhichModel::VoxCPM1_5 => "voxcpm1.5", WhichModel::GlmASRNano2512 => "glm-asr-nano-2512", WhichModel::FunASRNano2512 => "fun-asr-nano-2512", + WhichModel::GlmOCR => "glm-ocr", } } @@ -275,7 +276,7 @@ fn which_model_to_owner(which_model: WhichModel) -> &'static str { WhichModel::PaddleOCRVL => "PaddlePaddle", WhichModel::RMBG2_0 => "AI-ModelScope", WhichModel::VoxCPM | WhichModel::VoxCPM1_5 => "OpenBMB", - WhichModel::GlmASRNano2512 => "ZhipuAI", + WhichModel::GlmASRNano2512 | WhichModel::GlmOCR => "ZhipuAI", WhichModel::FunASRNano2512 => "FunAudioLLM", } } diff --git a/src/exec/glm_ocr.rs b/src/exec/glm_ocr.rs new file mode 100644 index 0000000..273d397 --- /dev/null +++ b/src/exec/glm_ocr.rs @@ -0,0 +1,69 @@ +//! Glm-OCR exec implementation for CLI `run` subcommand +use std::time::Instant; + +use anyhow::{Ok, Result}; + +use crate::exec::ExecModel; +use crate::models::{GenerateModel, glm_ocr::generate::GlmOcrGenerateModel}; + +pub struct GlmOcrExec; + +impl ExecModel for GlmOcrExec { + fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { + let url = &input[0]; + let input_url = if url.starts_with("http://") + || url.starts_with("https://") + || url.starts_with("file://") + { + url.clone() + } else { + format!("file://{}", url) + }; + + let i_start = Instant::now(); + let mut model = GlmOcrGenerateModel::init(weight_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let message = format!( + r#"{{ + "model": "glm-ocr", + "messages": [ + {{ + "role": "user", + "content": [ + {{ + "type": "image_url", + "image_url": {{ + "url": "{}" + }} + }}, + {{ + "type": "text", + "text": "Text Recognition:" + }} + ] + }} + ], + "max_tokens": 1024 + }}"#, + input_url + ); + + let mes = serde_json::from_str(&message)?; + + let i_start = Instant::now(); + let result = model.generate(mes)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + + println!("Result: {:?}", result); + + if let Some(out) = output { + std::fs::write(out, format!("{:?}", result))?; + println!("Output saved to: {}", out); + } + + Ok(()) + } +} \ No newline at end of file diff --git a/src/exec/mod.rs b/src/exec/mod.rs index aac7164..1e2c08e 100644 --- a/src/exec/mod.rs +++ b/src/exec/mod.rs @@ -6,6 +6,7 @@ pub mod deepseek_ocr; pub mod fun_asr_nano; pub mod glm_asr_nano; +pub mod glm_ocr; pub mod hunyuan_ocr; pub mod minicpm4; pub mod paddleocr_vl; diff --git a/src/main.rs b/src/main.rs index 95d8179..b5aec59 100644 --- a/src/main.rs +++ b/src/main.rs @@ -219,6 +219,7 @@ fn run_list(args: ListArgs) -> anyhow::Result<()> { WhichModel::VoxCPM1_5, WhichModel::GlmASRNano2512, WhichModel::FunASRNano2512, + WhichModel::GlmOCR, ]; if args.json { @@ -466,6 +467,10 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> { use aha::exec::fun_asr_nano::FunASRNanoExec; FunASRNanoExec::run(&input, output.as_deref(), &weight_path)?; } + WhichModel::GlmOCR => { + use aha::exec::glm_ocr::GlmOcrExec; + GlmOcrExec::run(&input, output.as_deref(), &weight_path)?; + } } Ok(()) diff --git a/src/models/glm_ocr/config.rs b/src/models/glm_ocr/config.rs new file mode 100644 index 0000000..769af59 --- /dev/null +++ b/src/models/glm_ocr/config.rs @@ -0,0 +1,369 @@ +use candle_nn::Activation; +use serde::Deserialize; + +/// 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 +} + +/// 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) +} + +/// Top-level configuration for GLM-OCR multimodal model. +#[derive(Debug, Clone, PartialEq, Deserialize, Default)] +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)] +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 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, +} + +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)] +pub struct GlmOcrPreprocessorConfig { + /// Mean values for image normalization (per channel). + #[serde(default)] + 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) +} diff --git a/src/models/glm_ocr/generate.rs b/src/models/glm_ocr/generate.rs new file mode 100644 index 0000000..79e7f61 --- /dev/null +++ b/src/models/glm_ocr/generate.rs @@ -0,0 +1,450 @@ +//! GLM-OCR Inference and Generation +use anyhow::{Result, anyhow}; +use aha_openai_dive::v1::resources::chat::{ + ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, +}; +use candle_core::{DType, Device, IndexOp, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::utils::apply_repeat_penalty; +use rocket::async_stream::stream; +use rocket::futures::Stream; + +use crate::{ + chat_template::ChatTemplate, + models::{ + GenerateModel, + glm_ocr::{ + config::{GlmOcrConfig, GlmOcrGenerationConfig}, + model::GlmOcrModel, + processor::GlmOcrProcessor, + }, + }, + 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, + }, +}; + +pub struct GlmOcrGenerateModel<'a> { + chat_template: ChatTemplate<'a>, + tokenizer: TokenizerModel, + processor: GlmOcrProcessor, + model: GlmOcrModel, + device: Device, + eos_token_ids: Vec, + generation_config: GlmOcrGenerationConfig, + model_name: String, + image_token_id: u32, + image_start_token_id: u32, + image_end_token_id: u32, + patch_size: usize, + temporal_patch_size: usize, + spatial_merge_size: usize, +} + +impl<'a> GlmOcrGenerateModel<'a> { + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + 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 device = get_device(device); + let cfg_dtype = if cfg.torch_dtype.is_empty() { + "bfloat16" + } else { + &cfg.torch_dtype + }; + 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 + }; + + #[cfg(debug_assertions)] + { + eprintln!("GLM-OCR Config Debug:"); + eprintln!(" text_config.hidden_size: {}", cfg.text_config.hidden_size); + eprintln!( + " text_config.num_attention_heads: {}", + cfg.text_config.num_attention_heads + ); + eprintln!( + " text_config.num_key_value_heads: {}", + cfg.text_config.num_key_value_heads + ); + eprintln!( + " text_config.head_dim: {}", + cfg.text_config.head_dim.unwrap_or_else(|| { + // Integer division, panics if num_attention_heads is 0 (like Python) + cfg.text_config.hidden_size / cfg.text_config.num_attention_heads + }) + ); + eprintln!( + " text_config.mrope_section: {:?}", + cfg.text_config.mrope_section + ); + eprintln!( + " Calculated head_dim: {}", + cfg.text_config.hidden_size / cfg.text_config.num_attention_heads + ); + } + + 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)? }; + let model = GlmOcrModel::new(vb, cfg.clone())?; + let generation_config_path = path.to_string() + "/generation_config.json"; + let generation_config: GlmOcrGenerationConfig = + serde_json::from_slice(&std::fs::read(generation_config_path)?)?; + + Ok(Self { + chat_template, + tokenizer, + processor, + model, + device, + eos_token_ids, + generation_config, + model_name: "glm-ocr".to_string(), + image_token_id: cfg.image_token_id, + image_start_token_id: cfg.image_start_token_id, + image_end_token_id: cfg.image_end_token_id, + patch_size: cfg.vision_config.patch_size, + temporal_patch_size: cfg.vision_config.temporal_patch_size, + spatial_merge_size: cfg.vision_config.spatial_merge_size, + }) + } +} + +impl<'a> GenerateModel for GlmOcrGenerateModel<'a> { + 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); + + // Extract image path and prompt from messages + let image_urls = extract_image_url(&mes); + let image_path = image_urls + .first() + .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()); + + #[cfg(debug_assertions)] + { + eprintln!("[GLM-OCR] ===== DEBUG START ====="); + eprintln!("[GLM-OCR] Image: {}", image_path); + eprintln!("[GLM-OCR] Prompt: {}", prompt); + eprintln!("[GLM-OCR] Device: {:?}", self.device); + eprintln!("[GLM-OCR] Temperature: {:?}", temperature); + eprintln!("[GLM-OCR] Top_p: {:?}", top_p); + eprintln!("[GLM-OCR] Max tokens: {}", mes.max_tokens.unwrap_or(512)); + } + + let processed = self.processor.process_info( + image_path, + &prompt, + &self.tokenizer, + self.image_token_id, + self.image_start_token_id, + self.image_end_token_id, + self.patch_size, + self.temporal_patch_size, + self.spatial_merge_size, + )?; + + if std::env::var("GLM_DEBUG").is_ok() { + eprintln!("[GLM-OCR] ===== AFTER PROCESS_INFO ====="); + eprintln!("[GLM-OCR] input_ids shape: {:?}", processed.input_ids.shape()); + let input_ids_vec = processed.input_ids.squeeze(0).unwrap().to_vec1::().unwrap(); + eprintln!("[GLM-OCR] input_ids (first 20): {:?}", &input_ids_vec[..20.min(input_ids_vec.len())]); + eprintln!("[GLM-OCR] pixel_values shape: {:?}", processed.pixel_values.shape()); + eprintln!("[GLM-OCR] image_mask shape: {:?}", processed.image_mask.shape()); + eprintln!("[GLM-OCR] grid_thw: {:?}", processed.grid_thw); + } + + let mut input_ids = processed.input_ids; + let pixel_values = Some(processed.pixel_values); + let image_grid_thw = Some(processed.grid_thw); + let image_mask = Some(processed.image_mask); + let mut seqlen_offset = 0; + let mut seq_len = input_ids.dim(1)?; + let mut generate = Vec::new(); + let sample_len = mes.max_tokens.unwrap_or(512); + + #[cfg(debug_assertions)] + { + eprintln!("[GLM-OCR] ===== START GENERATION ====="); + eprintln!("[GLM-OCR] Initial seq_len: {}", seq_len); + eprintln!("[GLM-OCR] eos_token_ids: {:?}", self.eos_token_ids); + } + + for _ in 0..sample_len { + let is_first_pass = seqlen_offset == 0; + let logits = self.model.forward( + &input_ids, + if is_first_pass { + pixel_values.as_ref() + } else { + None + }, + if is_first_pass { + image_grid_thw.as_ref() + } else { + None + }, + if is_first_pass { + image_mask.as_ref() + } else { + None + }, + 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 next_token = logit_processor.sample(&logits)?; + + generate.push(next_token); + if self.eos_token_ids.contains(&next_token) { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; + } + + 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)); + Ok(response) + } + + fn generate_stream( + &mut self, + mes: ChatCompletionParameters, + ) -> Result< + Box< + dyn Stream> + + Send + + Unpin + + '_, + >, + > { + // 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); + + // Extract image path and prompt from messages + let image_urls = extract_image_url(&mes); + let image_path = image_urls + .first() + .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 processed = self.processor.process_info( + image_path, + &prompt, + &self.tokenizer, + self.image_token_id, + self.image_start_token_id, + self.image_end_token_id, + self.patch_size, + self.temporal_patch_size, + self.spatial_merge_size, + )?; + + let mut input_ids = processed.input_ids; + let pixel_values = Some(processed.pixel_values); + let image_grid_thw = Some(processed.grid_thw); + let image_mask = Some(processed.image_mask); + let mut seqlen_offset = 0; + let mut seq_len = input_ids.dim(1)?; + let sample_len = mes.max_tokens.unwrap_or(512); + + let stream = stream! { + let mut generated: Vec = Vec::new(); + let mut error_tokens = Vec::new(); + for _ in 0..sample_len { + let is_first_pass = seqlen_offset == 0; + let logits = self.model.forward( + &input_ids, + if is_first_pass { pixel_values.as_ref() } else { None }, + if is_first_pass { image_grid_thw.as_ref() } else { None }, + if is_first_pass { image_mask.as_ref() } else { None }, + 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); + + let mut decode_ids = Vec::new(); + if !error_tokens.is_empty() { + decode_ids.extend_from_slice(&error_tokens); + } + decode_ids.push(next_token); + + let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("decode error: {e}")))?; + if decoded_token.contains("�") { + error_tokens.push(next_token); + if error_tokens.len() > 3 { + error_tokens.clear(); + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device).map_err(|e| anyhow!(format!("tensor error: {e}")))?; + continue; + } + error_tokens.clear(); + + let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); + yield Ok(chunk); + + if self.eos_token_ids.contains(&next_token) { + break; + } + seqlen_offset += seq_len; + seq_len = 1; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device).map_err(|e| anyhow!(format!("tensor error: {e}")))?; + } + self.model.clear_kv_cache(); + }; + + Ok(Box::new(Box::pin(stream))) + } +} + +/// 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/mod.rs b/src/models/glm_ocr/mod.rs new file mode 100644 index 0000000..8b1baf7 --- /dev/null +++ b/src/models/glm_ocr/mod.rs @@ -0,0 +1,4 @@ +pub mod config; +pub mod generate; +pub mod model; +pub mod processor; diff --git a/src/models/glm_ocr/model.rs b/src/models/glm_ocr/model.rs new file mode 100644 index 0000000..62a9dfd --- /dev/null +++ b/src/models/glm_ocr/model.rs @@ -0,0 +1,1801 @@ +//! GLM-OCR Model Implementation +//! +//! A multimodal vision-language model for OCR tasks that integrates: +//! - Vision Encoder: Processes images via patch embedding and transformer blocks +//! - Projector: Maps vision features to the language model's embedding space +//! - Language Model: Causal transformer decoder for text generation +//! +//! # Architecture +//! +//! ```text +//! Image → [Patch Embed] → [Vision Transformer] → [Spatial Merge] → [Projector] +//! ↓ +//! Text → [Token Embed] → [Decoder Layers with M-RoPE] ← [Feature Fusion] +//! ↓ +//! [LM Head] → Output Tokens +//! ``` +//! +//! Key features: +//! - M-RoPE (Multimodal Rotary Position Embedding) for unified 1D text and 3D vision positions +//! - Spatial merge to reduce visual token count before feeding to LLM +//! - KV-cache support for efficient autoregressive generation + +use anyhow::Result; +use candle_core::{D, DType, IndexOp, Tensor}; +use candle_nn::{ + Activation, Conv2d, Conv2dConfig, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, + conv2d, embedding, layer_norm, linear, linear_no_bias, rms_norm, +}; + +use crate::{ + models::{ + common::GateUpDownMLP, + glm_ocr::config::{ + GlmOcrConfig, GlmOcrProjectorConfig, 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}, + }, +}; + +/// Print tensor statistics in the same format as Python's `stats()` helper in compare_intermediate.py. +/// Gated by environment variable `GLM_INTERMEDIATE=1`. +fn tensor_stats(name: &str, t: &Tensor) { + if std::env::var("GLM_INTERMEDIATE").is_err() { + return; + } + let Ok(t_f32) = t.to_dtype(DType::F32) else { return }; + let Ok(flat) = t_f32.flatten_all() else { return }; + let n = flat.elem_count(); + if n == 0 { + eprintln!("[RS] {name}: shape={:?} EMPTY", t.shape()); + return; + } + let Ok(vals) = flat.to_vec1::() else { return }; + let mean = vals.iter().copied().sum::() / n as f32; + let variance = vals.iter().map(|v| (v - mean).powi(2)).sum::() / n as f32; + let std = variance.sqrt(); + let min = vals.iter().copied().fold(f32::INFINITY, f32::min); + let max = vals.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let n_show = 8.min(n); + let first: Vec = vals[..n_show].iter().map(|v| format!("{v:.4}")).collect(); + eprintln!( + "[RS] {name}: shape={:?} mean={mean:.6} std={std:.6} min={min:.6} max={max:.6}", + t.shape() + ); + eprintln!("[RS] {name}: first{n_show}={first:?}"); +} + +// ============================================================================ +// 1. GlmOcrRMSNorm +// ============================================================================ + +// Python: @use_kernel_forward_from_hub("RMSNorm") +// class GlmOcrRMSNorm(nn.Module): +// def __init__(self, hidden_size, eps: float = 1e-6) -> None: +// """ +// GlmOcrRMSNorm is equivalent to T5LayerNorm +// """ +// super().__init__() +// self.weight = nn.Parameter(torch.ones(hidden_size)) +// self.variance_epsilon = eps + +// def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: +// input_dtype = hidden_states.dtype +// hidden_states = hidden_states.to(torch.float32) +// variance = hidden_states.pow(2).mean(-1, keepdim=True) +// hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) +// return self.weight * hidden_states.to(input_dtype) + +// def extra_repr(self): +// return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" +pub struct GlmOcrRMSNorm(RmsNorm); + +impl GlmOcrRMSNorm { + pub fn new(vb: VarBuilder, hidden_size: usize, eps: f64) -> Result { + let rms = rms_norm(hidden_size, eps, vb)?; + Ok(Self(rms)) + } + pub fn forward(&self, xs: &Tensor) -> Result { + Ok(self.0.forward(xs)?) + } + pub fn extra_repr(&self) -> String { + "GlmOcrRMSNorm".to_string() + } +} + +// ============================================================================ +// 2. GlmOcrVisionMlp +// ============================================================================ + +// Python reference (transformers commit 4854dbf9): +// class GlmOcrVisionMlp(nn.Module): +// def __init__(self, config, bias=False): +// self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=bias) +// self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=bias) +// self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=bias) +// self.act_fn = ACT2FN[config.hidden_act] +// def forward(self, hidden_state): +// return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) +pub struct GlmOcrVisionMlp(GateUpDownMLP); + +impl GlmOcrVisionMlp { + pub fn new(vb: VarBuilder, config: &GlmOcrVisionConfig) -> Result { + let mlp = GateUpDownMLP::new( + vb, + config.hidden_size, + config.intermediate_size, + config.hidden_act, + config.attention_bias, + Some("gate_proj"), + Some("up_proj"), + Some("down_proj"), + )?; + Ok(Self(mlp)) + } + + pub fn forward(&self, hidden_state: &Tensor) -> Result { + Ok(self.0.forward(hidden_state)?) + } +} + +// ============================================================================ +// 3. eager_attention_forward (+ repeat_kv from utils) +// ============================================================================ + +// Python eager_attention_forward implementation +// Reference: py-glm-ocr/glm_ocr/modeling_glm_ocr.py +fn eager_attention_forward( + query_states: &Tensor, + key_states: &Tensor, + value_states: &Tensor, + num_key_value_groups: Option, + attention_mask: Option<&Tensor>, + scaling: f64, + dropout: f64, +) -> Result<(Tensor, Tensor)> { + // Attention matrix size info — enable with GLM_DEBUG=1. + if std::env::var("GLM_DEBUG").is_ok() { + let dims = query_states.dims(); + let (heads, q_len, k_len) = (dims[1], dims[2], key_states.dims()[2]); + let attn_mb = (heads * q_len * k_len * 4) as f64 / 1_048_576.0; + eprintln!( + "[GLM-OCR OOM-DBG] eager_attention_forward: heads={heads} q_len={q_len} k_len={k_len} \ + attn_weights={attn_mb:.1}MB (f32)" + ); + } + let key_states = match num_key_value_groups { + Some(g) => repeat_kv(key_states.clone(), g)?.contiguous()?, + None => key_states.clone(), + }; + let value_states = match num_key_value_groups { + Some(g) => repeat_kv(value_states.clone(), g)?.contiguous()?, + None => value_states.clone(), + }; + let query_states = query_states.contiguous()?; + let key_states = key_states.contiguous()?; + let value_states = value_states.contiguous()?; + + let output = { + #[cfg(feature = "flash-attn")] + { + // Flash attention: causal iff attention_mask is present. + let q = query_states.transpose(1, 2)?; + let k = key_states.transpose(1, 2)?; + let v = value_states.transpose(1, 2)?; + 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 + } + #[cfg(not(feature = "flash-attn"))] + { + // Chunked Q-attention: process Q in blocks so the attention matrix + // [batch, heads, CHUNK, k_len] stays bounded regardless of q_len. + // Peak memory per chunk: CHUNK × k_len × heads × 4 bytes (f32 softmax). + // Mathematically equivalent to full attention. + const CHUNK_SIZE: usize = 512; + let q_len = query_states.dim(2)?; + 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 start = 0; + while start < q_len { + let len = CHUNK_SIZE.min(q_len - start); + let q_chunk = query_states.narrow(2, start, len)?; + 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())?)? + } + }; + let attn = candle_nn::ops::softmax_last_dim(&attn.to_dtype(DType::F32)?)? + .to_dtype(query_states.dtype())?; + chunks.push(attn.matmul(&value_states)?); + start += len; + } + Tensor::cat(&chunks, 2)? // [batch, heads, q_len, head_dim] + } else { + let attn = (query_states.matmul(&k_t)? * scaling)?; + let attn = match attention_mask { + None => attn, + Some(mask) => attn.broadcast_add(&mask.to_dtype(attn.dtype())?)?, + }; + let attn = candle_nn::ops::softmax_last_dim(&attn.to_dtype(DType::F32)?)? + .to_dtype(query_states.dtype())?; + candle_nn::ops::dropout(&attn, dropout as f32)?.matmul(&value_states)? + }; + // [batch, heads, q_len, head_dim] -> [batch, q_len, heads, head_dim] + raw.transpose(1, 2)?.contiguous()? + } + }; + + // output layout: [batch, q_len, heads, head_dim] + let placeholder = Tensor::zeros((0,), query_states.dtype(), query_states.device())?; + Ok((output, placeholder)) +} + +// ============================================================================ +// 5. GlmOcrTextAttention +// ============================================================================ + +// Python: class GlmOcrTextAttention(nn.Module): +// def __init__(self, config): +// self.q_proj = nn.Linear(...) +// self.k_proj = nn.Linear(...) +// self.v_proj = nn.Linear(...) +// self.o_proj = nn.Linear(...) +// def forward(self, hidden_states, attention_mask, position_ids, past_key_value, use_cache): +// # QKV projection, attention computation, output projection +// return attn_output, attn_weights +pub struct GlmOcrTextAttention { + q_proj: Linear, + k_proj: Linear, + v_proj: Linear, + o_proj: Linear, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + scaling: f64, + kv_cache: Option<(Tensor, Tensor)>, +} + +impl GlmOcrTextAttention { + pub fn new( + vb: VarBuilder, + config: &GlmOcrTextConfig, + _layer_idx: Option, + ) -> Result { + 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 num_kv_groups = config.num_attention_heads / config.num_key_value_heads; + + let scaling = 1.0 / (head_dim as f64).sqrt(); + + let q_proj = linear_no_bias( + config.hidden_size, + config.num_attention_heads * head_dim, + vb.pp("q_proj"), + )?; + let k_proj = linear_no_bias( + config.hidden_size, + config.num_key_value_heads * head_dim, + vb.pp("k_proj"), + )?; + let v_proj = linear_no_bias( + config.hidden_size, + config.num_key_value_heads * head_dim, + vb.pp("v_proj"), + )?; + let o_proj = linear_no_bias( + config.num_attention_heads * head_dim, + config.hidden_size, + vb.pp("o_proj"), + )?; + + Ok(Self { + q_proj, + k_proj, + v_proj, + o_proj, + num_heads: config.num_attention_heads, + num_kv_heads: config.num_key_value_heads, + num_kv_groups, + head_dim, + scaling, + kv_cache: None, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + position_embeddings: (&Tensor, &Tensor), + attention_mask: Option<&Tensor>, + ) -> Result<(Tensor, Tensor)> { + let (bs, q_len, _) = xs.dims3()?; + + let query_states = self.q_proj.forward(xs)?; + let key_states = self.k_proj.forward(xs)?; + let value_states = self.v_proj.forward(xs)?; + + let query_states = query_states + .reshape((bs, q_len, self.num_heads, self.head_dim))? + .transpose(1, 2)?; + let key_states = key_states + .reshape((bs, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + let value_states = value_states + .reshape((bs, q_len, self.num_kv_heads, self.head_dim))? + .transpose(1, 2)?; + + let (cos, sin) = position_embeddings; + let (query_states, key_states) = + glm_ocr_apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?; + + // Python: if past_key_values is not None: + // cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + // key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) + let (key_states, value_states) = match &self.kv_cache { + None => (key_states, value_states), + Some((prev_k, prev_v)) => { + let key_states = Tensor::cat(&[prev_k, &key_states], 2)?; + let value_states = Tensor::cat(&[prev_v, &value_states], 2)?; + (key_states, value_states) + } + }; + self.kv_cache = Some((key_states.clone(), value_states.clone())); + + // Python: dropout=0.0 if not self.training else self.attention_dropout + // Rust is inference-only, so always 0.0 + let (attn_output, attn_weights) = eager_attention_forward( + &query_states, + &key_states, + &value_states, + Some(self.num_kv_groups), + attention_mask, + self.scaling, + 0.0, + )?; + + let attn_output = attn_output.reshape((bs, q_len, ()))?; + Ok((self.o_proj.forward(&attn_output)?, attn_weights)) + } + + pub fn clear_kv_cache(&mut self) { + self.kv_cache = None; + } +} + +// ============================================================================ +// 7. GlmOcrVisionRotaryEmbedding +// ============================================================================ + +// Python: class GlmOcrVisionRotaryEmbedding(nn.Module): +// def __init__(self, dim, theta=10000.0): +// inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) +// def forward(self, seqlen): +// seq = torch.arange(seqlen, device=self.inv_freq.device) +// return torch.outer(seq, self.inv_freq) # (seqlen, dim/2) +pub struct GlmOcrVisionRotaryEmbedding { + inv_freq: Tensor, +} + +impl GlmOcrVisionRotaryEmbedding { + pub fn new(dim: usize, theta: f32, device: &candle_core::Device, dtype: DType) -> Result { + let inv_freq: Vec = (0..dim) + .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)?; + Ok(Self { inv_freq }) + } + + pub fn forward(&self, seqlen: usize) -> Result { + // Python: freqs = torch.outer(seq, self.inv_freq) -> (seqlen, dim/4) + let seq = Tensor::arange(0f32, seqlen as f32, self.inv_freq.device())?; + let seq = seq.to_dtype(self.inv_freq.dtype())?; + let freqs = seq.unsqueeze(1)?.matmul(&self.inv_freq.unsqueeze(0)?)?; + Ok(freqs) + } + + /// Python: GlmOcrVisionModel.rot_pos_emb(self, grid_thw) + /// pos_ids = [] + /// for t, h, w in grid_thw: + /// hpos_ids = arange(h).unsqueeze(1).expand(-1, w) + /// hpos_ids = hpos_ids.reshape(h//sms, sms, w//sms, sms).permute(0,2,1,3).flatten() + /// wpos_ids = arange(w).unsqueeze(0).expand(h, -1) + /// wpos_ids = wpos_ids.reshape(h//sms, sms, w//sms, sms).permute(0,2,1,3).flatten() + /// pos_ids.append(stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + /// pos_ids = cat(pos_ids, dim=0) # (total, 2) + /// max_grid_size = grid_thw[:, 1:].max() + /// rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) # (max_grid_size, dim/4) + /// rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) # (total, dim/2) + /// return rotary_pos_emb, pos_ids + pub fn rot_pos_emb( + &self, + grid_thw: &[(usize, usize, usize)], + spatial_merge_size: usize, + ) -> Result<(Tensor, Tensor)> { + let sms = spatial_merge_size; + let mut all_hpos: Vec = Vec::new(); + let mut all_wpos: Vec = Vec::new(); + let mut max_grid_size: usize = 0; + + for &(t, h, w) in grid_thw { + max_grid_size = max_grid_size.max(h).max(w); + + // Generate position indices for each patch BEFORE spatial merge + // The vision encoder processes all patches, then merges them + // Python: for each (t, h, w), generate h*w position indices + for _ in 0..t { + for hi in 0..h { + for wi in 0..w { + // Apply spatial merge rearrangement + // Python: hpos_ids = hpos_ids.reshape(h//sms, sms, w//sms, sms).permute(0,2,1,3).flatten() + let _hb = hi / sms; + 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 + all_hpos.push(hi as u32); + all_wpos.push(wi as u32); + } + } + } + } + + let total_seq = all_hpos.len(); + let freqs_full = self.forward(max_grid_size)?; // (max_grid_size, dim/4) + + // Python: rotary_pos_emb_full[pos_ids].flatten(1) + // pos_ids is (total, 2) with [h_idx, w_idx] entries + // rotary_pos_emb_full[pos_ids] -> (total, 2, dim/4) -> flatten(1) -> (total, dim/2) + // This is equivalent to cat(freqs[h_indices], freqs[w_indices], dim=-1) + let h_indices = Tensor::from_vec(all_hpos, (total_seq,), self.inv_freq.device())?; + let w_indices = Tensor::from_vec(all_wpos, (total_seq,), self.inv_freq.device())?; + let h_freqs = freqs_full.index_select(&h_indices, 0)?; // (total_seq, dim/4) + let w_freqs = freqs_full.index_select(&w_indices, 0)?; // (total_seq, dim/4) + + // Concatenate h and w freqs: (total_seq, dim/2) + let rotary_pos_emb = Tensor::cat(&[&h_freqs, &w_freqs], 1)?; + + // Python (in GlmOcrVisionModel.forward): + // emb = cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + // position_embeddings = (emb.cos(), emb.sin()) + let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], 1)?; + let cos = emb.cos()?; + let sin = emb.sin()?; + + Ok((cos, sin)) + } +} + +// ============================================================================ +// 8. GlmOcrTextMLP +// ============================================================================ + +// Python: class GlmOcrTextMLP(nn.Module): +// def forward(self, hidden_states): +// up_states = self.gate_up_proj(hidden_states) +// gate, up_states = up_states.chunk(2, dim=-1) +// up_states = up_states * self.activation_fn(gate) +// return self.down_proj(up_states) +pub struct GlmOcrTextMLP { + gate_up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl GlmOcrTextMLP { + pub fn new(vb: VarBuilder, config: &GlmOcrTextConfig) -> Result { + let gate_up_proj = linear_no_bias( + config.hidden_size, + 2 * config.intermediate_size, + vb.pp("gate_up_proj"), + )?; + let down_proj = linear_no_bias( + config.intermediate_size, + config.hidden_size, + vb.pp("down_proj"), + )?; + Ok(Self { + gate_up_proj, + down_proj, + act_fn: config.hidden_act, + }) + } + + pub fn forward(&self, xs: &Tensor) -> Result { + let up_states = self.gate_up_proj.forward(xs)?; + let dim = up_states.dims().len() - 1; + let chunks = up_states.chunk(2, dim)?; + let gate = &chunks[0]; + let up = &chunks[1]; + let up_states = up.broadcast_mul(&self.act_fn.forward(gate)?)?; + Ok(self.down_proj.forward(&up_states)?) + } +} + +// ============================================================================ +// 9. GlmOcrTextDecoderLayer +// ============================================================================ + +// Python: class GlmOcrTextDecoderLayer(nn.Module): +// def forward(self, hidden_states, position_embeddings): +// hidden_states = self.input_layernorm(hidden_states) +// hidden_states, _ = self.self_attn(hidden_states, position_embeddings) +// hidden_states = self.post_self_attn_layernorm(hidden_states) +// hidden_states = residual + hidden_states +// residual = hidden_states +// hidden_states = self.post_attention_layernorm(hidden_states) +// hidden_states = self.mlp(hidden_states) +// hidden_states = self.post_mlp_layernorm(hidden_states) +// return residual + hidden_states +pub struct GlmOcrTextDecoderLayer { + self_attn: GlmOcrTextAttention, + mlp: GlmOcrTextMLP, + input_layernorm: GlmOcrRMSNorm, + post_attention_layernorm: GlmOcrRMSNorm, + post_self_attn_layernorm: GlmOcrRMSNorm, + post_mlp_layernorm: GlmOcrRMSNorm, +} + +impl GlmOcrTextDecoderLayer { + pub fn new(vb: VarBuilder, config: &GlmOcrTextConfig, layer_idx: usize) -> Result { + let self_attn = GlmOcrTextAttention::new(vb.pp("self_attn"), config, Some(layer_idx))?; + let mlp = GlmOcrTextMLP::new(vb.pp("mlp"), config)?; + let input_layernorm = GlmOcrRMSNorm::new( + vb.pp("input_layernorm"), + config.hidden_size, + config.rms_norm_eps, + )?; + let post_attention_layernorm = GlmOcrRMSNorm::new( + vb.pp("post_attention_layernorm"), + config.hidden_size, + config.rms_norm_eps, + )?; + let post_self_attn_layernorm = GlmOcrRMSNorm::new( + vb.pp("post_self_attn_layernorm"), + config.hidden_size, + config.rms_norm_eps, + )?; + let post_mlp_layernorm = GlmOcrRMSNorm::new( + vb.pp("post_mlp_layernorm"), + config.hidden_size, + config.rms_norm_eps, + )?; + + Ok(Self { + self_attn, + mlp, + input_layernorm, + post_attention_layernorm, + post_self_attn_layernorm, + post_mlp_layernorm, + }) + } + + pub fn forward( + &mut self, + xs: &Tensor, + position_embeddings: (&Tensor, &Tensor), + attention_mask: Option<&Tensor>, + ) -> Result { + let residual = xs.clone(); + let xs = self.input_layernorm.forward(xs)?; + let (xs, _attn_weights) = + self.self_attn + .forward(&xs, position_embeddings, attention_mask)?; + let xs = self.post_self_attn_layernorm.forward(&xs)?; + let xs = residual.add(&xs)?; + + let residual = xs.clone(); + let xs = self.post_attention_layernorm.forward(&xs)?; + let xs = self.mlp.forward(&xs)?; + let xs = self.post_mlp_layernorm.forward(&xs)?; + Ok(xs.add(&residual)?) + } + + pub fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache(); + } +} + +// ============================================================================ +// 10. rotate_half + apply_rotary_pos_emb_vision (in position_embed/rope.rs) +// ============================================================================ + +// ============================================================================ +// 11. GlmOcrVisionAttention +// ============================================================================ + +// Python reference (transformers): +// class GlmOcrVisionAttention(nn.Module): +// def __init__(self, config): +// self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.attention_bias) +// self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias) +// self.q_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps) +// self.k_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps) +// def forward(self, hidden_states, position_embeddings): +// q, k, v = self.qkv(hidden_states).reshape(seq_len, 3, num_heads, -1).permute(1,0,2,3).unbind(0) +// query_states = self.q_norm(query_states) +// key_states = self.k_norm(key_states) +// q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) +// attn_output = attention(q, k, v) +// return self.proj(attn_output.reshape(seq_len, -1)) +pub struct GlmOcrVisionAttention { + num_heads: usize, + head_dim: usize, + scaling: f64, + qkv: Linear, + proj: Linear, + q_norm: GlmOcrRMSNorm, + k_norm: GlmOcrRMSNorm, +} + +impl GlmOcrVisionAttention { + pub fn new(vb: VarBuilder, config: &GlmOcrVisionConfig) -> Result { + let head_dim = config.hidden_size / config.num_heads; + let scaling = 1.0 / (head_dim as f64).sqrt(); + let qkv = linear(config.hidden_size, config.hidden_size * 3, vb.pp("qkv"))?; + let proj = linear(config.hidden_size, config.hidden_size, vb.pp("proj"))?; + let q_norm = GlmOcrRMSNorm::new(vb.pp("q_norm"), head_dim, config.rms_norm_eps)?; + let k_norm = GlmOcrRMSNorm::new(vb.pp("k_norm"), head_dim, config.rms_norm_eps)?; + + Ok(Self { + num_heads: config.num_heads, + head_dim, + scaling, + qkv, + proj, + q_norm, + k_norm, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + position_embeddings: Option<(&Tensor, &Tensor)>, + ) -> Result { + let (seq_len, _) = xs.dims2()?; + let qkv = self.qkv.forward(xs)?; + let qkv = qkv + .reshape((seq_len, 3, self.num_heads, self.head_dim))? + .permute((1, 0, 2, 3))?; + + let q = qkv.i(0)?; + let k = qkv.i(1)?; + let v = qkv.i(2)?; + + let q = self.q_norm.forward(&q)?; + let k = self.k_norm.forward(&k)?; + + let (cos, sin) = if let Some((cos, sin)) = position_embeddings { + (cos, sin) + } else { + return Err(anyhow::anyhow!( + "Position embeddings required for vision attention" + )); + }; + + let (q, k) = apply_rotary_pos_emb_vision(&q, &k, cos, sin)?; + + let q = q.transpose(0, 1)?.unsqueeze(0)?; + let k = k.transpose(0, 1)?.unsqueeze(0)?; + let v = v.transpose(0, 1)?.unsqueeze(0)?; + + // Vision attention always uses num_key_value_groups = 1 + let (attn_output, _attn_weights) = + eager_attention_forward(&q, &k, &v, Some(1), None, self.scaling, 0.0)?; + let attn_output = attn_output.reshape((seq_len, ()))?; + Ok(self.proj.forward(&attn_output)?) + } + + pub fn forward_with_params( + &self, + xs: &Tensor, + _cu_seqlens: &Tensor, + _rotary_pos_emb: Option<&Tensor>, + position_embeddings: Option<(&Tensor, &Tensor)>, + ) -> Result { + let (seq_len, _) = xs.dims2()?; + let qkv = self.qkv.forward(xs)?; + let qkv = qkv + .reshape((seq_len, 3, self.num_heads, self.head_dim))? + .permute((1, 0, 2, 3))?; + + let q = qkv.i(0)?; + let k = qkv.i(1)?; + let v = qkv.i(2)?; + + let q = self.q_norm.forward(&q)?; + let k = self.k_norm.forward(&k)?; + + let (cos, sin) = if let Some((cos, sin)) = position_embeddings { + (cos, sin) + } else { + return Err(anyhow::anyhow!( + "Position embeddings required for vision attention" + )); + }; + + let (q, k) = apply_rotary_pos_emb_vision(&q, &k, cos, sin)?; + + let q = q.transpose(0, 1)?.unsqueeze(0)?; + let k = k.transpose(0, 1)?.unsqueeze(0)?; + let v = v.transpose(0, 1)?.unsqueeze(0)?; + + // Vision attention always uses num_key_value_groups = 1 + let (attn_output, _attn_weights) = + eager_attention_forward(&q, &k, &v, Some(1), None, self.scaling, 0.0)?; + let attn_output = attn_output.reshape((seq_len, ()))?; + Ok(self.proj.forward(&attn_output)?) + } +} + +// ============================================================================ +// 12. GlmOcrVisionBlock +// ============================================================================ + +// Python: class GlmOcrVisionBlock(GradientCheckpointingLayer): +// def __init__(self, config): +// self.norm1 = GlmOcrRMSNorm(...) +// self.attn = GlmOcrVisionAttention(...) +// self.norm2 = GlmOcrRMSNorm(...) +// self.mlp = GlmOcrVisionMlp(...) +// def forward(self, x): +// x = x + self.attn(self.norm1(x)) +// x = x + self.mlp(self.norm2(x)) +// return x +pub struct GlmOcrVisionBlock { + norm1: GlmOcrRMSNorm, + norm2: GlmOcrRMSNorm, + attn: GlmOcrVisionAttention, + mlp: GlmOcrVisionMlp, +} + +impl GlmOcrVisionBlock { + pub fn new(vb: VarBuilder, config: &GlmOcrVisionConfig) -> Result { + let norm1 = GlmOcrRMSNorm::new(vb.pp("norm1"), config.hidden_size, config.rms_norm_eps)?; + let attn = GlmOcrVisionAttention::new(vb.pp("attn"), config)?; + let norm2 = GlmOcrRMSNorm::new(vb.pp("norm2"), config.hidden_size, config.rms_norm_eps)?; + let mlp = GlmOcrVisionMlp::new(vb.pp("mlp"), config)?; + + Ok(Self { + norm1, + norm2, + attn, + mlp, + }) + } + + pub fn forward( + &self, + xs: &Tensor, + cu_seqlens: &Tensor, + rotary_pos_emb: Option<&Tensor>, + position_embeddings: Option<(&Tensor, &Tensor)>, + ) -> Result { + let residual = xs.clone(); + let xs = self.norm1.forward(xs)?; + let xs = + self.attn + .forward_with_params(&xs, cu_seqlens, rotary_pos_emb, position_embeddings)?; + let xs = residual.add(&xs)?; + + let residual = xs.clone(); + let xs = self.norm2.forward(&xs)?; + let xs = self.mlp.forward(&xs)?; + Ok(xs.add(&residual)?) + } +} + +// ============================================================================ +// 13. GlmOcrVisionPatchMerger +// ============================================================================ + +// Python reference (transformers commit 4854dbf9): +// class GlmOcrVisionPatchMerger(nn.Module): +// def __init__(self, dim, context_dim, hidden_act, bias=False): +// # dim = out_hidden_size (1536) +// # context_dim = intermediate_size (4096) NOT out_hidden_size * in_channels! +// self.proj = nn.Linear(dim, dim, bias=bias) +// self.post_projection_norm = LayerNorm(dim) +// self.gate_proj = nn.Linear(dim, context_dim, bias=bias) +// self.up_proj = nn.Linear(dim, context_dim, bias=bias) +// self.down_proj = nn.Linear(context_dim, dim, bias=bias) +// def forward(self, x): +// x = self.proj(x) +// x = self.post_projection_norm(x) +// x = GELU(x) # fixed activation, not config.hidden_act +// x = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) +// return x +// +// FIX: context_dim should be intermediate_size (4096), not out_hidden_size * in_channels (1536*3=4608) +pub struct GlmOcrVisionPatchMerger { + proj: Linear, + post_projection_norm: LayerNorm, + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, + act_fn: Activation, +} + +impl GlmOcrVisionPatchMerger { + pub fn new(vb: VarBuilder, config: &GlmOcrVisionConfig) -> Result { + // NOTE: The checkpoint stores proj.weight as [out_features, in_features] = [1536, 1536] + // PyTorch Linear computes: output = input @ weight.T + // Candle Linear computes: output = weight.matmul(input) which is equivalent to input @ weight.T + // So the weight should be loaded correctly. + let proj = linear_no_bias( + config.out_hidden_size, + config.out_hidden_size, + vb.pp("proj"), + )?; + + let post_projection_norm = layer_norm( + config.out_hidden_size, + config.rms_norm_eps, + vb.pp("post_projection_norm"), + )?; + + // Patch merger MLP: uses out_hidden_size * in_channels as intermediate dim + // Checkpoint shape: [4608, 1536] = [out_hidden_size * in_channels, out_hidden_size] + 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"), + )?; + + Ok(Self { + proj, + post_projection_norm, + gate_proj, + up_proj, + down_proj, + act_fn: config.hidden_act, + }) + } + + pub fn forward(&self, hidden_state: &Tensor) -> Result { + let mut hidden_state = self.proj.forward(hidden_state)?; + hidden_state = self.post_projection_norm.forward(&hidden_state)?; + // Python: x = self.act1(x) where act1 = nn.GELU() - fixed GELU, not config activation + hidden_state = hidden_state.gelu()?; + + let gate = self.gate_proj.forward(&hidden_state)?; + let gate = self.act_fn.forward(&gate)?; + let up = self.up_proj.forward(&hidden_state)?; + let result = gate.broadcast_mul(&up)?; + + Ok(self.down_proj.forward(&result)?) + } +} + +// ============================================================================ +// 14. GlmOcrVisionPatchEmbed +// ============================================================================ + +// Python: class GlmOcrVisionPatchEmbed(nn.Module): +// def __init__(self, config): +// self.proj = nn.Conv3d(...) +// def forward(self, x): +// x = x.view(-1, C, T, P, P) +// x = self.proj(x).view(-1, embed_dim) +// return x +pub struct GlmOcrVisionPatchEmbed { + patch_size: usize, + temporal_patch_size: usize, + in_channels: usize, + #[allow(dead_code)] embed_dim: usize, + proj: Linear, +} + +impl GlmOcrVisionPatchEmbed { + pub fn new(vb: VarBuilder, config: &GlmOcrVisionConfig) -> Result { + let patch_dim = + config.in_channels * config.temporal_patch_size * config.patch_size * config.patch_size; + + let weight = vb + .get( + ( + config.hidden_size, + config.in_channels, + config.temporal_patch_size, + config.patch_size, + config.patch_size, + ), + "proj.weight", + )? + .reshape((config.hidden_size, patch_dim))?; + + let bias = vb.get(config.hidden_size, "proj.bias").ok(); + + let proj = candle_nn::Linear::new(weight, bias); + + Ok(Self { + patch_size: config.patch_size, + temporal_patch_size: config.temporal_patch_size, + in_channels: config.in_channels, + embed_dim: config.hidden_size, + proj, + }) + } + + pub fn forward(&self, pixel_values: &Tensor) -> Result { + // pixel_values can be either: + // 1. Flattened patches: [num_patches, patch_dim] (Python format) + // 2. Standard image: [batch, C, H, W] (traditional format) + + let rank = pixel_values.rank(); + + if rank == 2 { + // Flattened patches format: [num_patches, patch_dim] + // Directly project to hidden_size + let hidden_states = self.proj.forward(pixel_values)?; + Ok(hidden_states) + } else { + // Standard image format: [batch, C, H, W] + let (batch, _c, h, w) = pixel_values.dims4()?; + + // Support non-square images + let patches_h = h / self.patch_size; + let patches_w = w / self.patch_size; + let num_patches = patches_h * patches_w; + + // Reshape: (batch, C, H, W) -> (batch, patches_h, patch_size, patches_w, patch_size, C) + let pv = pixel_values.reshape(( + batch, + patches_h, + self.patch_size, + patches_w, + self.patch_size, + self.in_channels, + ))?; + + // Permute: (batch, patches_h, patches_w, C, patch_size, patch_size) + let pv = pv.permute((0, 1, 3, 5, 2, 4))?; + + // Reshape: (batch * num_patches, C * patch_size * patch_size) + let pv = pv.reshape(( + batch * num_patches, + self.in_channels * self.patch_size * self.patch_size, + ))?; + + // Add temporal dimension + let pv = pv.unsqueeze(1)?; + let ones_shape: Vec = vec![1, self.temporal_patch_size]; + let pv = pv.broadcast_mul(&Tensor::ones(ones_shape, pv.dtype(), pv.device())?)?; + let pv = pv.reshape(( + batch * num_patches, + self.in_channels * self.temporal_patch_size * self.patch_size * self.patch_size, + ))?; + + let hidden_states = self.proj.forward(&pv)?; + Ok(hidden_states) + } + } +} + +// ============================================================================ +// 15. GlmOcrVisionModel +// ============================================================================ + +// Python: class GlmOcrVisionModel(GlmOcrPreTrainedModel): +// def __init__(self, config): +// self.patch_embed = GlmOcrVisionPatchEmbed(config) +// self.rotary_pos_emb = GlmOcrVisionRotaryEmbedding(...) +// self.blocks = nn.ModuleList([GlmOcrVisionBlock(config) for _ in range(config.depth)]) +// self.merger = GlmOcrVisionPatchMerger(...) +// self.downsample = nn.Conv2d(...) +// self.post_layernorm = GlmOcrRMSNorm(...) +// def forward(self, hidden_states, grid_thw): +// hidden_states = self.patch_embed(hidden_states) +// # ... apply rotary pos emb, blocks, merger, downsample +// return hidden_states +pub struct GlmOcrVisionModel { + patch_embed: GlmOcrVisionPatchEmbed, + rotary_pos_emb: GlmOcrVisionRotaryEmbedding, + blocks: Vec, + merger: GlmOcrVisionPatchMerger, + downsample: Conv2d, + post_layernorm: GlmOcrRMSNorm, + config: GlmOcrVisionConfig, +} + +impl GlmOcrVisionModel { + pub fn new(vb: VarBuilder, config: &GlmOcrVisionConfig) -> Result { + 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 mut blocks = Vec::new(); + let depth = config.depth; + for i in 0..depth { + let block = GlmOcrVisionBlock::new(vb.pp("blocks").pp(i), config)?; + blocks.push(block); + } + + let merger = GlmOcrVisionPatchMerger::new(vb.pp("merger"), config)?; + + let downsample = conv2d( + config.hidden_size, + config.out_hidden_size, + config.spatial_merge_size, + Conv2dConfig { + stride: config.spatial_merge_size, + ..Default::default() + }, + vb.pp("downsample"), + )?; + + let post_layernorm = GlmOcrRMSNorm::new( + vb.pp("post_layernorm"), + config.hidden_size, + config.rms_norm_eps, + )?; + + Ok(Self { + patch_embed, + rotary_pos_emb, + blocks, + merger, + downsample, + post_layernorm, + config: config.clone(), + }) + } + + pub fn forward(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result { + let mut hidden_states = self.patch_embed.forward(pixel_values)?; + + // Parse grid_thw - may be shape (3,) or (N, 3) + let grid_thw_parsed = if grid_thw.dims().len() == 1 { + let t = grid_thw.i(0)?.to_dtype(DType::F32)?.to_scalar::()? as usize; + let h = grid_thw.i(1)?.to_dtype(DType::F32)?.to_scalar::()? as usize; + let w = grid_thw.i(2)?.to_dtype(DType::F32)?.to_scalar::()? as usize; + vec![(t, h, w)] + } else { + let grid_thw = grid_thw.to_dtype(DType::F32)?; + let n = grid_thw.dim(0)?; + let mut result = Vec::new(); + for i in 0..n { + let row = grid_thw.i(i)?; + let t = row.i(0)?.to_scalar::()? as usize; + let h = row.i(1)?.to_scalar::()? as usize; + let w = row.i(2)?.to_scalar::()? as usize; + result.push((t, h, w)); + } + result + }; + + tensor_stats("after_patch_embed", &hidden_states); + + // Compute rotary embeddings matching Python rot_pos_emb exactly + let (cos, sin) = self + .rotary_pos_emb + .rot_pos_emb(&grid_thw_parsed, self.config.spatial_merge_size)?; + + tensor_stats("vision_cos", &cos); + tensor_stats("vision_sin", &sin); + + let rotary_pos_emb = Tensor::cat(&[&cos, &sin], D::Minus1)?; + let position_embeddings = (&cos, &sin); + + // Compute cu_seqlens + let mut cu_seqlens_values: Vec = vec![0]; + let mut cumsum: i32 = 0; + for (t, h, w) in &grid_thw_parsed { + let spatial_patches = (h * w) as i32; + for _ in 0..*t { + cumsum += spatial_patches; + cu_seqlens_values.push(cumsum); + } + } + let cu_seqlens = Tensor::from_slice( + &cu_seqlens_values, + &[cu_seqlens_values.len()], + hidden_states.device(), + )?; + + let num_blocks = self.blocks.len(); + if std::env::var("GLM_DEBUG").is_ok() { + let n_patches = hidden_states.dim(0).unwrap_or(0); + let hidden_mb = (hidden_states.elem_count() * 2) as f64 / 1_048_576.0; // bf16 + let attn_mb = (self.config.num_heads * n_patches * n_patches * 4) as f64 / 1_048_576.0; + eprintln!( + "[GLM-OCR OOM-DBG] Vision encoder: n_patches={n_patches} \ + hidden={hidden_mb:.1}MB per-layer attn_weights={attn_mb:.1}MB (f32) \ + num_blocks={num_blocks}" + ); + } + for (i, block) in self.blocks.iter().enumerate() { + hidden_states = block.forward( + &hidden_states, + &cu_seqlens, + Some(&rotary_pos_emb), + Some(position_embeddings), + )?; + if i == 0 { + tensor_stats("after_vision_block[0]", &hidden_states); + } + if i == num_blocks - 1 { + tensor_stats(&format!("after_vision_block[{i}] (last)"), &hidden_states); + } + } + + let hidden_states = self.post_layernorm.forward(&hidden_states)?; + tensor_stats("after_vision_post_layernorm", &hidden_states); + + let sms = self.config.spatial_merge_size; + let hidden_dim = hidden_states.dim(hidden_states.dims().len() - 1)?; + + // Python: hidden_states.view(-1, sms, sms, hidden_dim).permute(0, 3, 1, 2) + // Input: [2816, 1024] where 2816 = grid_h * grid_w = 44 * 64 + // After reshape: [704, 2, 2, 1024] where 704 = 2816 / 4 + // After permute: [704, 1024, 2, 2] + // After downsample: [704, 1536, 1, 1] + // After reshape: [704, 1536] + 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] + + tensor_stats("after_downsample", &hidden_states); + + let merged = self.merger.forward(&hidden_states)?; + tensor_stats("after_merger (vision features)", &merged); + + let merged = merged.unsqueeze(0)?; + Ok(merged) + } +} + +// ============================================================================ +// GlmOcrProjector (Rust-specific, no Python equivalent) +// ============================================================================ + +// Rust-specific: Projects vision features to language model space +// (Not present in Python - integrated in GlmOcrVisionModel forward) +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(); + + 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, + }) + } + + pub fn forward(&self, image_features: &Tensor) -> Result { + let projected = self.proj.forward(image_features)?; + Ok(self.norm.forward(&projected)?) + } +} + +// ============================================================================ +// 16. GlmOcrTextRotaryEmbedding +// ============================================================================ + +// Python: class GlmOcrTextRotaryEmbedding(nn.Module): +// def forward(self, x, position_ids): # position_ids: (3, bs, seq_len) +// inv_freq_expanded = self.inv_freq[None, None, :, None].expand(3, bs, -1, 1) +// position_ids_expanded = position_ids[:, :, None, :].float() +// freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3) +// freqs = self.apply_mrope(freqs, self.mrope_section) +// emb = torch.cat((freqs, freqs), dim=-1) +// return emb.cos() * self.attention_scaling, emb.sin() * self.attention_scaling +// def apply_mrope(self, freqs, mrope_section): +// chunks = freqs.split(mrope_section, dim=-1) +// return torch.cat([chunk[i % 3] for i, chunk in enumerate(chunks)], dim=-1) +pub struct GlmOcrTextRotaryEmbedding { + inv_freq: Tensor, + mrope_section: Vec, +} + +impl GlmOcrTextRotaryEmbedding { + pub fn new( + config: &GlmOcrTextConfig, + device: &candle_core::Device, + dtype: DType, + ) -> Result { + let rope_theta = config.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 inv_freq: Vec = (0..dim) + .step_by(2) + .map(|i| 1.0 / (rope_theta as f64).powf(i as f64 / dim as f64) as f32) + .collect(); + 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, + }) + } + + fn apply_mrope(&self, freqs: &Tensor) -> Result { + // freqs: (3, bs, seq_len, head_dim/2) + // Split by mrope_section and select from different axes + let section = &self.mrope_section; + let mut chunks = Vec::new(); + let mut offset = 0; + for &s in section.iter() { + let chunk = freqs.narrow(D::Minus1, offset, s)?; + chunks.push(chunk); + offset += s; + } + // Select chunk[i % 3] from axis 0 + let mut result_parts = Vec::new(); + for (i, chunk) in chunks.iter().enumerate() { + let selected = chunk.i(i % 3)?; // (bs, seq_len, section_size) + result_parts.push(selected); + } + Ok(Tensor::cat(&result_parts, D::Minus1)?) + } + + /// Compute cos/sin from explicit 3D position IDs (used for prefill with image tokens). + /// position_ids: (3, bs, seq_len) — axis 0 = temporal, 1 = height, 2 = width. + 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)?; // head_dim/2 + + // 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) + let inv_freq = inv_freq.broadcast_as((3, bs, inv_freq_len, 1))?; + let inv_freq = inv_freq.to_dtype(DType::F32)?.contiguous()?; + + // position_ids: (3, bs, seq_len) -> (3, bs, 1, seq_len) + let pos_expanded = position_ids + .unsqueeze(D::Minus2)? + .to_dtype(DType::F32)? + .contiguous()?; + + // freqs = inv_freq @ pos_expanded -> (3, bs, hd/2, seq_len) -> T -> (3, bs, seq_len, hd/2) + let freqs = inv_freq.matmul(&pos_expanded)?.transpose(2, 3)?; + + // Apply M-RoPE section selection + let freqs = self.apply_mrope(&freqs)?; // (bs, seq_len, hd/2) + + let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?.contiguous()?; + Ok(( + emb.cos()?.to_dtype(self.inv_freq.dtype())?, + emb.sin()?.to_dtype(self.inv_freq.dtype())?, + )) + } + + pub fn forward( + &self, + seq_len: usize, + seqlen_offset: usize, + device: &candle_core::Device, + ) -> Result<(Tensor, Tensor)> { + // For text-only (no image), all 3 axes have the same position IDs + let positions = Tensor::arange( + seqlen_offset as f32, + (seqlen_offset + seq_len) as f32, + device, + )? + .to_dtype(self.inv_freq.dtype())?; + + // position_ids: (3, 1, seq_len) + let positions = positions.unsqueeze(0)?; // (1, seq_len) + let positions_3d = positions.unsqueeze(0)?.expand((3, 1, seq_len))?; // (3, 1, seq_len) + + // inv_freq: (1, head_dim/2) -> (1, 1, head_dim/2, 1) -> (3, 1, head_dim/2, 1) + let inv_freq = self.inv_freq.unsqueeze(0)?.unsqueeze(D::Minus1)?; // (1, 1, hd/2, 1) + let inv_freq = inv_freq.broadcast_as((3, 1, self.inv_freq.dim(1)?, 1))?; // (3, 1, hd/2, 1) + let inv_freq = inv_freq.to_dtype(DType::F32)?.contiguous()?; + + // position_ids: (3, 1, 1, seq_len) + let positions_expanded = positions_3d + .unsqueeze(D::Minus2)? + .to_dtype(DType::F32)? + .contiguous()?; + + // freqs = inv_freq @ positions -> (3, 1, hd/2, seq_len) -> transpose -> (3, 1, seq_len, hd/2) + let freqs = inv_freq.matmul(&positions_expanded)?.transpose(2, 3)?; + + // Apply M-RoPE + let freqs = self.apply_mrope(&freqs)?; // (1, seq_len, hd/2) + + // Double: emb = cat(freqs, freqs) -> (1, seq_len, head_dim) + let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?.contiguous()?; + let cos = emb.cos()?; + let sin = emb.sin()?; + + Ok(( + cos.to_dtype(self.inv_freq.dtype())?, + sin.to_dtype(self.inv_freq.dtype())?, + )) + } +} + +// ============================================================================ +// 17. GlmOcrTextModel +// ============================================================================ + +// Python: class GlmOcrTextModel(GlmOcrPreTrainedModel): +// def __init__(self, config): +// self.embed_tokens = nn.Embedding(...) +// self.layers = nn.ModuleList([GlmOcrTextDecoderLayer(config) for _ in range(...)]) +// self.norm = GlmOcrRMSNorm(...) +// self.rotary_emb = GlmOcrTextRotaryEmbedding(...) +// def forward(self, input_ids, ...): +// hidden_states = self.embed_tokens(input_ids) +// # ... apply layers, rotary emb, return hidden_states +pub struct GlmOcrTextModel { + embed_tokens: Embedding, + layers: Vec, + norm: GlmOcrRMSNorm, + lm_head: Linear, + rotary_emb: GlmOcrTextRotaryEmbedding, + config: GlmOcrTextConfig, + spatial_merge_size: usize, + /// max_mrope_position + 1 after prefill (stored for decode-pass position computation) + next_mrope_pos: usize, + /// Number of tokens in the prefill pass + prefill_seq_len: usize, +} + +impl GlmOcrTextModel { + 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(); + for i in 0..config.num_hidden_layers { + let layer = GlmOcrTextDecoderLayer::new(vb.pp("layers").pp(i), &config, i)?; + layers.push(layer); + } + + let norm = GlmOcrRMSNorm::new(vb.pp("norm"), config.hidden_size, config.rms_norm_eps)?; + + // lm_head.weight lives at the checkpoint root (not under model.language_model) + let root_vb = vb.root(); + let lm_head = linear_no_bias(config.hidden_size, config.vocab_size, root_vb.pp("lm_head"))?; + + let rotary_emb = GlmOcrTextRotaryEmbedding::new(&config, vb.device(), vb.dtype())?; + + Ok(Self { + embed_tokens, + layers, + norm, + lm_head, + rotary_emb, + config, + spatial_merge_size, + next_mrope_pos: 0, + prefill_seq_len: 0, + }) + } + + /// Compute 3D M-RoPE position IDs matching Python's GlmOcrModel.get_rope_index(). + /// + /// For image tokens: each gets (t, h, w) grid coordinates. + /// For text tokens: sequential positions on all 3 axes. + /// Positions continue from where the previous group left off. + /// + /// Returns tensor of shape (3, 1, seq_len) containing [temporal, height, width] position IDs. + fn compute_mrope_position_ids( + &mut self, + image_mask: &Tensor, + grid_thw: &Tensor, + seq_len: usize, + device: &candle_core::Device, + ) -> Result { + // Parse original grid dimensions (before spatial merge) + let t_dim = grid_thw.i(0)?.to_dtype(DType::F32)?.to_scalar::()? as usize; + let h_dim = grid_thw.i(1)?.to_dtype(DType::F32)?.to_scalar::()? as usize; + let w_dim = grid_thw.i(2)?.to_dtype(DType::F32)?.to_scalar::()? as usize; + + // Merged grid dimensions (what the LLM sees as image tokens) + let llm_grid_t = t_dim; + let llm_grid_h = h_dim / self.spatial_merge_size; + let llm_grid_w = w_dim / self.spatial_merge_size; + 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 mut t_ids: Vec = Vec::with_capacity(seq_len); + let mut h_ids: Vec = Vec::with_capacity(seq_len); + let mut w_ids: Vec = Vec::with_capacity(seq_len); + + let mut st_idx: i64 = 0; // start index for the next group + let mut i = 0usize; + + while i < seq_len { + let is_img = mask_vec[i] == 1; + let start = i; + while i < seq_len && (mask_vec[i] == 1) == is_img { + i += 1; + } + let run_len = i - start; + + if is_img { + // Assign 3D (t, h, w) positions for merged image grid + assert_eq!( + run_len, num_image_tokens, + "image token count mismatch: mask={}, grid={}", + run_len, num_image_tokens + ); + for ti in 0..llm_grid_t { + for hi in 0..llm_grid_h { + for wi in 0..llm_grid_w { + t_ids.push(ti as i64 + st_idx); + h_ids.push(hi as i64 + st_idx); + w_ids.push(wi as i64 + st_idx); + } + } + } + // Next group starts after max(t, h, w) + 1 + let max_offset = (llm_grid_t as i64 - 1) + .max(llm_grid_h as i64 - 1) + .max(llm_grid_w as i64 - 1); + st_idx += max_offset + 1; + } else { + // Sequential positions for text tokens + for j in 0..run_len { + let pos = st_idx + j as i64; + t_ids.push(pos); + h_ids.push(pos); + w_ids.push(pos); + } + st_idx += run_len as i64; + } + } + + // st_idx is now max_mrope_pos + 1; store for decode passes + self.next_mrope_pos = st_idx as usize; + self.prefill_seq_len = seq_len; + + if std::env::var("GLM_DEBUG").is_ok() { + eprintln!( + "[M-RoPE] prefill: seq_len={}, next_mrope_pos={}", + seq_len, self.next_mrope_pos + ); + } + + let t_t = Tensor::from_vec(t_ids, (1, seq_len), device)?; + let h_t = Tensor::from_vec(h_ids, (1, seq_len), device)?; + let w_t = Tensor::from_vec(w_ids, (1, seq_len), device)?; + Ok(Tensor::stack(&[&t_t, &h_t, &w_t], 0)?) // (3, 1, seq_len) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + image_features: Option<&Tensor>, + image_mask: Option<&Tensor>, + image_grid_thw: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + let (bs, seq_len) = input_ids.dims2()?; + let mut inputs_embeds = self.embed_tokens.forward(input_ids)?; + tensor_stats("embed_tokens output", &inputs_embeds); + + #[cfg(debug_assertions)] + eprintln!( + "LanguageModel forward: bs={}, seq_len={}, head_dim={}", + bs, + seq_len, + self.config + .head_dim + .unwrap_or_else(|| self.config.hidden_size / self.config.num_attention_heads) + ); + + // Merge image features into embeddings at image token positions + if let (Some(img_feats), Some(img_mask)) = (image_features, image_mask) { + // img_feats: (1, num_features, hidden_size) + // img_mask: (1, seq_len) with 1s at image_token positions + let img_mask_bool = img_mask.squeeze(0)?.to_dtype(DType::U8)?.to_vec1::()?; + let _hidden_size = inputs_embeds.dim(2)?; + + // Collect image token indices + let image_indices: Vec = img_mask_bool + .iter() + .enumerate() + .filter(|&(_, &v)| v == 1) + .map(|(i, _)| i) + .collect(); + + let num_features = img_feats.dim(1)?; + let num_to_replace = image_indices.len().min(num_features); + + if std::env::var("GLM_DEBUG").is_ok() { + eprintln!("[LM] Image indices count: {}, num_features: {}, num_to_replace: {}", + image_indices.len(), num_features, num_to_replace); + if !image_indices.is_empty() { + eprintln!("[LM] First image index: {}, Last image index: {}", + image_indices.first().unwrap(), image_indices.last().unwrap()); + } + } + + // Replace embeddings at image positions with image features + // Build the merged embeddings by copying + 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 pos = 0; + for &img_pos in image_indices.iter().take(num_to_replace) { + 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; + pos = img_pos + 1; + } + if pos < seq_len { + embeds_vec.push(embeds_flat.narrow(0, pos, seq_len - pos)?); + } + + let refs: Vec<&Tensor> = embeds_vec.iter().collect(); + inputs_embeds = Tensor::cat(&refs, 0)?.unsqueeze(0)?; + if std::env::var("GLM_DEBUG").is_ok() { + eprintln!("[LM] inputs_embeds after image injection shape: {:?}", inputs_embeds.shape()); + if seq_len > 710 { + let text_embed = inputs_embeds.i((0, 710, ..))?; + let text_mean = text_embed.to_dtype(DType::F32)?.mean_all()?.to_scalar::()?; + eprintln!("[LM] Text token 710 embedding mean: {:.6}", text_mean); + } + } + } + tensor_stats("inputs_embeds after image injection", &inputs_embeds); + + let attention_mask = if seq_len > 1 { + Some(prepare_causal_attention_mask( + bs, + seq_len, + seqlen_offset, + input_ids.device(), + )?) + } else { + None + }; + + 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())?; + self.prefill_seq_len = seq_len; + self.rotary_emb.forward_with_position_ids(&pos_ids)? + } else { + // Pure text prefill: all three axes have sequential positions + self.next_mrope_pos = seq_len; + self.prefill_seq_len = seq_len; + self.rotary_emb.forward(seq_len, 0, input_ids.device())? + } + } else { + // Decode pass: single token, mrope position = next_mrope_pos + decode_step + let decode_pos = self.next_mrope_pos + (seqlen_offset - self.prefill_seq_len); + self.rotary_emb.forward(1, decode_pos, input_ids.device())? + }; + tensor_stats("text_rotary cos", &cos); + tensor_stats("text_rotary sin", &sin); + + let num_layers = self.layers.len(); + let mut hidden_states = inputs_embeds; + for (i, layer) in self.layers.iter_mut().enumerate() { + hidden_states = layer.forward(&hidden_states, (&cos, &sin), attention_mask.as_ref())?; + if i == 0 { + tensor_stats("after_text_layer[0]", &hidden_states); + } + if i == num_layers - 1 { + tensor_stats(&format!("after_text_layer[{i}] (last)"), &hidden_states); + } + } + + hidden_states = self.norm.forward(&hidden_states)?; + tensor_stats("after_final_norm", &hidden_states); + let logits = self.lm_head.forward(&hidden_states)?; + + // Log top-5 logits at last position (only on first pass) + if seqlen_offset == 0 && std::env::var("GLM_INTERMEDIATE").is_ok() { + let last_logits = logits.i((0, seq_len - 1, ..))?.to_dtype(DType::F32)?; + tensor_stats("logits at last position", &last_logits); + if let Ok(vals) = last_logits.to_vec1::() { + let mut indexed: Vec<(usize, f32)> = vals.iter().copied().enumerate().collect(); + indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + eprintln!("[RS] Top-5 logit token_ids: {:?}", + indexed[..5.min(indexed.len())].iter().map(|(i, v)| format!("id={i} val={v:.4}")).collect::>()); + } + } + + Ok(logits) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache(); + } + } +} + +// ============================================================================ +// 18. GlmOcrModel (corresponds to GlmOcrForConditionalGeneration in Python) +// ============================================================================ + +// Python: class GlmOcrForConditionalGeneration(GlmOcrPreTrainedModel, GenerationMixin): +// def __init__(self, config): +// self.model = GlmOcrModel(config) # Contains visual + language_model +// self.lm_head = nn.Linear(...) +// def forward(self, input_ids, pixel_values, ...): +// # Vision encoder -> language model -> lm_head +// return logits +pub struct GlmOcrModel { + vision_encoder: GlmOcrVisionModel, + language_model: GlmOcrTextModel, +} + +impl GlmOcrModel { + pub fn new(vb: VarBuilder, config: GlmOcrConfig) -> Result { + let vision_encoder = + GlmOcrVisionModel::new(vb.pp("model").pp("visual"), &config.vision_config)?; + let language_model = GlmOcrTextModel::new( + vb.pp("model").pp("language_model"), + config.text_config, + config.vision_config.spatial_merge_size, + )?; + + Ok(Self { + vision_encoder, + language_model, + }) + } + + pub fn forward( + &mut self, + input_ids: &Tensor, + pixel_values: Option<&Tensor>, + image_grid_thw: Option<&Tensor>, + image_mask: Option<&Tensor>, + seqlen_offset: usize, + ) -> Result { + #[cfg(debug_assertions)] + { + eprintln!("[GLM-OCR Model] ===== FORWARD START ====="); + eprintln!("[GLM-OCR Model] input_ids shape: {:?}", input_ids.shape()); + eprintln!("[GLM-OCR Model] seqlen_offset: {}", seqlen_offset); + if let Some(pv) = pixel_values { + eprintln!("[GLM-OCR Model] pixel_values shape: {:?}", pv.shape()); + } + if let Some(g) = image_grid_thw { + eprintln!("[GLM-OCR Model] grid_thw shape: {:?}", g.shape()); + } + if let Some(m) = image_mask { + eprintln!("[GLM-OCR Model] image_mask shape: {:?}", m.shape()); + let mask_sum = m.sum_all()?.to_scalar::()?; + eprintln!("[GLM-OCR Model] image_mask sum (num image tokens): {}", mask_sum); + } + } + + let image_features = if let Some(pixels) = pixel_values { + let grid_thw = if let Some(grid) = image_grid_thw { + grid.clone() + } else { + Tensor::new( + &[ + 1u32, + (pixels.dim(0)? / 44) as u32, // Approximate + (pixels.dim(1)? / 44) as u32, + ], + input_ids.device(), + )? + }; + + if std::env::var("GLM_DEBUG").is_ok() { + let pv_mb = (pixels.elem_count() * 2) as f64 / 1_048_576.0; // bf16 + eprintln!( + "[GLM-OCR OOM-DBG] pixel_values shape={:?} size={pv_mb:.1}MB", + pixels.shape() + ); + eprintln!("[GLM-OCR OOM-DBG] grid_thw={:?}", grid_thw.to_vec1::()); + } + + let vision_output = self.vision_encoder.forward(pixels, &grid_thw)?; + + if std::env::var("GLM_DEBUG").is_ok() { + eprintln!( + "[GLM-OCR Model] vision_output shape: {:?}", + vision_output.shape() + ); + let vis_mean = vision_output.to_dtype(candle_core::DType::F32)?.mean_all()?.to_scalar::()?; + eprintln!("[GLM-OCR Model] vision_output mean: {:.6}", vis_mean); + } + + Some(vision_output) + } else { + None + }; + + let result = self.language_model.forward( + input_ids, + image_features.as_ref(), + image_mask, + image_grid_thw, + seqlen_offset, + ); + + #[cfg(debug_assertions)] + { + eprintln!("[GLM-OCR Model] ===== FORWARD END ====="); + if let Ok(ref r) = result { + eprintln!("[GLM-OCR Model] output shape: {:?}", r.shape()); + } + } + + result + } + + pub fn clear_kv_cache(&mut self) { + self.language_model.clear_kv_cache(); + } +} diff --git a/src/models/glm_ocr/processor.rs b/src/models/glm_ocr/processor.rs new file mode 100644 index 0000000..c9ed381 --- /dev/null +++ b/src/models/glm_ocr/processor.rs @@ -0,0 +1,344 @@ +use anyhow::Result; +use candle_core::{DType, Device, Tensor}; + +use super::config::GlmOcrPreprocessorConfig; +use crate::tokenizer::TokenizerModel; +use crate::utils::img_utils::get_image; + +/// GLM-OCR Processor for image and text preprocessing. +/// +/// Matches Python's Glm46VImageProcessor behavior: +/// - Uses smart_resize to compute target dimensions +/// - Outputs flattened patches format [num_patches, patch_dim] +pub struct GlmOcrProcessor { + image_mean: Vec, + image_std: Vec, + shortest_edge: usize, // min_pixels in Python + longest_edge: usize, // max_pixels in Python + patch_size: usize, + merge_size: usize, + temporal_patch_size: usize, + device: Device, + dtype: DType, +} + +pub struct ProcessedImage { + 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 image_mask: Tensor, + pub grid_thw: Tensor, +} + +impl GlmOcrProcessor { + pub fn new(path: &str, device: &Device, dtype: DType) -> Result { + 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 + ) + }; + + Ok(Self { + image_mean, + image_std, + shortest_edge, + longest_edge, + patch_size, + merge_size, + temporal_patch_size: 2, // 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(): + /// 1. Resize using smart_resize + /// 2. Normalize + /// 3. Reshape into flattened patches [num_patches, patch_dim] + /// + /// Output format: [grid_t * grid_h * grid_w, channels * temporal_patch_size * patch_size * patch_size] + /// 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); + if std::env::var("GLM_DEBUG").is_ok() { + let grid_h = target_h / self.patch_size; + let grid_w = target_w / self.patch_size; + let n_patches = grid_h * grid_w; + let pv_elems = n_patches * 3 * self.temporal_patch_size * self.patch_size * self.patch_size; + let pv_mb = (pv_elems * 2) as f64 / 1_048_576.0; // bf16 + let attn_mb = (16 * n_patches * n_patches * 4) as f64 / 1_048_576.0; // 16 heads, f32 + eprintln!( + "[GLM-OCR OOM-DBG] Image: orig={orig_w}x{orig_h} -> target={target_w}x{target_h} \ + grid={grid_w}x{grid_h} n_patches={n_patches} \ + pixel_values={pv_mb:.1}MB vision_attn_per_layer={attn_mb:.1}MB" + ); + } + + // Resize image + let img = img.resize_exact( + target_w as u32, + target_h as u32, + image::imageops::FilterType::Lanczos3, + ); + + // Convert to RGB and normalize + let img = img.to_rgb8(); + let pixels: Vec = img + .pixels() + .flat_map(|p| { + vec![ + p[0] as f32 / 255.0, + p[1] as f32 / 255.0, + p[2] as f32 / 255.0, + ] + }) + .collect(); + + // Reshape to [H, W, 3] + let tensor = Tensor::from_vec(pixels, (target_h, target_w, 3), &self.device)?; + + // Normalize + 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; // Python uses temporal_patch_size=2 even for images + + // Reshape: [H, W, 3] -> [grid_h, patch_size, grid_w, patch_size, 3] + let tensor = tensor.reshape(( + grid_h, patch_size, + grid_w, patch_size, + channels, + ))?; + + // Permute to: [grid_h, grid_w, patch_size, patch_size, channels] + let tensor = tensor.permute((0, 2, 1, 3, 4))?; + + // Permute to put channels first: [grid_h, grid_w, channels, patch_size, patch_size] + // Python's patch_embed.forward does: view(-1, C, T, P, P) so we need (C, T, P_h, P_w) order + let tensor = tensor.permute((0, 1, 4, 2, 3))?; + + // Reshape to: [num_patches, channels, patch_size, patch_size] + let num_patches = grid_h * grid_w; + let tensor = tensor.reshape((num_patches, channels, patch_size, patch_size))?; + + // Add temporal dimension after C: [num_patches, channels, 1, patch_size, patch_size] + let tensor = tensor.unsqueeze(2)?; + // Repeat T times along temporal dim 2: [num_patches, channels, temporal_patch_size, patch_size, patch_size] + let tensor = tensor.repeat((1, 1, temporal_patch_size, 1, 1))?; + + // Flatten to: [num_patches, channels * temporal_patch_size * patch_size * patch_size] + // This gives (C, T, P_h, P_w) order matching Python's patch_embed input + let patch_dim = channels * temporal_patch_size * patch_size * patch_size; + let tensor = tensor.reshape((num_patches, patch_dim))?; + + // Convert to model dtype + let tensor = tensor.to_dtype(self.dtype)?; + + Ok(ProcessedImage { + pixel_values: tensor, + grid_h, + grid_w, + }) + } + + /// Process image and text for multimodal input. + /// + /// # Arguments + /// * `image_path` - Path to input image + /// * `prompt` - Text prompt/question about the image + /// * `tokenizer` - Tokenizer for text encoding + /// * `image_token_id` - Token ID for image content placeholders + /// * `image_start_token_id` - Token ID marking start of image region + /// * `image_end_token_id` - Token ID marking end of image region + /// * `patch_size` - Vision encoder patch size (default: 14) + /// * `temporal_patch_size` - Temporal patch size (default: 2, unused for images) + /// * `spatial_merge_size` - Spatial merge factor (default: 2) + /// + /// # Returns + /// ProcessedInput containing: + /// - input_ids: Combined image placeholder + text token IDs + /// - pixel_values: Flattened patches tensor [num_patches, patch_dim] + /// - image_mask: Boolean mask for image token positions + /// - grid_thw: (temporal, height, width) grid dimensions for RoPE + pub fn process_info( + &self, + image_path: &str, + prompt: &str, + tokenizer: &TokenizerModel, + image_token_id: u32, + image_start_token_id: u32, + image_end_token_id: u32, + _patch_size: usize, + _temporal_patch_size: usize, + spatial_merge_size: usize, + ) -> Result { + let processed_image = self.process_image(image_path)?; + let pixel_values = processed_image.pixel_values; + let grid_h = processed_image.grid_h; + let grid_w = processed_image.grid_w; + + // After spatial merge, each spatial_merge_size x spatial_merge_size block becomes 1 token + let merged_h = grid_h / spatial_merge_size; + let merged_w = grid_w / spatial_merge_size; + let num_image_tokens = merged_h * merged_w; + + // GLM-OCR format: [gMASK] <|user|> \n <|begin_of_image|> <|image|>*N <|end_of_image|> text <|assistant|> \n + // Special token IDs: + // 59248 = [gMASK] + // 59250 = + // 59253 = <|user|> + // 59256 = <|begin_of_image|> + // 59280 = <|image|> + // 59257 = <|end_of_image|> + // 59254 = <|assistant|> + + // Build input_ids following Python format + let mut input_ids_vec = Vec::new(); + + // 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|> + } + input_ids_vec.push(image_end_token_id); // <|end_of_image|> + + // Text prompt (without special tokens - they're already added) + let text_ids = tokenizer.text_encode_vec(prompt.to_string(), false)?; + input_ids_vec.extend(text_ids); + + // Generation prompt: <|assistant|> \n + input_ids_vec.push(59254); // <|assistant|> + input_ids_vec.push(10); // newline + + let input_ids = Tensor::from_vec( + input_ids_vec.clone(), + (1, input_ids_vec.len()), + &self.device, + )?; + + // Create image mask (1s at image token positions, 0 elsewhere) + // Image tokens start after header (4 tokens) + start token (1 token) = index 5 + let mut image_mask_vec = vec![0u32; input_ids_vec.len()]; + let image_start_idx = 5; // After [gMASK, sop, user, newline, begin_image] + for i in 0..num_image_tokens { + image_mask_vec[image_start_idx + i] = 1; + } + let image_mask = Tensor::from_vec(image_mask_vec, (1, input_ids_vec.len()), &self.device)?; + + // Compute grid_thw for RoPE + // For images: grid_t = 1 + let grid_thw = Tensor::from_vec( + vec![1u32, grid_h as u32, grid_w as u32], + (3,), + &self.device, + )?; + + Ok(ProcessedInput { + input_ids, + pixel_values, + image_mask, + grid_thw, + }) + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 741d629..7435dbc 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -5,6 +5,7 @@ pub mod deepseek_ocr; pub mod feature_extractor; pub mod fun_asr_nano; pub mod glm_asr_nano; +pub mod glm_ocr; pub mod hunyuan_ocr; pub mod mask_gct; pub mod minicpm4; @@ -27,7 +28,7 @@ use rocket::futures::Stream; use crate::models::{ deepseek_ocr::generate::DeepseekOCRGenerateModel, fun_asr_nano::generate::FunAsrNanoGenerateModel, - glm_asr_nano::generate::GlmAsrNanoGenerateModel, + glm_asr_nano::generate::GlmAsrNanoGenerateModel, glm_ocr::generate::GlmOcrGenerateModel, hunyuan_ocr::generate::HunyuanOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel, paddleocr_vl::generate::PaddleOCRVLGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, qwen3::generate::Qwen3GenerateModel, qwen3_5::generate::Qwen3_5GenerateModel, @@ -81,6 +82,8 @@ pub enum WhichModel { GlmASRNano2512, #[value(name = "fun-asr-nano-2512", hide = true)] FunASRNano2512, + #[value(name = "glm-ocr", hide = true)] + GlmOCR, } impl WhichModel { @@ -109,6 +112,7 @@ impl WhichModel { WhichModel::VoxCPM1_5 => "OpenBMB/VoxCPM1.5", WhichModel::GlmASRNano2512 => "ZhipuAI/GLM-ASR-Nano-2512", WhichModel::FunASRNano2512 => "FunAudioLLM/Fun-ASR-Nano-2512", + WhichModel::GlmOCR => "ZhipuAI/GLM-OCR", } } @@ -128,7 +132,7 @@ impl WhichModel { | WhichModel::Qwen3_5_4B | WhichModel::Qwen3_5_9B => "vlm", // OCR models - WhichModel::DeepSeekOCR | WhichModel::HunyuanOCR | WhichModel::PaddleOCRVL => "ocr", + WhichModel::DeepSeekOCR | WhichModel::HunyuanOCR | WhichModel::GlmOCR | WhichModel::PaddleOCRVL => "ocr", // ASR models WhichModel::Qwen3ASR0_6B | WhichModel::Qwen3ASR1_7B @@ -169,6 +173,7 @@ pub enum ModelInstance<'a> { VoxCPM(Box), GlmASRNano(GlmAsrNanoGenerateModel<'a>), FunASRNano(FunAsrNanoGenerateModel), + GlmOCR(GlmOcrGenerateModel<'a>), } impl<'a> GenerateModel for ModelInstance<'a> { @@ -187,6 +192,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::VoxCPM(model) => model.generate(mes), ModelInstance::GlmASRNano(model) => model.generate(mes), ModelInstance::FunASRNano(model) => model.generate(mes), + ModelInstance::GlmOCR(model) => model.generate(mes), } } @@ -215,6 +221,7 @@ impl<'a> GenerateModel for ModelInstance<'a> { ModelInstance::VoxCPM(model) => model.generate_stream(mes), ModelInstance::GlmASRNano(model) => model.generate_stream(mes), ModelInstance::FunASRNano(model) => model.generate_stream(mes), + ModelInstance::GlmOCR(model) => model.generate_stream(mes), } } } @@ -309,6 +316,10 @@ pub fn load_model(model_type: WhichModel, path: &str) -> Result { + let model = GlmOcrGenerateModel::init(path, None, None)?; + ModelInstance::GlmOCR(model) + } }; Ok(model) } diff --git a/src/position_embed/rope.rs b/src/position_embed/rope.rs index c5013fc..8026fc7 100644 --- a/src/position_embed/rope.rs +++ b/src/position_embed/rope.rs @@ -159,6 +159,105 @@ pub fn glm_asr_apply_rotary_pos_emb( Ok((q_embed, k_embed)) } +/// Interleaved rotation used by GLM-OCR text decoder. +/// +/// Python `rotate_half_llm`: +/// x1 = x[..., 0::2] # even indices +/// x2 = x[..., 1::2] # odd indices +/// return stack((-x2, x1), dim=-1).flatten(-2) +/// # e.g. [q0,q1,q2,q3] → [-q1, q0, -q3, q2] +/// +/// Each adjacent pair (x_{2i}, x_{2i+1}) is rotated to (-x_{2i+1}, x_{2i}). +/// This is the correct counterpart to `repeat_interleave(2)` style cos/sin. +fn rotate_half_llm(x: &Tensor) -> Result { + let last_dim = x.dim(D::Minus1)?; + let half = last_dim / 2; + // Reshape (..., D) → (..., D/2, 2) so each row is one adjacent pair + let mut pair_shape = x.dims().to_vec(); + let rank = pair_shape.len(); + pair_shape[rank - 1] = half; + pair_shape.push(2); + 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 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) + // Flatten last two dims back to D: [-x1, x0, -x3, x2, ...] + Ok(result_pairs.reshape(x.dims().to_vec())?) +} + +pub fn glm_ocr_apply_rotary_pos_emb( + q: &Tensor, + k: &Tensor, + cos: &Tensor, + sin: &Tensor, +) -> Result<(Tensor, Tensor)> { + // GLM-OCR applies rotary to only the first rotary_dim of head_dim + // cos/sin: (bs, seq_len, head_dim) - already doubled via cat(freqs, freqs) + // q/k: (bs, n_head, seq_len, head_dim) + // Python: unsqueeze_dim=1 + // - rank 2 (seq_len, head_dim) -> unsqueeze(1) -> (seq_len, 1, head_dim) + // - rank 3 (bs, seq_len, head_dim) -> unsqueeze(1) -> (bs, 1, seq_len, head_dim) + let mut cos = cos.clone(); + let mut sin = sin.clone(); + + cos = cos.unsqueeze(1)?; // (seq_len, head_dim) -> (seq_len, 1, head_dim) + sin = sin.unsqueeze(1)?; + + // Python: cos = cos[..., :cos.shape[-1]//2].repeat_interleave(2, dim=-1) + // Take first half and interleave each element + let full_dim = cos.dim(D::Minus1)?; + let half_dim = full_dim / 2; + let cos_half = cos.narrow(D::Minus1, 0, half_dim)?; + let sin_half = sin.narrow(D::Minus1, 0, half_dim)?; + // repeat_interleave(2, dim=-1): [a,b,c] -> [a,a,b,b,c,c] + let cos_interleaved = cos_half + .unsqueeze(D::Minus1)? + .broadcast_mul(&Tensor::ones( + &[1, 1, 1, 1, 2], + cos_half.dtype(), + cos_half.device(), + )?)? + .reshape(cos.shape())?; + let sin_interleaved = sin_half + .unsqueeze(D::Minus1)? + .broadcast_mul(&Tensor::ones( + &[1, 1, 1, 1, 2], + sin_half.dtype(), + sin_half.device(), + )?)? + .reshape(sin.shape())?; + + let cos = cos_interleaved.to_dtype(q.dtype())?; + let sin = sin_interleaved.to_dtype(q.dtype())?; + + let rotary_dim = cos.dim(D::Minus1)?; + // Split q/k into rotary and pass-through portions + let q_rot = q.narrow(D::Minus1, 0, rotary_dim)?; + let q_pass = q.narrow(D::Minus1, rotary_dim, q.dim(D::Minus1)? - rotary_dim)?; + let k_rot = k.narrow(D::Minus1, 0, rotary_dim)?; + let k_pass = k.narrow(D::Minus1, rotary_dim, k.dim(D::Minus1)? - rotary_dim)?; + + // Apply rotary: q_rot * cos + rotate_half_llm(q_rot) * sin + // Must use interleaved rotate_half_llm (not split-half rotate_half) because + // cos/sin use repeat_interleave(2) format: [c0,c0,c1,c1,...]. + // rotate_half_llm rotates adjacent pairs (q_{2i},q_{2i+1}) → (-q_{2i+1}, q_{2i}), + // which is the correct counterpart for this cos/sin format. + let q_embed = q_rot + .broadcast_mul(&cos)? + .add(&rotate_half_llm(&q_rot)?.broadcast_mul(&sin)?)?; + let k_embed = k_rot + .broadcast_mul(&cos)? + .add(&rotate_half_llm(&k_rot)?.broadcast_mul(&sin)?)?; + + // Concatenate rotary and pass-through portions + let q_embed = Tensor::cat(&[&q_embed, &q_pass], D::Minus1)?; + let k_embed = Tensor::cat(&[&k_embed, &k_pass], D::Minus1)?; + Ok((q_embed, k_embed)) +} + pub fn roformer_rotate(x: &Tensor) -> Result { let dims = x.dims(); let last_dim = dims From 1fe06110ff2b4c422ef26a986b5e74ae353acb25 Mon Sep 17 00:00:00 2001 From: jason Date: Fri, 6 Mar 2026 11:25:44 -0500 Subject: [PATCH 2/2] more glm-ocr tweaks --- src/models/glm_ocr/generate.rs | 57 ---- src/models/glm_ocr/model.rs | 512 ++------------------------------ src/models/glm_ocr/processor.rs | 54 +--- 3 files changed, 33 insertions(+), 590 deletions(-) diff --git a/src/models/glm_ocr/generate.rs b/src/models/glm_ocr/generate.rs index 79e7f61..49d79c5 100644 --- a/src/models/glm_ocr/generate.rs +++ b/src/models/glm_ocr/generate.rs @@ -124,35 +124,6 @@ impl<'a> GlmOcrGenerateModel<'a> { dtype }; - #[cfg(debug_assertions)] - { - eprintln!("GLM-OCR Config Debug:"); - eprintln!(" text_config.hidden_size: {}", cfg.text_config.hidden_size); - eprintln!( - " text_config.num_attention_heads: {}", - cfg.text_config.num_attention_heads - ); - eprintln!( - " text_config.num_key_value_heads: {}", - cfg.text_config.num_key_value_heads - ); - eprintln!( - " text_config.head_dim: {}", - cfg.text_config.head_dim.unwrap_or_else(|| { - // Integer division, panics if num_attention_heads is 0 (like Python) - cfg.text_config.hidden_size / cfg.text_config.num_attention_heads - }) - ); - eprintln!( - " text_config.mrope_section: {:?}", - cfg.text_config.mrope_section - ); - eprintln!( - " Calculated head_dim: {}", - cfg.text_config.hidden_size / cfg.text_config.num_attention_heads - ); - } - 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)? }; @@ -213,17 +184,6 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> { // Get prompt text from messages let prompt = extract_text_from_messages(&mes).unwrap_or_else(|| "Extract all text from this image.".to_string()); - #[cfg(debug_assertions)] - { - eprintln!("[GLM-OCR] ===== DEBUG START ====="); - eprintln!("[GLM-OCR] Image: {}", image_path); - eprintln!("[GLM-OCR] Prompt: {}", prompt); - eprintln!("[GLM-OCR] Device: {:?}", self.device); - eprintln!("[GLM-OCR] Temperature: {:?}", temperature); - eprintln!("[GLM-OCR] Top_p: {:?}", top_p); - eprintln!("[GLM-OCR] Max tokens: {}", mes.max_tokens.unwrap_or(512)); - } - let processed = self.processor.process_info( image_path, &prompt, @@ -236,16 +196,6 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> { self.spatial_merge_size, )?; - if std::env::var("GLM_DEBUG").is_ok() { - eprintln!("[GLM-OCR] ===== AFTER PROCESS_INFO ====="); - eprintln!("[GLM-OCR] input_ids shape: {:?}", processed.input_ids.shape()); - let input_ids_vec = processed.input_ids.squeeze(0).unwrap().to_vec1::().unwrap(); - eprintln!("[GLM-OCR] input_ids (first 20): {:?}", &input_ids_vec[..20.min(input_ids_vec.len())]); - eprintln!("[GLM-OCR] pixel_values shape: {:?}", processed.pixel_values.shape()); - eprintln!("[GLM-OCR] image_mask shape: {:?}", processed.image_mask.shape()); - eprintln!("[GLM-OCR] grid_thw: {:?}", processed.grid_thw); - } - let mut input_ids = processed.input_ids; let pixel_values = Some(processed.pixel_values); let image_grid_thw = Some(processed.grid_thw); @@ -255,13 +205,6 @@ impl<'a> GenerateModel for GlmOcrGenerateModel<'a> { let mut generate = Vec::new(); let sample_len = mes.max_tokens.unwrap_or(512); - #[cfg(debug_assertions)] - { - eprintln!("[GLM-OCR] ===== START GENERATION ====="); - eprintln!("[GLM-OCR] Initial seq_len: {}", seq_len); - eprintln!("[GLM-OCR] eos_token_ids: {:?}", self.eos_token_ids); - } - for _ in 0..sample_len { let is_first_pass = seqlen_offset == 0; let logits = self.model.forward( diff --git a/src/models/glm_ocr/model.rs b/src/models/glm_ocr/model.rs index 62a9dfd..d4fe3d7 100644 --- a/src/models/glm_ocr/model.rs +++ b/src/models/glm_ocr/model.rs @@ -1,24 +1,4 @@ //! GLM-OCR Model Implementation -//! -//! A multimodal vision-language model for OCR tasks that integrates: -//! - Vision Encoder: Processes images via patch embedding and transformer blocks -//! - Projector: Maps vision features to the language model's embedding space -//! - Language Model: Causal transformer decoder for text generation -//! -//! # Architecture -//! -//! ```text -//! Image → [Patch Embed] → [Vision Transformer] → [Spatial Merge] → [Projector] -//! ↓ -//! Text → [Token Embed] → [Decoder Layers with M-RoPE] ← [Feature Fusion] -//! ↓ -//! [LM Head] → Output Tokens -//! ``` -//! -//! Key features: -//! - M-RoPE (Multimodal Rotary Position Embedding) for unified 1D text and 3D vision positions -//! - Spatial merge to reduce visual token count before feeding to LLM -//! - KV-cache support for efficient autoregressive generation use anyhow::Result; use candle_core::{D, DType, IndexOp, Tensor}; @@ -40,57 +20,6 @@ use crate::{ }, }; -/// Print tensor statistics in the same format as Python's `stats()` helper in compare_intermediate.py. -/// Gated by environment variable `GLM_INTERMEDIATE=1`. -fn tensor_stats(name: &str, t: &Tensor) { - if std::env::var("GLM_INTERMEDIATE").is_err() { - return; - } - let Ok(t_f32) = t.to_dtype(DType::F32) else { return }; - let Ok(flat) = t_f32.flatten_all() else { return }; - let n = flat.elem_count(); - if n == 0 { - eprintln!("[RS] {name}: shape={:?} EMPTY", t.shape()); - return; - } - let Ok(vals) = flat.to_vec1::() else { return }; - let mean = vals.iter().copied().sum::() / n as f32; - let variance = vals.iter().map(|v| (v - mean).powi(2)).sum::() / n as f32; - let std = variance.sqrt(); - let min = vals.iter().copied().fold(f32::INFINITY, f32::min); - let max = vals.iter().copied().fold(f32::NEG_INFINITY, f32::max); - let n_show = 8.min(n); - let first: Vec = vals[..n_show].iter().map(|v| format!("{v:.4}")).collect(); - eprintln!( - "[RS] {name}: shape={:?} mean={mean:.6} std={std:.6} min={min:.6} max={max:.6}", - t.shape() - ); - eprintln!("[RS] {name}: first{n_show}={first:?}"); -} - -// ============================================================================ -// 1. GlmOcrRMSNorm -// ============================================================================ - -// Python: @use_kernel_forward_from_hub("RMSNorm") -// class GlmOcrRMSNorm(nn.Module): -// def __init__(self, hidden_size, eps: float = 1e-6) -> None: -// """ -// GlmOcrRMSNorm is equivalent to T5LayerNorm -// """ -// super().__init__() -// self.weight = nn.Parameter(torch.ones(hidden_size)) -// self.variance_epsilon = eps - -// def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: -// input_dtype = hidden_states.dtype -// hidden_states = hidden_states.to(torch.float32) -// variance = hidden_states.pow(2).mean(-1, keepdim=True) -// hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) -// return self.weight * hidden_states.to(input_dtype) - -// def extra_repr(self): -// return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" pub struct GlmOcrRMSNorm(RmsNorm); impl GlmOcrRMSNorm { @@ -106,19 +35,6 @@ impl GlmOcrRMSNorm { } } -// ============================================================================ -// 2. GlmOcrVisionMlp -// ============================================================================ - -// Python reference (transformers commit 4854dbf9): -// class GlmOcrVisionMlp(nn.Module): -// def __init__(self, config, bias=False): -// self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=bias) -// self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=bias) -// self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=bias) -// self.act_fn = ACT2FN[config.hidden_act] -// def forward(self, hidden_state): -// return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) pub struct GlmOcrVisionMlp(GateUpDownMLP); impl GlmOcrVisionMlp { @@ -141,12 +57,6 @@ impl GlmOcrVisionMlp { } } -// ============================================================================ -// 3. eager_attention_forward (+ repeat_kv from utils) -// ============================================================================ - -// Python eager_attention_forward implementation -// Reference: py-glm-ocr/glm_ocr/modeling_glm_ocr.py fn eager_attention_forward( query_states: &Tensor, key_states: &Tensor, @@ -156,16 +66,6 @@ fn eager_attention_forward( scaling: f64, dropout: f64, ) -> Result<(Tensor, Tensor)> { - // Attention matrix size info — enable with GLM_DEBUG=1. - if std::env::var("GLM_DEBUG").is_ok() { - let dims = query_states.dims(); - let (heads, q_len, k_len) = (dims[1], dims[2], key_states.dims()[2]); - let attn_mb = (heads * q_len * k_len * 4) as f64 / 1_048_576.0; - eprintln!( - "[GLM-OCR OOM-DBG] eager_attention_forward: heads={heads} q_len={q_len} k_len={k_len} \ - attn_weights={attn_mb:.1}MB (f32)" - ); - } let key_states = match num_key_value_groups { Some(g) => repeat_kv(key_states.clone(), g)?.contiguous()?, None => key_states.clone(), @@ -182,9 +82,10 @@ fn eager_attention_forward( #[cfg(feature = "flash-attn")] { // Flash attention: causal iff attention_mask is present. - let q = query_states.transpose(1, 2)?; - let k = key_states.transpose(1, 2)?; - let v = value_states.transpose(1, 2)?; + // Explicit contiguous() ensures proper memory layout for flash_attn kernel + let q = query_states.transpose(1, 2)?.contiguous()?; + 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 } @@ -194,6 +95,7 @@ fn eager_attention_forward( // [batch, heads, CHUNK, k_len] stays bounded regardless of q_len. // Peak memory per chunk: CHUNK × k_len × heads × 4 bytes (f32 softmax). // Mathematically equivalent to full attention. + // CHUNK_SIZE=512 is empirically optimal for most hardware (CPU/GPU balance) const CHUNK_SIZE: usize = 512; let q_len = query_states.dim(2)?; let k_t = key_states.transpose(D::Minus2, D::Minus1)?.contiguous()?; @@ -211,8 +113,15 @@ fn eager_attention_forward( attn.broadcast_add(&mask.narrow(2, start, len)?.to_dtype(attn.dtype())?)? } }; - let attn = candle_nn::ops::softmax_last_dim(&attn.to_dtype(DType::F32)?)? - .to_dtype(query_states.dtype())?; + // Softmax computation: Optimize dtype conversions for CPU (which uses F32) + let attn = if query_states.dtype() == DType::F32 { + candle_nn::ops::softmax_last_dim(&attn)? + } else { + candle_nn::ops::softmax_last_dim(&attn.to_dtype(DType::F32)?)? + .to_dtype(query_states.dtype())? + }; + // Apply dropout uniformly across chunked and non-chunked paths for consistency + let attn = candle_nn::ops::dropout(&attn, dropout as f32)?; chunks.push(attn.matmul(&value_states)?); start += len; } @@ -223,9 +132,16 @@ fn eager_attention_forward( None => attn, Some(mask) => attn.broadcast_add(&mask.to_dtype(attn.dtype())?)?, }; - let attn = candle_nn::ops::softmax_last_dim(&attn.to_dtype(DType::F32)?)? - .to_dtype(query_states.dtype())?; - candle_nn::ops::dropout(&attn, dropout as f32)?.matmul(&value_states)? + // Softmax computation: Same optimization as chunked path + let attn = if query_states.dtype() == DType::F32 { + candle_nn::ops::softmax_last_dim(&attn)? + } else { + candle_nn::ops::softmax_last_dim(&attn.to_dtype(DType::F32)?)? + .to_dtype(query_states.dtype())? + }; + // Apply dropout uniformly (now consistent across both paths) + let attn = candle_nn::ops::dropout(&attn, dropout as f32)?; + attn.matmul(&value_states)? }; // [batch, heads, q_len, head_dim] -> [batch, q_len, heads, head_dim] raw.transpose(1, 2)?.contiguous()? @@ -237,19 +153,6 @@ fn eager_attention_forward( Ok((output, placeholder)) } -// ============================================================================ -// 5. GlmOcrTextAttention -// ============================================================================ - -// Python: class GlmOcrTextAttention(nn.Module): -// def __init__(self, config): -// self.q_proj = nn.Linear(...) -// self.k_proj = nn.Linear(...) -// self.v_proj = nn.Linear(...) -// self.o_proj = nn.Linear(...) -// def forward(self, hidden_states, attention_mask, position_ids, past_key_value, use_cache): -// # QKV projection, attention computation, output projection -// return attn_output, attn_weights pub struct GlmOcrTextAttention { q_proj: Linear, k_proj: Linear, @@ -338,9 +241,6 @@ impl GlmOcrTextAttention { let (query_states, key_states) = glm_ocr_apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?; - // Python: if past_key_values is not None: - // cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - // key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) let (key_states, value_states) = match &self.kv_cache { None => (key_states, value_states), Some((prev_k, prev_v)) => { @@ -372,16 +272,6 @@ impl GlmOcrTextAttention { } } -// ============================================================================ -// 7. GlmOcrVisionRotaryEmbedding -// ============================================================================ - -// Python: class GlmOcrVisionRotaryEmbedding(nn.Module): -// def __init__(self, dim, theta=10000.0): -// inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) -// def forward(self, seqlen): -// seq = torch.arange(seqlen, device=self.inv_freq.device) -// return torch.outer(seq, self.inv_freq) # (seqlen, dim/2) pub struct GlmOcrVisionRotaryEmbedding { inv_freq: Tensor, } @@ -405,19 +295,6 @@ impl GlmOcrVisionRotaryEmbedding { Ok(freqs) } - /// Python: GlmOcrVisionModel.rot_pos_emb(self, grid_thw) - /// pos_ids = [] - /// for t, h, w in grid_thw: - /// hpos_ids = arange(h).unsqueeze(1).expand(-1, w) - /// hpos_ids = hpos_ids.reshape(h//sms, sms, w//sms, sms).permute(0,2,1,3).flatten() - /// wpos_ids = arange(w).unsqueeze(0).expand(h, -1) - /// wpos_ids = wpos_ids.reshape(h//sms, sms, w//sms, sms).permute(0,2,1,3).flatten() - /// pos_ids.append(stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) - /// pos_ids = cat(pos_ids, dim=0) # (total, 2) - /// max_grid_size = grid_thw[:, 1:].max() - /// rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) # (max_grid_size, dim/4) - /// rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) # (total, dim/2) - /// return rotary_pos_emb, pos_ids pub fn rot_pos_emb( &self, grid_thw: &[(usize, usize, usize)], @@ -431,14 +308,10 @@ impl GlmOcrVisionRotaryEmbedding { for &(t, h, w) in grid_thw { max_grid_size = max_grid_size.max(h).max(w); - // Generate position indices for each patch BEFORE spatial merge - // The vision encoder processes all patches, then merges them - // Python: for each (t, h, w), generate h*w position indices for _ in 0..t { for hi in 0..h { for wi in 0..w { // Apply spatial merge rearrangement - // Python: hpos_ids = hpos_ids.reshape(h//sms, sms, w//sms, sms).permute(0,2,1,3).flatten() let _hb = hi / sms; let _si = hi % sms; let _wb = wi / sms; @@ -457,10 +330,6 @@ impl GlmOcrVisionRotaryEmbedding { let total_seq = all_hpos.len(); let freqs_full = self.forward(max_grid_size)?; // (max_grid_size, dim/4) - // Python: rotary_pos_emb_full[pos_ids].flatten(1) - // pos_ids is (total, 2) with [h_idx, w_idx] entries - // rotary_pos_emb_full[pos_ids] -> (total, 2, dim/4) -> flatten(1) -> (total, dim/2) - // This is equivalent to cat(freqs[h_indices], freqs[w_indices], dim=-1) let h_indices = Tensor::from_vec(all_hpos, (total_seq,), self.inv_freq.device())?; let w_indices = Tensor::from_vec(all_wpos, (total_seq,), self.inv_freq.device())?; let h_freqs = freqs_full.index_select(&h_indices, 0)?; // (total_seq, dim/4) @@ -469,9 +338,6 @@ impl GlmOcrVisionRotaryEmbedding { // Concatenate h and w freqs: (total_seq, dim/2) let rotary_pos_emb = Tensor::cat(&[&h_freqs, &w_freqs], 1)?; - // Python (in GlmOcrVisionModel.forward): - // emb = cat((rotary_pos_emb, rotary_pos_emb), dim=-1) - // position_embeddings = (emb.cos(), emb.sin()) let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], 1)?; let cos = emb.cos()?; let sin = emb.sin()?; @@ -480,16 +346,6 @@ impl GlmOcrVisionRotaryEmbedding { } } -// ============================================================================ -// 8. GlmOcrTextMLP -// ============================================================================ - -// Python: class GlmOcrTextMLP(nn.Module): -// def forward(self, hidden_states): -// up_states = self.gate_up_proj(hidden_states) -// gate, up_states = up_states.chunk(2, dim=-1) -// up_states = up_states * self.activation_fn(gate) -// return self.down_proj(up_states) pub struct GlmOcrTextMLP { gate_up_proj: Linear, down_proj: Linear, @@ -526,21 +382,6 @@ impl GlmOcrTextMLP { } } -// ============================================================================ -// 9. GlmOcrTextDecoderLayer -// ============================================================================ - -// Python: class GlmOcrTextDecoderLayer(nn.Module): -// def forward(self, hidden_states, position_embeddings): -// hidden_states = self.input_layernorm(hidden_states) -// hidden_states, _ = self.self_attn(hidden_states, position_embeddings) -// hidden_states = self.post_self_attn_layernorm(hidden_states) -// hidden_states = residual + hidden_states -// residual = hidden_states -// hidden_states = self.post_attention_layernorm(hidden_states) -// hidden_states = self.mlp(hidden_states) -// hidden_states = self.post_mlp_layernorm(hidden_states) -// return residual + hidden_states pub struct GlmOcrTextDecoderLayer { self_attn: GlmOcrTextAttention, mlp: GlmOcrTextMLP, @@ -611,28 +452,6 @@ impl GlmOcrTextDecoderLayer { } } -// ============================================================================ -// 10. rotate_half + apply_rotary_pos_emb_vision (in position_embed/rope.rs) -// ============================================================================ - -// ============================================================================ -// 11. GlmOcrVisionAttention -// ============================================================================ - -// Python reference (transformers): -// class GlmOcrVisionAttention(nn.Module): -// def __init__(self, config): -// self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.attention_bias) -// self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias) -// self.q_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps) -// self.k_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps) -// def forward(self, hidden_states, position_embeddings): -// q, k, v = self.qkv(hidden_states).reshape(seq_len, 3, num_heads, -1).permute(1,0,2,3).unbind(0) -// query_states = self.q_norm(query_states) -// key_states = self.k_norm(key_states) -// q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) -// attn_output = attention(q, k, v) -// return self.proj(attn_output.reshape(seq_len, -1)) pub struct GlmOcrVisionAttention { num_heads: usize, head_dim: usize, @@ -744,20 +563,6 @@ impl GlmOcrVisionAttention { } } -// ============================================================================ -// 12. GlmOcrVisionBlock -// ============================================================================ - -// Python: class GlmOcrVisionBlock(GradientCheckpointingLayer): -// def __init__(self, config): -// self.norm1 = GlmOcrRMSNorm(...) -// self.attn = GlmOcrVisionAttention(...) -// self.norm2 = GlmOcrRMSNorm(...) -// self.mlp = GlmOcrVisionMlp(...) -// def forward(self, x): -// x = x + self.attn(self.norm1(x)) -// x = x + self.mlp(self.norm2(x)) -// return x pub struct GlmOcrVisionBlock { norm1: GlmOcrRMSNorm, norm2: GlmOcrRMSNorm, @@ -801,28 +606,6 @@ impl GlmOcrVisionBlock { } } -// ============================================================================ -// 13. GlmOcrVisionPatchMerger -// ============================================================================ - -// Python reference (transformers commit 4854dbf9): -// class GlmOcrVisionPatchMerger(nn.Module): -// def __init__(self, dim, context_dim, hidden_act, bias=False): -// # dim = out_hidden_size (1536) -// # context_dim = intermediate_size (4096) NOT out_hidden_size * in_channels! -// self.proj = nn.Linear(dim, dim, bias=bias) -// self.post_projection_norm = LayerNorm(dim) -// self.gate_proj = nn.Linear(dim, context_dim, bias=bias) -// self.up_proj = nn.Linear(dim, context_dim, bias=bias) -// self.down_proj = nn.Linear(context_dim, dim, bias=bias) -// def forward(self, x): -// x = self.proj(x) -// x = self.post_projection_norm(x) -// x = GELU(x) # fixed activation, not config.hidden_act -// x = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) -// return x -// -// FIX: context_dim should be intermediate_size (4096), not out_hidden_size * in_channels (1536*3=4608) pub struct GlmOcrVisionPatchMerger { proj: Linear, post_projection_norm: LayerNorm, @@ -834,10 +617,6 @@ pub struct GlmOcrVisionPatchMerger { impl GlmOcrVisionPatchMerger { pub fn new(vb: VarBuilder, config: &GlmOcrVisionConfig) -> Result { - // NOTE: The checkpoint stores proj.weight as [out_features, in_features] = [1536, 1536] - // PyTorch Linear computes: output = input @ weight.T - // Candle Linear computes: output = weight.matmul(input) which is equivalent to input @ weight.T - // So the weight should be loaded correctly. let proj = linear_no_bias( config.out_hidden_size, config.out_hidden_size, @@ -850,8 +629,6 @@ impl GlmOcrVisionPatchMerger { vb.pp("post_projection_norm"), )?; - // Patch merger MLP: uses out_hidden_size * in_channels as intermediate dim - // Checkpoint shape: [4608, 1536] = [out_hidden_size * in_channels, out_hidden_size] let context_dim = config.out_hidden_size * config.in_channels; let gate_proj = linear_no_bias( config.out_hidden_size, @@ -882,7 +659,6 @@ impl GlmOcrVisionPatchMerger { pub fn forward(&self, hidden_state: &Tensor) -> Result { let mut hidden_state = self.proj.forward(hidden_state)?; hidden_state = self.post_projection_norm.forward(&hidden_state)?; - // Python: x = self.act1(x) where act1 = nn.GELU() - fixed GELU, not config activation hidden_state = hidden_state.gelu()?; let gate = self.gate_proj.forward(&hidden_state)?; @@ -894,17 +670,6 @@ impl GlmOcrVisionPatchMerger { } } -// ============================================================================ -// 14. GlmOcrVisionPatchEmbed -// ============================================================================ - -// Python: class GlmOcrVisionPatchEmbed(nn.Module): -// def __init__(self, config): -// self.proj = nn.Conv3d(...) -// def forward(self, x): -// x = x.view(-1, C, T, P, P) -// x = self.proj(x).view(-1, embed_dim) -// return x pub struct GlmOcrVisionPatchEmbed { patch_size: usize, temporal_patch_size: usize, @@ -945,27 +710,18 @@ impl GlmOcrVisionPatchEmbed { } pub fn forward(&self, pixel_values: &Tensor) -> Result { - // pixel_values can be either: - // 1. Flattened patches: [num_patches, patch_dim] (Python format) - // 2. Standard image: [batch, C, H, W] (traditional format) - let rank = pixel_values.rank(); if rank == 2 { - // Flattened patches format: [num_patches, patch_dim] - // Directly project to hidden_size let hidden_states = self.proj.forward(pixel_values)?; Ok(hidden_states) } else { - // Standard image format: [batch, C, H, W] let (batch, _c, h, w) = pixel_values.dims4()?; - // Support non-square images let patches_h = h / self.patch_size; let patches_w = w / self.patch_size; let num_patches = patches_h * patches_w; - // Reshape: (batch, C, H, W) -> (batch, patches_h, patch_size, patches_w, patch_size, C) let pv = pixel_values.reshape(( batch, patches_h, @@ -975,16 +731,13 @@ impl GlmOcrVisionPatchEmbed { self.in_channels, ))?; - // Permute: (batch, patches_h, patches_w, C, patch_size, patch_size) let pv = pv.permute((0, 1, 3, 5, 2, 4))?; - // Reshape: (batch * num_patches, C * patch_size * patch_size) let pv = pv.reshape(( batch * num_patches, self.in_channels * self.patch_size * self.patch_size, ))?; - // Add temporal dimension let pv = pv.unsqueeze(1)?; let ones_shape: Vec = vec![1, self.temporal_patch_size]; let pv = pv.broadcast_mul(&Tensor::ones(ones_shape, pv.dtype(), pv.device())?)?; @@ -999,22 +752,6 @@ impl GlmOcrVisionPatchEmbed { } } -// ============================================================================ -// 15. GlmOcrVisionModel -// ============================================================================ - -// Python: class GlmOcrVisionModel(GlmOcrPreTrainedModel): -// def __init__(self, config): -// self.patch_embed = GlmOcrVisionPatchEmbed(config) -// self.rotary_pos_emb = GlmOcrVisionRotaryEmbedding(...) -// self.blocks = nn.ModuleList([GlmOcrVisionBlock(config) for _ in range(config.depth)]) -// self.merger = GlmOcrVisionPatchMerger(...) -// self.downsample = nn.Conv2d(...) -// self.post_layernorm = GlmOcrRMSNorm(...) -// def forward(self, hidden_states, grid_thw): -// hidden_states = self.patch_embed(hidden_states) -// # ... apply rotary pos emb, blocks, merger, downsample -// return hidden_states pub struct GlmOcrVisionModel { patch_embed: GlmOcrVisionPatchEmbed, rotary_pos_emb: GlmOcrVisionRotaryEmbedding, @@ -1093,20 +830,13 @@ impl GlmOcrVisionModel { result }; - tensor_stats("after_patch_embed", &hidden_states); - - // Compute rotary embeddings matching Python rot_pos_emb exactly let (cos, sin) = self .rotary_pos_emb .rot_pos_emb(&grid_thw_parsed, self.config.spatial_merge_size)?; - tensor_stats("vision_cos", &cos); - tensor_stats("vision_sin", &sin); - let rotary_pos_emb = Tensor::cat(&[&cos, &sin], D::Minus1)?; let position_embeddings = (&cos, &sin); - // Compute cu_seqlens let mut cu_seqlens_values: Vec = vec![0]; let mut cumsum: i32 = 0; for (t, h, w) in &grid_thw_parsed { @@ -1122,67 +852,34 @@ impl GlmOcrVisionModel { hidden_states.device(), )?; - let num_blocks = self.blocks.len(); - if std::env::var("GLM_DEBUG").is_ok() { - let n_patches = hidden_states.dim(0).unwrap_or(0); - let hidden_mb = (hidden_states.elem_count() * 2) as f64 / 1_048_576.0; // bf16 - let attn_mb = (self.config.num_heads * n_patches * n_patches * 4) as f64 / 1_048_576.0; - eprintln!( - "[GLM-OCR OOM-DBG] Vision encoder: n_patches={n_patches} \ - hidden={hidden_mb:.1}MB per-layer attn_weights={attn_mb:.1}MB (f32) \ - num_blocks={num_blocks}" - ); - } - for (i, block) in self.blocks.iter().enumerate() { + for block in self.blocks.iter() { hidden_states = block.forward( &hidden_states, &cu_seqlens, Some(&rotary_pos_emb), Some(position_embeddings), )?; - if i == 0 { - tensor_stats("after_vision_block[0]", &hidden_states); - } - if i == num_blocks - 1 { - tensor_stats(&format!("after_vision_block[{i}] (last)"), &hidden_states); - } } let hidden_states = self.post_layernorm.forward(&hidden_states)?; - tensor_stats("after_vision_post_layernorm", &hidden_states); let sms = self.config.spatial_merge_size; let hidden_dim = hidden_states.dim(hidden_states.dims().len() - 1)?; - // Python: hidden_states.view(-1, sms, sms, hidden_dim).permute(0, 3, 1, 2) - // Input: [2816, 1024] where 2816 = grid_h * grid_w = 44 * 64 - // After reshape: [704, 2, 2, 1024] where 704 = 2816 / 4 - // After permute: [704, 1024, 2, 2] - // After downsample: [704, 1536, 1, 1] - // After reshape: [704, 1536] 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] - - tensor_stats("after_downsample", &hidden_states); let merged = self.merger.forward(&hidden_states)?; - tensor_stats("after_merger (vision features)", &merged); let merged = merged.unsqueeze(0)?; Ok(merged) } } -// ============================================================================ -// GlmOcrProjector (Rust-specific, no Python equivalent) -// ============================================================================ - -// Rust-specific: Projects vision features to language model space -// (Not present in Python - integrated in GlmOcrVisionModel forward) pub struct GlmOcrProjector { #[allow(dead_code)] query_embed: Option, proj: Linear, @@ -1224,21 +921,6 @@ impl GlmOcrProjector { } } -// ============================================================================ -// 16. GlmOcrTextRotaryEmbedding -// ============================================================================ - -// Python: class GlmOcrTextRotaryEmbedding(nn.Module): -// def forward(self, x, position_ids): # position_ids: (3, bs, seq_len) -// inv_freq_expanded = self.inv_freq[None, None, :, None].expand(3, bs, -1, 1) -// position_ids_expanded = position_ids[:, :, None, :].float() -// freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3) -// freqs = self.apply_mrope(freqs, self.mrope_section) -// emb = torch.cat((freqs, freqs), dim=-1) -// return emb.cos() * self.attention_scaling, emb.sin() * self.attention_scaling -// def apply_mrope(self, freqs, mrope_section): -// chunks = freqs.split(mrope_section, dim=-1) -// return torch.cat([chunk[i % 3] for i, chunk in enumerate(chunks)], dim=-1) pub struct GlmOcrTextRotaryEmbedding { inv_freq: Tensor, mrope_section: Vec, @@ -1300,12 +982,10 @@ impl GlmOcrTextRotaryEmbedding { } Ok(Tensor::cat(&result_parts, D::Minus1)?) } - - /// Compute cos/sin from explicit 3D position IDs (used for prefill with image tokens). - /// position_ids: (3, bs, seq_len) — axis 0 = temporal, 1 = height, 2 = width. - 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)?; // head_dim/2 + 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) @@ -1378,19 +1058,6 @@ impl GlmOcrTextRotaryEmbedding { } } -// ============================================================================ -// 17. GlmOcrTextModel -// ============================================================================ - -// Python: class GlmOcrTextModel(GlmOcrPreTrainedModel): -// def __init__(self, config): -// self.embed_tokens = nn.Embedding(...) -// self.layers = nn.ModuleList([GlmOcrTextDecoderLayer(config) for _ in range(...)]) -// self.norm = GlmOcrRMSNorm(...) -// self.rotary_emb = GlmOcrTextRotaryEmbedding(...) -// def forward(self, input_ids, ...): -// hidden_states = self.embed_tokens(input_ids) -// # ... apply layers, rotary emb, return hidden_states pub struct GlmOcrTextModel { embed_tokens: Embedding, layers: Vec, @@ -1417,7 +1084,6 @@ impl GlmOcrTextModel { let norm = GlmOcrRMSNorm::new(vb.pp("norm"), config.hidden_size, config.rms_norm_eps)?; - // lm_head.weight lives at the checkpoint root (not under model.language_model) let root_vb = vb.root(); let lm_head = linear_no_bias(config.hidden_size, config.vocab_size, root_vb.pp("lm_head"))?; @@ -1436,13 +1102,6 @@ impl GlmOcrTextModel { }) } - /// Compute 3D M-RoPE position IDs matching Python's GlmOcrModel.get_rope_index(). - /// - /// For image tokens: each gets (t, h, w) grid coordinates. - /// For text tokens: sequential positions on all 3 axes. - /// Positions continue from where the previous group left off. - /// - /// Returns tensor of shape (3, 1, seq_len) containing [temporal, height, width] position IDs. fn compute_mrope_position_ids( &mut self, image_mask: &Tensor, @@ -1516,13 +1175,6 @@ impl GlmOcrTextModel { self.next_mrope_pos = st_idx as usize; self.prefill_seq_len = seq_len; - if std::env::var("GLM_DEBUG").is_ok() { - eprintln!( - "[M-RoPE] prefill: seq_len={}, next_mrope_pos={}", - seq_len, self.next_mrope_pos - ); - } - let t_t = Tensor::from_vec(t_ids, (1, seq_len), device)?; let h_t = Tensor::from_vec(h_ids, (1, seq_len), device)?; let w_t = Tensor::from_vec(w_ids, (1, seq_len), device)?; @@ -1539,17 +1191,6 @@ impl GlmOcrTextModel { ) -> Result { let (bs, seq_len) = input_ids.dims2()?; let mut inputs_embeds = self.embed_tokens.forward(input_ids)?; - tensor_stats("embed_tokens output", &inputs_embeds); - - #[cfg(debug_assertions)] - eprintln!( - "LanguageModel forward: bs={}, seq_len={}, head_dim={}", - bs, - seq_len, - self.config - .head_dim - .unwrap_or_else(|| self.config.hidden_size / self.config.num_attention_heads) - ); // Merge image features into embeddings at image token positions if let (Some(img_feats), Some(img_mask)) = (image_features, image_mask) { @@ -1569,15 +1210,6 @@ impl GlmOcrTextModel { let num_features = img_feats.dim(1)?; let num_to_replace = image_indices.len().min(num_features); - if std::env::var("GLM_DEBUG").is_ok() { - eprintln!("[LM] Image indices count: {}, num_features: {}, num_to_replace: {}", - image_indices.len(), num_features, num_to_replace); - if !image_indices.is_empty() { - eprintln!("[LM] First image index: {}, Last image index: {}", - image_indices.first().unwrap(), image_indices.last().unwrap()); - } - } - // Replace embeddings at image positions with image features // Build the merged embeddings by copying let embeds_flat = inputs_embeds.squeeze(0)?; // (seq_len, hidden_size) @@ -1599,16 +1231,7 @@ impl GlmOcrTextModel { let refs: Vec<&Tensor> = embeds_vec.iter().collect(); inputs_embeds = Tensor::cat(&refs, 0)?.unsqueeze(0)?; - if std::env::var("GLM_DEBUG").is_ok() { - eprintln!("[LM] inputs_embeds after image injection shape: {:?}", inputs_embeds.shape()); - if seq_len > 710 { - let text_embed = inputs_embeds.i((0, 710, ..))?; - let text_mean = text_embed.to_dtype(DType::F32)?.mean_all()?.to_scalar::()?; - eprintln!("[LM] Text token 710 embedding mean: {:.6}", text_mean); - } - } } - tensor_stats("inputs_embeds after image injection", &inputs_embeds); let attention_mask = if seq_len > 1 { Some(prepare_causal_attention_mask( @@ -1638,37 +1261,15 @@ impl GlmOcrTextModel { let decode_pos = self.next_mrope_pos + (seqlen_offset - self.prefill_seq_len); self.rotary_emb.forward(1, decode_pos, input_ids.device())? }; - tensor_stats("text_rotary cos", &cos); - tensor_stats("text_rotary sin", &sin); - let num_layers = self.layers.len(); let mut hidden_states = inputs_embeds; - for (i, layer) in self.layers.iter_mut().enumerate() { + for layer in self.layers.iter_mut() { hidden_states = layer.forward(&hidden_states, (&cos, &sin), attention_mask.as_ref())?; - if i == 0 { - tensor_stats("after_text_layer[0]", &hidden_states); - } - if i == num_layers - 1 { - tensor_stats(&format!("after_text_layer[{i}] (last)"), &hidden_states); - } } hidden_states = self.norm.forward(&hidden_states)?; - tensor_stats("after_final_norm", &hidden_states); let logits = self.lm_head.forward(&hidden_states)?; - // Log top-5 logits at last position (only on first pass) - if seqlen_offset == 0 && std::env::var("GLM_INTERMEDIATE").is_ok() { - let last_logits = logits.i((0, seq_len - 1, ..))?.to_dtype(DType::F32)?; - tensor_stats("logits at last position", &last_logits); - if let Ok(vals) = last_logits.to_vec1::() { - let mut indexed: Vec<(usize, f32)> = vals.iter().copied().enumerate().collect(); - indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - eprintln!("[RS] Top-5 logit token_ids: {:?}", - indexed[..5.min(indexed.len())].iter().map(|(i, v)| format!("id={i} val={v:.4}")).collect::>()); - } - } - Ok(logits) } @@ -1679,17 +1280,6 @@ impl GlmOcrTextModel { } } -// ============================================================================ -// 18. GlmOcrModel (corresponds to GlmOcrForConditionalGeneration in Python) -// ============================================================================ - -// Python: class GlmOcrForConditionalGeneration(GlmOcrPreTrainedModel, GenerationMixin): -// def __init__(self, config): -// self.model = GlmOcrModel(config) # Contains visual + language_model -// self.lm_head = nn.Linear(...) -// def forward(self, input_ids, pixel_values, ...): -// # Vision encoder -> language model -> lm_head -// return logits pub struct GlmOcrModel { vision_encoder: GlmOcrVisionModel, language_model: GlmOcrTextModel, @@ -1719,24 +1309,6 @@ impl GlmOcrModel { image_mask: Option<&Tensor>, seqlen_offset: usize, ) -> Result { - #[cfg(debug_assertions)] - { - eprintln!("[GLM-OCR Model] ===== FORWARD START ====="); - eprintln!("[GLM-OCR Model] input_ids shape: {:?}", input_ids.shape()); - eprintln!("[GLM-OCR Model] seqlen_offset: {}", seqlen_offset); - if let Some(pv) = pixel_values { - eprintln!("[GLM-OCR Model] pixel_values shape: {:?}", pv.shape()); - } - if let Some(g) = image_grid_thw { - eprintln!("[GLM-OCR Model] grid_thw shape: {:?}", g.shape()); - } - if let Some(m) = image_mask { - eprintln!("[GLM-OCR Model] image_mask shape: {:?}", m.shape()); - let mask_sum = m.sum_all()?.to_scalar::()?; - eprintln!("[GLM-OCR Model] image_mask sum (num image tokens): {}", mask_sum); - } - } - let image_features = if let Some(pixels) = pixel_values { let grid_thw = if let Some(grid) = image_grid_thw { grid.clone() @@ -1751,26 +1323,8 @@ impl GlmOcrModel { )? }; - if std::env::var("GLM_DEBUG").is_ok() { - let pv_mb = (pixels.elem_count() * 2) as f64 / 1_048_576.0; // bf16 - eprintln!( - "[GLM-OCR OOM-DBG] pixel_values shape={:?} size={pv_mb:.1}MB", - pixels.shape() - ); - eprintln!("[GLM-OCR OOM-DBG] grid_thw={:?}", grid_thw.to_vec1::()); - } - let vision_output = self.vision_encoder.forward(pixels, &grid_thw)?; - if std::env::var("GLM_DEBUG").is_ok() { - eprintln!( - "[GLM-OCR Model] vision_output shape: {:?}", - vision_output.shape() - ); - let vis_mean = vision_output.to_dtype(candle_core::DType::F32)?.mean_all()?.to_scalar::()?; - eprintln!("[GLM-OCR Model] vision_output mean: {:.6}", vis_mean); - } - Some(vision_output) } else { None @@ -1784,14 +1338,6 @@ impl GlmOcrModel { seqlen_offset, ); - #[cfg(debug_assertions)] - { - eprintln!("[GLM-OCR Model] ===== FORWARD END ====="); - if let Ok(ref r) = result { - eprintln!("[GLM-OCR Model] output shape: {:?}", r.shape()); - } - } - result } diff --git a/src/models/glm_ocr/processor.rs b/src/models/glm_ocr/processor.rs index c9ed381..ceabe56 100644 --- a/src/models/glm_ocr/processor.rs +++ b/src/models/glm_ocr/processor.rs @@ -149,20 +149,7 @@ impl GlmOcrProcessor { // Use smart_resize to compute target dimensions let (target_h, target_w) = self.smart_resize(orig_h, orig_w); - if std::env::var("GLM_DEBUG").is_ok() { - let grid_h = target_h / self.patch_size; - let grid_w = target_w / self.patch_size; - let n_patches = grid_h * grid_w; - let pv_elems = n_patches * 3 * self.temporal_patch_size * self.patch_size * self.patch_size; - let pv_mb = (pv_elems * 2) as f64 / 1_048_576.0; // bf16 - let attn_mb = (16 * n_patches * n_patches * 4) as f64 / 1_048_576.0; // 16 heads, f32 - eprintln!( - "[GLM-OCR OOM-DBG] Image: orig={orig_w}x{orig_h} -> target={target_w}x{target_h} \ - grid={grid_w}x{grid_h} n_patches={n_patches} \ - pixel_values={pv_mb:.1}MB vision_attn_per_layer={attn_mb:.1}MB" - ); - } - + // Resize image let img = img.resize_exact( target_w as u32, @@ -183,10 +170,8 @@ impl GlmOcrProcessor { }) .collect(); - // Reshape to [H, W, 3] let tensor = Tensor::from_vec(pixels, (target_h, target_w, 3), &self.device)?; - // Normalize 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)?; @@ -196,37 +181,25 @@ impl GlmOcrProcessor { let grid_w = target_w / self.patch_size; let patch_size = self.patch_size; let channels = 3; - let temporal_patch_size = 2; // Python uses temporal_patch_size=2 even for images + let temporal_patch_size = 2; - // Reshape: [H, W, 3] -> [grid_h, patch_size, grid_w, patch_size, 3] let tensor = tensor.reshape(( grid_h, patch_size, grid_w, patch_size, channels, ))?; - // Permute to: [grid_h, grid_w, patch_size, patch_size, channels] - let tensor = tensor.permute((0, 2, 1, 3, 4))?; + let tensor = tensor.permute((0, 2, 4, 1, 3))?; - // Permute to put channels first: [grid_h, grid_w, channels, patch_size, patch_size] - // Python's patch_embed.forward does: view(-1, C, T, P, P) so we need (C, T, P_h, P_w) order - let tensor = tensor.permute((0, 1, 4, 2, 3))?; - - // Reshape to: [num_patches, channels, patch_size, patch_size] let num_patches = grid_h * grid_w; let tensor = tensor.reshape((num_patches, channels, patch_size, patch_size))?; - // Add temporal dimension after C: [num_patches, channels, 1, patch_size, patch_size] let tensor = tensor.unsqueeze(2)?; - // Repeat T times along temporal dim 2: [num_patches, channels, temporal_patch_size, patch_size, patch_size] let tensor = tensor.repeat((1, 1, temporal_patch_size, 1, 1))?; - // Flatten to: [num_patches, channels * temporal_patch_size * patch_size * patch_size] - // This gives (C, T, P_h, P_w) order matching Python's patch_embed input let patch_dim = channels * temporal_patch_size * patch_size * patch_size; let tensor = tensor.reshape((num_patches, patch_dim))?; - // Convert to model dtype let tensor = tensor.to_dtype(self.dtype)?; Ok(ProcessedImage { @@ -236,25 +209,7 @@ impl GlmOcrProcessor { }) } - /// Process image and text for multimodal input. - /// - /// # Arguments - /// * `image_path` - Path to input image - /// * `prompt` - Text prompt/question about the image - /// * `tokenizer` - Tokenizer for text encoding - /// * `image_token_id` - Token ID for image content placeholders - /// * `image_start_token_id` - Token ID marking start of image region - /// * `image_end_token_id` - Token ID marking end of image region - /// * `patch_size` - Vision encoder patch size (default: 14) - /// * `temporal_patch_size` - Temporal patch size (default: 2, unused for images) - /// * `spatial_merge_size` - Spatial merge factor (default: 2) - /// - /// # Returns - /// ProcessedInput containing: - /// - input_ids: Combined image placeholder + text token IDs - /// - pixel_values: Flattened patches tensor [num_patches, patch_dim] - /// - image_mask: Boolean mask for image token positions - /// - grid_thw: (temporal, height, width) grid dimensions for RoPE + /// Process image and text for multimodal input pub fn process_info( &self, image_path: &str, @@ -327,7 +282,6 @@ impl GlmOcrProcessor { let image_mask = Tensor::from_vec(image_mask_vec, (1, input_ids_vec.len()), &self.device)?; // Compute grid_thw for RoPE - // For images: grid_t = 1 let grid_thw = Tensor::from_vec( vec![1u32, grid_h as u32, grid_w as u32], (3,),