diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f45f70..7a99cbd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -93,6 +93,12 @@ aims to follow [Semantic Versioning](https://semver.org/spec/v2.0.0.html) - **Explore polish** — the Tag Cloud now uses the in-app tooltip instead of the native browser one (and reads "1 image", not "1 images"), and the folder-scope dropdown no longer opens behind the cluster cards. +- **AI tagging no longer freezes the app** — tagging now runs inference in small + GPU micro-batches with a brief yield between them, instead of one wide batch + that monopolised the GPU (and with it the whole UI) for seconds at a time. The + app stays responsive while tagging runs, throughput is steadier (the old wide + batches caused periodic slowdowns), and the first batch after launch starts + faster. ## [0.1.1] — 2026-06-23 diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index 7f04378..ba09c05 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -1248,13 +1248,16 @@ fn process_tagging_batch( // Exclude actively-indexing folders for the same reason as the other // workers: don't compete with a running scan. + let batch_started_at = Instant::now(); let mut excluded_folders = paused_folder_ids("tagging"); excluded_folders.extend(active_indexing_folders()); let batch_size = crate::tagger::tagger_batch_size(app_data_dir); + let claim_started_at = Instant::now(); let jobs = with_db_write_lock(|| { let mut conn = pool.get()?; db::claim_tagging_jobs(&mut conn, &excluded_folders, batch_size) })?; + let claim_elapsed = claim_started_at.elapsed(); if jobs.is_empty() { return Ok(false); @@ -1288,21 +1291,41 @@ fn process_tagging_batch( .as_mut() .expect("tagger should be initialized before tagging batch processing"); - let tag_results = jobs + // Resolve each job's source image (AVIF can't be decoded directly, so it + // falls back to its thumbnail), then tag the batch in small chunks (below). + let source_paths: Vec = jobs .iter() .map(|job| { - let source_path = if is_avif_path(Path::new(&job.path)) { - job.thumbnail_path.as_deref().unwrap_or(&job.path) + if is_avif_path(Path::new(&job.path)) { + PathBuf::from(job.thumbnail_path.as_deref().unwrap_or(&job.path)) } else { - &job.path - }; - ( - job.clone(), - tagger_ref.run(Path::new(source_path), tagger::DEFAULT_MAX_TAGS), - ) + PathBuf::from(&job.path) + } }) - .collect::>(); + .collect(); + // Tag in small micro-batches instead of one wide forward pass. On a shared + // GPU every DirectML dispatch blocks the WebView2 compositor for its whole + // duration, so a 16-wide batch freezes the UI for seconds; small chunks keep + // each GPU lock short (and bound peak decode memory) while a brief yield + // between them lets the UI and other workers grab the GPU/CPU. The model is + // compute-bound here, so this costs almost no throughput. + let infer_started_at = Instant::now(); + let mut outputs = Vec::with_capacity(source_paths.len()); + let mut chunks = source_paths.chunks(tagger::TAGGER_INFER_CHUNK).peekable(); + while let Some(chunk) = chunks.next() { + outputs.extend(tagger_ref.run_batch(chunk, tagger::DEFAULT_MAX_TAGS)); + if chunks.peek().is_some() { + std::thread::sleep(std::time::Duration::from_millis( + tagger::TAGGER_INFER_YIELD_MS, + )); + } + } + let infer_elapsed = infer_started_at.elapsed(); + + let tag_results = jobs.iter().cloned().zip(outputs).collect::>(); + + let write_started_at = Instant::now(); let updated_images = with_db_write_lock(|| { let mut conn = pool.get()?; let tx = conn.transaction()?; @@ -1354,6 +1377,7 @@ fn process_tagging_batch( db::requeue_tagging_jobs(&conn, &image_ids) }); })?; + let write_elapsed = write_started_at.elapsed(); if !updated_images.is_empty() { let folder_ids = updated_images @@ -1369,6 +1393,15 @@ fn process_tagging_batch( emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); } + log::info!( + "Tagging batch timing: {} items, claim {:?}, tag {:?}, write {:?}, total {:?}", + jobs.len(), + claim_elapsed, + infer_elapsed, + write_elapsed, + batch_started_at.elapsed() + ); + Ok(true) } diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index 52214b6..8451d36 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -4,6 +4,7 @@ use image::{imageops::FilterType, DynamicImage, ImageReader}; use ort::ep; use ort::session::{builder::GraphOptimizationLevel, Session}; use ort::value::Tensor; +use rayon::prelude::*; use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; @@ -37,6 +38,19 @@ const RATING_CATEGORY: u8 = 9; pub const DEFAULT_THRESHOLD: f32 = 0.35; pub const DEFAULT_MAX_TAGS: usize = 30; +/// How many images are fed to the GPU in a single forward pass, regardless of +/// how many the worker claims from the DB at once. The WD model is compute- +/// bound on a shared GPU, so a wide batch buys little throughput but holds the +/// GPU (and therefore the WebView2 compositor) hostage for its whole duration, +/// freezing the UI. Small chunks keep each DirectML dispatch short — lower this +/// for a smoother UI under heavy tagging, raise it for marginally higher +/// throughput on a dedicated GPU. +pub const TAGGER_INFER_CHUNK: usize = 4; + +/// Brief pause between inference chunks so the compositor and other workers can +/// claim the GPU/CPU between dispatches. Negligible against per-chunk inference. +pub const TAGGER_INFER_YIELD_MS: u64 = 40; + /// Set to `true` by `set_tagger_acceleration` so the tagging worker loop /// knows to drop its cached `WdTagger` and reload with the new EP. pub static TAGGER_SESSION_DIRTY: AtomicBool = AtomicBool::new(false); @@ -456,33 +470,34 @@ impl WdTagger { let session = create_tagger_session(&model_path, acceleration)?; - // Determine the input spatial size from the ONNX model graph. - // WD v3 models use (1, H, W, 3) where H == W, typically 448. - let input_size = { + // Determine the input spatial size (and batch axis, for diagnostics) from + // the ONNX model graph. WD v3 models use (N, H, W, 3) with H == W, + // typically 448, and a dynamic batch dimension (reported as -1/0) which + // batched inference relies on. + let (input_size, batch_axis) = { let inputs = session.inputs(); - let dim = inputs + inputs .first() .and_then(|inp| { if let ort::value::ValueType::Tensor { shape, .. } = inp.dtype() { - shape - .get(1) - .and_then(|&d| if d > 0 { Some(d as usize) } else { None }) + let size = shape.get(1).copied().filter(|&d| d > 0).map(|d| d as usize); + Some((size.unwrap_or(448), shape.first().copied())) } else { None } }) - .unwrap_or(448); - dim + .unwrap_or((448, None)) }; let labels = load_labels(&labels_path)?; log::info!( - "WD tagger loaded in {:?} ({} labels, input {}x{}, {:?} acceleration)", + "WD tagger loaded in {:?} ({} labels, input {}x{}, batch axis {:?}, {:?} acceleration)", started_at.elapsed(), labels.len(), input_size, input_size, + batch_axis, acceleration, ); @@ -494,14 +509,96 @@ impl WdTagger { }) } - pub fn run(&mut self, image_path: &Path, max_tags: usize) -> Result { - let started_at = Instant::now(); + /// Tag a batch of images in a single forward pass. Returns one result per + /// input path, in the same order. Per-image decode failures — and, if the + /// batched inference itself fails (e.g. a model pinned to batch=1), a + /// per-image inference fallback — are reflected in the individual `Result`s. + pub fn run_batch( + &mut self, + image_paths: &[PathBuf], + max_tags: usize, + ) -> Vec> { + if image_paths.is_empty() { + return Vec::new(); + } - let image_array = preprocess_image(image_path, self.input_size)?; - let batch_size = 1usize; + let started_at = Instant::now(); + let input_size = self.input_size; + let stride = input_size * input_size * 3; + + // Decode + preprocess every image in parallel, keeping results aligned to + // the inputs so a failure stays attached to its own path. + let preprocessed: Vec>> = image_paths + .par_iter() + .map(|path| preprocess_image(path, input_size)) + .collect(); + + // Pack the successfully-decoded images into one contiguous buffer for a + // single batched inference, remembering which input slot each came from. + let mut batch_slots: Vec = Vec::new(); + let mut batch_pixels: Vec = Vec::with_capacity(image_paths.len() * stride); + for (i, result) in preprocessed.iter().enumerate() { + if let Ok(pixels) = result { + batch_slots.push(i); + batch_pixels.extend_from_slice(pixels); + } + } + + // Seed each slot with its decode error; inference overwrites the ones + // that decoded. The placeholder error never survives a decoded slot. + let mut results: Vec> = preprocessed + .into_iter() + .map(|r| match r { + Ok(_) => Err(anyhow::anyhow!( + "tagging inference did not run for this image" + )), + Err(error) => Err(error), + }) + .collect(); + + if batch_slots.is_empty() { + return results; // nothing decoded + } + + match self.infer_batch(&batch_pixels, batch_slots.len(), max_tags) { + Ok(outputs) => { + for (&slot, output) in batch_slots.iter().zip(outputs) { + results[slot] = Ok(output); + } + } + Err(batch_error) => { + log::warn!( + "WD tagger batch inference failed for {} images, falling back to per-image: {batch_error}", + batch_slots.len() + ); + for (k, &slot) in batch_slots.iter().enumerate() { + let pixels = batch_pixels[k * stride..(k + 1) * stride].to_vec(); + results[slot] = self.infer_one(pixels, max_tags); + } + } + } + + log::debug!( + "WD tagger batch: {} images ({} decoded) in {:?}", + image_paths.len(), + batch_slots.len(), + started_at.elapsed(), + ); + + results + } + + /// Run `count` already-preprocessed images (contiguous `[count, H, W, 3]` + /// pixels) through the model in one forward pass. + fn infer_batch( + &mut self, + pixels: &[f32], + count: usize, + max_tags: usize, + ) -> Result> { let input = Tensor::from_array(( - [batch_size, self.input_size, self.input_size, 3usize], - image_array.into_boxed_slice(), + [count, self.input_size, self.input_size, 3usize], + pixels.to_vec().into_boxed_slice(), )) .map_err(|error| anyhow::anyhow!("{error}"))?; @@ -511,23 +608,48 @@ impl WdTagger { .run(ort::inputs! { input_name.as_str() => input }) .map_err(|error| anyhow::anyhow!("{error}"))?; - let (_, probabilities) = outputs[0] + let (_, probs) = outputs[0] .try_extract_tensor::() .map_err(|error| anyhow::anyhow!("{error}"))?; - let probs: &[f32] = probabilities; - - if probs.len() != self.labels.len() { + let n_labels = self.labels.len(); + if probs.len() != count * n_labels { anyhow::bail!( - "Model output length {} does not match label count {}", + "Model output length {} does not match {} images x {} labels", probs.len(), - self.labels.len() + count, + n_labels ); } - // Collect rating scores (category 9) - pick the argmax as the rating. - let rating = self - .labels + // Borrow disjoint fields (not all of `self`) so this doesn't conflict + // with the mutable session borrow `outputs` still holds. + let labels = &self.labels; + let threshold = self.threshold; + Ok(probs + .chunks_exact(n_labels) + .map(|row| Self::tags_from_probs(labels, threshold, row, max_tags)) + .collect()) + } + + /// Run a single preprocessed image through the model. + fn infer_one(&mut self, image_array: Vec, max_tags: usize) -> Result { + self.infer_batch(&image_array, 1, max_tags)? + .into_iter() + .next() + .ok_or_else(|| anyhow::anyhow!("tagger produced no output for image")) + } + + /// Convert one image's class probabilities into its rating + sorted tags. + /// Associated (not `&self`) so callers can hold a disjoint session borrow. + fn tags_from_probs( + labels: &[TagEntry], + threshold: f32, + probs: &[f32], + max_tags: usize, + ) -> TaggerOutput { + // Rating (category 9): pick the argmax label as the rating. + let rating = labels .iter() .zip(probs.iter()) .filter(|(entry, _)| entry.category == RATING_CATEGORY) @@ -535,14 +657,13 @@ impl WdTagger { .map(|(entry, _)| entry.name.clone()) .unwrap_or_else(|| "general".to_string()); - // Collect general + character tags above threshold, sorted by confidence. - let mut tags: Vec = self - .labels + // General + character tags above threshold, sorted by confidence. + let mut tags: Vec = labels .iter() .zip(probs.iter()) .filter(|(entry, prob)| { (entry.category == GENERAL_CATEGORY || entry.category == CHARACTER_CATEGORY) - && **prob >= self.threshold + && **prob >= threshold }) .map(|(entry, prob)| TagResult { tag: entry.name.clone(), @@ -553,16 +674,7 @@ impl WdTagger { tags.sort_by(|a, b| b.confidence.total_cmp(&a.confidence)); tags.truncate(max_tags); - log::info!( - "WD tagger: {} tags (threshold={}, rating={}) in {:?} for {}", - tags.len(), - self.threshold, - rating, - started_at.elapsed(), - image_path.display(), - ); - - Ok(TaggerOutput { tags, rating }) + TaggerOutput { tags, rating } } }