1#![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
2
3use std::{any::Any, num::NonZeroUsize, sync::Arc};
4
5use candle_core::{Device, Result, Tensor};
6use image::{imageops::FilterType, DynamicImage, GenericImageView};
7use mistralrs_vision::{ApplyTransforms, Normalize, Rescale, ToTensorNoNorm, Transforms};
8use tokenizers::Tokenizer;
9use tracing::warn;
10
11use crate::{
12 device_map::DeviceMapper,
13 pipeline::{
14 text_models_inputs_processor::{
15 self, get_completion_input, get_prompt_input, PagedAttentionMeta,
16 },
17 InputProcessorOutput, InputsProcessor, InputsProcessorType, MessagesAction, Processor,
18 },
19 sequence::Sequence,
20 vision_models::{
21 image_processor::{ImagePreProcessor, PreprocessedImages},
22 preprocessor_config::{PreProcessorConfig, ToFilter},
23 processor_config::ProcessorConfig,
24 ModelInputs,
25 },
26};
27
28use super::Mistral3SpecificArgs;
29
30const PLACEHOLDER: &str = "<placeholder>";
31
32struct Mistral3ImageProcessor {
33 image_break_token: String,
34 image_end_token: String,
35 image_token: String,
36 patch_size: usize,
37 spatial_merge_size: usize,
38}
39
40pub struct Mistral3Processor {
41 image_break_token: String,
42 image_end_token: String,
43 image_token: String,
44 patch_size: usize,
45 spatial_merge_size: usize,
46}
47
48impl Mistral3Processor {
49 pub fn new(processor_config: ProcessorConfig) -> Self {
50 Self {
51 image_break_token: processor_config.image_break_token.unwrap().clone(),
52 image_end_token: processor_config.image_end_token.unwrap().clone(),
53 image_token: processor_config.image_token.unwrap().clone(),
54 patch_size: processor_config.patch_size.unwrap(),
55 spatial_merge_size: processor_config.spatial_merge_size.unwrap(),
56 }
57 }
58}
59
60impl Processor for Mistral3Processor {
61 fn inputs_processor(&self) -> Arc<dyn InputsProcessor> {
62 Arc::new(Mistral3ImageProcessor {
63 image_break_token: self.image_break_token.clone(),
64 image_end_token: self.image_end_token.clone(),
65 image_token: self.image_token.clone(),
66 patch_size: self.patch_size,
67 spatial_merge_size: self.spatial_merge_size,
68 })
69 }
70
71 fn get_special_tokens(&self) -> &[&'static str] {
72 &[]
73 }
74
75 fn template_action(&self) -> MessagesAction {
76 MessagesAction::Keep
77 }
78}
79
80impl InputsProcessor for Mistral3ImageProcessor {
81 fn get_type(&self) -> InputsProcessorType {
82 InputsProcessorType::Vision
83 }
84 fn process_inputs(
85 &self,
86 tokenizer: Option<Arc<Tokenizer>>,
87 input_seqs: &mut [&mut Sequence],
88 is_prompt: bool,
89 is_xlora: bool,
90 device: &Device,
91 no_kv_cache: bool,
92 last_n_context_len: Option<(usize, usize)>,
93 return_raw_logits: bool,
94 other_config: Option<Arc<dyn Any>>,
95 mut paged_attn_metadata: Option<PagedAttentionMeta<'_>>,
96 prompt_chunksize: Option<NonZeroUsize>,
97 mapper: Option<&dyn DeviceMapper>,
98 ) -> Box<dyn Iterator<Item = anyhow::Result<InputProcessorOutput>>> {
99 if is_xlora {
100 return Box::new(std::iter::once(Err(anyhow::Error::msg(
101 "Cannot make inputs for X-LoRA vision model.",
102 ))));
103 }
104 if no_kv_cache {
105 return Box::new(std::iter::once(Err(anyhow::Error::msg(
106 "Vision model must have kv cache.",
107 ))));
108 }
109 if prompt_chunksize.is_some() {
111 warn!("`prompt_chunksize` is set. Mistral3 does not support prompt batching.");
112 }
113 let Some(tokenizer) = tokenizer else {
114 return Box::new(std::iter::once(Err(anyhow::Error::msg(
115 "Idefics3ImageProcessor requires a specified tokenizer.",
116 ))));
117 };
118
119 let config = other_config.expect("Need a PreProcessorConfig config.");
120 let config: &PreProcessorConfig = config.downcast_ref().expect("Downcast failed.");
121
122 let has_images = input_seqs.iter().all(|seq| seq.has_images());
123
124 let (pixel_values, image_sizes) = if has_images {
125 let mut pixel_values_accum = Vec::new();
126 let mut image_sizes_accum = Vec::new();
127
128 for seq in input_seqs.iter_mut() {
129 let PreprocessedImages {
130 pixel_values,
131 pixel_attention_mask: _,
132 image_sizes: _,
133 num_img_tokens: _,
134 aspect_ratio_ids: _,
135 aspect_ratio_mask: _,
136 num_tiles: _,
137 image_grid_thw: _,
138 video_grid_thw: _,
139 rows: _,
140 cols: _,
141 pixel_values_list: _,
142 tgt_sizes: _,
143 image_sizes_all,
144 num_crops: _,
145 } = self
146 .preprocess(
147 seq.take_images()
148 .expect("Need to have images by this point."),
149 vec![],
150 config,
151 device,
152 (usize::MAX, usize::MAX), )
154 .expect("Preprocessing failed");
155 let image_sizes_all = image_sizes_all.unwrap();
156
157 pixel_values_accum.push(pixel_values.clone());
159 image_sizes_accum.extend_from_slice(&image_sizes_all);
160
161 let mut prompt = tokenizer
162 .decode(seq.get_toks(), false)
163 .expect("Detokenization failed!");
164
165 let mut image_sizes_all_iter = image_sizes_all.into_iter();
166 let mut replace_strings = Vec::new();
167 while prompt.contains(&self.image_token) {
168 let (height, width) = image_sizes_all_iter.next().unwrap();
169 let num_height_tokens =
170 (height as usize) / (self.patch_size * self.spatial_merge_size);
171 let num_width_tokens =
172 (width as usize) / (self.patch_size * self.spatial_merge_size);
173
174 let mut replace_tokens = vec![
175 [
176 vec![self.image_token.clone(); num_width_tokens],
177 vec![self.image_break_token.clone()],
178 ]
179 .concat();
180 num_height_tokens
181 ]
182 .into_iter()
183 .flatten()
184 .collect::<Vec<_>>();
185
186 *replace_tokens.last_mut().unwrap() = self.image_end_token.clone();
187
188 replace_strings.push(replace_tokens.join(""));
189 prompt = prompt.replace(&self.image_token, PLACEHOLDER);
190 }
191
192 while prompt.contains(PLACEHOLDER) {
193 let replace_str = replace_strings.pop().unwrap();
194 prompt = prompt.replace(PLACEHOLDER, &replace_str);
195 }
196
197 seq.set_initial_prompt(prompt.clone());
198 let toks = tokenizer
199 .encode_fast(prompt, false)
200 .expect("Detokenization failed!");
201
202 let ids = toks.get_ids().to_vec();
203 seq.set_toks_and_reallocate(ids, paged_attn_metadata.as_mut());
204 }
205
206 (
207 Some(Tensor::cat(&pixel_values_accum, 0).unwrap()),
208 Some(image_sizes_accum),
209 )
210 } else {
211 (None, None)
212 };
213
214 let text_models_inputs_processor::InnerInputProcessorOutput {
215 inputs:
216 text_models_inputs_processor::InputMetadata {
217 input,
218 positions,
219 context_lens,
220 position_ids,
221 paged_attn_meta,
222 flash_meta,
223 },
224 seq_indices,
225 } = if is_prompt {
226 get_prompt_input(
227 input_seqs
228 .iter()
229 .map(|seq| seq.get_toks().to_vec())
230 .collect::<Vec<_>>(),
231 input_seqs,
232 device,
233 last_n_context_len,
234 return_raw_logits,
235 paged_attn_metadata.as_mut(),
236 None, mapper,
238 )
239 .nth(0)
240 .unwrap()
241 .unwrap()
242 } else {
243 get_completion_input(
244 input_seqs
245 .iter()
246 .map(|seq| seq.get_toks().to_vec())
247 .collect::<Vec<_>>(),
248 input_seqs,
249 device,
250 no_kv_cache,
251 last_n_context_len,
252 return_raw_logits,
253 paged_attn_metadata.as_mut(),
254 None, mapper,
256 )
257 .nth(0)
258 .unwrap()
259 .unwrap()
260 };
261
262 let inputs: Box<dyn Any> = Box::new(ModelInputs {
263 input_ids: input,
264 seqlen_offsets: positions,
265 context_lens,
266 position_ids,
267 pixel_values,
268 model_specific_args: Box::new(Mistral3SpecificArgs { image_sizes }),
269 paged_attn_meta,
270 flash_meta,
271 });
272 Box::new(std::iter::once(Ok(InputProcessorOutput {
273 inputs,
274 seq_indices,
275 })))
276 }
277}
278
279impl Mistral3ImageProcessor {
280 #[allow(clippy::too_many_arguments)]
281 fn resize(
282 &self,
283 image: &DynamicImage,
284 mut height: usize,
285 mut width: usize,
286 max_height: usize,
287 max_width: usize,
288 patch_size: usize,
289 filter: FilterType,
290 ) -> DynamicImage {
291 let ratio = (height as f64 / max_height as f64).max(width as f64 / max_width as f64);
292 if ratio > 1. {
293 height = (height as f64 / ratio).floor() as usize;
294 width = (width as f64 / ratio).floor() as usize;
295 }
296
297 let num_height_tokens = (height - 1) / patch_size + 1;
298 let num_width_tokens = (width - 1) / patch_size + 1;
299
300 let resize_height = num_height_tokens * patch_size;
301 let resize_width = num_width_tokens * patch_size;
302
303 image.resize_exact(resize_width as u32, resize_height as u32, filter)
304 }
305}
306
307impl ImagePreProcessor for Mistral3ImageProcessor {
308 #[allow(clippy::excessive_precision)]
309 const DEFAULT_MEAN: [f64; 3] = [0.48145466, 0.4578275, 0.40821073];
310 #[allow(clippy::excessive_precision)]
311 const DEFAULT_STD: [f64; 3] = [0.26862954, 0.26130258, 0.27577711];
312
313 fn preprocess(
315 &self,
316 mut images: Vec<DynamicImage>,
317 videos: Vec<Vec<DynamicImage>>,
318 config: &PreProcessorConfig,
319 device: &Device,
320 (_bs, _max_num_images): (usize, usize),
321 ) -> Result<PreprocessedImages> {
322 assert!(videos.is_empty());
323
324 let do_resize = config.do_resize.unwrap();
325 let do_rescale = config.do_rescale.unwrap();
326 let rescale_factor = config.rescale_factor.unwrap();
327 let do_normalize = config.do_normalize.unwrap();
328 let image_mean = config.image_mean.unwrap_or(Self::DEFAULT_MEAN);
329 let image_std = config.image_std.unwrap_or(Self::DEFAULT_STD);
330 let do_convert_rgb = config.do_convert_rgb.unwrap_or(true);
331 let patch_size = config.patch_size.unwrap();
332 let size = config.size.as_ref().unwrap();
333 let resample = config.resampling.to_filter()?;
334
335 let default_to_square = config.default_to_square.unwrap();
336 assert!(default_to_square);
337
338 let mut pixel_values = Vec::new();
339 let mut image_sizes = Vec::new();
340
341 let (max_height, max_width) = if size.contains_key("longest_edge") {
342 (size["longest_edge"] as usize, size["longest_edge"] as usize)
343 } else if size.contains_key("height") && size.contains_key("width") {
344 (size["height"] as usize, size["width"] as usize)
345 } else {
346 candle_core::bail!("Size must be a map of `longest_edge` or `height` and `width`.");
347 };
348
349 for image in images.iter_mut() {
350 let (width, height) = image.dimensions();
351
352 if do_convert_rgb {
354 *image = DynamicImage::ImageRgb8(image.to_rgb8());
355 }
356
357 if do_resize {
358 *image = self.resize(
359 image,
360 height as usize,
361 width as usize,
362 max_height,
363 max_width,
364 patch_size,
365 resample,
366 );
367 }
368
369 let (width, height) = image.dimensions();
370
371 image_sizes.push((height, width));
372 }
373
374 images = mistralrs_vision::pad_to_max_image_size(images);
375
376 for image in images.iter_mut() {
377 let transforms = Transforms {
378 input: &ToTensorNoNorm,
379 inner_transforms: &[
380 &do_rescale.then_some(Rescale {
381 factor: Some(rescale_factor),
382 }),
383 &do_normalize.then(|| Normalize {
384 mean: image_mean.to_vec(),
385 std: image_std.to_vec(),
386 }),
387 ],
388 };
389
390 let image = image.apply(transforms, device)?;
391 pixel_values.push(image.unsqueeze(0)?);
392 }
393
394 Ok(PreprocessedImages {
395 pixel_values: Tensor::cat(&pixel_values, 0)?,
396 pixel_attention_mask: None,
397 image_sizes: None,
398 num_img_tokens: None,
399 aspect_ratio_ids: None,
400 aspect_ratio_mask: None,
401 num_tiles: None,
402 image_grid_thw: None,
403 video_grid_thw: None,
404 rows: None,
405 cols: None,
406 pixel_values_list: None,
407 tgt_sizes: None,
408 image_sizes_all: Some(image_sizes),
409 num_crops: None,
410 })
411 }
412}