perf(workers): scaled embedding preprocessing and idle-only backoff

- 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.
This commit is contained in:
2026-06-12 09:00:25 +01:00
parent a4486547e8
commit 334ac54e00
3 changed files with 81 additions and 38 deletions
+10 -7
View File
@@ -191,9 +191,11 @@ fn resolve_device() -> Result<Device> {
} }
fn load_image(path: &Path, image_size: usize) -> Result<Tensor> { fn load_image(path: &Path, image_size: usize) -> Result<Tensor> {
let image = image::ImageReader::open(path)? // Scaled decode: CLIP only needs image_size² pixels, so decoding a large
.with_guessed_format()? // JPEG at full resolution is wasted work. Cover mode keeps the shortest
.decode()?; // 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( let image = image.resize_to_fill(
image_size as u32, image_size as u32,
image_size as u32, image_size as u32,
@@ -208,10 +210,11 @@ fn load_image(path: &Path, image_size: usize) -> Result<Tensor> {
} }
fn load_images(paths: &[PathBuf], image_size: usize) -> Result<Tensor> { fn load_images(paths: &[PathBuf], image_size: usize) -> Result<Tensor> {
let mut images = Vec::with_capacity(paths.len()); use rayon::prelude::*;
for path in paths { let images = paths
images.push(load_image(path, image_size)?); .par_iter()
} .map(|path| load_image(path, image_size))
.collect::<Result<Vec<_>>>()?;
Ok(Tensor::stack(&images, 0)?) Ok(Tensor::stack(&images, 0)?)
} }
+30 -16
View File
@@ -231,10 +231,16 @@ pub fn start_embedding_worker(app: AppHandle, pool: DbPool) {
let mut embedder: Option<ClipImageEmbedder> = None; let mut embedder: Option<ClipImageEmbedder> = None;
println!("Embedding worker started."); println!("Embedding worker started.");
loop { loop {
if let Err(error) = process_embedding_batch(&app, &pool, &mut embedder) { // Only back off when the queue is empty (or errored); while jobs
eprintln!("Embedding worker error: {}", error); // 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."); println!("Tagging worker: acceleration setting changed — resetting session.");
tagger_instance = None; tagger_instance = None;
} }
if let Err(error) = // Only back off when the queue is empty (or errored); while jobs
process_tagging_batch(&app, &pool, &app_data_dir, &mut tagger_instance) // are pending, claim the next batch immediately.
{ match process_tagging_batch(&app, &pool, &app_data_dir, &mut tagger_instance) {
eprintln!("Tagging worker error: {}", error); Ok(true) => {}
tagger_instance = None; 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(()) Ok(())
} }
/// Returns `Ok(true)` if a batch was claimed and processed, `Ok(false)` if
/// the queue was empty.
fn process_embedding_batch( fn process_embedding_batch(
app: &AppHandle, app: &AppHandle,
pool: &DbPool, pool: &DbPool,
embedder: &mut Option<ClipImageEmbedder>, embedder: &mut Option<ClipImageEmbedder>,
) -> Result<()> { ) -> Result<bool> {
let batch_started_at = Instant::now(); let batch_started_at = Instant::now();
let claim_started_at = Instant::now(); let claim_started_at = Instant::now();
let paused_folders = paused_folder_ids("embedding"); let paused_folders = paused_folder_ids("embedding");
@@ -774,7 +786,7 @@ fn process_embedding_batch(
let claim_elapsed = claim_started_at.elapsed(); let claim_elapsed = claim_started_at.elapsed();
if jobs.is_empty() { if jobs.is_empty() {
return Ok(()); return Ok(false);
} }
if embedder.is_none() { if embedder.is_none() {
@@ -902,7 +914,7 @@ fn process_embedding_batch(
EMBEDDING_BATCH_SIZE, claim_elapsed, infer_elapsed, write_elapsed, batch_elapsed EMBEDDING_BATCH_SIZE, claim_elapsed, infer_elapsed, write_elapsed, batch_elapsed
); );
Ok(()) Ok(true)
} }
fn process_caption_batch( fn process_caption_batch(
@@ -1007,14 +1019,16 @@ fn process_caption_batch(
Ok(()) 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( fn process_tagging_batch(
app: &AppHandle, app: &AppHandle,
pool: &DbPool, pool: &DbPool,
app_data_dir: &Path, app_data_dir: &Path,
tagger_instance: &mut Option<WdTagger>, tagger_instance: &mut Option<WdTagger>,
) -> Result<()> { ) -> Result<bool> {
if !tagger::tagger_model_status(app_data_dir).ready { if !tagger::tagger_model_status(app_data_dir).ready {
return Ok(()); return Ok(false);
} }
let paused_folders = paused_folder_ids("tagging"); let paused_folders = paused_folder_ids("tagging");
@@ -1025,7 +1039,7 @@ fn process_tagging_batch(
})?; })?;
if jobs.is_empty() { if jobs.is_empty() {
return Ok(()); return Ok(false);
} }
if tagger_instance.is_none() { if tagger_instance.is_none() {
@@ -1133,7 +1147,7 @@ fn process_tagging_batch(
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true); emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
} }
Ok(()) Ok(true)
} }
fn active_indexing_folders() -> HashSet<i64> { fn active_indexing_folders() -> HashSet<i64> {
+41 -15
View File
@@ -62,13 +62,26 @@ pub fn generate_image_thumbnail(image_path: &Path, cache_dir: &Path) -> Result<G
} }
/// Decodes an image for thumbnailing, with EXIF orientation already applied. /// Decodes an image for thumbnailing, with EXIF orientation already applied.
fn decode_for_thumbnail(image_path: &Path) -> Result<image::DynamicImage> {
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), /// JPEGs take a fast path that decodes at reduced resolution (DCT scaling),
/// which skips most of the IDCT/upsampling work for large photos. Anything /// which skips most of the IDCT/upsampling work for large photos. Anything
/// that path can't handle falls back to the `image` crate. /// that path can't handle falls back to the `image` crate at full size.
fn decode_for_thumbnail(image_path: &Path) -> Result<image::DynamicImage> { ///
/// `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<image::DynamicImage> {
if is_jpeg(image_path) { 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); 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")) .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 /// Returns `None` on any failure (corrupt file, CMYK without conversion
/// support, etc.) so the caller can fall back to the generic decoder. /// support, etc.) so the caller can fall back to the generic decoder.
/// mozjpeg reports errors by panicking, hence the `catch_unwind`. /// mozjpeg reports errors by panicking, hence the `catch_unwind`.
fn decode_jpeg_scaled(image_path: &Path) -> Option<image::DynamicImage> { fn decode_jpeg_scaled(image_path: &Path, target: u32, cover: bool) -> Option<image::DynamicImage> {
let path = image_path.to_path_buf(); let path = image_path.to_path_buf();
let decoded = std::panic::catch_unwind(move || -> std::io::Result<(Vec<u8>, u32, u32)> { let decoded = std::panic::catch_unwind(move || -> std::io::Result<(Vec<u8>, u32, u32)> {
let mut decompress = mozjpeg::Decompress::new_path(&path)?; let mut decompress = mozjpeg::Decompress::new_path(&path)?;
let (width, height) = decompress.size(); 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 mut started = decompress.rgb()?;
let (out_width, out_height) = (started.width() as u32, started.height() as u32); let (out_width, out_height) = (started.width() as u32, started.height() as u32);
let pixels = started.read_scanlines::<u8>()?; let pixels = started.read_scanlines::<u8>()?;
@@ -115,11 +133,11 @@ fn decode_jpeg_scaled(image_path: &Path) -> Option<image::DynamicImage> {
Some(img) Some(img)
} }
/// Smallest numerator (of /8) that keeps the longest side at or above /// Smallest numerator (of /8) that keeps `reference` at or above `target`,
/// `THUMB_SIZE`, so the subsequent resize never upscales. /// so the subsequent resize never upscales.
fn scale_numerator(max_dimension: usize) -> u8 { fn scale_numerator(reference: usize, target: u32) -> u8 {
for numerator in 1..=8u8 { 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; return numerator;
} }
} }
@@ -248,10 +266,12 @@ mod tests {
#[test] #[test]
fn scale_numerator_picks_smallest_sufficient() { fn scale_numerator_picks_smallest_sufficient() {
assert_eq!(scale_numerator(6000), 1); assert_eq!(scale_numerator(6000, THUMB_SIZE), 1);
assert_eq!(scale_numerator(640), 4); assert_eq!(scale_numerator(640, THUMB_SIZE), 4);
assert_eq!(scale_numerator(320), 8); assert_eq!(scale_numerator(320, THUMB_SIZE), 8);
assert_eq!(scale_numerator(100), 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] #[test]
@@ -264,9 +284,15 @@ mod tests {
img.save(&src_path).unwrap(); img.save(&src_path).unwrap();
// 1600 max dim -> numerator 2 -> 400x300 decode output // 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)); 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 cache = dir.join("cache");
let _ = std::fs::remove_file(thumb_path(&cache, &src_path.to_string_lossy())); let _ = std::fs::remove_file(thumb_path(&cache, &src_path.to_string_lossy()));
let result = generate_image_thumbnail(&src_path, &cache).unwrap(); let result = generate_image_thumbnail(&src_path, &cache).unwrap();