diff --git a/src/api/mod.rs b/src/api/mod.rs index c5b30ce..73f0451 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -261,6 +261,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", } } @@ -283,7 +284,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 4fa6fe4..f408503 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..49d79c5 --- /dev/null +++ b/src/models/glm_ocr/generate.rs @@ -0,0 +1,393 @@ +//! 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 + }; + + 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()); + + 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 mut generate = Vec::new(); + let sample_len = mes.max_tokens.unwrap_or(512); + + 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..d4fe3d7 --- /dev/null +++ b/src/models/glm_ocr/model.rs @@ -0,0 +1,1347 @@ +//! GLM-OCR Model Implementation + +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}, + }, +}; + +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() + } +} + +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)?) + } +} + +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)> { + 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. + // 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 + } + #[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. + // 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()?; + + 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())?)? + } + }; + // 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; + } + 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())?)?, + }; + // 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()? + } + }; + + // output layout: [batch, q_len, heads, head_dim] + let placeholder = Tensor::zeros((0,), query_states.dtype(), query_states.device())?; + Ok((output, placeholder)) +} + +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)?; + + 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; + } +} + +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) + } + + 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); + + for _ in 0..t { + for hi in 0..h { + for wi in 0..w { + // Apply spatial merge rearrangement + 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) + + 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)?; + + let emb = Tensor::cat(&[&rotary_pos_emb, &rotary_pos_emb], 1)?; + let cos = emb.cos()?; + let sin = emb.sin()?; + + Ok((cos, sin)) + } +} + +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)?) + } +} + +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(); + } +} + +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)?) + } +} + +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)?) + } +} + +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 { + 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"), + )?; + + 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)?; + 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)?) + } +} + +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 { + let rank = pixel_values.rank(); + + if rank == 2 { + let hidden_states = self.proj.forward(pixel_values)?; + Ok(hidden_states) + } else { + let (batch, _c, h, w) = pixel_values.dims4()?; + + let patches_h = h / self.patch_size; + let patches_w = w / self.patch_size; + let num_patches = patches_h * patches_w; + + let pv = pixel_values.reshape(( + batch, + patches_h, + self.patch_size, + patches_w, + self.patch_size, + self.in_channels, + ))?; + + let pv = pv.permute((0, 1, 3, 5, 2, 4))?; + + let pv = pv.reshape(( + batch * num_patches, + self.in_channels * self.patch_size * self.patch_size, + ))?; + + 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) + } + } +} + +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 + }; + + let (cos, sin) = self + .rotary_pos_emb + .rot_pos_emb(&grid_thw_parsed, self.config.spatial_merge_size)?; + + let rotary_pos_emb = Tensor::cat(&[&cos, &sin], D::Minus1)?; + let position_embeddings = (&cos, &sin); + + 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(), + )?; + + for block in self.blocks.iter() { + hidden_states = block.forward( + &hidden_states, + &cu_seqlens, + Some(&rotary_pos_emb), + Some(position_embeddings), + )?; + } + + let hidden_states = self.post_layernorm.forward(&hidden_states)?; + + let sms = self.config.spatial_merge_size; + let hidden_dim = hidden_states.dim(hidden_states.dims().len() - 1)?; + + let total_patches = hidden_states.dim(0)?; // 2816 + let merged_patches = total_patches / (sms * sms); // 704 + let hidden_states = hidden_states.reshape((merged_patches, sms, sms, hidden_dim))?; + let hidden_states = hidden_states.permute((0, 3, 1, 2))?; // [704, 1024, 2, 2] + let hidden_states = self.downsample.forward(&hidden_states)?; // [704, 1536, 1, 1] + let hidden_states = hidden_states.reshape((merged_patches, self.config.out_hidden_size))?; // [704, 1536] + + let merged = self.merger.forward(&hidden_states)?; + + let merged = merged.unsqueeze(0)?; + Ok(merged) + } +} + +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)?) + } +} + +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)?) + } + + 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)?; + + // 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())?, + )) + } +} + +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)?; + + 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, + }) + } + + 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; + + 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)?; + + // 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); + + // 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)?; + } + + 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())? + }; + + let mut hidden_states = inputs_embeds; + for layer in self.layers.iter_mut() { + hidden_states = layer.forward(&hidden_states, (&cos, &sin), attention_mask.as_ref())?; + } + + hidden_states = self.norm.forward(&hidden_states)?; + let logits = self.lm_head.forward(&hidden_states)?; + + Ok(logits) + } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache(); + } + } +} + +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 { + 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(), + )? + }; + + let vision_output = self.vision_encoder.forward(pixels, &grid_thw)?; + + Some(vision_output) + } else { + None + }; + + let result = self.language_model.forward( + input_ids, + image_features.as_ref(), + image_mask, + image_grid_thw, + seqlen_offset, + ); + + result + } + + pub fn clear_kv_cache(&mut self) { + 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..ceabe56 --- /dev/null +++ b/src/models/glm_ocr/processor.rs @@ -0,0 +1,298 @@ +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); + + // 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(); + + let tensor = Tensor::from_vec(pixels, (target_h, target_w, 3), &self.device)?; + + let mean = Tensor::new(self.image_mean.clone(), &self.device)?.reshape((1, 1, 3))?; + let std = Tensor::new(self.image_std.clone(), &self.device)?.reshape((1, 1, 3))?; + let tensor = tensor.broadcast_sub(&mean)?.broadcast_div(&std)?; + + // Now reshape into flattened patches like Python + let grid_h = target_h / self.patch_size; + let grid_w = target_w / self.patch_size; + let patch_size = self.patch_size; + let channels = 3; + let temporal_patch_size = 2; + + let tensor = tensor.reshape(( + grid_h, patch_size, + grid_w, patch_size, + channels, + ))?; + + let tensor = tensor.permute((0, 2, 4, 1, 3))?; + + let num_patches = grid_h * grid_w; + let tensor = tensor.reshape((num_patches, channels, patch_size, patch_size))?; + + let tensor = tensor.unsqueeze(2)?; + let tensor = tensor.repeat((1, 1, temporal_patch_size, 1, 1))?; + + let patch_dim = channels * temporal_patch_size * patch_size * patch_size; + let tensor = tensor.reshape((num_patches, patch_dim))?; + + let tensor = tensor.to_dtype(self.dtype)?; + + Ok(ProcessedImage { + pixel_values: tensor, + grid_h, + grid_w, + }) + } + + /// Process image and text for multimodal input + 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 + 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