From aec47e47afcb5fd4277642343cfec5b8e240ceed Mon Sep 17 00:00:00 2001 From: "273265088@qq.com" <273265088@qq.com> Date: Thu, 26 Mar 2026 17:42:58 +0800 Subject: [PATCH] =?UTF-8?q?=E7=BB=99=20glm-ocr=20=E6=B7=BB=E5=8A=A0=20gguf?= =?UTF-8?q?=20=E5=92=8C=20onnx=20=E6=A0=BC=E5=BC=8F=E6=8E=A8=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 7 + README.zh-CN.md | 7 + docs/supported-models.md | 2 + docs/supported-models.zh-CN.md | 2 + src/exec/glm_ocr.rs | 154 +++-- src/main.rs | 72 ++ src/models/common/gguf.rs | 4 + src/models/core/artifact.rs | 5 + src/models/core/registry.rs | 16 +- src/models/glm_ocr/generate.rs | 1002 +++++++++++++++++++++++++++- src/models/glm_ocr/mod.rs | 1 + src/models/glm_ocr/onnx.rs | 847 +++++++++++++++++++++++ src/models/glm_ocr/processor.rs | 24 + tests/test_glm_ocr_multi_format.rs | 220 ++++++ tests/test_load_spec.rs | 29 + 15 files changed, 2315 insertions(+), 77 deletions(-) create mode 100644 src/models/glm_ocr/onnx.rs create mode 100644 tests/test_glm_ocr_multi_format.rs diff --git a/README.md b/README.md index d83b85a..16df83c 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ aha is a high-performance, cross-platform AI inference engine built with Rust an ## Changelog ### v0.2.3 (2026-03-18) - add DeepSeek-OCR-2 +- add GLM-OCR gguf and onnx local loading ### 2026-03-17 - add PaddleOCR-VL1.5 model @@ -106,6 +107,12 @@ aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format gguf --gg # Run local all-MiniLM-L6-v2 embedding (ONNX) aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format onnx --onnx-path D:\model_download\all-MiniLM-L6-v2\onnx --tokenizer-dir D:\model_download\all-MiniLM-L6-v2 +# Run local GLM-OCR (GGUF) +aha run -m glm-ocr -i .\assets\img\ocr_test1.png --artifact-format gguf --gguf-path D:\model_download\GLM-OCR-GGUF + +# Run local GLM-OCR (ONNX) +aha run -m glm-ocr -i .\assets\img\ocr_test1.png --artifact-format onnx --onnx-path D:\model_download\GLM-OCR-ONNX --tokenizer-dir D:\model_download\GLM-OCR-ONNX + # Start service only (model already downloaded) aha serv -m qwen3asr-0.6b -p 10100 diff --git a/README.zh-CN.md b/README.zh-CN.md index cc31391..0467746 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -27,6 +27,7 @@ aha 是一款基于 Rust 和 Candle 框架构建的高性能跨平台 AI 推理 ## 更新日志 ### v0.2.3 (2026-03-18) - 新增 DeepSeek-OCR-2 +- 增加 GLM-OCR 的 GGUF 和 ONNX 本地加载 ### 2026-03-17 - 新增 PaddleOCR-VL1.5 模型 @@ -105,6 +106,12 @@ aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format gguf --gg # 本地运行 all-MiniLM-L6-v2 向量模型（ONNX） aha run -m all-minilm-l6-v2 -i "Rust embedding test" --artifact-format onnx --onnx-path D:\model_download\all-MiniLM-L6-v2\onnx --tokenizer-dir D:\model_download\all-MiniLM-L6-v2 +# 本地运行 GLM-OCR（GGUF） +aha run -m glm-ocr -i .\assets\img\ocr_test1.png --artifact-format gguf --gguf-path D:\model_download\GLM-OCR-GGUF + +# 本地运行 GLM-OCR（ONNX） +aha run -m glm-ocr -i .\assets\img\ocr_test1.png --artifact-format onnx --onnx-path D:\model_download\GLM-OCR-ONNX --tokenizer-dir D:\model_download\GLM-OCR-ONNX + # 仅启动服务（模型已下载） aha serv -m qwen3asr-0.6b -p 10100 diff --git a/docs/supported-models.md b/docs/supported-models.md index 96478c9..48fe0f5 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -63,6 +63,8 @@ aha supports a growing collection of state-of-the-art AI models across multiple | **DeepSeek-OCR** | Multi | Scene text | Natural images | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) | | **GLM-OCR** | 8 | Scene text | complex document | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) | +GLM-OCR local artifacts: `safetensors`, `gguf`, `onnx` + ## Speech Recognition (ASR) | Model | Parameters | Language | Real-time | Speed | License | diff --git a/docs/supported-models.zh-CN.md b/docs/supported-models.zh-CN.md index d35308f..98afd4a 100644 --- a/docs/supported-models.zh-CN.md +++ b/docs/supported-models.zh-CN.md @@ -63,6 +63,8 @@ aha 支持多个领域的最先进 AI 模型集合。 | **DeepSeek-OCR** | 多语言 | 场景文字 | 自然图像 | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) | | **GLM-OCR** | 8 | 场景文字 | 复杂文档 | [MIT](https://huggingface.co/datasets/choosealicense/licenses/blob/main/markdown/mit.md) | +GLM-OCR 本地制品格式：`safetensors`、`gguf`、`onnx` + ## 语音识别 (ASR) | 模型 | 参数量 | 语言 | 实时 | 速度 | 开源协议 | diff --git a/src/exec/glm_ocr.rs b/src/exec/glm_ocr.rs index 4505071..2ec8b9b 100644 --- a/src/exec/glm_ocr.rs +++ b/src/exec/glm_ocr.rs @@ -1,56 +1,84 @@ //! Glm-OCR exec implementation for CLI `run` subcommand -use std::time::Instant; +use std::{path::Path, time::Instant}; -use anyhow::{Ok, Result}; +use anyhow::{Ok, Result, anyhow}; +use serde_json::json; use crate::exec::ExecModel; -use crate::models::{GenerateModel, glm_ocr::generate::GlmOcrGenerateModel}; +use crate::models::{GenerateModel, LoadSpec, 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) - }; +fn resolve_input_url(input: &str) -> Result { + if input.starts_with("http://") || input.starts_with("https://") || input.starts_with("file://") + { + return Ok(input.to_string()); + } + + let path = Path::new(input); + let path = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir()?.join(path) + }; + let path = path + .canonicalize() + .map_err(|e| anyhow!("failed to resolve input path {}: {e}", path.display()))?; + url::Url::from_file_path(&path) + .map(|url| url.to_string()) + .map_err(|_| { + anyhow!( + "failed to convert input path to file url: {}", + path.display() + ) + }) +} + +fn build_request( + input_url: &str, + max_tokens: Option, +) -> Result { + Ok(serde_json::from_value(json!({ + "model": "glm-ocr", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": input_url + } + }, + { + "type": "text", + "text": "Text Recognition:" + } + ] + } + ], + "max_tokens": max_tokens.unwrap_or(256) + }))?) +} + +impl GlmOcrExec { + pub fn run_with_spec( + input: &[String], + output: Option<&str>, + spec: &LoadSpec, + max_tokens: Option, + ) -> Result<()> { + let url = input + .first() + .ok_or_else(|| anyhow!("glm-ocr run requires an input image path or url"))?; + let input_url = resolve_input_url(url)?; let i_start = Instant::now(); - let mut model = GlmOcrGenerateModel::init(weight_path, None, None)?; + let mut model = GlmOcrGenerateModel::init_from_spec(spec, 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 mes = build_request(&input_url, max_tokens)?; let i_start = Instant::now(); let result = model.generate(mes)?; @@ -67,3 +95,47 @@ impl ExecModel for GlmOcrExec { Ok(()) } } + +impl ExecModel for GlmOcrExec { + fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> { + let spec = LoadSpec::for_safetensors(crate::models::WhichModel::GlmOCR, weight_path); + Self::run_with_spec(input, output, &spec, None) + } +} + +#[cfg(test)] +mod tests { + use super::{build_request, resolve_input_url}; + use aha_openai_dive::v1::resources::chat::{ + ChatMessage, ChatMessageContent, ChatMessageContentPart, + }; + use anyhow::Result; + + #[test] + fn resolve_input_url_converts_relative_windows_path_to_file_url() -> Result<()> { + let url = resolve_input_url(r".\assets\img\ocr_test1.png")?; + assert!(url.starts_with("file:///")); + assert!(url.contains("assets/img/ocr_test1.png")); + Ok(()) + } + + #[test] + fn build_request_accepts_file_url_without_json_escape_issues() -> Result<()> { + let mes = build_request("file:///D:/model_download/ocr_test1.png", Some(32))?; + let ChatMessage::User { content, .. } = &mes.messages[0] else { + panic!("expected user message"); + }; + let ChatMessageContent::ContentPart(parts) = content else { + panic!("expected multipart content"); + }; + let ChatMessageContentPart::Image(image) = &parts[0] else { + panic!("expected image content part"); + }; + assert_eq!( + image.image_url.url, + "file:///D:/model_download/ocr_test1.png" + ); + assert_eq!(mes.max_tokens, Some(32)); + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs index e76357c..cd0f0df 100644 --- a/src/main.rs +++ b/src/main.rs @@ -236,6 +236,10 @@ struct RunArgs { #[arg(short, long)] output: Option, + /// Maximum number of tokens to generate + #[arg(long)] + max_tokens: Option, + /// Local model weight path (defaults to ~/.aha/{model_id} if not specified) #[arg(long)] weight_path: Option, @@ -467,6 +471,11 @@ fn run_target_model_with_spec(args: &RunArgs, spec: &LoadSpec) -> anyhow::Result Qwen3_5Exec::run_with_spec(&args.input, args.output.as_deref(), spec)?; Ok(true) } + WhichModel::GlmOCR => { + use aha::exec::glm_ocr::GlmOcrExec; + GlmOcrExec::run_with_spec(&args.input, args.output.as_deref(), spec, args.max_tokens)?; + Ok(true) + } _ => Ok(false), } } @@ -1017,6 +1026,69 @@ mod tests { ); } + #[test] + fn parse_glm_ocr_run_gguf_flags() { + let cli = Cli::try_parse_from([ + "aha", + "run", + "--model", + "glm-ocr", + "--input", + "ocr.png", + "--artifact-format", + "gguf", + "--gguf-path", + "D:\\model_download\\GLM-OCR-GGUF", + "--max-tokens", + "8", + ]) + .expect("run args should parse"); + + let Some(Commands::Run(args)) = cli.command else { + panic!("expected run subcommand"); + }; + assert!(matches!(args.artifact_format, Some(ArtifactArg::Gguf))); + assert_eq!(args.model, WhichModel::GlmOCR); + assert_eq!( + args.gguf_path.as_deref(), + Some("D:\\model_download\\GLM-OCR-GGUF") + ); + assert_eq!(args.max_tokens, Some(8)); + } + + #[test] + fn parse_glm_ocr_run_onnx_flags() { + let cli = Cli::try_parse_from([ + "aha", + "run", + "--model", + "glm-ocr", + "--input", + "ocr.png", + "--artifact-format", + "onnx", + "--onnx-path", + "D:\\model_download\\GLM-OCR-ONNX", + "--tokenizer-dir", + "D:\\model_download\\GLM-OCR-ONNX", + ]) + .expect("run args should parse"); + + let Some(Commands::Run(args)) = cli.command else { + panic!("expected run subcommand"); + }; + assert!(matches!(args.artifact_format, Some(ArtifactArg::Onnx))); + assert_eq!(args.model, WhichModel::GlmOCR); + assert_eq!( + args.onnx_path.as_deref(), + Some("D:\\model_download\\GLM-OCR-ONNX") + ); + assert_eq!( + args.tokenizer_dir.as_deref(), + Some("D:\\model_download\\GLM-OCR-ONNX") + ); + } + #[test] fn parse_serv_onnx_flags() { let cli = Cli::try_parse_from([ diff --git a/src/models/common/gguf.rs b/src/models/common/gguf.rs index a0d92f6..7980976 100644 --- a/src/models/common/gguf.rs +++ b/src/models/common/gguf.rs @@ -86,6 +86,10 @@ impl Gguf { &self.ct.metadata } + pub fn has_tensor(&self, name: &str) -> bool { + self.ct.tensor_infos.contains_key(name) + } + pub fn tensor(&mut self, name: &str) -> Result { Ok(self.ct.tensor(&mut self.reader, name, &self.device)?) } diff --git a/src/models/core/artifact.rs b/src/models/core/artifact.rs index dcc386a..fcc599a 100644 --- a/src/models/core/artifact.rs +++ b/src/models/core/artifact.rs @@ -33,6 +33,11 @@ pub fn supported_artifacts(model: WhichModel) -> &'static [ArtifactKind] { ArtifactKind::Gguf, ArtifactKind::Onnx, ], + WhichModel::GlmOCR => &[ + ArtifactKind::Safetensors, + ArtifactKind::Gguf, + ArtifactKind::Onnx, + ], WhichModel::Qwen3_0_6B => &[ ArtifactKind::Safetensors, ArtifactKind::Gguf, diff --git a/src/models/core/registry.rs b/src/models/core/registry.rs index 87584a6..c60ea12 100644 --- a/src/models/core/registry.rs +++ b/src/models/core/registry.rs @@ -1,9 +1,10 @@ use anyhow::Result; use crate::models::{ - ModelInstance, WhichModel, all_minilm_l6_v2::generate::AllMiniLML6V2Model, load_model_legacy, - qwen3::generate::Qwen3GenerateModel, qwen3_5::generate::Qwen3_5GenerateModel, - qwen3_embedding::generate::Qwen3EmbeddingModel, qwen3_reranker::generate::Qwen3RerankerModel, + ModelInstance, WhichModel, all_minilm_l6_v2::generate::AllMiniLML6V2Model, + glm_ocr::generate::GlmOcrGenerateModel, load_model_legacy, qwen3::generate::Qwen3GenerateModel, + qwen3_5::generate::Qwen3_5GenerateModel, qwen3_embedding::generate::Qwen3EmbeddingModel, + qwen3_reranker::generate::Qwen3RerankerModel, }; use super::artifact::LoadSpec; @@ -11,6 +12,7 @@ use super::artifact::LoadSpec; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ModelLoaderFamily { AllMiniLML6V2, + GlmOCR, Qwen3, Qwen3Embedding, Qwen3Reranker, @@ -21,6 +23,7 @@ enum ModelLoaderFamily { fn resolve_model_loader_family(model: WhichModel) -> ModelLoaderFamily { match model { WhichModel::AllMiniLML6V2 => ModelLoaderFamily::AllMiniLML6V2, + WhichModel::GlmOCR => ModelLoaderFamily::GlmOCR, WhichModel::Qwen3_0_6B => ModelLoaderFamily::Qwen3, WhichModel::Qwen3Embedding0_6B | WhichModel::Qwen3Embedding4B @@ -49,6 +52,9 @@ pub fn load_model_from_spec<'a>(spec: &LoadSpec) -> Result> { ModelLoaderFamily::AllMiniLML6V2 => { ModelInstance::AllMiniLML6V2(AllMiniLML6V2Model::init_from_spec(spec, None, None)?) } + ModelLoaderFamily::GlmOCR => { + ModelInstance::GlmOCR(GlmOcrGenerateModel::init_from_spec(spec, None, None)?) + } ModelLoaderFamily::Qwen3 => { ModelInstance::Qwen3(Qwen3GenerateModel::init_from_spec(spec, None, None)?) } @@ -86,6 +92,10 @@ mod tests { resolve_model_loader_family(WhichModel::Qwen3_0_6B), ModelLoaderFamily::Qwen3 ); + assert_eq!( + resolve_model_loader_family(WhichModel::GlmOCR), + ModelLoaderFamily::GlmOCR + ); assert_eq!( resolve_model_loader_family(WhichModel::Qwen3Embedding0_6B), ModelLoaderFamily::Qwen3Embedding diff --git a/src/models/glm_ocr/generate.rs b/src/models/glm_ocr/generate.rs index 2f072a7..399bd4d 100644 --- a/src/models/glm_ocr/generate.rs +++ b/src/models/glm_ocr/generate.rs @@ -1,4 +1,6 @@ //! GLM-OCR Inference and Generation +use std::{collections::HashMap, path::Path}; + use aha_openai_dive::v1::resources::chat::{ ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse, }; @@ -9,12 +11,20 @@ use rocket::async_stream::stream; use rocket::futures::Stream; use crate::{ - // chat_template::ChatTemplate, models::{ GenerateModel, + artifact::{ArtifactKind, LoadSpec}, + common::{ + gguf::{Gguf, load_gguf_file, load_text_bootstrap_from_gguf}, + onnx::resolve_tokenizer_dir, + }, glm_ocr::{ - config::{GlmOcrConfig, GlmOcrGenerationConfig}, + config::{ + GlmOcrConfig, GlmOcrGenerationConfig, GlmOcrRopeParameters, GlmOcrTextConfig, + GlmOcrVisionConfig, + }, model::GlmOcrModel, + onnx::GlmOcrOnnxBackend, processor::GlmOcrProcessor, }, }, @@ -25,11 +35,20 @@ use crate::{ }, }; +const DEFAULT_GLM_OCR_SHORTEST_EDGE: usize = 12_544; +const DEFAULT_GLM_OCR_LONGEST_EDGE: usize = 9_633_792; +const DEFAULT_GLM_OCR_IMAGE_TOKEN_ID: u32 = 59_280; +const DEFAULT_GLM_OCR_IMAGE_START_TOKEN_ID: u32 = 59_256; +const DEFAULT_GLM_OCR_IMAGE_END_TOKEN_ID: u32 = 59_257; +const DEFAULT_GLM_OCR_VIDEO_TOKEN_ID: u32 = 59_281; +const DEFAULT_GLM_OCR_VIDEO_START_TOKEN_ID: u32 = 59_258; +const DEFAULT_GLM_OCR_VIDEO_END_TOKEN_ID: u32 = 59_259; + pub struct GlmOcrGenerateModel { - // chat_template: ChatTemplate<'a>, tokenizer: TokenizerModel, processor: GlmOcrProcessor, - model: GlmOcrModel, + model: Option, + onnx_backend: Option, device: Device, eos_token_ids: Vec, model_name: String, @@ -42,8 +61,40 @@ pub struct GlmOcrGenerateModel { } impl GlmOcrGenerateModel { + pub fn init_from_spec( + spec: &LoadSpec, + device: Option<&Device>, + dtype: Option, + ) -> Result { + match spec.resolved_artifact() { + ArtifactKind::Safetensors => { + let path = + spec.paths.weight_dir.as_deref().ok_or_else(|| { + anyhow!("weight_path is required for glm-ocr safetensors") + })?; + Self::init(path, device, dtype) + } + ArtifactKind::Gguf => { + let gguf_path = spec + .paths + .gguf_path + .as_deref() + .ok_or_else(|| anyhow!("gguf_path is required for glm-ocr gguf"))?; + Self::init_from_gguf(gguf_path, spec.paths.mmproj_path.as_deref(), device, dtype) + } + ArtifactKind::Onnx => { + let onnx_path = spec + .paths + .onnx_path + .as_deref() + .ok_or_else(|| anyhow!("onnx_path is required for glm-ocr onnx"))?; + Self::init_from_onnx(onnx_path, spec.paths.tokenizer_dir.as_deref()) + } + ArtifactKind::Auto => unreachable!("artifact kind should be resolved before init"), + } + } + 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 cfg: GlmOcrConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; @@ -59,10 +110,10 @@ impl GlmOcrGenerateModel { serde_json::from_slice(&std::fs::read(generation_config_path)?)?; Ok(Self { - // chat_template, tokenizer, processor, - model, + model: Some(model), + onnx_backend: None, device, eos_token_ids: generation_config.eos_token_id.clone(), model_name: "glm-ocr".to_string(), @@ -74,6 +125,116 @@ impl GlmOcrGenerateModel { spatial_merge_size: cfg.vision_config.spatial_merge_size, }) } + + pub fn init_from_onnx(onnx_path: &str, tokenizer_dir: Option<&str>) -> Result { + let tokenizer_dir = resolve_tokenizer_dir( + onnx_path, + tokenizer_dir, + &[ + "tokenizer.json", + "config.json", + "generation_config.json", + "preprocessor_config.json", + ], + )?; + let base_path = tokenizer_dir.to_string_lossy().to_string(); + let tokenizer = TokenizerModel::init(&base_path)?; + let cfg: GlmOcrConfig = + serde_json::from_slice(&std::fs::read(tokenizer_dir.join("config.json"))?)?; + let generation_config: GlmOcrGenerationConfig = serde_json::from_slice(&std::fs::read( + tokenizer_dir.join("generation_config.json"), + )?)?; + let processor = GlmOcrProcessor::new(&base_path, &Device::Cpu, DType::F32)?; + let onnx_backend = + GlmOcrOnnxBackend::load(onnx_path, cfg.vision_config.spatial_merge_size)?; + let model_name = tokenizer_dir + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("glm-ocr") + .to_string(); + + Ok(Self { + tokenizer, + processor, + model: None, + onnx_backend: Some(onnx_backend), + device: Device::Cpu, + eos_token_ids: generation_config.eos_token_id, + model_name, + 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, + }) + } + + pub fn init_from_gguf( + gguf_path: &str, + mmproj_path: Option<&str>, + device: Option<&Device>, + dtype: Option, + ) -> Result { + let (model_file, mmproj_file) = resolve_glm_ocr_gguf_files(gguf_path, mmproj_path)?; + let device = get_device(device); + let dtype = dtype.unwrap_or(DType::F32); + let bootstrap = + load_text_bootstrap_from_gguf(&model_file, Some(false), Some(false), Some(false))?; + let mut model_gguf = load_gguf_file(&model_file, &device)?; + let mut mmproj_gguf = load_gguf_file(&mmproj_file, &device)?; + let cfg = build_glm_ocr_gguf_config(&mut model_gguf, &mut mmproj_gguf)?; + let processor = build_glm_ocr_gguf_processor(&mut mmproj_gguf, &device, dtype)?; + let tensors = load_glm_ocr_gguf_tensors(&mut model_gguf, &mut mmproj_gguf, &device, dtype)?; + let vb = VarBuilder::from_tensors(tensors, dtype, &device); + let model = GlmOcrModel::new(vb, cfg.clone())?; + let eos_token_ids = build_glm_ocr_gguf_eos_tokens(&mut model_gguf)?; + let image_token_id = resolve_special_token_id( + &bootstrap.tokenizer, + "<|image|>", + DEFAULT_GLM_OCR_IMAGE_TOKEN_ID, + ); + let image_start_token_id = resolve_special_token_id( + &bootstrap.tokenizer, + "<|begin_of_image|>", + DEFAULT_GLM_OCR_IMAGE_START_TOKEN_ID, + ); + let image_end_token_id = resolve_special_token_id( + &bootstrap.tokenizer, + "<|end_of_image|>", + DEFAULT_GLM_OCR_IMAGE_END_TOKEN_ID, + ); + let model_name = Path::new(&model_file) + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("glm-ocr") + .to_string(); + + Ok(Self { + tokenizer: bootstrap.tokenizer, + processor, + model: Some(model), + onnx_backend: None, + device, + eos_token_ids, + model_name, + image_token_id, + image_start_token_id, + 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, + }) + } + + fn clear_runtime_cache(&mut self) { + if let Some(model) = self.model.as_mut() { + model.clear_kv_cache(); + } + if let Some(onnx_backend) = self.onnx_backend.as_mut() { + onnx_backend.clear_cache(); + } + } } impl GenerateModel for GlmOcrGenerateModel { @@ -81,15 +242,13 @@ impl GenerateModel for GlmOcrGenerateModel { let seed = mes.seed.unwrap_or(34562) as u64; let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed); - // Extract image path and prompt from messages let image_urls = extract_image_url(&mes); let image_path = image_urls .first() .ok_or_else(|| anyhow!("No image provided"))?; - // Get prompt text from messages let mut prompt = extract_user_text(&mes)?; - if prompt.chars().count() == 0 { + if prompt.is_empty() { prompt = "Extract all text from this image.".to_string() } @@ -105,6 +264,62 @@ impl GenerateModel for GlmOcrGenerateModel { self.spatial_merge_size, )?; + if let Some(onnx_backend) = self.onnx_backend.as_mut() { + let mut current_ids = processed.input_ids.squeeze(0)?.to_vec1::()?; + let pixel_values = Some(processed.pixel_values); + let image_grid_thw = Some(processed.grid_thw); + let image_mask = Some(processed.image_mask); + let prompt_tokens = current_ids.len() as u32; + let sample_len = mes.max_tokens.unwrap_or(512); + let mut position_start = 0usize; + let mut generate = Vec::new(); + + for _ in 0..sample_len { + let logits = onnx_backend.forward_logits( + ¤t_ids, + if position_start == 0 { + image_mask.as_ref() + } else { + None + }, + if position_start == 0 { + pixel_values.as_ref() + } else { + None + }, + if position_start == 0 { + image_grid_thw.as_ref() + } else { + None + }, + position_start, + )?; + let vocab = logits.len(); + let logits = Tensor::from_vec(logits, vocab, &self.device)?; + let next_token = logit_processor.sample(&logits)?; + generate.push(next_token); + if self.eos_token_ids.contains(&next_token) { + break; + } + position_start += current_ids.len(); + current_ids = vec![next_token]; + } + + self.clear_runtime_cache(); + let num_token = generate.len() as u32; + let res = self.tokenizer.token_decode(generate)?; + return Ok(build_completion_response( + res, + &self.model_name, + Some(num_token), + Some(prompt_tokens), + )); + } + + let model = self + .model + .as_mut() + .ok_or_else(|| anyhow!("glm-ocr native runtime is not initialized"))?; let mut input_ids = processed.input_ids; let pixel_values = Some(processed.pixel_values); let image_grid_thw = Some(processed.grid_thw); @@ -117,7 +332,7 @@ impl GenerateModel for GlmOcrGenerateModel { for _ in 0..sample_len { let is_first_pass = seqlen_offset == 0; - let logits = self.model.forward( + let logits = model.forward( &input_ids, if is_first_pass { pixel_values.as_ref() @@ -148,12 +363,15 @@ impl GenerateModel for GlmOcrGenerateModel { input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?; } - self.model.clear_kv_cache(); + self.clear_runtime_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), Some(prompt_tokens)); - Ok(response) + Ok(build_completion_response( + res, + &self.model_name, + Some(num_token), + Some(prompt_tokens), + )) } fn generate_stream( @@ -170,15 +388,13 @@ impl GenerateModel for GlmOcrGenerateModel { let seed = mes.seed.unwrap_or(34562) as u64; let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p, None, seed); - // Extract image path and prompt from messages let image_urls = extract_image_url(&mes); let image_path = image_urls .first() .ok_or_else(|| anyhow!("No image provided"))?; - // Get prompt text from messages let mut prompt = extract_user_text(&mes)?; - if prompt.chars().count() == 0 { + if prompt.is_empty() { prompt = "Extract all text from this image.".to_string() } @@ -194,6 +410,80 @@ impl GenerateModel for GlmOcrGenerateModel { self.spatial_merge_size, )?; + if let Some(onnx_backend) = self.onnx_backend.as_mut() { + let tokenizer = &self.tokenizer; + let model_name = self.model_name.clone(); + let eos_token_ids = self.eos_token_ids.clone(); + let device = self.device.clone(); + let mut current_ids = processed.input_ids.squeeze(0)?.to_vec1::()?; + let mut position_start = 0usize; + let sample_len = mes.max_tokens.unwrap_or(512); + let mut pixel_values = Some(processed.pixel_values); + let mut image_grid_thw = Some(processed.grid_thw); + let mut image_mask = Some(processed.image_mask); + + let stream = stream! { + let mut error_tokens = Vec::new(); + for _ in 0..sample_len { + let logits = onnx_backend.forward_logits( + ¤t_ids, + image_mask.as_ref(), + pixel_values.as_ref(), + image_grid_thw.as_ref(), + position_start, + )?; + let vocab = logits.len(); + let logits = Tensor::from_vec(logits, vocab, &device)?; + let next_token = logit_processor.sample(&logits)?; + + let mut decode_ids = Vec::new(); + if !error_tokens.is_empty() { + decode_ids.extend_from_slice(&error_tokens); + } + decode_ids.push(next_token); + + let decoded_token = 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(); + } + position_start += current_ids.len(); + current_ids = vec![next_token]; + pixel_values = None; + image_grid_thw = None; + image_mask = None; + continue; + } + error_tokens.clear(); + + let chunk = build_completion_chunk_response(decoded_token, &model_name, None, None); + yield Ok(chunk); + + if eos_token_ids.contains(&next_token) { + break; + } + position_start += current_ids.len(); + current_ids = vec![next_token]; + pixel_values = None; + image_grid_thw = None; + image_mask = None; + } + onnx_backend.clear_cache(); + }; + return Ok(Box::new(Box::pin(stream))); + } + + let model = self + .model + .as_mut() + .ok_or_else(|| anyhow!("glm-ocr native runtime is not initialized"))?; + let tokenizer = &self.tokenizer; + let model_name = self.model_name.clone(); + let eos_token_ids = self.eos_token_ids.clone(); + let device = self.device.clone(); let mut input_ids = processed.input_ids; let pixel_values = Some(processed.pixel_values); let image_grid_thw = Some(processed.grid_thw); @@ -203,21 +493,27 @@ impl GenerateModel for GlmOcrGenerateModel { let sample_len = mes.max_tokens.unwrap_or(512); let stream = stream! { - let mut generated: Vec = Vec::new(); let mut error_tokens = Vec::new(); + let mut pixel_values = pixel_values.as_ref(); + let image_grid_thw = image_grid_thw.as_ref(); + let mut image_mask = image_mask.as_ref(); for _ in 0..sample_len { - let is_first_pass = seqlen_offset == 0; - let logits = self.model.forward( + let logits = 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 }, + pixel_values, + image_grid_thw, + image_mask, 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 = 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 next_token = logit_processor.sample(&logits).map_err(|e| anyhow!(format!("sample error: {e}")))?; - generated.push(next_token); + let next_token = logit_processor + .sample(&logits) + .map_err(|e| anyhow!(format!("sample error: {e}")))?; let mut decode_ids = Vec::new(); if !error_tokens.is_empty() { @@ -225,7 +521,9 @@ impl GenerateModel for GlmOcrGenerateModel { } decode_ids.push(next_token); - let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("decode error: {e}")))?; + let decoded_token = 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 { @@ -233,24 +531,662 @@ impl GenerateModel for GlmOcrGenerateModel { } 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}")))?; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &device) + .map_err(|e| anyhow!(format!("tensor error: {e}")))?; + pixel_values = None; + image_mask = None; continue; } error_tokens.clear(); - let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None); + let chunk = build_completion_chunk_response(decoded_token, &model_name, None, None); yield Ok(chunk); - if self.eos_token_ids.contains(&next_token) { + if 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}")))?; + input_ids = Tensor::from_vec(vec![next_token], (1, 1), &device) + .map_err(|e| anyhow!(format!("tensor error: {e}")))?; + pixel_values = None; + image_mask = None; } - self.model.clear_kv_cache(); + model.clear_kv_cache(); }; Ok(Box::new(Box::pin(stream))) } } + +fn resolve_glm_ocr_gguf_files( + gguf_path: &str, + mmproj_path: Option<&str>, +) -> Result<(String, String)> { + fn find_gguf_file(dir: &Path, marker: &str, negate: bool) -> Result { + let mut matches = std::fs::read_dir(dir)? + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + let file_name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or_default(); + path.is_file() + && path + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case("gguf")) + && if negate { + !file_name.contains(marker) + } else { + file_name.contains(marker) + } + }) + .collect::>(); + matches.sort(); + matches + .into_iter() + .next() + .map(|path| path.to_string_lossy().to_string()) + .ok_or_else(|| { + anyhow!( + "unable to locate gguf component {} under {}", + marker, + dir.display() + ) + }) + } + + let model_path = Path::new(gguf_path); + if !model_path.exists() { + return Err(anyhow!("gguf model path not found: {}", gguf_path)); + } + + let model_file = if model_path.is_dir() { + find_gguf_file(model_path, "mmproj", true)? + } else { + let file_name = model_path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or_default(); + if !file_name.ends_with(".gguf") || file_name.contains("mmproj") { + return Err(anyhow!("glm-ocr gguf model file is invalid: {}", gguf_path)); + } + model_path.to_string_lossy().to_string() + }; + + let mmproj_file = if let Some(mmproj_path) = mmproj_path { + let mmproj = Path::new(mmproj_path); + if !mmproj.exists() { + return Err(anyhow!("glm-ocr mmproj path not found: {}", mmproj_path)); + } + if mmproj.is_dir() { + find_gguf_file(mmproj, "mmproj", false)? + } else { + mmproj.to_string_lossy().to_string() + } + } else { + let search_dir = Path::new(&model_file) + .parent() + .ok_or_else(|| anyhow!("glm-ocr gguf parent directory not found"))?; + find_gguf_file(search_dir, "mmproj", false)? + }; + + Ok((model_file, mmproj_file)) +} + +fn resolve_special_token_id(tokenizer: &TokenizerModel, token: &str, fallback: u32) -> u32 { + tokenizer.tokenizer.token_to_id(token).unwrap_or(fallback) +} + +fn build_glm_ocr_gguf_eos_tokens( + gguf: &mut Gguf, +) -> Result> { + let mut eos = vec![gguf.get_matedata("tokenizer.ggml.eos_token_id")?.to_u32()?]; + if let Ok(eot) = gguf.get_matedata("tokenizer.ggml.eot_token_id") + && let Ok(eot) = eot.to_u32() + && !eos.contains(&eot) + { + eos.push(eot); + } + Ok(eos) +} + +fn build_glm_ocr_gguf_config< + R1: std::io::Read + std::io::Seek, + R2: std::io::Read + std::io::Seek, +>( + model_gguf: &mut Gguf, + mmproj_gguf: &mut Gguf, +) -> Result { + let rope_sections = model_gguf + .get_matedata("glm4.rope.dimension_sections")? + .to_vec()? + .iter() + .map(|value| { + value + .to_i32() + .map(|value| value as usize) + .or_else(|_| value.to_u32().map(|value| value as usize)) + }) + .collect::>>()?; + let image_size = mmproj_gguf + .get_matedata("clip.vision.image_size")? + .to_u32()? as usize; + let patch_size = mmproj_gguf + .get_matedata("clip.vision.patch_size")? + .to_u32()? as usize; + + Ok(GlmOcrConfig { + model_type: "glm_ocr".to_string(), + vision_config: GlmOcrVisionConfig { + depth: mmproj_gguf + .get_matedata("clip.vision.block_count")? + .to_u32()? as usize, + hidden_size: mmproj_gguf + .get_matedata("clip.vision.embedding_length")? + .to_u32()? as usize, + hidden_act: candle_nn::Activation::Silu, + attention_bias: true, + num_heads: mmproj_gguf + .get_matedata("clip.vision.attention.head_count")? + .to_u32()? as usize, + in_channels: 3, + image_size, + patch_size, + rms_norm_eps: mmproj_gguf + .get_matedata("clip.vision.attention.layer_norm_epsilon")? + .to_f32()? as f64, + spatial_merge_size: 2, + temporal_patch_size: 2, + out_hidden_size: mmproj_gguf + .get_matedata("clip.vision.projection_dim")? + .to_u32()? as usize, + intermediate_size: mmproj_gguf + .get_matedata("clip.vision.feed_forward_length")? + .to_u32()? as usize, + initializer_range: 0.02, + rope_theta: 10_000.0, + }, + text_config: GlmOcrTextConfig { + vocab_size: model_gguf + .get_matedata("tokenizer.ggml.tokens")? + .to_vec()? + .len(), + hidden_size: model_gguf.get_matedata("glm4.embedding_length")?.to_u32()? as usize, + intermediate_size: model_gguf + .get_matedata("glm4.feed_forward_length")? + .to_u32()? as usize, + num_hidden_layers: model_gguf.get_matedata("glm4.block_count")?.to_u32()? as usize, + num_attention_heads: model_gguf + .get_matedata("glm4.attention.head_count")? + .to_u32()? as usize, + num_key_value_heads: model_gguf + .get_matedata("glm4.attention.head_count_kv")? + .to_u32()? as usize, + head_dim: Some( + model_gguf + .get_matedata("glm4.attention.key_length")? + .to_u32()? as usize, + ), + max_position_embeddings: model_gguf.get_matedata("glm4.context_length")?.to_u32()? + as usize, + rms_norm_eps: model_gguf + .get_matedata("glm4.attention.layer_norm_rms_epsilon")? + .to_f32()? as f64, + hidden_act: candle_nn::Activation::Silu, + use_cache: true, + rope_parameters: GlmOcrRopeParameters { + rope_type: "default".to_string(), + mrope_section: rope_sections, + partial_rotary_factor: 1.0, + rope_theta: model_gguf.get_matedata("glm4.rope.freq_base")?.to_f32()?, + }, + eos_token_id: build_glm_ocr_gguf_eos_tokens(model_gguf)?, + dtype: "float32".to_string(), + }, + image_token_id: DEFAULT_GLM_OCR_IMAGE_TOKEN_ID, + video_token_id: DEFAULT_GLM_OCR_VIDEO_TOKEN_ID, + image_start_token_id: DEFAULT_GLM_OCR_IMAGE_START_TOKEN_ID, + image_end_token_id: DEFAULT_GLM_OCR_IMAGE_END_TOKEN_ID, + video_start_token_id: DEFAULT_GLM_OCR_VIDEO_START_TOKEN_ID, + video_end_token_id: DEFAULT_GLM_OCR_VIDEO_END_TOKEN_ID, + }) +} + +fn build_glm_ocr_gguf_processor( + mmproj_gguf: &mut Gguf, + device: &Device, + dtype: DType, +) -> Result { + let image_mean = mmproj_gguf + .get_matedata("clip.vision.image_mean")? + .to_vec()? + .iter() + .map(|value| value.to_f32()) + .collect::>>()?; + let image_std = mmproj_gguf + .get_matedata("clip.vision.image_std")? + .to_vec()? + .iter() + .map(|value| value.to_f32()) + .collect::>>()?; + let patch_size = mmproj_gguf + .get_matedata("clip.vision.patch_size")? + .to_u32()? as usize; + + Ok(GlmOcrProcessor::from_params( + image_mean, + image_std, + DEFAULT_GLM_OCR_SHORTEST_EDGE, + DEFAULT_GLM_OCR_LONGEST_EDGE, + patch_size, + 2, + 2, + device, + dtype, + )) +} + +fn load_glm_ocr_gguf_tensors< + R1: std::io::Read + std::io::Seek, + R2: std::io::Read + std::io::Seek, +>( + model_gguf: &mut Gguf, + mmproj_gguf: &mut Gguf, + device: &Device, + dtype: DType, +) -> Result> { + let mut tensors = HashMap::new(); + + insert_gguf_tensor( + model_gguf, + &mut tensors, + "token_embd.weight", + "model.language_model.embed_tokens.weight", + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + "output.weight", + "lm_head.weight", + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + "output_norm.weight", + "model.language_model.norm.weight", + device, + dtype, + )?; + + let num_layers = model_gguf.get_matedata("glm4.block_count")?.to_u32()? as usize; + for idx in 0..num_layers { + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.attn_q.weight"), + &format!("model.language_model.layers.{idx}.self_attn.q_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.attn_k.weight"), + &format!("model.language_model.layers.{idx}.self_attn.k_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.attn_v.weight"), + &format!("model.language_model.layers.{idx}.self_attn.v_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.attn_output.weight"), + &format!("model.language_model.layers.{idx}.self_attn.o_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.attn_norm.weight"), + &format!("model.language_model.layers.{idx}.input_layernorm.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.post_attention_norm.weight"), + &format!("model.language_model.layers.{idx}.post_self_attn_layernorm.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.ffn_norm.weight"), + &format!("model.language_model.layers.{idx}.post_attention_layernorm.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.post_ffw_norm.weight"), + &format!("model.language_model.layers.{idx}.post_mlp_layernorm.weight"), + device, + dtype, + )?; + insert_text_gate_up_gguf_tensor( + model_gguf, + &mut tensors, + idx, + &format!("model.language_model.layers.{idx}.mlp.gate_up_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + model_gguf, + &mut tensors, + &format!("blk.{idx}.ffn_down.weight"), + &format!("model.language_model.layers.{idx}.mlp.down_proj.weight"), + device, + dtype, + )?; + } + + let patch_weight_0 = take_gguf_tensor(mmproj_gguf, "v.patch_embd.weight", device, dtype)?; + let patch_weight_1 = take_gguf_tensor(mmproj_gguf, "v.patch_embd.weight.1", device, dtype)?; + let patch_weight = Tensor::cat( + &[&patch_weight_0.unsqueeze(2)?, &patch_weight_1.unsqueeze(2)?], + 2, + )?; + tensors.insert( + "model.visual.patch_embed.proj.weight".to_string(), + patch_weight, + ); + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "v.patch_embd.bias", + "model.visual.patch_embed.proj.bias", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "v.post_ln.weight", + "model.visual.post_layernorm.weight", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.patch_merger.weight", + "model.visual.downsample.weight", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.patch_merger.bias", + "model.visual.downsample.bias", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.model.fc.weight", + "model.visual.merger.proj.weight", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.post_norm.weight", + "model.visual.merger.post_projection_norm.weight", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.post_norm.bias", + "model.visual.merger.post_projection_norm.bias", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.gate.weight", + "model.visual.merger.gate_proj.weight", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.up.weight", + "model.visual.merger.up_proj.weight", + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + "mm.down.weight", + "model.visual.merger.down_proj.weight", + device, + dtype, + )?; + + let vision_depth = mmproj_gguf + .get_matedata("clip.vision.block_count")? + .to_u32()? as usize; + for idx in 0..vision_depth { + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ln1.weight"), + &format!("model.visual.blocks.{idx}.norm1.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ln2.weight"), + &format!("model.visual.blocks.{idx}.norm2.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.attn_qkv.weight"), + &format!("model.visual.blocks.{idx}.attn.qkv.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.attn_qkv.bias"), + &format!("model.visual.blocks.{idx}.attn.qkv.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.attn_out.weight"), + &format!("model.visual.blocks.{idx}.attn.proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.attn_out.bias"), + &format!("model.visual.blocks.{idx}.attn.proj.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.attn_q_norm.weight"), + &format!("model.visual.blocks.{idx}.attn.q_norm.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.attn_k_norm.weight"), + &format!("model.visual.blocks.{idx}.attn.k_norm.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ffn_gate.weight"), + &format!("model.visual.blocks.{idx}.mlp.gate_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ffn_gate.bias"), + &format!("model.visual.blocks.{idx}.mlp.gate_proj.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ffn_up.weight"), + &format!("model.visual.blocks.{idx}.mlp.up_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ffn_up.bias"), + &format!("model.visual.blocks.{idx}.mlp.up_proj.bias"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ffn_down.weight"), + &format!("model.visual.blocks.{idx}.mlp.down_proj.weight"), + device, + dtype, + )?; + insert_gguf_tensor( + mmproj_gguf, + &mut tensors, + &format!("v.blk.{idx}.ffn_down.bias"), + &format!("model.visual.blocks.{idx}.mlp.down_proj.bias"), + device, + dtype, + )?; + } + + Ok(tensors) +} + +fn take_gguf_tensor( + gguf: &mut Gguf, + name: &str, + device: &Device, + dtype: DType, +) -> Result { + let tensor = gguf + .get_dequantized(name) + .map_err(|err| anyhow!("failed to load gguf tensor {}: {}", name, err))? + .to_device(device) + .map_err(|err| anyhow!("failed to move gguf tensor {}: {}", name, err))?; + tensor + .to_dtype(dtype) + .map_err(|err| anyhow!("failed to convert gguf tensor {}: {}", name, err)) +} + +fn insert_gguf_tensor( + gguf: &mut Gguf, + tensors: &mut HashMap, + gguf_name: &str, + target_name: &str, + device: &Device, + dtype: DType, +) -> Result<()> { + let tensor = take_gguf_tensor(gguf, gguf_name, device, dtype)?; + tensors.insert(target_name.to_string(), tensor); + Ok(()) +} + +fn insert_combined_gguf_tensors( + gguf: &mut Gguf, + tensors: &mut HashMap, + gate_name: &str, + up_name: &str, + target_name: &str, + device: &Device, + dtype: DType, +) -> Result<()> { + let gate = take_gguf_tensor(gguf, gate_name, device, dtype)?; + let up = take_gguf_tensor(gguf, up_name, device, dtype)?; + let tensor = Tensor::cat(&[&gate, &up], 0)?; + tensors.insert(target_name.to_string(), tensor); + Ok(()) +} + +fn insert_text_gate_up_gguf_tensor( + gguf: &mut Gguf, + tensors: &mut HashMap, + layer_idx: usize, + target_name: &str, + device: &Device, + dtype: DType, +) -> Result<()> { + let fused_name = format!("blk.{layer_idx}.ffn_up.weight"); + let gate_name = format!("blk.{layer_idx}.ffn_gate.weight"); + + if gguf.has_tensor(&gate_name) { + return insert_combined_gguf_tensors( + gguf, + tensors, + &gate_name, + &fused_name, + target_name, + device, + dtype, + ); + } + + insert_gguf_tensor(gguf, tensors, &fused_name, target_name, device, dtype) +} diff --git a/src/models/glm_ocr/mod.rs b/src/models/glm_ocr/mod.rs index 8b1baf7..757655d 100644 --- a/src/models/glm_ocr/mod.rs +++ b/src/models/glm_ocr/mod.rs @@ -1,4 +1,5 @@ pub mod config; pub mod generate; pub mod model; +pub mod onnx; pub mod processor; diff --git a/src/models/glm_ocr/onnx.rs b/src/models/glm_ocr/onnx.rs new file mode 100644 index 0000000..95913f8 --- /dev/null +++ b/src/models/glm_ocr/onnx.rs @@ -0,0 +1,847 @@ +use anyhow::{Result, anyhow}; +use candle_core::Tensor; + +#[cfg(feature = "onnx-runtime")] +use candle_core::{DType, IndexOp}; + +#[cfg(feature = "onnx-runtime")] +use std::path::{Path, PathBuf}; + +#[cfg(feature = "onnx-runtime")] +use half::f16; +#[cfg(feature = "onnx-runtime")] +use ndarray::{Array, IxDyn}; + +#[cfg(feature = "onnx-runtime")] +use crate::models::common::onnx::create_session; + +#[cfg(feature = "onnx-runtime")] +pub struct GlmOcrOnnxCacheEntry { + pub name: String, + pub dims: Vec, + pub data: Vec, +} + +#[cfg(feature = "onnx-runtime")] +#[derive(Clone)] +struct OnnxInputDescriptor { + name: String, + shape: Vec, + kind: Option, +} + +#[cfg(feature = "onnx-runtime")] +#[derive(Clone, Copy)] +enum OnnxTensorKind { + Bool, + I32, + I64, + F16, + F32, +} + +#[cfg(feature = "onnx-runtime")] +fn map_tensor_kind(ty: ort::value::TensorElementType) -> Option { + match ty { + ort::value::TensorElementType::Bool => Some(OnnxTensorKind::Bool), + ort::value::TensorElementType::Int32 => Some(OnnxTensorKind::I32), + ort::value::TensorElementType::Int64 => Some(OnnxTensorKind::I64), + ort::value::TensorElementType::Float16 => Some(OnnxTensorKind::F16), + ort::value::TensorElementType::Float32 => Some(OnnxTensorKind::F32), + _ => None, + } +} + +#[cfg(feature = "onnx-runtime")] +pub struct GlmOcrOnnxBackend { + embed_session: ort::session::Session, + decoder_session: ort::session::Session, + vision_session: ort::session::Session, + decoder_input_descriptors: Vec, + decoder_output_names: Vec, + vision_input_descriptors: Vec, + vision_output_names: Vec, + cache_values: Vec, + spatial_merge_size: usize, + next_mrope_pos: usize, + prefill_seq_len: usize, +} + +#[cfg(feature = "onnx-runtime")] +fn find_onnx_component_file(path: &str, marker: &str) -> Result { + let model_path = Path::new(path); + if !model_path.exists() { + return Err(anyhow!("onnx model path not found: {}", path)); + } + + if model_path.is_file() { + let file_name = model_path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or_default(); + if file_name.contains(marker) && file_name.ends_with(".onnx") { + return Ok(model_path.to_path_buf()); + } + } + + let search_root = if model_path.is_dir() { + model_path.to_path_buf() + } else { + model_path + .parent() + .ok_or_else(|| anyhow!("onnx component parent directory not found for {}", path))? + .to_path_buf() + }; + + let mut stack = vec![search_root]; + let mut matches = Vec::new(); + while let Some(current) = stack.pop() { + for entry in std::fs::read_dir(¤t)? { + let entry = entry?; + let entry_path = entry.path(); + if entry_path.is_dir() { + stack.push(entry_path); + continue; + } + let file_name = entry_path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or_default(); + if file_name.contains(marker) && file_name.ends_with(".onnx") { + matches.push(entry_path); + } + } + } + matches.sort(); + matches.into_iter().next().ok_or_else(|| { + anyhow!( + "unable to locate onnx component {} under {}", + marker, + model_path.display() + ) + }) +} + +#[cfg(feature = "onnx-runtime")] +impl GlmOcrOnnxBackend { + pub fn load(onnx_path: &str, spatial_merge_size: usize) -> Result { + let embed_file = find_onnx_component_file(onnx_path, "embed_tokens")?; + let decoder_file = find_onnx_component_file(onnx_path, "decoder_model_merged")?; + let vision_file = find_onnx_component_file(onnx_path, "vision_encoder")?; + let embed_bundle = create_session(&embed_file.to_string_lossy(), None)?; + let decoder_bundle = create_session(&decoder_file.to_string_lossy(), None)?; + let vision_bundle = create_session(&vision_file.to_string_lossy(), None)?; + + let decoder_input_descriptors = decoder_bundle + .session + .inputs() + .iter() + .map(|input| { + let (shape, kind) = match input.dtype() { + ort::value::ValueType::Tensor { ty, shape, .. } => ( + shape.iter().copied().collect::>(), + map_tensor_kind(*ty), + ), + _ => (Vec::new(), None), + }; + OnnxInputDescriptor { + name: input.name().to_string(), + shape, + kind, + } + }) + .collect::>(); + let vision_input_descriptors = vision_bundle + .session + .inputs() + .iter() + .map(|input| { + let (shape, kind) = match input.dtype() { + ort::value::ValueType::Tensor { ty, shape, .. } => ( + shape.iter().copied().collect::>(), + map_tensor_kind(*ty), + ), + _ => (Vec::new(), None), + }; + OnnxInputDescriptor { + name: input.name().to_string(), + shape, + kind, + } + }) + .collect::>(); + + Ok(Self { + embed_session: embed_bundle.session, + decoder_session: decoder_bundle.session, + vision_session: vision_bundle.session, + decoder_input_descriptors, + decoder_output_names: decoder_bundle.output_names, + vision_input_descriptors, + vision_output_names: vision_bundle.output_names, + cache_values: Vec::new(), + spatial_merge_size, + next_mrope_pos: 0, + prefill_seq_len: 0, + }) + } + + pub fn clear_cache(&mut self) { + self.cache_values.clear(); + self.next_mrope_pos = 0; + self.prefill_seq_len = 0; + } + + pub fn forward_logits( + &mut self, + input_ids: &[u32], + image_mask: Option<&Tensor>, + pixel_values: Option<&Tensor>, + image_grid_thw: Option<&Tensor>, + position_start: usize, + ) -> Result> { + if input_ids.is_empty() { + return Err(anyhow!("glm-ocr onnx input_ids cannot be empty")); + } + + let (mut embed_data, hidden_size) = self.embed_input_ids(input_ids)?; + if let Some(pixel_values) = pixel_values { + let image_mask = + image_mask.ok_or_else(|| anyhow!("glm-ocr onnx image_mask is required"))?; + let image_grid_thw = + image_grid_thw.ok_or_else(|| anyhow!("glm-ocr onnx image_grid_thw is required"))?; + self.apply_vision_embeds( + &mut embed_data, + hidden_size, + image_mask, + pixel_values, + image_grid_thw, + )?; + } + + let position_ids = if self.cache_values.is_empty() { + if let (Some(mask), Some(grid_thw)) = (image_mask, image_grid_thw) { + self.compute_prefill_position_ids(mask, grid_thw, input_ids.len())? + } else { + self.next_mrope_pos = input_ids.len(); + self.prefill_seq_len = input_ids.len(); + build_text_position_ids(0, input_ids.len()) + } + } else { + let decode_pos = self + .next_mrope_pos + .saturating_add(position_start.saturating_sub(self.prefill_seq_len)); + build_text_position_ids(decode_pos, input_ids.len()) + }; + + let mut decoder_inputs = Vec::with_capacity(self.decoder_input_descriptors.len()); + for desc in &self.decoder_input_descriptors { + let value = match desc.name.as_str() { + "inputs_embeds" => build_float_input( + desc, + vec![1_i64, input_ids.len() as i64, hidden_size as i64], + &embed_data, + )?, + "attention_mask" => { + let attention_mask = self.build_attention_mask(input_ids.len())?; + build_i64_like_input( + desc, + vec![1_i64, attention_mask.len() as i64], + &attention_mask, + )? + } + "position_ids" => build_i64_like_input( + desc, + vec![3_i64, 1_i64, input_ids.len() as i64], + &position_ids, + )?, + "past_sequence_length" => build_i64_scalar_input(desc, self.past_seq_len() as i64)?, + name if name.starts_with("past_key_values.") => { + if let Some(cache) = self.cache_values.iter().find(|entry| entry.name == name) { + ort::value::Tensor::from_array((cache.dims.clone(), cache.data.clone()))? + .into_dyn() + } else { + build_zero_cache_input(desc)? + } + } + _ => build_zero_input(desc)?, + }; + decoder_inputs.push((desc.name.clone(), value)); + } + + let decoder_outputs = self.decoder_session.run(decoder_inputs)?; + let logits_value = decoder_outputs + .get("logits") + .or_else(|| { + self.decoder_output_names + .first() + .and_then(|name| decoder_outputs.get(name)) + }) + .ok_or_else(|| anyhow!("glm-ocr onnx output logits not found"))?; + let logits = extract_last_logits(logits_value)?; + + let mut new_cache_values = Vec::new(); + for desc in &self.decoder_input_descriptors { + let name = &desc.name; + if !name.starts_with("past_key_values.") { + continue; + } + let present_name = name.replace("past_key_values.", "present."); + let value = decoder_outputs + .get(&present_name) + .ok_or_else(|| anyhow!("missing glm-ocr onnx output {}", present_name))?; + let (shape, data) = value.try_extract_tensor::()?; + new_cache_values.push(GlmOcrOnnxCacheEntry { + name: name.clone(), + dims: shape.iter().copied().collect::>(), + data: data.to_vec(), + }); + } + self.cache_values = new_cache_values; + Ok(logits) + } + + fn embed_input_ids(&mut self, input_ids: &[u32]) -> Result<(Vec, usize)> { + let outputs = self.embed_session.run(vec![( + "input_ids".to_string(), + ort::value::Tensor::from_array(( + vec![1_i64, input_ids.len() as i64], + input_ids.iter().map(|id| *id as i64).collect::>(), + ))? + .into_dyn(), + )])?; + let embed_value = outputs + .get("inputs_embeds") + .ok_or_else(|| anyhow!("glm-ocr onnx output inputs_embeds not found"))?; + + if let Ok((shape, embed_data)) = embed_value.try_extract_tensor::() { + if shape.len() != 3 || shape[0] != 1 { + return Err(anyhow!( + "unexpected glm-ocr onnx inputs_embeds shape: {}", + shape + )); + } + return Ok((embed_data.to_vec(), shape[2] as usize)); + } + if let Ok((shape, embed_data)) = embed_value.try_extract_tensor::() { + if shape.len() != 3 || shape[0] != 1 { + return Err(anyhow!( + "unexpected glm-ocr onnx inputs_embeds shape: {}", + shape + )); + } + return Ok(( + embed_data + .iter() + .map(|value| value.to_f32()) + .collect::>(), + shape[2] as usize, + )); + } + Err(anyhow!( + "glm-ocr onnx inputs_embeds output must be a f32/f16 tensor" + )) + } + + fn apply_vision_embeds( + &mut self, + embed_data: &mut [f32], + hidden_size: usize, + image_mask: &Tensor, + pixel_values: &Tensor, + image_grid_thw: &Tensor, + ) -> Result<()> { + let (vision_embeds, vision_rows, vision_hidden) = + self.run_vision_encoder(pixel_values, image_grid_thw)?; + if vision_hidden != hidden_size { + return Err(anyhow!( + "glm-ocr onnx vision hidden size mismatch: vision={}, text={}", + vision_hidden, + hidden_size + )); + } + + let image_positions = image_mask + .squeeze(0)? + .to_dtype(DType::U8)? + .to_vec1::()? + .into_iter() + .enumerate() + .filter_map(|(idx, value)| (value == 1).then_some(idx)) + .collect::>(); + + if image_positions.len() != vision_rows { + return Err(anyhow!( + "glm-ocr onnx image token/vision embed mismatch: image_tokens={}, vision_embeds={}", + image_positions.len(), + vision_rows + )); + } + + for (row_idx, token_idx) in image_positions.into_iter().enumerate() { + let src_start = row_idx * hidden_size; + let dst_start = token_idx * hidden_size; + let src_end = src_start + hidden_size; + let dst_end = dst_start + hidden_size; + embed_data[dst_start..dst_end].copy_from_slice(&vision_embeds[src_start..src_end]); + } + Ok(()) + } + + fn run_vision_encoder( + &mut self, + pixel_values: &Tensor, + image_grid_thw: &Tensor, + ) -> Result<(Vec, usize, usize)> { + let raw_pixel_values_shape = pixel_values + .dims() + .iter() + .map(|dim| *dim as i64) + .collect::>(); + let pixel_values_data = pixel_values + .flatten_all()? + .to_dtype(DType::F32)? + .to_vec1::()?; + let raw_image_grid_shape = image_grid_thw + .dims() + .iter() + .map(|dim| *dim as i64) + .collect::>(); + let image_grid_data = image_grid_thw + .flatten_all()? + .to_vec1::()? + .into_iter() + .map(|value| value as i64) + .collect::>(); + + let mut vision_inputs = Vec::with_capacity(self.vision_input_descriptors.len()); + for desc in &self.vision_input_descriptors { + let value = match desc.name.as_str() { + "pixel_values" => { + let pixel_values_shape = + align_onnx_input_shape(desc, raw_pixel_values_shape.clone()); + build_float_input(desc, pixel_values_shape, &pixel_values_data)? + } + "image_grid_thw" => { + let image_grid_shape = + align_onnx_input_shape(desc, raw_image_grid_shape.clone()); + build_i64_like_input(desc, image_grid_shape, &image_grid_data)? + } + _ => build_zero_input(desc)?, + }; + vision_inputs.push((desc.name.clone(), value)); + } + + let outputs = self.vision_session.run(vision_inputs)?; + let output_value = outputs + .get("image_embeds") + .or_else(|| outputs.get("vision_embeds")) + .or_else(|| { + self.vision_output_names + .first() + .and_then(|name| outputs.get(name)) + }) + .ok_or_else(|| anyhow!("glm-ocr onnx vision output not found"))?; + + if let Ok((shape, values)) = output_value.try_extract_tensor::() { + let shape_vec = shape.iter().copied().collect::>(); + return extract_vision_output(shape_vec.as_slice(), values.to_vec()); + } + if let Ok((shape, values)) = output_value.try_extract_tensor::() { + let shape_vec = shape.iter().copied().collect::>(); + let values = values + .iter() + .map(|value| value.to_f32()) + .collect::>(); + return extract_vision_output(shape_vec.as_slice(), values); + } + Err(anyhow!( + "glm-ocr onnx vision output must be a f32/f16 tensor" + )) + } + + fn past_seq_len(&self) -> usize { + self.cache_values + .iter() + .find(|entry| entry.name.ends_with(".key")) + .map(|entry| entry.dims.get(2).copied().unwrap_or_default() as usize) + .unwrap_or(0) + } + + fn build_attention_mask(&self, current_seq_len: usize) -> Result> { + let total_len = self.past_seq_len().saturating_add(current_seq_len); + if total_len == 0 { + return Err(anyhow!("glm-ocr onnx attention length cannot be zero")); + } + Ok(vec![1_i64; total_len]) + } + + fn compute_prefill_position_ids( + &mut self, + image_mask: &Tensor, + grid_thw: &Tensor, + seq_len: usize, + ) -> Result> { + 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; + + 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; + + let mask_vec = image_mask + .squeeze(0)? + .to_dtype(DType::U8)? + .to_vec1::()?; + let mut t_ids = Vec::with_capacity(seq_len); + let mut h_ids = Vec::with_capacity(seq_len); + let mut w_ids = Vec::with_capacity(seq_len); + let mut st_idx: i64 = 0; + 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 { + if run_len != num_image_tokens { + return Err(anyhow!( + "glm-ocr onnx 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); + } + } + } + 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 { + 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; + } + } + + self.next_mrope_pos = st_idx as usize; + self.prefill_seq_len = seq_len; + + let mut out = Vec::with_capacity(seq_len * 3); + out.extend_from_slice(&t_ids); + out.extend_from_slice(&h_ids); + out.extend_from_slice(&w_ids); + Ok(out) + } +} + +#[cfg(feature = "onnx-runtime")] +fn build_text_position_ids(position_start: usize, seq_len: usize) -> Vec { + let positions = (position_start..position_start + seq_len) + .map(|idx| idx as i64) + .collect::>(); + let mut ids = Vec::with_capacity(seq_len * 3); + for _ in 0..3 { + ids.extend_from_slice(&positions); + } + ids +} + +#[cfg(feature = "onnx-runtime")] +fn align_onnx_input_shape(desc: &OnnxInputDescriptor, actual_shape: Vec) -> Vec { + if desc.shape.is_empty() || desc.shape.len() <= actual_shape.len() { + return actual_shape; + } + + let mut aligned = actual_shape; + while aligned.len() < desc.shape.len() { + aligned.insert(0, 1); + } + aligned +} + +#[cfg(feature = "onnx-runtime")] +fn build_i64_scalar_input(desc: &OnnxInputDescriptor, value: i64) -> Result { + match desc.kind { + Some(OnnxTensorKind::I32) => { + Ok(ort::value::Tensor::from_array((vec![1_i64], vec![value as i32]))?.into_dyn()) + } + Some(OnnxTensorKind::I64) | None => { + Ok(ort::value::Tensor::from_array((vec![1_i64], vec![value]))?.into_dyn()) + } + _ => Err(anyhow!( + "unsupported glm-ocr onnx scalar input dtype for {}", + desc.name + )), + } +} + +#[cfg(feature = "onnx-runtime")] +fn build_i64_like_input( + desc: &OnnxInputDescriptor, + shape: Vec, + data: &[i64], +) -> Result { + match desc.kind { + Some(OnnxTensorKind::I32) => Ok(ort::value::Tensor::from_array(( + shape, + data.iter().map(|value| *value as i32).collect::>(), + ))? + .into_dyn()), + Some(OnnxTensorKind::I64) | None => { + Ok(ort::value::Tensor::from_array((shape, data.to_vec()))?.into_dyn()) + } + Some(OnnxTensorKind::Bool) => Ok(ort::value::Tensor::from_array(( + shape, + data.iter().map(|value| *value != 0).collect::>(), + ))? + .into_dyn()), + _ => Err(anyhow!( + "unsupported glm-ocr onnx integer-like input dtype for {}", + desc.name + )), + } +} + +#[cfg(feature = "onnx-runtime")] +fn build_float_input( + desc: &OnnxInputDescriptor, + shape: Vec, + data: &[f32], +) -> Result { + match desc.kind { + Some(OnnxTensorKind::F16) => Ok(ort::value::Tensor::from_array(( + shape, + data.iter() + .map(|value| f16::from_f32(*value)) + .collect::>(), + ))? + .into_dyn()), + Some(OnnxTensorKind::F32) | None => { + Ok(ort::value::Tensor::from_array((shape, data.to_vec()))?.into_dyn()) + } + _ => Err(anyhow!( + "unsupported glm-ocr onnx float input dtype for {}", + desc.name + )), + } +} + +#[cfg(feature = "onnx-runtime")] +fn build_zero_cache_input(desc: &OnnxInputDescriptor) -> Result { + let shape = desc + .shape + .iter() + .enumerate() + .map(|(idx, dim)| { + if *dim >= 0 { + *dim + } else if idx == 2 { + 0 + } else { + 1 + } + }) + .collect::>(); + build_zero_input_with_shape(desc, shape) +} + +#[cfg(feature = "onnx-runtime")] +fn build_zero_input(desc: &OnnxInputDescriptor) -> Result { + let shape = desc + .shape + .iter() + .map(|dim| if *dim < 0 { 1 } else { *dim }) + .collect::>(); + build_zero_input_with_shape(desc, shape) +} + +#[cfg(feature = "onnx-runtime")] +fn build_zero_input_with_shape( + desc: &OnnxInputDescriptor, + shape: Vec, +) -> Result { + let kind = desc + .kind + .ok_or_else(|| anyhow!("unsupported glm-ocr onnx input dtype for {}", desc.name))?; + let elem_count = shape.iter().try_fold(1_usize, |acc, dim| { + if *dim < 0 { + Err(anyhow!( + "cannot resolve dynamic onnx shape for input {}: {:?}", + desc.name, + shape + )) + } else { + Ok(acc.saturating_mul(*dim as usize)) + } + })?; + let has_zero_dim = shape.contains(&0); + let make_ndarray = || { + let dims = shape.iter().map(|dim| *dim as usize).collect::>(); + IxDyn(&dims) + }; + let value = match kind { + OnnxTensorKind::Bool => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![false; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![false; elem_count]))?.into_dyn() + } + } + OnnxTensorKind::I32 => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![0_i32; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![0_i32; elem_count]))?.into_dyn() + } + } + OnnxTensorKind::I64 => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![0_i64; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![0_i64; elem_count]))?.into_dyn() + } + } + OnnxTensorKind::F16 => { + if has_zero_dim { + let arr = + Array::from_shape_vec(make_ndarray(), vec![f16::from_f32(0.0); elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![f16::from_f32(0.0); elem_count]))? + .into_dyn() + } + } + OnnxTensorKind::F32 => { + if has_zero_dim { + let arr = Array::from_shape_vec(make_ndarray(), vec![0_f32; elem_count])?; + ort::value::Tensor::from_array(arr)?.into_dyn() + } else { + ort::value::Tensor::from_array((shape, vec![0_f32; elem_count]))?.into_dyn() + } + } + }; + Ok(value) +} + +#[cfg(feature = "onnx-runtime")] +fn extract_last_logits(value: &ort::value::DynValue) -> Result> { + if let Ok((shape, values)) = value.try_extract_tensor::() { + return extract_last_logits_from_shape( + shape.iter().copied().collect::>().as_slice(), + values, + ); + } + if let Ok((shape, values)) = value.try_extract_tensor::() { + let values = values + .iter() + .map(|value| value.to_f32()) + .collect::>(); + return extract_last_logits_from_shape( + shape.iter().copied().collect::>().as_slice(), + &values, + ); + } + Err(anyhow!( + "glm-ocr onnx logits output must be a f32/f16 tensor" + )) +} + +#[cfg(feature = "onnx-runtime")] +fn extract_last_logits_from_shape(shape: &[i64], values: &[f32]) -> Result> { + match shape { + [1, seq_len, vocab_size] => { + let seq_len = *seq_len as usize; + let vocab_size = *vocab_size as usize; + let start = seq_len + .checked_sub(1) + .ok_or_else(|| anyhow!("glm-ocr onnx logits sequence is empty"))? + * vocab_size; + Ok(values[start..start + vocab_size].to_vec()) + } + [seq_len, vocab_size] => { + let seq_len = *seq_len as usize; + let vocab_size = *vocab_size as usize; + let start = seq_len + .checked_sub(1) + .ok_or_else(|| anyhow!("glm-ocr onnx logits sequence is empty"))? + * vocab_size; + Ok(values[start..start + vocab_size].to_vec()) + } + _ => Err(anyhow!("unexpected glm-ocr onnx logits shape: {:?}", shape)), + } +} + +#[cfg(feature = "onnx-runtime")] +fn extract_vision_output(shape: &[i64], values: Vec) -> Result<(Vec, usize, usize)> { + match shape { + [rows, hidden] => Ok((values, *rows as usize, *hidden as usize)), + [1, rows, hidden] => Ok((values, *rows as usize, *hidden as usize)), + _ => Err(anyhow!( + "unexpected glm-ocr onnx vision output shape: {:?}", + shape + )), + } +} + +#[cfg(all(test, feature = "onnx-runtime"))] +mod tests { + use super::{OnnxInputDescriptor, OnnxTensorKind, align_onnx_input_shape}; + + #[test] + fn align_onnx_input_shape_prepends_batch_dimension_when_descriptor_has_higher_rank() { + let desc = OnnxInputDescriptor { + name: "image_grid_thw".to_string(), + shape: vec![-1, 3], + kind: Some(OnnxTensorKind::I64), + }; + assert_eq!(align_onnx_input_shape(&desc, vec![3]), vec![1, 3]); + } +} + +#[cfg(not(feature = "onnx-runtime"))] +pub struct GlmOcrOnnxBackend; + +#[cfg(not(feature = "onnx-runtime"))] +impl GlmOcrOnnxBackend { + pub fn load(_onnx_path: &str, _spatial_merge_size: usize) -> Result { + Err(anyhow!( + "onnx runtime support is not enabled; rebuild with --features onnx-runtime" + )) + } + + pub fn clear_cache(&mut self) {} + + pub fn forward_logits( + &mut self, + _input_ids: &[u32], + _image_mask: Option<&Tensor>, + _pixel_values: Option<&Tensor>, + _image_grid_thw: Option<&Tensor>, + _position_start: usize, + ) -> Result> { + Err(anyhow!( + "onnx runtime support is not enabled; rebuild with --features onnx-runtime" + )) + } +} diff --git a/src/models/glm_ocr/processor.rs b/src/models/glm_ocr/processor.rs index fa324df..c0aa1ad 100644 --- a/src/models/glm_ocr/processor.rs +++ b/src/models/glm_ocr/processor.rs @@ -63,6 +63,30 @@ impl GlmOcrProcessor { }) } + pub fn from_params( + image_mean: Vec, + image_std: Vec, + shortest_edge: usize, + longest_edge: usize, + patch_size: usize, + merge_size: usize, + temporal_patch_size: usize, + device: &Device, + dtype: DType, + ) -> Self { + Self { + image_mean, + image_std, + shortest_edge, + longest_edge, + patch_size, + merge_size, + temporal_patch_size, + device: device.clone(), + dtype, + } + } + /// Process image for vision encoder. /// /// Matches Python's Glm46VImageProcessor._preprocess(): diff --git a/tests/test_glm_ocr_multi_format.rs b/tests/test_glm_ocr_multi_format.rs new file mode 100644 index 0000000..7666339 --- /dev/null +++ b/tests/test_glm_ocr_multi_format.rs @@ -0,0 +1,220 @@ +use std::path::{Path, PathBuf}; + +use aha::models::{ + ArtifactKind, GenerateModel, LoadSpec, ModelPaths, WhichModel, + glm_ocr::generate::GlmOcrGenerateModel, +}; +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::{Context, Result, anyhow}; + +const GLM_OCR_GGUF_DIR: &str = r"D:\model_download\GLM-OCR-GGUF"; +const GLM_OCR_ONNX_DIR: &str = r"D:\model_download\GLM-OCR-ONNX"; + +fn require_existing_dir(path: &str) -> Result<()> { + let dir = Path::new(path); + if !dir.exists() { + return Err(anyhow!("model dir not found: {}", path)); + } + if !dir.is_dir() { + return Err(anyhow!("path is not a directory: {}", path)); + } + Ok(()) +} + +fn first_file_with_extension(dir: &str, extension: &str) -> Result { + require_existing_dir(dir)?; + let mut matches = std::fs::read_dir(dir)? + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + path.is_file() + && path + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case(extension)) + }) + .collect::>(); + matches.sort(); + matches + .into_iter() + .next() + .ok_or_else(|| anyhow!("no .{} file found in {}", extension, dir)) +} + +fn first_file_with_extension_recursive(dir: &str, extension: &str) -> Result { + require_existing_dir(dir)?; + + let mut stack = vec![PathBuf::from(dir)]; + let mut matches = Vec::new(); + while let Some(current) = stack.pop() { + for entry in std::fs::read_dir(¤t)? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + stack.push(path); + continue; + } + if path + .extension() + .is_some_and(|ext| ext.eq_ignore_ascii_case(extension)) + { + matches.push(path); + } + } + } + + matches.sort(); + matches + .into_iter() + .next() + .ok_or_else(|| anyhow!("no .{} file found (recursive) in {}", extension, dir)) +} + +fn ocr_test_image_url() -> Result { + let path = std::env::current_dir()? + .join("assets") + .join("img") + .join("ocr_test1.png"); + let path = path.canonicalize()?; + url::Url::from_file_path(&path) + .map(|url| url.to_string()) + .map_err(|_| anyhow!("failed to build file url for {}", path.display())) +} + +#[test] +fn glm_ocr_gguf_files_can_load() -> Result<()> { + let gguf_path = first_file_with_extension(GLM_OCR_GGUF_DIR, "gguf")?; + let metadata = std::fs::metadata(&gguf_path) + .with_context(|| format!("failed to read gguf metadata: {}", gguf_path.display()))?; + if metadata.len() == 0 { + return Err(anyhow!("gguf file is empty: {}", gguf_path.display())); + } + Ok(()) +} + +#[test] +fn glm_ocr_onnx_files_can_load() -> Result<()> { + let onnx_path = first_file_with_extension_recursive(GLM_OCR_ONNX_DIR, "onnx")?; + let metadata = std::fs::metadata(&onnx_path) + .with_context(|| format!("failed to read onnx metadata: {}", onnx_path.display()))?; + if metadata.len() == 0 { + return Err(anyhow!("onnx file is empty: {}", onnx_path.display())); + } + Ok(()) +} + +#[test] +fn glm_ocr_load_spec_accepts_gguf() { + let spec = LoadSpec { + model: WhichModel::GlmOCR, + artifact: ArtifactKind::Gguf, + paths: ModelPaths { + gguf_path: Some(GLM_OCR_GGUF_DIR.to_string()), + ..Default::default() + }, + }; + spec.validate() + .expect("glm-ocr should accept gguf artifact"); +} + +#[test] +fn glm_ocr_load_spec_accepts_onnx() { + let spec = LoadSpec { + model: WhichModel::GlmOCR, + artifact: ArtifactKind::Onnx, + paths: ModelPaths { + onnx_path: Some(GLM_OCR_ONNX_DIR.to_string()), + tokenizer_dir: Some(GLM_OCR_ONNX_DIR.to_string()), + ..Default::default() + }, + }; + spec.validate() + .expect("glm-ocr should accept onnx artifact"); +} + +#[test] +fn glm_ocr_gguf_init_from_spec_can_init() -> Result<()> { + if let Err(err) = require_existing_dir(GLM_OCR_GGUF_DIR) { + println!("skip glm-ocr gguf init test: {err}"); + return Ok(()); + } + + let spec = LoadSpec { + model: WhichModel::GlmOCR, + artifact: ArtifactKind::Gguf, + paths: ModelPaths { + gguf_path: Some(GLM_OCR_GGUF_DIR.to_string()), + ..Default::default() + }, + }; + GlmOcrGenerateModel::init_from_spec(&spec, None, None)?; + Ok(()) +} + +#[test] +fn glm_ocr_onnx_init_from_spec_can_init() -> Result<()> { + if let Err(err) = require_existing_dir(GLM_OCR_ONNX_DIR) { + println!("skip glm-ocr onnx init test: {err}"); + return Ok(()); + } + + let spec = LoadSpec { + model: WhichModel::GlmOCR, + artifact: ArtifactKind::Onnx, + paths: ModelPaths { + onnx_path: Some(GLM_OCR_ONNX_DIR.to_string()), + tokenizer_dir: Some(GLM_OCR_ONNX_DIR.to_string()), + ..Default::default() + }, + }; + GlmOcrGenerateModel::init_from_spec(&spec, None, None)?; + Ok(()) +} + +#[test] +#[ignore = "manual real-model inference smoke test"] +fn glm_ocr_gguf_generate_smoke() -> Result<()> { + if let Err(err) = require_existing_dir(GLM_OCR_GGUF_DIR) { + println!("skip glm-ocr gguf generate smoke test: {err}"); + return Ok(()); + } + + let spec = LoadSpec { + model: WhichModel::GlmOCR, + artifact: ArtifactKind::Gguf, + paths: ModelPaths { + gguf_path: Some(GLM_OCR_GGUF_DIR.to_string()), + ..Default::default() + }, + }; + let ocr_test_image = ocr_test_image_url()?; + let message = format!( + r#"{{ + "model": "glm-ocr", + "messages": [ + {{ + "role": "user", + "content": [ + {{ + "type": "image", + "image_url": {{ + "url": "{ocr_test_image}" + }} + }}, + {{ + "type": "text", + "text": "Text Recognition:" + }} + ] + }} + ], + "max_tokens": 1 + }}"# + ); + let mes: ChatCompletionParameters = serde_json::from_str(&message)?; + let mut model = GlmOcrGenerateModel::init_from_spec(&spec, None, None)?; + let res = model.generate(mes)?; + if let Some(usage) = res.usage { + assert!(usage.total_tokens > 0); + } + Ok(()) +} diff --git a/tests/test_load_spec.rs b/tests/test_load_spec.rs index 2a659dd..0c275e7 100644 --- a/tests/test_load_spec.rs +++ b/tests/test_load_spec.rs @@ -48,6 +48,35 @@ fn load_spec_all_minilm_accepts_gguf() { assert!(spec.validate().is_ok()); } +#[test] +fn load_spec_glm_ocr_accepts_gguf() { + let spec = LoadSpec { + model: WhichModel::GlmOCR, + artifact: ArtifactKind::Gguf, + paths: ModelPaths { + gguf_path: Some("D:/model_download/GLM-OCR-GGUF".to_string()), + ..Default::default() + }, + }; + + assert!(spec.validate().is_ok()); +} + +#[test] +fn load_spec_glm_ocr_accepts_onnx() { + let spec = LoadSpec { + model: WhichModel::GlmOCR, + artifact: ArtifactKind::Onnx, + paths: ModelPaths { + onnx_path: Some("D:/model_download/GLM-OCR-ONNX".to_string()), + tokenizer_dir: Some("D:/model_download/GLM-OCR-ONNX".to_string()), + ..Default::default() + }, + }; + + assert!(spec.validate().is_ok()); +} + #[test] fn load_spec_gguf_requires_gguf_path() { let spec = LoadSpec {