From 334ac54e003505ff53e43d99d933db7a68b73178 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Fri, 12 Jun 2026 09:00:25 +0100 Subject: [PATCH] perf(workers): scaled embedding preprocessing and idle-only backoff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Embedding and tagging workers now sleep their backoff interval (500ms/750ms) only when the queue is empty, the model is not ready, or a batch errored — previously every batch paid the sleep even with work pending. - CLIP preprocessing no longer decodes originals at full resolution: decode_for_thumbnail is generalized into decode_image_scaled with a cover mode (shortest side >= target) and the embedder decodes JPEGs at the smallest DCT scale covering its 224px fill-crop, in parallel via rayon. The forward pass, not image decode, is now the dominant embedding cost. - EXIF orientation is now applied before embedding, so rotated photos embed the way they are displayed (previously embedded sideways). Existing stored embeddings are unaffected. --- src-tauri/src/embedder.rs | 17 +++++++----- src-tauri/src/indexer.rs | 46 ++++++++++++++++++++----------- src-tauri/src/thumbnail.rs | 56 ++++++++++++++++++++++++++++---------- 3 files changed, 81 insertions(+), 38 deletions(-) diff --git a/src-tauri/src/embedder.rs b/src-tauri/src/embedder.rs index 3a977c7..f2a7551 100644 --- a/src-tauri/src/embedder.rs +++ b/src-tauri/src/embedder.rs @@ -191,9 +191,11 @@ fn resolve_device() -> Result { } fn load_image(path: &Path, image_size: usize) -> Result { - let image = image::ImageReader::open(path)? - .with_guessed_format()? - .decode()?; + // Scaled decode: CLIP only needs image_size² pixels, so decoding a large + // JPEG at full resolution is wasted work. Cover mode keeps the shortest + // side at or above image_size for the fill-crop below. Also applies EXIF + // orientation, so rotated photos embed the way they are displayed. + let image = crate::thumbnail::decode_image_scaled(path, image_size as u32, true)?; let image = image.resize_to_fill( image_size as u32, image_size as u32, @@ -208,10 +210,11 @@ fn load_image(path: &Path, image_size: usize) -> Result { } fn load_images(paths: &[PathBuf], image_size: usize) -> Result { - let mut images = Vec::with_capacity(paths.len()); - for path in paths { - images.push(load_image(path, image_size)?); - } + use rayon::prelude::*; + let images = paths + .par_iter() + .map(|path| load_image(path, image_size)) + .collect::>>()?; Ok(Tensor::stack(&images, 0)?) } diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index 35583c4..440dd7e 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -231,10 +231,16 @@ pub fn start_embedding_worker(app: AppHandle, pool: DbPool) { let mut embedder: Option = None; println!("Embedding worker started."); loop { - if let Err(error) = process_embedding_batch(&app, &pool, &mut embedder) { - eprintln!("Embedding worker error: {}", error); + // Only back off when the queue is empty (or errored); while jobs + // are pending, claim the next batch immediately. + match process_embedding_batch(&app, &pool, &mut embedder) { + Ok(true) => {} + Ok(false) => std::thread::sleep(std::time::Duration::from_millis(500)), + Err(error) => { + eprintln!("Embedding worker error: {}", error); + std::thread::sleep(std::time::Duration::from_millis(500)); + } } - std::thread::sleep(std::time::Duration::from_millis(500)); } }); } @@ -270,13 +276,17 @@ pub fn start_tagging_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) println!("Tagging worker: acceleration setting changed — resetting session."); tagger_instance = None; } - if let Err(error) = - process_tagging_batch(&app, &pool, &app_data_dir, &mut tagger_instance) - { - eprintln!("Tagging worker error: {}", error); - tagger_instance = None; + // Only back off when the queue is empty (or errored); while jobs + // are pending, claim the next batch immediately. + match process_tagging_batch(&app, &pool, &app_data_dir, &mut tagger_instance) { + Ok(true) => {} + Ok(false) => std::thread::sleep(std::time::Duration::from_millis(750)), + Err(error) => { + eprintln!("Tagging worker error: {}", error); + tagger_instance = None; + std::thread::sleep(std::time::Duration::from_millis(750)); + } } - std::thread::sleep(std::time::Duration::from_millis(750)); } }); } @@ -759,11 +769,13 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo Ok(()) } +/// Returns `Ok(true)` if a batch was claimed and processed, `Ok(false)` if +/// the queue was empty. fn process_embedding_batch( app: &AppHandle, pool: &DbPool, embedder: &mut Option, -) -> Result<()> { +) -> Result { let batch_started_at = Instant::now(); let claim_started_at = Instant::now(); let paused_folders = paused_folder_ids("embedding"); @@ -774,7 +786,7 @@ fn process_embedding_batch( let claim_elapsed = claim_started_at.elapsed(); if jobs.is_empty() { - return Ok(()); + return Ok(false); } if embedder.is_none() { @@ -902,7 +914,7 @@ fn process_embedding_batch( EMBEDDING_BATCH_SIZE, claim_elapsed, infer_elapsed, write_elapsed, batch_elapsed ); - Ok(()) + Ok(true) } fn process_caption_batch( @@ -1007,14 +1019,16 @@ fn process_caption_batch( Ok(()) } +/// Returns `Ok(true)` if a batch was claimed and processed, `Ok(false)` if +/// the queue was empty or the model is not ready. fn process_tagging_batch( app: &AppHandle, pool: &DbPool, app_data_dir: &Path, tagger_instance: &mut Option, -) -> Result<()> { +) -> Result { if !tagger::tagger_model_status(app_data_dir).ready { - return Ok(()); + return Ok(false); } let paused_folders = paused_folder_ids("tagging"); @@ -1025,7 +1039,7 @@ fn process_tagging_batch( })?; if jobs.is_empty() { - return Ok(()); + return Ok(false); } if tagger_instance.is_none() { @@ -1133,7 +1147,7 @@ fn process_tagging_batch( emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); } - Ok(()) + Ok(true) } fn active_indexing_folders() -> HashSet { diff --git a/src-tauri/src/thumbnail.rs b/src-tauri/src/thumbnail.rs index ce74bfb..8b48765 100644 --- a/src-tauri/src/thumbnail.rs +++ b/src-tauri/src/thumbnail.rs @@ -62,13 +62,26 @@ pub fn generate_image_thumbnail(image_path: &Path, cache_dir: &Path) -> Result Result { + decode_image_scaled(image_path, THUMB_SIZE, false) +} + +/// Decodes an image at reduced resolution, with EXIF orientation applied. /// /// JPEGs take a fast path that decodes at reduced resolution (DCT scaling), /// which skips most of the IDCT/upsampling work for large photos. Anything -/// that path can't handle falls back to the `image` crate. -fn decode_for_thumbnail(image_path: &Path) -> Result { +/// that path can't handle falls back to the `image` crate at full size. +/// +/// `cover` selects which side must stay at or above `target`: `false` keeps +/// the longest side (for aspect-fit consumers), `true` keeps the shortest +/// side (for fill-crop consumers like the CLIP preprocessor). +pub fn decode_image_scaled( + image_path: &Path, + target: u32, + cover: bool, +) -> Result { if is_jpeg(image_path) { - if let Some(img) = decode_jpeg_scaled(image_path) { + if let Some(img) = decode_jpeg_scaled(image_path, target, cover) { return Ok(img); } } @@ -87,16 +100,21 @@ fn is_jpeg(path: &Path) -> bool { .is_some_and(|ext| ext.eq_ignore_ascii_case("jpg") || ext.eq_ignore_ascii_case("jpeg")) } -/// Decodes a JPEG at the smallest DCT scale that still covers `THUMB_SIZE`. +/// Decodes a JPEG at the smallest DCT scale that still covers `target`. /// Returns `None` on any failure (corrupt file, CMYK without conversion /// support, etc.) so the caller can fall back to the generic decoder. /// mozjpeg reports errors by panicking, hence the `catch_unwind`. -fn decode_jpeg_scaled(image_path: &Path) -> Option { +fn decode_jpeg_scaled(image_path: &Path, target: u32, cover: bool) -> Option { let path = image_path.to_path_buf(); let decoded = std::panic::catch_unwind(move || -> std::io::Result<(Vec, u32, u32)> { let mut decompress = mozjpeg::Decompress::new_path(&path)?; let (width, height) = decompress.size(); - decompress.scale(scale_numerator(width.max(height))); + let reference = if cover { + width.min(height) + } else { + width.max(height) + }; + decompress.scale(scale_numerator(reference, target)); let mut started = decompress.rgb()?; let (out_width, out_height) = (started.width() as u32, started.height() as u32); let pixels = started.read_scanlines::()?; @@ -115,11 +133,11 @@ fn decode_jpeg_scaled(image_path: &Path) -> Option { Some(img) } -/// Smallest numerator (of /8) that keeps the longest side at or above -/// `THUMB_SIZE`, so the subsequent resize never upscales. -fn scale_numerator(max_dimension: usize) -> u8 { +/// Smallest numerator (of /8) that keeps `reference` at or above `target`, +/// so the subsequent resize never upscales. +fn scale_numerator(reference: usize, target: u32) -> u8 { for numerator in 1..=8u8 { - if max_dimension * numerator as usize / 8 >= THUMB_SIZE as usize { + if reference * numerator as usize / 8 >= target as usize { return numerator; } } @@ -248,10 +266,12 @@ mod tests { #[test] fn scale_numerator_picks_smallest_sufficient() { - assert_eq!(scale_numerator(6000), 1); - assert_eq!(scale_numerator(640), 4); - assert_eq!(scale_numerator(320), 8); - assert_eq!(scale_numerator(100), 8); + assert_eq!(scale_numerator(6000, THUMB_SIZE), 1); + assert_eq!(scale_numerator(640, THUMB_SIZE), 4); + assert_eq!(scale_numerator(320, THUMB_SIZE), 8); + assert_eq!(scale_numerator(100, THUMB_SIZE), 8); + // Cover mode reference: shortest side must reach 224 for CLIP. + assert_eq!(scale_numerator(1200, 224), 2); } #[test] @@ -264,9 +284,15 @@ mod tests { img.save(&src_path).unwrap(); // 1600 max dim -> numerator 2 -> 400x300 decode output - let decoded = decode_jpeg_scaled(&src_path).expect("fast path should handle plain JPEG"); + let decoded = decode_jpeg_scaled(&src_path, THUMB_SIZE, false) + .expect("fast path should handle plain JPEG"); assert_eq!((decoded.width(), decoded.height()), (400, 300)); + // Cover mode: shortest side (1200) must stay >= 224 -> numerator 2 + let covered = decode_jpeg_scaled(&src_path, 224, true) + .expect("fast path should handle plain JPEG"); + assert_eq!((covered.width(), covered.height()), (400, 300)); + let cache = dir.join("cache"); let _ = std::fs::remove_file(thumb_path(&cache, &src_path.to_string_lossy())); let result = generate_image_thumbnail(&src_path, &cache).unwrap();