fastembed/image_embedding/
impl.rs

1#[cfg(feature = "hf-hub")]
2use hf_hub::api::sync::ApiRepo;
3use image::DynamicImage;
4use ndarray::{Array3, ArrayView3};
5use ort::{session::Session, value::Value};
6#[cfg(feature = "hf-hub")]
7use std::path::PathBuf;
8use std::{io::Cursor, path::Path};
9
10use crate::{
11    common::{init_session_builder, normalize},
12    models::image_embedding::models_list,
13    Embedding, ImageEmbeddingModel, ModelInfo,
14};
15use anyhow::anyhow;
16#[cfg(feature = "hf-hub")]
17use anyhow::Context;
18
19#[cfg(feature = "hf-hub")]
20use super::ImageInitOptions;
21use super::{
22    init::{ImageInitOptionsUserDefined, UserDefinedImageEmbeddingModel},
23    utils::{Compose, Transform, TransformData},
24    ImageEmbedding, DEFAULT_BATCH_SIZE,
25};
26
27impl ImageEmbedding {
28    /// Try to generate a new ImageEmbedding Instance
29    ///
30    /// Uses the highest level of Graph optimization
31    ///
32    /// Uses the total number of CPUs available as the number of intra-threads
33    #[cfg(feature = "hf-hub")]
34    pub fn try_new(options: ImageInitOptions) -> anyhow::Result<Self> {
35        let ImageInitOptions {
36            model_name,
37            execution_providers,
38            cache_dir,
39            show_download_progress,
40            intra_threads,
41        } = options;
42
43        let model_repo = ImageEmbedding::retrieve_model(
44            model_name.clone(),
45            cache_dir.clone(),
46            show_download_progress,
47        )?;
48
49        let preprocessor_file = model_repo
50            .get("preprocessor_config.json")
51            .context("Failed to retrieve preprocessor_config.json")?;
52        let preprocessor = Compose::from_file(preprocessor_file)?;
53
54        let model_file_name = ImageEmbedding::get_model_info(&model_name).model_file;
55        let model_file_reference = model_repo
56            .get(&model_file_name)
57            .context(format!("Failed to retrieve {}", model_file_name))?;
58
59        let session = init_session_builder(execution_providers, intra_threads)?
60            .commit_from_file(model_file_reference)?;
61
62        Ok(Self::new(preprocessor, session))
63    }
64
65    /// Create a ImageEmbedding instance from model files provided by the user.
66    ///
67    /// This can be used for 'bring your own' embedding models
68    pub fn try_new_from_user_defined(
69        model: UserDefinedImageEmbeddingModel,
70        options: ImageInitOptionsUserDefined,
71    ) -> anyhow::Result<Self> {
72        let ImageInitOptionsUserDefined {
73            execution_providers,
74            intra_threads,
75        } = options;
76
77        let preprocessor = Compose::from_bytes(model.preprocessor_file)?;
78
79        let session = init_session_builder(execution_providers, intra_threads)?
80            .commit_from_memory(&model.onnx_file)?;
81
82        Ok(Self::new(preprocessor, session))
83    }
84
85    /// Private method to return an instance
86    fn new(preprocessor: Compose, session: Session) -> Self {
87        Self {
88            preprocessor,
89            session,
90        }
91    }
92
93    /// Return the ImageEmbedding model's directory from cache or remote retrieval
94    #[cfg(feature = "hf-hub")]
95    fn retrieve_model(
96        model: ImageEmbeddingModel,
97        cache_dir: PathBuf,
98        show_download_progress: bool,
99    ) -> anyhow::Result<ApiRepo> {
100        use crate::common::pull_from_hf;
101
102        pull_from_hf(model.to_string(), cache_dir, show_download_progress)
103    }
104
105    /// Retrieve a list of supported models
106    pub fn list_supported_models() -> Vec<ModelInfo<ImageEmbeddingModel>> {
107        models_list()
108    }
109
110    /// Get ModelInfo from ImageEmbeddingModel
111    pub fn get_model_info(model: &ImageEmbeddingModel) -> ModelInfo<ImageEmbeddingModel> {
112        ImageEmbedding::list_supported_models()
113            .into_iter()
114            .find(|m| &m.model == model)
115            .expect("Model not found in supported models list. This is a bug - please report it.")
116    }
117
118    /// Method to generate image embeddings for a Vec of image bytes
119    pub fn embed_bytes(
120        &mut self,
121        images: &[&[u8]],
122        batch_size: Option<usize>,
123    ) -> anyhow::Result<Vec<Embedding>> {
124        let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
125        anyhow::ensure!(batch_size > 0, "batch_size must be greater than 0");
126
127        let output = images
128            .chunks(batch_size)
129            .map(|batch| {
130                // Encode the texts in the batch
131                let inputs = batch
132                    .iter()
133                    .map(|img| {
134                        image::ImageReader::new(Cursor::new(img))
135                            .with_guessed_format()?
136                            .decode()
137                            .map_err(|err| anyhow!("image decode: {}", err))
138                    })
139                    .collect::<Result<_, _>>()?;
140
141                self.embed_images(inputs)
142            })
143            .collect::<anyhow::Result<Vec<_>>>()?
144            .into_iter()
145            .flatten()
146            .collect();
147
148        Ok(output)
149    }
150
151    /// Method to generate image embeddings for a collection of image paths.
152    ///
153    /// Accepts anything that can be referenced as a slice of elements implementing
154    /// [`AsRef<Path>`], such as `Vec<String>`, `Vec<PathBuf>`, `&[&str]`, or `&[&Path]`.
155    pub fn embed<S: AsRef<Path> + Send + Sync>(
156        &mut self,
157        images: impl AsRef<[S]>,
158        batch_size: Option<usize>,
159    ) -> anyhow::Result<Vec<Embedding>> {
160        let images = images.as_ref();
161        // Determine the batch size, default if not specified
162        let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
163        anyhow::ensure!(batch_size > 0, "batch_size must be greater than 0");
164
165        let output = images
166            .chunks(batch_size)
167            .map(|batch| {
168                // Encode the texts in the batch
169                let inputs = batch
170                    .iter()
171                    .map(|img| {
172                        image::ImageReader::open(img)?
173                            .decode()
174                            .map_err(|err| anyhow!("image decode: {}", err))
175                    })
176                    .collect::<Result<_, _>>()?;
177
178                self.embed_images(inputs)
179            })
180            .collect::<anyhow::Result<Vec<_>>>()?
181            .into_iter()
182            .flatten()
183            .collect();
184
185        Ok(output)
186    }
187
188    /// Embed DynamicImages
189    pub fn embed_images(&mut self, imgs: Vec<DynamicImage>) -> anyhow::Result<Vec<Embedding>> {
190        let inputs = imgs
191            .into_iter()
192            .map(|img| {
193                let pixels = self.preprocessor.transform(TransformData::Image(img))?;
194                match pixels {
195                    TransformData::NdArray(array) => Ok(array),
196                    _ => Err(anyhow!("Preprocessor configuration error!")),
197                }
198            })
199            .collect::<anyhow::Result<Vec<Array3<f32>>>>()?;
200
201        // Extract the batch size
202        let inputs_view: Vec<ArrayView3<f32>> = inputs.iter().map(|img| img.view()).collect();
203        let pixel_values_array = ndarray::stack(ndarray::Axis(0), &inputs_view)?;
204
205        let input_name = self.session.inputs()[0].name().to_string();
206        let session_inputs = ort::inputs![
207            input_name => Value::from_array(pixel_values_array)?,
208        ];
209
210        let outputs = self.session.run(session_inputs)?;
211
212        // Try to get the only output key
213        // If multiple, then default to few known keys `image_embeds` and `last_hidden_state`
214        let last_hidden_state_key = match outputs.len() {
215            1 => vec![outputs
216                .keys()
217                .next()
218                .ok_or_else(|| anyhow!("Expected one output but found none"))?],
219            _ => vec!["image_embeds", "last_hidden_state"],
220        };
221
222        // Extract tensor and handle different dimensionalities
223        let (shape, data) = last_hidden_state_key
224            .iter()
225            .find_map(|&key| {
226                outputs
227                    .get(key)
228                    .and_then(|v| v.try_extract_tensor::<f32>().ok())
229            })
230            .ok_or_else(|| anyhow!("Could not extract tensor from any known output key"))?;
231        let shape: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
232        let output_array = ndarray::ArrayViewD::from_shape(shape.as_slice(), data)?;
233
234        let embeddings = match output_array.ndim() {
235            3 => {
236                // For 3D output [batch_size, sequence_length, hidden_size]
237                // Take only the first token, sequence_length[0] (CLS token), embedding
238                // and return [batch_size, hidden_size]
239                (0..output_array.shape()[0])
240                    .map(|batch_idx| {
241                        let cls_embedding = output_array
242                            .slice(ndarray::s![batch_idx, 0, ..])
243                            .to_owned()
244                            .to_vec();
245                        normalize(&cls_embedding)
246                    })
247                    .collect()
248            }
249            2 => {
250                // For 2D output [batch_size, hidden_size]
251                output_array
252                    .outer_iter()
253                    .map(|row| {
254                        row.as_slice()
255                            .ok_or_else(|| anyhow!("Failed to convert array row to slice"))
256                            .map(normalize)
257                    })
258                    .collect::<anyhow::Result<Vec<_>>>()?
259            }
260            _ => {
261                return Err(anyhow!(
262                    "Unexpected output tensor shape: {:?}",
263                    output_array.shape()
264                ))
265            }
266        };
267
268        Ok(embeddings)
269    }
270}
fastembed/image_embedding/impl.rs

fastembed/image_embedding/
impl.rs
