add lfm2vl

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<String>,
     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<u32>,
     #[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<u32>,
+    pub rope_theta: Option<f32>,
+    pub rope_parameters: Option<RopeParameters>,
     pub torch_dtype: Option<String>,
     pub dtype: Option<String>,
     pub use_cache: bool,
@@ -43,6 +44,12 @@ pub struct Lfm2Config {
     pub tie_embedding: Option<bool>,
 }
 
+#[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,

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 {

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<f64>,
+    pub image_std: Vec<f64>,
+    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,
+}
+

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<DType>) -> Result<Self> {
+        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(())
+    }
+}

src/models/lfm2vl/mod.rs

@@ -0,0 +1,4 @@
+pub mod config;
+pub mod generate;
+pub mod model;
+pub mod processor;

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<Self> {
+        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: "<image>".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<DynamicImage>, 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<DynamicImage>, 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<DynamicImage>,
+    ) -> Result<(
+        Tensor,
+        Tensor,
+        Tensor,
+        Vec<usize>,
+        Vec<usize>,
+        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<usize>,
+        num_rows_list: Vec<usize>,
+        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))
+    }
+}

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;

src/utils/img_utils.rs

@@ -110,6 +110,7 @@ pub fn extract_images(mes: &ChatCompletionParameters) -> Result<Vec<DynamicImage
     img_url_vec.par_iter().map(|url| get_image(url)).collect()
 }
 
+/// return vec<(grid_width, grid_height)>
 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<DynamicImage> {
+    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,