From b881dfcd8d66c5ce35125f0b057d0b8f49896ebd Mon Sep 17 00:00:00 2001 From: jhqxxx <18280426169@163.com> Date: Fri, 27 Mar 2026 11:16:31 +0800 Subject: [PATCH] add lfm2vl --- src/models/lfm2/config.rs | 19 +- src/models/lfm2/model.rs | 9 +- src/models/lfm2vl/config.rs | 82 ++++++++ src/models/lfm2vl/generate.rs | 73 +++++++ src/models/lfm2vl/mod.rs | 4 + src/models/lfm2vl/model.rs | 0 src/models/lfm2vl/processor.rs | 337 +++++++++++++++++++++++++++++++++ src/models/mod.rs | 1 + src/utils/img_utils.rs | 75 ++++++-- tests/config_tests.rs | 18 +- tests/test_lfm2vl.rs | 54 ++++++ 11 files changed, 642 insertions(+), 30 deletions(-) create mode 100644 src/models/lfm2vl/config.rs create mode 100644 src/models/lfm2vl/generate.rs create mode 100644 src/models/lfm2vl/mod.rs create mode 100644 src/models/lfm2vl/model.rs create mode 100644 src/models/lfm2vl/processor.rs create mode 100644 tests/test_lfm2vl.rs diff --git a/src/models/lfm2/config.rs b/src/models/lfm2/config.rs index 2d97e31..1748525 100644 --- a/src/models/lfm2/config.rs +++ b/src/models/lfm2/config.rs @@ -1,6 +1,6 @@ use anyhow::{Result, anyhow}; -use serde::{Deserialize, Serialize}; -#[derive(Debug, PartialEq, Deserialize, Serialize)] + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] pub struct Lfm2Config { pub architectures: Vec, pub block_auto_adjust_ff_dim: bool, @@ -13,7 +13,7 @@ pub struct Lfm2Config { pub block_out_init_scale: f64, pub block_use_swiglu: bool, pub block_use_xavier_init: bool, - pub bos_token_id: u32, + pub bos_token_id: Option, #[serde[rename="conv_L_cache"]] pub conv_l_cache: usize, pub conv_bias: bool, @@ -33,8 +33,9 @@ pub struct Lfm2Config { pub num_heads: usize, pub num_hidden_layers: usize, pub num_key_value_heads: usize, - pub pad_token_id: u32, - pub rope_theta: f32, + pub pad_token_id: Option, + pub rope_theta: Option, + pub rope_parameters: Option, pub torch_dtype: Option, pub dtype: Option, pub use_cache: bool, @@ -43,6 +44,12 @@ pub struct Lfm2Config { pub tie_embedding: Option, } +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct RopeParameters { + pub rope_theta: f32, + pub rope_type: String, +} + impl Lfm2Config { pub fn full_attn_idx2layer_type(&mut self) { if self.layer_types.is_none() @@ -81,7 +88,7 @@ impl Lfm2Config { } } -#[derive(Debug, PartialEq, Deserialize, Serialize)] +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] pub struct Lfm2GenerateConfig { pub bos_token_id: u32, pub eos_token_id: u32, diff --git a/src/models/lfm2/model.rs b/src/models/lfm2/model.rs index 8dd13c7..7e25a33 100644 --- a/src/models/lfm2/model.rs +++ b/src/models/lfm2/model.rs @@ -217,7 +217,14 @@ impl Lfm2Decoder { layers.push(layer); } let dim = config.hidden_size / config.num_attention_heads; - let pos_emb = RoPE::new(dim, config.rope_theta, vb.device())?; + let theta_base = if let Some(theta) = config.rope_theta { + theta + } else if let Some(param) = &config.rope_parameters { + param.rope_theta + } else { + 1000000.0 + }; + let pos_emb = RoPE::new(dim, theta_base, vb.device())?; let embedding_norm = rms_norm(config.hidden_size, config.norm_eps, vb.pp("embedding_norm"))?; Ok(Self { diff --git a/src/models/lfm2vl/config.rs b/src/models/lfm2vl/config.rs new file mode 100644 index 0000000..6d6872e --- /dev/null +++ b/src/models/lfm2vl/config.rs @@ -0,0 +1,82 @@ +use candle_nn::Activation; + +use crate::models::lfm2::config::Lfm2Config; + + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Lfm2VLConfig { + pub do_image_splitting: bool, + pub downsample_factor: usize, + pub dtype: String, + pub encoder_patch_size: usize, + pub image_token_id: u32, + pub max_image_tokens: usize, + pub max_pixels_tolerance: f64, + pub max_tiles: usize, + pub min_image_tokens: usize, + pub min_tiles: usize, + pub model_type: String, + pub projector_bias: bool, + pub projector_hidden_act: Activation, + pub projector_hidden_size: usize, + pub projector_use_layernorm: bool, + pub text_config: Lfm2Config, + pub tile_size: usize, + pub use_image_special_tokens: bool, + pub use_thumbnail: bool, + pub vision_config: Lfm2VLVisionConfig, +} + + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Lfm2VLVisionConfig { + pub attention_dropout: f64, + pub dtype: String, + pub hidden_act: Activation, + pub hidden_size: usize, + pub intermediate_size: usize, + pub layer_norm_eps: f64, + pub model_type: String, + pub num_attention_heads: u32, + pub num_channels: u32, + pub num_hidden_layers: usize, + pub num_patches: usize, + pub patch_size: usize, + pub vision_use_head: bool, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Lfm2ImageConfig { + pub do_image_splitting: bool, + pub do_normalize: bool, + pub do_pad: bool, + pub do_rescale: bool, + pub do_resize: bool, + pub downsample_factor: usize, + pub encoder_patch_size: usize, + pub image_mean: Vec, + pub image_std: Vec, + pub max_image_tokens: usize, + pub max_num_patches: usize, + pub max_pixels_tolerance: f64, + pub max_tiles: usize, + pub min_image_tokens: usize, + pub min_tiles: usize, + pub resample: usize, + pub rescale_factor: f64, + pub size: Size, + pub tile_size: usize, + pub use_thumbnail: bool, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Size { + pub height: usize, + pub width: usize, +} + +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct Lfm2ProcessorConfig { + pub image_processor: Lfm2ImageConfig, +} + diff --git a/src/models/lfm2vl/generate.rs b/src/models/lfm2vl/generate.rs new file mode 100644 index 0000000..cc3ba78 --- /dev/null +++ b/src/models/lfm2vl/generate.rs @@ -0,0 +1,73 @@ +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device}; + +use crate::{ + chat_template::ChatTemplate, + models::{ + lfm2::config::Lfm2GenerateConfig, + lfm2vl::{config::Lfm2VLConfig, processor::Lfm2VLProcessor}, + }, + tokenizer::TokenizerModel, + utils::{find_type_files, get_device, get_dtype, get_logit_processor}, +}; + +pub struct Lfm2VLGenerateModel<'a> { + chat_template: ChatTemplate<'a>, + tokenizer: TokenizerModel, + device: Device, + // model: Lfm2VLModel, + processor: Lfm2VLProcessor, + eos_token_id: u32, + model_name: String, +} +impl<'a> Lfm2VLGenerateModel<'a> { + pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { + let chat_template = ChatTemplate::init(path)?; + let tokenizer = TokenizerModel::init(path)?; + let device = get_device(device); + let gen_cfg_path = path.to_string() + "/generation_config.json"; + let gen_cfg: Lfm2GenerateConfig = serde_json::from_slice(&std::fs::read(gen_cfg_path)?)?; + let cfg_path = path.to_string() + "/config.json"; + let cfg: Lfm2VLConfig = serde_json::from_slice(&std::fs::read(cfg_path)?)?; + + let model_path = find_type_files(path, "safetensors")?; + let dtype = get_dtype(dtype, &cfg.dtype); + // let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_path, dtype, &device)? }; + // let model = Lfm2Model::new(vb, &cfg)?; + let processor = Lfm2VLProcessor::new(path, dtype, &device)?; + let eos_token_id = gen_cfg.eos_token_id; + let model_name = std::path::Path::new(path) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or("lfm2.5-vl") + .to_string(); + Ok(Self { + chat_template, + tokenizer, + device, + // model, + processor, + eos_token_id, + model_name, + }) + } + + pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> { + let mes_render = self.chat_template.apply_chat_template(&mes)?; + let mut logits = get_logit_processor( + mes.temperature, + mes.top_p, + None, + mes.seed.unwrap_or(34562) as u64, + ); + let (pixel_values, pixel_attention_mask, spatial_shapes, text) = + self.processor.process_info(&mes, &mes_render)?; + let input_ids = self.tokenizer.text_encode(text, &self.device)?; + println!("pixel_values: {}", pixel_values); + println!("pixel_attention_mask: {}", pixel_attention_mask); + println!("spatial_shapes: {}", spatial_shapes); + println!("input_ids: {}", input_ids); + Ok(()) + } +} diff --git a/src/models/lfm2vl/mod.rs b/src/models/lfm2vl/mod.rs new file mode 100644 index 0000000..ce9b2a0 --- /dev/null +++ b/src/models/lfm2vl/mod.rs @@ -0,0 +1,4 @@ +pub mod config; +pub mod generate; +pub mod model; +pub mod processor; \ No newline at end of file diff --git a/src/models/lfm2vl/model.rs b/src/models/lfm2vl/model.rs new file mode 100644 index 0000000..e69de29 diff --git a/src/models/lfm2vl/processor.rs b/src/models/lfm2vl/processor.rs new file mode 100644 index 0000000..327020e --- /dev/null +++ b/src/models/lfm2vl/processor.rs @@ -0,0 +1,337 @@ +use crate::{ + models::lfm2vl::config::{Lfm2ImageConfig, Lfm2ProcessorConfig}, + utils::{ + img_utils::{ + crop_img, extract_images, find_closest_aspect_ratio, generate_target_ratios_sorted, + img_smart_resize, img_transform, + }, + round_by_factor, + }, +}; +use aha_openai_dive::v1::resources::chat::ChatCompletionParameters; +use anyhow::Result; +use candle_core::{DType, Device, Tensor}; +use image::DynamicImage; + +pub struct Lfm2VLProcessor { + dtype: DType, + device: Device, + image_config: Lfm2ImageConfig, + max_num_patches: usize, + total_factor: u32, + max_pixel_num: usize, + smart_resize_min_pixels: usize, + smart_resize_max_pixels: usize, + target_ratios: Vec<(u32, u32)>, + img_mean: Tensor, + img_std: Tensor, + tokens_per_tile: usize, + image_token: String, + // image_token_id: u32, + image_start_token: String, + image_end_token: String, + image_thumbnail_token: String, +} + +impl Lfm2VLProcessor { + pub fn new(path: &str, dtype: DType, device: &Device) -> Result { + let path = path.to_string(); + assert!( + std::path::Path::new(&path).exists(), + "model path file not exists" + ); + let processor_cfg_path = path + "/processor_config.json"; + let processor_cfg: Lfm2ProcessorConfig = + serde_json::from_slice(&std::fs::read(processor_cfg_path)?)?; + let image_config = processor_cfg.image_processor; + // 256 + let max_thumbnail_image_patches = + image_config.max_image_tokens * image_config.downsample_factor.pow(2); + // 1024 + let tile_size_patches = if image_config.do_image_splitting { + (image_config.tile_size / image_config.encoder_patch_size).pow(2) + } else { + 1 + }; + // 1024 + let max_num_patches = max_thumbnail_image_patches.max(tile_size_patches); + let total_factor = + (image_config.encoder_patch_size * image_config.downsample_factor) as u32; + let token_pixels = + image_config.encoder_patch_size.pow(2) * image_config.downsample_factor.pow(2); + let max_pixel_num = ((image_config.max_image_tokens * token_pixels) as f64 + * image_config.max_pixels_tolerance) as usize; + + let smart_resize_min_pixels = image_config.min_image_tokens * token_pixels; + let smart_resize_max_pixels = image_config.max_image_tokens * token_pixels; + let target_ratios = generate_target_ratios_sorted( + image_config.min_tiles as u32, + image_config.max_tiles as u32, + ); + let img_mean = + Tensor::from_slice(&image_config.image_mean, (3, 1, 1), device)?.to_dtype(dtype)?; + let img_std = + Tensor::from_slice(&image_config.image_std, (3, 1, 1), device)?.to_dtype(dtype)?; + + let tokens_per_tile = (image_config.tile_size + / image_config.encoder_patch_size + / image_config.downsample_factor) + .pow(2); + + Ok(Self { + dtype, + device: device.clone(), + image_config, + max_num_patches, + total_factor, + max_pixel_num, + smart_resize_min_pixels, + smart_resize_max_pixels, + target_ratios, + img_mean, + img_std, + tokens_per_tile, + image_token: "".to_string(), + // image_token_id: 396, + image_start_token: "<|image_start|>".to_string(), + image_end_token: "<|image_end|>".to_string(), + image_thumbnail_token: "<|img_thumbnail|>".to_string(), + }) + } + + fn is_image_too_large(&self, height: u32, width: u32) -> bool { + let h_bar = self + .image_config + .encoder_patch_size + .max(round_by_factor(height, self.total_factor) as usize); + let w_bar = self + .image_config + .encoder_patch_size + .max(round_by_factor(width, self.total_factor) as usize); + h_bar * w_bar > self.max_pixel_num + } + + fn get_grid_layout(&self, height: u32, width: u32) -> (u32, u32) { + let aspect_ratio = width as f64 / height as f64; + let (grid_width, grid_height) = find_closest_aspect_ratio( + aspect_ratio, + &self.target_ratios, + width, + height, + self.image_config.tile_size as u32, + ); + (grid_width, grid_height) + } + + fn crop_image_to_patches( + &self, + img: &DynamicImage, + height: u32, + width: u32, + new_height: u32, + new_width: u32, + ) -> Result<(Vec, usize, usize)> { + let (grid_width, grid_height) = self.get_grid_layout(height, width); + let mut processed_images = crop_img( + img, + grid_height, + grid_width, + self.image_config.tile_size as u32, + ); + if self.image_config.use_thumbnail && processed_images.len() != 1 { + let thumbnail_img = img.resize_exact( + new_width, + new_height, + image::imageops::FilterType::CatmullRom, + ); + processed_images.push(thumbnail_img); + } + Ok((processed_images, grid_width as usize, grid_height as usize)) + } + + fn resize_and_split( + &self, + img: &DynamicImage, + ) -> Result<(Vec, usize, usize, u32, u32)> { + let height = img.height(); + let width = img.width(); + let is_image_large = self.is_image_too_large(height, width); + + let (new_height, new_width) = img_smart_resize( + height, + width, + self.total_factor as u32, + self.smart_resize_min_pixels as u32, + self.smart_resize_max_pixels as u32, + )?; + let (images, num_cols, num_rows) = if is_image_large && self.image_config.do_image_splitting + { + self.crop_image_to_patches(img, height, width, new_height, new_width)? + } else { + let img = img.resize_exact( + new_width, + new_height, + image::imageops::FilterType::CatmullRom, + ); + (vec![img], 1, 1) + }; + + Ok((images, num_cols, num_rows, new_height, new_width)) + } + + pub fn process_imgs( + &self, + imgs: Vec, + ) -> Result<( + Tensor, + Tensor, + Tensor, + Vec, + Vec, + Vec<(u32, u32)>, + )> { + let patch_size = self.image_config.encoder_patch_size; + let mut images_list = vec![]; + let mut images_mask_list = vec![]; + let mut processed_spatial_shapes = vec![]; + let mut num_cols_list = vec![]; + let mut num_rows_list = vec![]; + let mut image_size_list = vec![]; + for img in &imgs { + // img: 过大的切分，返回切块图像 + reshape图像 + // 小的直接reshape + let (imgs, num_cols, num_rows, new_height, new_width) = self.resize_and_split(img)?; + num_cols_list.push(num_cols); + num_rows_list.push(num_rows); + image_size_list.push((new_height, new_width)); + for img in imgs { + // img-> tensor + let img_t = img_transform( + &img, + &self.img_mean, + &self.img_std, + &self.device, + self.dtype, + )?; + + // 图像嵌入 -> (seq_len, embedding_size) + let (c, h, w) = img_t.dims3()?; + let num_patches_height = h / patch_size; + let num_patches_width = w / patch_size; + let patched_image = img_t.reshape(( + c, + num_patches_height, + patch_size, + num_patches_width, + patch_size, + ))?; + // (c, num_patches_height, patch_size, num_patches_width, patch_size) + // -> (num_patches_height, num_patches_width, patch_size, patch_size, c) + let patched_image = patched_image.permute((1, 3, 2, 4, 0))?; + let patched_image = + patched_image.reshape((num_patches_height * num_patches_width, ()))?; + + // padding + let curren_length = patched_image.dim(0)?; + let padding_length = self.max_num_patches - curren_length; + let (patched_image, pixel_mask) = if self.image_config.do_pad && padding_length > 0 + { + let mut pixel_mask = Tensor::ones(curren_length, DType::U32, &self.device)?; + let padding_image = patched_image.pad_with_zeros(0, 0, padding_length)?; + let pad = Tensor::zeros(padding_length, DType::U32, &self.device)?; + pixel_mask = Tensor::cat(&[&pixel_mask, &pad], 0)?; + (padding_image, pixel_mask) + } else { + let pixel_mask = Tensor::ones(curren_length, DType::U32, &self.device)?; + (patched_image, pixel_mask) + }; + images_list.push(patched_image); + images_mask_list.push(pixel_mask); + processed_spatial_shapes + .push(vec![num_patches_height as u32, num_patches_width as u32]); + } + } + let pixel_values = Tensor::stack(&images_list, 0)?; + let pixel_attention_mask = Tensor::stack(&images_mask_list, 0)?; + let spatial_shapes = Tensor::new(processed_spatial_shapes, &self.device)?; + Ok(( + pixel_values, + pixel_attention_mask, + spatial_shapes, + num_cols_list, + num_rows_list, + image_size_list, + )) + } + + fn build_image_tokens(&self, rows: usize, cols: usize, tokens_for_image: usize) -> String { + let mut parts = "".to_string(); + parts += &self.image_start_token; + if rows > 1 && cols > 1 { + for row in 0..rows { + for col in 0..cols { + parts += &format!("<|img_row_{}_col_{}|>", row + 1, col + 1); + parts += &(self.image_token.repeat(self.tokens_per_tile)); + } + } + if self.image_config.use_thumbnail { + parts += &self.image_thumbnail_token; + parts += &(self.image_token.repeat(tokens_for_image)); + } + } else { + parts += &(self.image_token.repeat(tokens_for_image)); + } + parts += &self.image_end_token; + parts + } + + fn expand_text_with_placeholders( + &self, + text: &str, + num_cols_list: Vec, + num_rows_list: Vec, + image_size_list: Vec<(u32, u32)>, + ) -> String { + let text_parts: Vec<&str> = text.split(&self.image_token).collect(); + let mut result_parts = "".to_string(); + for i in 0..num_cols_list.len() { + result_parts += text_parts[i]; + let rows = num_rows_list[i]; + let cols = num_cols_list[i]; + let image_size = image_size_list[i]; + let (h, w) = image_size; + let tokens_for_image = (h as usize + / self.image_config.encoder_patch_size + / self.image_config.downsample_factor) + * (w as usize + / self.image_config.encoder_patch_size + / self.image_config.downsample_factor); + let sub_str = self.build_image_tokens(rows, cols, tokens_for_image); + result_parts += &sub_str; + } + if text_parts.len() > num_cols_list.len() { + result_parts += text_parts[text_parts.len() - 1]; + } + result_parts + } + + pub fn process_info( + &self, + messages: &ChatCompletionParameters, + text: &str, + ) -> Result<(Tensor, Tensor, Tensor, String)> { + let imgs = extract_images(messages)?; + let ( + pixel_values, + pixel_attention_mask, + spatial_shapes, + num_cols_list, + num_rows_list, + image_size_list, + ) = self.process_imgs(imgs)?; + let text = + self.expand_text_with_placeholders(text, num_cols_list, num_rows_list, image_size_list); + + Ok((pixel_values, pixel_attention_mask, spatial_shapes, text)) + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 625784e..c15debb 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -8,6 +8,7 @@ pub mod glm_asr_nano; pub mod glm_ocr; pub mod hunyuan_ocr; pub mod lfm2; +pub mod lfm2vl; pub mod mask_gct; pub mod minicpm4; pub mod paddleocr_vl; diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs index bd0d5c0..0e334d7 100644 --- a/src/utils/img_utils.rs +++ b/src/utils/img_utils.rs @@ -110,6 +110,7 @@ pub fn extract_images(mes: &ChatCompletionParameters) -> Result pub fn generate_target_ratios_sorted(min_num: u32, max_num: u32) -> Vec<(u32, u32)> { let mut target_ratios = HashSet::new(); @@ -130,6 +131,7 @@ pub fn generate_target_ratios_sorted(min_num: u32, max_num: u32) -> Vec<(u32, u3 sorted_ratios } +/// return (grid_width, grid_height) pub fn find_closest_aspect_ratio( aspect_ratio: f64, target_ratios: &[(u32, u32)], @@ -160,6 +162,34 @@ pub fn find_closest_aspect_ratio( best_ratio } +pub fn crop_img( + image: &DynamicImage, + grid_height: u32, + grid_width: u32, + image_size: u32, +) -> Vec { + let target_width = image_size * grid_width; + let target_height = image_size * grid_height; + let blocks = grid_width * grid_height; + let mut resized_img = image.resize_exact( + target_width, + target_height, + image::imageops::FilterType::CatmullRom, + ); + let mut processed_images = Vec::new(); + for i in 0..blocks { + // Calculate box coordinates + let x1 = (i % grid_width) * image_size; + let y1 = (i / grid_width) * image_size; + + // Crop the image + let split_img = resized_img.crop(x1, y1, image_size, image_size); + processed_images.push(split_img); + } + assert_eq!(processed_images.len() as u32, blocks); + processed_images +} + pub fn dynamic_preprocess( image: &DynamicImage, min_num: u32, @@ -179,26 +209,32 @@ pub fn dynamic_preprocess( orig_height, image_size, ); - let target_width = image_size * target_aspect_ratio.0; - let target_height = image_size * target_aspect_ratio.1; - let blocks = target_aspect_ratio.0 * target_aspect_ratio.1; - let mut resized_img = image.resize_exact( - target_width, - target_height, - image::imageops::FilterType::CatmullRom, - ); - let mut processed_images = Vec::new(); - let grid_width = target_width / image_size; - for i in 0..blocks { - // Calculate box coordinates - let x1 = (i % grid_width) * image_size; - let y1 = (i / grid_width) * image_size; + // let target_width = image_size * target_aspect_ratio.0; + // let target_height = image_size * target_aspect_ratio.1; + // let blocks = target_aspect_ratio.0 * target_aspect_ratio.1; + // let mut resized_img = image.resize_exact( + // target_width, + // target_height, + // image::imageops::FilterType::CatmullRom, + // ); + // let mut processed_images = Vec::new(); + // let grid_width = target_aspect_ratio.0; + // for i in 0..blocks { + // // Calculate box coordinates + // let x1 = (i % grid_width) * image_size; + // let y1 = (i / grid_width) * image_size; - // Crop the image - let split_img = resized_img.crop(x1, y1, image_size, image_size); - processed_images.push(split_img); - } - assert_eq!(processed_images.len() as u32, blocks); + // // Crop the image + // let split_img = resized_img.crop(x1, y1, image_size, image_size); + // processed_images.push(split_img); + // } + // assert_eq!(processed_images.len() as u32, blocks); + let mut processed_images = crop_img( + image, + target_aspect_ratio.1, + target_aspect_ratio.0, + image_size, + ); if use_thumbnail && processed_images.len() != 1 { let thumbnail_img = image.resize_exact( @@ -257,6 +293,7 @@ pub fn img_transform( Ok(img_tensor) } +/// return (height, width) pub fn img_smart_resize( img_h: u32, img_w: u32, diff --git a/tests/config_tests.rs b/tests/config_tests.rs index e78645f..62ca94b 100644 --- a/tests/config_tests.rs +++ b/tests/config_tests.rs @@ -1,8 +1,5 @@ use aha::models::{ - deepseek_ocr::config::DeepseekOCRConfig, hunyuan_ocr::config::HunYuanVLConfig, - lfm2::config::Lfm2Config, minicpm4::config::MiniCPM4Config, - paddleocr_vl::config::PaddleOCRVLConfig, qwen2_5vl::config::Qwen2_5VLConfig, - qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig, + deepseek_ocr::config::DeepseekOCRConfig, hunyuan_ocr::config::HunYuanVLConfig, lfm2::config::Lfm2Config, lfm2vl::config::{Lfm2ProcessorConfig, Lfm2VLConfig}, minicpm4::config::MiniCPM4Config, paddleocr_vl::config::PaddleOCRVLConfig, qwen2_5vl::config::Qwen2_5VLConfig, qwen3vl::config::Qwen3VLConfig, voxcpm::config::VoxCPMConfig }; use anyhow::Result; @@ -98,3 +95,16 @@ fn lfm2_config() -> Result<()> { println!("{:?}", config); Ok(()) } + +#[test] +fn lfm2vl_config() -> Result<()> { + // cargo test -F cuda --test config_tests lfm2vl_config -r -- --nocapture + let model_path = "/home/jhq/.aha/LiquidAI/LFM2.5-VL-1.6B/"; + let config_path = model_path.to_string() + "/config.json"; + let config: Lfm2VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + println!("{:?}", config); + let processor_config_path = model_path.to_string() + "/processor_config.json"; + let processor_config: Lfm2ProcessorConfig = serde_json::from_slice(&std::fs::read(processor_config_path)?)?; + println!("{:?}", processor_config); + Ok(()) +} \ No newline at end of file diff --git a/tests/test_lfm2vl.rs b/tests/test_lfm2vl.rs new file mode 100644 index 0000000..9665d5d --- /dev/null +++ b/tests/test_lfm2vl.rs @@ -0,0 +1,54 @@ +use std::time::Instant; + +use aha::{chat::ChatCompletionParameters, models::lfm2vl::generate::Lfm2VLGenerateModel}; +use anyhow::Result; +#[test] +fn lfm2vl_generate() -> Result<()> { + // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda --test test_lfm2vl lfm2vl_generate -r -- --nocapture + + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/LiquidAI/LFM2.5-VL-1.6B/", save_dir); + let message = r#" + { + "model": "lfm2vl", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "image_url": + { + "url": "file://./assets/img/ocr_test1.png" + } + }, + { + "type": "text", + "text": "请分析图片并提取所有可见文本内容，按从左到右、从上到下的布局，返回纯文本" + } + ] + } + ] + } + "#; + let mes: ChatCompletionParameters = serde_json::from_str(message)?; + let i_start = Instant::now(); + let mut model = Lfm2VLGenerateModel::init(&model_path, None, None)?; + let i_duration = i_start.elapsed(); + println!("Time elapsed in load model is: {:?}", i_duration); + + let i_start = Instant::now(); + let result = model.generate(mes)?; + let i_duration = i_start.elapsed(); + // println!("generate: \n {:?}", result); + // if let Some(usage) = &result.usage { + // let num_token = usage.total_tokens; + // let duration_secs = i_duration.as_secs_f64(); + // let tps = num_token as f64 / duration_secs; + // println!("Tokens per second (TPS): {:.2}", tps); + // } + // println!("Time elapsed in generate is: {:?}", i_duration); + + Ok(()) +} \ No newline at end of file