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

fastembed/image_embedding/
impl.rs
