From bee6adc61aa94065247e48d7c3892832cfc15074 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Tue, 7 Apr 2026 06:11:01 +0100 Subject: [PATCH 01/25] feat: add local AI captions and queue controls --- src-tauri/Cargo.lock | 78 ++++ src-tauri/Cargo.toml | 1 + src-tauri/src/captioner.rs | 620 +++++++++++++++++++++++++++++ src-tauri/src/commands.rs | 287 ++++++++++++- src-tauri/src/db.rs | 468 ++++++++++++++++++++-- src-tauri/src/indexer.rs | 262 ++++++++++-- src-tauri/src/lib.rs | 26 +- src-tauri/src/vector.rs | 151 ++++++- src/App.tsx | 4 + src/components/BackgroundTasks.tsx | 114 ++++-- src/components/Gallery.tsx | 28 +- src/components/Lightbox.tsx | 137 +++++++ src/components/SettingsModal.tsx | 354 ++++++++++++++++ src/components/TitleBar.tsx | 13 + src/store.ts | 271 +++++++++++-- 15 files changed, 2653 insertions(+), 161 deletions(-) create mode 100644 src-tauri/src/captioner.rs create mode 100644 src/components/SettingsModal.tsx diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 2871f9e..db7b209 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -2389,6 +2389,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "hmac-sha256" +version = "1.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f" + [[package]] name = "html5ever" version = "0.29.1" @@ -3059,6 +3065,12 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lzma-rust2" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1670343e58806300d87950e3401e820b519b9384281bbabfb15e3636689ffd69" + [[package]] name = "lzma-sys" version = "0.1.20" @@ -3134,6 +3146,16 @@ version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2532096657941c2fea9c289d370a250971c689d4f143798ff67113ec042024a5" +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" @@ -3262,6 +3284,21 @@ dependencies = [ "tempfile", ] +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "ndk" version = "0.9.0" @@ -3670,6 +3707,31 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "ort" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133" +dependencies = [ + "libloading 0.9.0", + "ndarray", + "ort-sys", + "smallvec", + "tracing", + "ureq", +] + +[[package]] +name = "ort-sys" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90" +dependencies = [ + "hmac-sha256", + "lzma-rust2", + "ureq", +] + [[package]] name = "pango" version = "0.18.3" @@ -3952,6 +4014,7 @@ dependencies = [ "hf-hub", "image", "log", + "ort", "r2d2", "r2d2_sqlite", "rayon", @@ -4060,6 +4123,15 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "portable-atomic-util" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -4421,6 +4493,12 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rayon" version = "1.11.0" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index b9daa8a..ce47032 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -44,3 +44,4 @@ candle-nn = { version = "0.10.2", features = ["cuda"] } candle-transformers = { version = "0.10.2", features = ["cuda"] } hf-hub = { version = "0.5.0", default-features = false, features = ["ureq", "native-tls"] } tokenizers = "0.22.1" +ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "download-binaries", "copy-dylibs", "load-dynamic", "api-24", "tls-native"] } diff --git a/src-tauri/src/captioner.rs b/src-tauri/src/captioner.rs new file mode 100644 index 0000000..112aa7c --- /dev/null +++ b/src-tauri/src/captioner.rs @@ -0,0 +1,620 @@ +use anyhow::Result; +use hf_hub::{api::sync::Api, Repo, RepoType}; +use image::{imageops::FilterType, ImageReader}; +use ort::session::SessionInputValue; +use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::{Shape, Tensor}; +use serde::Serialize; +use std::borrow::Cow; +use std::path::{Path, PathBuf}; +use tokenizers::Tokenizer; + +pub const FLORENCE_MODEL_ID: &str = "onnx-community/Florence-2-base-ft"; +pub const FLORENCE_CAPTION_MODEL_NAME: &str = "florence-2-base-ft-onnx-q4"; + +const REQUIRED_FILES: &[&str] = &[ + "config.json", + "generation_config.json", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "onnx/vision_encoder_fp16.onnx", + "onnx/encoder_model_q4.onnx", + "onnx/decoder_model_merged_q4.onnx", + "onnx/embed_tokens_fp16.onnx", +]; + +#[derive(Serialize)] +pub struct CaptionModelStatus { + pub model_id: &'static str, + pub model_name: &'static str, + pub local_dir: String, + pub ready: bool, + pub missing_files: Vec, +} + +#[derive(Clone, Serialize)] +pub struct CaptionModelProgress { + pub total_files: usize, + pub completed_files: usize, + pub current_file: Option, + pub done: bool, +} + +#[derive(Serialize)] +pub struct CaptionRuntimeProbe { + pub ready: bool, + pub tokenizer_vocab_size: usize, + pub sessions: Vec, +} + +#[derive(Serialize)] +pub struct CaptionRuntimeSessionProbe { + pub file: &'static str, + pub inputs: Vec, + pub outputs: Vec, +} + +#[derive(Serialize)] +pub struct CaptionVisionProbe { + pub input_shape: Vec, + pub output_shape: Vec, + pub output_values: usize, +} + +#[derive(Clone)] +struct TensorData { + shape: Vec, + values: Vec, +} + +pub struct FlorenceCaptioner { + tokenizer: Tokenizer, + vision_session: Session, + embed_session: Session, + encoder_session: Session, + decoder_session: Session, +} + +pub fn model_dir(app_data_dir: &Path) -> PathBuf { + app_data_dir.join("models").join("florence-2-base-ft") +} + +pub fn caption_model_status(app_data_dir: &Path) -> CaptionModelStatus { + let local_dir = model_dir(app_data_dir); + let missing_files = REQUIRED_FILES + .iter() + .filter(|file| !local_dir.join(file).exists()) + .map(|file| (*file).to_string()) + .collect::>(); + + CaptionModelStatus { + model_id: FLORENCE_MODEL_ID, + model_name: FLORENCE_CAPTION_MODEL_NAME, + local_dir: local_dir.to_string_lossy().to_string(), + ready: missing_files.is_empty(), + missing_files, + } +} + +pub fn prepare_caption_model_with_progress( + app_data_dir: &Path, + emit_progress: impl Fn(CaptionModelProgress), +) -> Result { + let local_dir = model_dir(app_data_dir); + std::fs::create_dir_all(&local_dir)?; + + let api = Api::new()?; + let repo = api.repo(Repo::new(FLORENCE_MODEL_ID.to_string(), RepoType::Model)); + let mut completed_files = REQUIRED_FILES + .iter() + .filter(|file| local_dir.join(file).exists()) + .count(); + + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: None, + done: completed_files == REQUIRED_FILES.len(), + }); + + for file in REQUIRED_FILES { + let destination = local_dir.join(file); + if destination.exists() { + continue; + } + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: Some((*file).to_string()), + done: false, + }); + if let Some(parent) = destination.parent() { + std::fs::create_dir_all(parent)?; + } + let cached = repo.get(file)?; + std::fs::copy(cached, destination)?; + completed_files += 1; + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: Some((*file).to_string()), + done: completed_files == REQUIRED_FILES.len(), + }); + } + + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: None, + done: true, + }); + + Ok(caption_model_status(app_data_dir)) +} + +pub fn delete_caption_model(app_data_dir: &Path) -> Result { + let local_dir = model_dir(app_data_dir); + if local_dir.exists() { + std::fs::remove_dir_all(&local_dir)?; + } + Ok(caption_model_status(app_data_dir)) +} + +pub fn probe_caption_runtime(app_data_dir: &Path) -> Result { + let status = caption_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "Florence-2 model is missing {} required file(s)", + status.missing_files.len() + ); + } + + let local_dir = model_dir(app_data_dir); + let tokenizer = + Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + + let sessions = [ + "onnx/vision_encoder_fp16.onnx", + "onnx/embed_tokens_fp16.onnx", + "onnx/encoder_model_q4.onnx", + "onnx/decoder_model_merged_q4.onnx", + ] + .into_iter() + .map(|file| probe_session(file, &local_dir.join(file))) + .collect::>>()?; + + Ok(CaptionRuntimeProbe { + ready: true, + tokenizer_vocab_size: tokenizer.get_vocab_size(false), + sessions, + }) +} + +pub fn probe_caption_vision(app_data_dir: &Path, image_path: &Path) -> Result { + let status = caption_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "Florence-2 model is missing {} required file(s)", + status.missing_files.len() + ); + } + + let local_dir = model_dir(app_data_dir); + let pixels = preprocess_image(image_path)?; + let input_shape = vec![1, 3, 768, 768]; + let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice())) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let mut session = create_session(&local_dir.join("onnx/vision_encoder_fp16.onnx"))?; + let outputs = session + .run(ort::inputs! { + "pixel_values" => input + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let (output_shape, output_values) = outputs[0] + .try_extract_tensor::() + .map_err(|error| anyhow::anyhow!("{error}"))?; + + Ok(CaptionVisionProbe { + input_shape, + output_shape: output_shape.to_vec(), + output_values: output_values.len(), + }) +} + +pub fn generate_caption(app_data_dir: &Path, image_path: &Path) -> Result { + let mut captioner = FlorenceCaptioner::new(app_data_dir)?; + captioner.generate(image_path) +} + +impl FlorenceCaptioner { + pub fn new(app_data_dir: &Path) -> Result { + let status = caption_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "Florence-2 model is missing {} required file(s)", + status.missing_files.len() + ); + } + + let local_dir = model_dir(app_data_dir); + let tokenizer = + Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + + let vision_session = create_session(&local_dir.join("onnx/vision_encoder_fp16.onnx"))?; + let embed_session = create_session(&local_dir.join("onnx/embed_tokens_fp16.onnx"))?; + let encoder_session = create_session(&local_dir.join("onnx/encoder_model_q4.onnx"))?; + let decoder_session = create_session(&local_dir.join("onnx/decoder_model_merged_q4.onnx"))?; + + Ok(Self { + tokenizer, + vision_session, + embed_session, + encoder_session, + decoder_session, + }) + } + + pub fn generate(&mut self, image_path: &Path) -> Result { + let image_features = run_vision_encoder(&mut self.vision_session, image_path)?; + let prompt_ids = self + .tokenizer + .encode("What does the image describe?", false) + .map_err(anyhow::Error::msg)? + .get_ids() + .iter() + .map(|id| i64::from(*id)) + .collect::>(); + let prompt_embeds = run_token_embedder(&mut self.embed_session, &prompt_ids)?; + let encoder_embeds = concatenate_sequence_embeddings(&prompt_embeds, &image_features)?; + let encoder_attention_mask = vec![1_i64; encoder_embeds.shape[1] as usize]; + let encoder_hidden_states = run_encoder( + &mut self.encoder_session, + &encoder_embeds, + &encoder_attention_mask, + )?; + + let generated_ids = run_decoder( + &mut self.decoder_session, + &mut self.embed_session, + &encoder_hidden_states, + &encoder_attention_mask, + )?; + + let generated_u32 = generated_ids + .into_iter() + .map(|id| id as u32) + .collect::>(); + let caption = self + .tokenizer + .decode(&generated_u32, true) + .map_err(anyhow::Error::msg)? + .trim() + .to_string(); + + if caption.is_empty() { + anyhow::bail!("Florence-2 generated an empty caption"); + } + + Ok(clean_caption(&caption)) + } +} + +fn probe_session(file: &'static str, path: &Path) -> Result { + let metadata = std::fs::metadata(path)?; + if metadata.len() == 0 { + anyhow::bail!("{} is empty", path.display()); + } + + let (inputs, outputs) = match file { + "onnx/vision_encoder_fp16.onnx" => ( + vec!["pixel_values".to_string()], + vec!["image_features".to_string()], + ), + "onnx/embed_tokens_fp16.onnx" => ( + vec!["input_ids".to_string()], + vec!["inputs_embeds".to_string()], + ), + "onnx/encoder_model_q4.onnx" => ( + vec!["attention_mask".to_string(), "inputs_embeds".to_string()], + vec!["last_hidden_state".to_string()], + ), + "onnx/decoder_model_merged_q4.onnx" => ( + vec![ + "encoder_attention_mask".to_string(), + "encoder_hidden_states".to_string(), + "inputs_embeds".to_string(), + "past_key_values".to_string(), + "use_cache_branch".to_string(), + ], + vec!["logits".to_string(), "present_key_values".to_string()], + ), + _ => (Vec::new(), Vec::new()), + }; + + Ok(CaptionRuntimeSessionProbe { + file, + inputs, + outputs, + }) +} + +fn run_vision_encoder(session: &mut Session, image_path: &Path) -> Result { + let pixels = preprocess_image(image_path)?; + let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice())) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let outputs = session + .run(ort::inputs! { + "pixel_values" => input + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + tensor_data(&outputs[0]) +} + +fn run_token_embedder(session: &mut Session, token_ids: &[i64]) -> Result { + let input_ids = Tensor::from_array(( + [1usize, token_ids.len()], + token_ids.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let outputs = session + .run(ort::inputs! { + "input_ids" => input_ids + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + tensor_data(&outputs[0]) +} + +fn run_encoder( + session: &mut Session, + inputs_embeds: &TensorData, + attention_mask: &[i64], +) -> Result { + let attention_mask_tensor = Tensor::from_array(( + [1usize, attention_mask.len()], + attention_mask.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let inputs_embeds_tensor = tensor_from_data(inputs_embeds)?; + let outputs = session + .run(ort::inputs! { + "attention_mask" => attention_mask_tensor, + "inputs_embeds" => inputs_embeds_tensor + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + tensor_data(&outputs[0]) +} + +fn run_decoder( + decoder_session: &mut Session, + embed_session: &mut Session, + encoder_hidden_states: &TensorData, + encoder_attention_mask: &[i64], +) -> Result> { + const DECODER_LAYERS: usize = 6; + const DECODER_HEADS: usize = 12; + const HEAD_DIM: usize = 64; + const DECODER_START_TOKEN_ID: i64 = 2; + const EOS_TOKEN_ID: i64 = 2; + const MAX_NEW_TOKENS: usize = 32; + + let mut generated = Vec::new(); + let mut next_input_id = DECODER_START_TOKEN_ID; + let mut past: Vec = Vec::new(); + let mut use_cache_branch = false; + + for _ in 0..MAX_NEW_TOKENS { + let inputs_embeds = run_token_embedder(embed_session, &[next_input_id])?; + let encoder_attention_mask_tensor = Tensor::from_array(( + [1usize, encoder_attention_mask.len()], + encoder_attention_mask.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let encoder_hidden_states_tensor = tensor_from_data(encoder_hidden_states)?; + let inputs_embeds_tensor = tensor_from_data(&inputs_embeds)?; + let use_cache_branch_tensor = + Tensor::from_array(([1usize], vec![use_cache_branch].into_boxed_slice())) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut inputs = ort::inputs! { + "encoder_attention_mask" => encoder_attention_mask_tensor, + "encoder_hidden_states" => encoder_hidden_states_tensor, + "inputs_embeds" => inputs_embeds_tensor, + "use_cache_branch" => use_cache_branch_tensor + }; + + if past.is_empty() { + for layer in 0..DECODER_LAYERS { + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.decoder.key"), + TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.decoder.value"), + TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.encoder.key"), + TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.encoder.value"), + TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), + )?; + } + } else { + for layer in 0..DECODER_LAYERS { + for cache_name in [ + "decoder.key", + "decoder.value", + "encoder.key", + "encoder.value", + ] { + let past_index = layer * 4 + + match cache_name { + "decoder.key" => 0, + "decoder.value" => 1, + "encoder.key" => 2, + "encoder.value" => 3, + _ => unreachable!(), + }; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.{cache_name}"), + past[past_index].clone(), + )?; + } + } + } + + let outputs = decoder_session + .run(inputs) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let logits = tensor_data(&outputs["logits"])?; + let token_id = argmax_last_token(&logits)?; + if token_id == EOS_TOKEN_ID { + break; + } + generated.push(token_id); + next_input_id = token_id; + + past.clear(); + for layer in 0..DECODER_LAYERS { + for cache_name in [ + "decoder.key", + "decoder.value", + "encoder.key", + "encoder.value", + ] { + past.push(tensor_data( + &outputs[format!("present.{layer}.{cache_name}").as_str()], + )?); + } + } + use_cache_branch = true; + } + + Ok(generated) +} + +fn create_session(path: &Path) -> Result { + let builder = Session::builder().map_err(|error| anyhow::anyhow!("{error}"))?; + let builder = builder + .with_optimization_level(GraphOptimizationLevel::Level3) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let mut builder = builder + .with_intra_threads(1) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let session = builder + .commit_from_file(path) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok(session) +} + +fn concatenate_sequence_embeddings(text: &TensorData, image: &TensorData) -> Result { + if text.shape.len() != 3 || image.shape.len() != 3 { + anyhow::bail!("Expected 3D text and image embeddings"); + } + if text.shape[0] != image.shape[0] || text.shape[2] != image.shape[2] { + anyhow::bail!("Text and image embedding dimensions do not match"); + } + + let text_tokens = text.shape[1] as usize; + let image_tokens = image.shape[1] as usize; + let dim = text.shape[2] as usize; + let mut values = Vec::with_capacity((text_tokens + image_tokens) * dim); + values.extend_from_slice(&text.values); + values.extend_from_slice(&image.values); + + Ok(TensorData { + shape: vec![text.shape[0], text.shape[1] + image.shape[1], text.shape[2]], + values, + }) +} + +fn tensor_data(value: &ort::value::DynValue) -> Result { + let (shape, values) = value + .try_extract_tensor::() + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok(TensorData { + shape: shape.to_vec(), + values: values.to_vec(), + }) +} + +fn tensor_from_data(data: &TensorData) -> Result> { + Tensor::from_array(( + Shape::new(data.shape.clone()), + data.values.clone().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}")) +} + +fn push_tensor_input( + inputs: &mut Vec<(Cow<'_, str>, SessionInputValue<'_>)>, + name: String, + data: TensorData, +) -> Result<()> { + inputs.push((Cow::Owned(name), tensor_from_data(&data)?.into())); + Ok(()) +} + +fn argmax_last_token(logits: &TensorData) -> Result { + if logits.shape.len() != 3 { + anyhow::bail!("Expected decoder logits to be 3D"); + } + let sequence_length = logits.shape[1] as usize; + let vocab_size = logits.shape[2] as usize; + if sequence_length == 0 || vocab_size == 0 { + anyhow::bail!("Decoder logits are empty"); + } + let start = (sequence_length - 1) * vocab_size; + let (index, _) = logits.values[start..start + vocab_size] + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.total_cmp(b)) + .ok_or_else(|| anyhow::anyhow!("Decoder logits are empty"))?; + Ok(index as i64) +} + +fn clean_caption(caption: &str) -> String { + caption + .trim() + .trim_matches(|ch| matches!(ch, '<' | '>' | '|' | ' ')) + .to_string() +} + +fn preprocess_image(image_path: &Path) -> Result> { + let image = ImageReader::open(image_path)?.decode()?.to_rgb8(); + let resized = image::imageops::resize(&image, 768, 768, FilterType::CatmullRom); + let mut pixel_values = vec![0.0f32; 3 * 768 * 768]; + let mean = [0.485f32, 0.456, 0.406]; + let std = [0.229f32, 0.224, 0.225]; + + for (x, y, pixel) in resized.enumerate_pixels() { + let x = x as usize; + let y = y as usize; + let base = y * 768 + x; + for channel in 0..3 { + let value = f32::from(pixel[channel]) / 255.0; + pixel_values[channel * 768 * 768 + base] = (value - mean[channel]) / std[channel]; + } + } + + Ok(pixel_values) +} + +impl TensorData { + fn zeros(shape: Vec) -> Self { + let values_len = shape.iter().map(|value| (*value).max(0) as usize).product(); + Self { + shape, + values: vec![0.0; values_len], + } + } +} diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index fb16359..d28e7d5 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -1,10 +1,11 @@ +use crate::captioner::{self, CaptionModelStatus, CaptionRuntimeProbe, CaptionVisionProbe}; use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord}; use crate::embedder; use crate::indexer; use crate::vector; use serde::{Deserialize, Serialize}; use std::path::PathBuf; -use tauri::{AppHandle, State}; +use tauri::{AppHandle, Emitter, Manager, State}; pub type DbState = DbPool; @@ -41,11 +42,46 @@ pub struct FindSimilarImagesParams { pub limit: Option, } +#[derive(Deserialize)] +pub struct DebugSimilarImagesParams { + pub image_id: i64, + pub limit: Option, +} + #[derive(Deserialize)] pub struct RetryFailedEmbeddingsParams { pub folder_id: i64, } +#[derive(Deserialize)] +pub struct SetGeneratedCaptionParams { + pub image_id: i64, + pub caption: String, + pub model: Option, +} + +#[derive(Deserialize)] +pub struct SuggestImageTagsParams { + pub image_id: i64, + pub limit: Option, +} + +#[derive(Deserialize)] +pub struct QueueCaptionJobsParams { + pub folder_id: Option, + pub image_id: Option, +} + +#[derive(Deserialize)] +pub struct ProbeCaptionImageParams { + pub image_id: i64, +} + +#[derive(Deserialize)] +pub struct GenerateCaptionParams { + pub image_id: i64, +} + #[derive(Deserialize)] pub struct SemanticSearchParams { pub query: String, @@ -127,8 +163,15 @@ pub async fn get_images( let favorites_only = params.favorites_only.unwrap_or(false); let embedding_failed_only = params.embedding_failed_only.unwrap_or(false); - let total = db::count_images(&conn, params.folder_id, search, media_kind, favorites_only, embedding_failed_only) - .map_err(|e| e.to_string())?; + let total = db::count_images( + &conn, + params.folder_id, + search, + media_kind, + favorites_only, + embedding_failed_only, + ) + .map_err(|e| e.to_string())?; let images = db::get_images( &conn, @@ -188,11 +231,45 @@ pub async fn find_similar_images( ) -> Result, String> { let conn = db.get().map_err(|e| e.to_string())?; let limit = params.limit.unwrap_or(32); - let image_ids = vector::find_similar_image_ids(&conn, params.image_id, limit) - .map_err(|e| e.to_string())?; + if !vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())? { + db::repair_embedding_consistency(&conn).map_err(|e| e.to_string())?; + return Ok(Vec::new()); + } + let image_ids = + vector::find_similar_image_ids(&conn, params.image_id, limit).map_err(|e| e.to_string())?; db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string()) } +#[derive(Serialize)] +pub struct SimilarImagesDebug { + pub image_id: i64, + pub vector_count: i64, + pub has_vector: bool, + pub similar_ids: Vec, +} + +#[tauri::command] +pub async fn debug_similar_images( + db: State<'_, DbState>, + params: DebugSimilarImagesParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let limit = params.limit.unwrap_or(32); + let vector_count = vector::count_image_vectors(&conn).map_err(|e| e.to_string())?; + let has_vector = vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())?; + let similar_ids = if has_vector { + vector::find_similar_image_ids(&conn, params.image_id, limit).map_err(|e| e.to_string())? + } else { + Vec::new() + }; + Ok(SimilarImagesDebug { + image_id: params.image_id, + vector_count, + has_vector, + similar_ids, + }) +} + #[tauri::command] pub async fn retry_failed_embeddings( db: State<'_, DbState>, @@ -211,7 +288,8 @@ pub async fn semantic_search_images( let conn = db.get().map_err(|e| e.to_string())?; let limit = params.limit.unwrap_or(64); - let ids = vector::search_image_ids_by_embedding(&conn, &embedding, limit).map_err(|e| e.to_string())?; + let ids = vector::search_image_ids_by_embedding(&conn, &embedding, limit) + .map_err(|e| e.to_string())?; let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?; if let Some(folder_id) = params.folder_id { @@ -227,6 +305,142 @@ pub async fn semantic_search_images( Ok(images) } +#[tauri::command] +pub async fn set_generated_caption( + db: State<'_, DbState>, + params: SetGeneratedCaptionParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let model = params.model.as_deref().unwrap_or("manual"); + db::update_generated_caption(&conn, params.image_id, ¶ms.caption, model) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn suggest_image_tags( + db: State<'_, DbState>, + params: SuggestImageTagsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::suggest_tags_from_caption(&conn, params.image_id, params.limit.unwrap_or(2)) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn get_caption_model_status(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(captioner::caption_model_status(&app_dir)) +} + +#[tauri::command] +pub async fn prepare_caption_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || { + let app = app.clone(); + captioner::prepare_caption_model_with_progress(&app_dir, move |progress| { + let _ = app.emit("caption-model-progress", progress); + }) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn delete_caption_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || captioner::delete_caption_model(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn probe_caption_runtime(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || captioner::probe_caption_runtime(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn probe_caption_image( + app: AppHandle, + db: State<'_, DbState>, + params: ProbeCaptionImageParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + let image_path = { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_image_by_id(&conn, params.image_id) + .map(|image| image.path) + .map_err(|e| e.to_string())? + }; + tauri::async_runtime::spawn_blocking(move || { + captioner::probe_caption_vision(&app_dir, std::path::Path::new(&image_path)) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn generate_caption_for_image( + app: AppHandle, + db: State<'_, DbState>, + params: GenerateCaptionParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + let image_path = { + let conn = db.get().map_err(|e| e.to_string())?; + let image = db::get_image_by_id(&conn, params.image_id).map_err(|e| e.to_string())?; + if image.media_kind != "image" { + return Err("AI captions can only be generated for images".to_string()); + } + image.path + }; + + let caption = tauri::async_runtime::spawn_blocking(move || { + captioner::generate_caption(&app_dir, std::path::Path::new(&image_path)) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|error| { + if let Ok(conn) = db.get() { + let _ = db::mark_caption_failed(&conn, params.image_id, &error.to_string()); + } + error.to_string() + })?; + + let conn = db.get().map_err(|e| e.to_string())?; + db::update_generated_caption( + &conn, + params.image_id, + &caption, + captioner::FLORENCE_CAPTION_MODEL_NAME, + ) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn queue_caption_jobs( + db: State<'_, DbState>, + params: QueueCaptionJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + match (params.folder_id, params.image_id) { + (_, Some(image_id)) => { + db::enqueue_caption_job(&conn, image_id).map_err(|e| e.to_string())?; + Ok(1) + } + (Some(folder_id), None) => { + db::enqueue_missing_caption_jobs_for_folder(&conn, folder_id).map_err(|e| e.to_string()) + } + (None, None) => db::enqueue_missing_caption_jobs(&conn).map_err(|e| e.to_string()), + } +} + #[derive(Serialize, Deserialize)] pub struct TagCloudEntry { pub count: usize, @@ -287,7 +501,10 @@ pub async fn get_tag_cloud( // Cache miss — run k-means let ids: Vec = embeddings_with_ids.iter().map(|(id, _)| *id).collect(); - let points: Vec> = embeddings_with_ids.into_iter().map(|(_, emb)| emb).collect(); + let points: Vec> = embeddings_with_ids + .into_iter() + .map(|(_, emb)| emb) + .collect(); let k = (n / 20).clamp(5, 30); let (centroids, cluster_counts, assignments) = kmeans_cosine(&points, k, 40); @@ -361,7 +578,10 @@ fn kmeans_cosine( let next = points .iter() .map(|p| { - let best_sim = centroids.iter().map(|c| dot(p, c)).fold(f32::NEG_INFINITY, f32::max); + let best_sim = centroids + .iter() + .map(|c| dot(p, c)) + .fold(f32::NEG_INFINITY, f32::max); 1.0 - best_sim // distance = 1 - cosine_similarity }) .enumerate() @@ -389,7 +609,9 @@ fn kmeans_cosine( changed = true; } } - if !changed { break; } + if !changed { + break; + } // Update step: mean of assigned points, then normalize let mut sums = vec![vec![0.0f32; dim]; k]; @@ -398,7 +620,9 @@ fn kmeans_cosine( sums[c].iter_mut().zip(p.iter()).for_each(|(s, v)| *s += v); counts[c] += 1; } - for (centroid, (sum, &count)) in centroids.iter_mut().zip(sums.iter_mut().zip(counts.iter())) { + for (centroid, (sum, &count)) in + centroids.iter_mut().zip(sums.iter_mut().zip(counts.iter())) + { if count > 0 { sum.iter_mut().for_each(|v| *v /= count as f32); normalize(sum); @@ -408,7 +632,9 @@ fn kmeans_cosine( } let mut counts = vec![0usize; k]; - for &a in &assignments { counts[a] += 1; } + for &a in &assignments { + counts[a] += 1; + } (centroids, counts, assignments) } @@ -438,24 +664,43 @@ pub async fn get_failed_embedding_images( } #[derive(Serialize)] -pub struct WorkerStates { +pub struct FolderWorkerStates { + pub folder_id: i64, pub thumbnail_paused: bool, pub metadata_paused: bool, pub embedding_paused: bool, + pub caption_paused: bool, } #[tauri::command] -pub async fn set_worker_paused(worker: String, paused: bool) -> Result<(), String> { - indexer::set_worker_paused(&worker, paused); +pub async fn set_worker_paused(worker: String, folder_id: i64, paused: bool) -> Result<(), String> { + indexer::set_worker_paused(&worker, folder_id, paused); Ok(()) } #[tauri::command] -pub async fn get_worker_states() -> Result { - let states = indexer::get_worker_paused_states(); - Ok(WorkerStates { - thumbnail_paused: states[0], - metadata_paused: states[1], - embedding_paused: states[2], - }) +pub async fn get_worker_states(folder_ids: Vec) -> Result, String> { + let states = indexer::get_worker_paused_states(&folder_ids); + Ok(folder_ids + .into_iter() + .map(|folder_id| { + let state = + states + .get(&folder_id) + .copied() + .unwrap_or(indexer::FolderWorkerPausedState { + thumbnail: false, + metadata: false, + embedding: false, + caption: false, + }); + FolderWorkerStates { + folder_id, + thumbnail_paused: state.thumbnail, + metadata_paused: state.metadata, + embedding_paused: state.embedding, + caption_paused: state.caption, + } + }) + .collect()) } diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 7fe7ead..b576687 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -57,6 +57,10 @@ pub struct ImageRecord { pub embedding_model: Option, pub embedding_updated_at: Option, pub embedding_error: Option, + pub generated_caption: Option, + pub caption_model: Option, + pub caption_updated_at: Option, + pub caption_error: Option, } #[allow(dead_code)] @@ -89,6 +93,13 @@ pub struct MetadataJob { pub path: String, } +#[derive(Debug, Clone)] +pub struct CaptionJob { + pub image_id: i64, + pub folder_id: i64, + pub path: String, +} + #[derive(Debug, Clone)] pub struct IndexedMediaEntry { pub id: i64, @@ -106,6 +117,9 @@ pub struct FolderJobProgress { pub embedding_pending: i64, pub embedding_ready: i64, pub embedding_failed: i64, + pub caption_pending: i64, + pub caption_ready: i64, + pub caption_failed: i64, } pub fn create_pool(db_path: &Path) -> Result { @@ -172,6 +186,15 @@ pub fn migrate(conn: &Connection) -> Result<()> { updated_at TEXT NOT NULL DEFAULT (datetime('now')) ); + CREATE TABLE IF NOT EXISTS caption_jobs ( + image_id INTEGER PRIMARY KEY REFERENCES images(id) ON DELETE CASCADE, + status TEXT NOT NULL DEFAULT 'pending', + attempts INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) + ); + CREATE TABLE IF NOT EXISTS tag_cloud_cache ( folder_scope TEXT PRIMARY KEY, image_ids_hash INTEGER NOT NULL, @@ -184,6 +207,7 @@ pub fn migrate(conn: &Connection) -> Result<()> { CREATE INDEX IF NOT EXISTS idx_embedding_jobs_status ON embedding_jobs(status); CREATE INDEX IF NOT EXISTS idx_thumbnail_jobs_status ON thumbnail_jobs(status); CREATE INDEX IF NOT EXISTS idx_metadata_jobs_status ON metadata_jobs(status); + CREATE INDEX IF NOT EXISTS idx_caption_jobs_status ON caption_jobs(status); ", )?; @@ -209,6 +233,10 @@ pub fn migrate(conn: &Connection) -> Result<()> { ensure_column(conn, "images", "audio_codec", "TEXT")?; ensure_column(conn, "images", "metadata_updated_at", "TEXT")?; ensure_column(conn, "images", "metadata_error", "TEXT")?; + ensure_column(conn, "images", "generated_caption", "TEXT")?; + ensure_column(conn, "images", "caption_model", "TEXT")?; + ensure_column(conn, "images", "caption_updated_at", "TEXT")?; + ensure_column(conn, "images", "caption_error", "TEXT")?; vector::migrate(conn)?; Ok(()) @@ -229,8 +257,8 @@ pub fn insert_folder(conn: &Connection, path: &str, name: &str) -> Result { pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { let id = conn.query_row( - "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22) + "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, generated_caption, caption_model, caption_updated_at, caption_error) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26) ON CONFLICT(path) DO UPDATE SET folder_id = excluded.folder_id, filename = excluded.filename, @@ -250,7 +278,11 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { embedding_status = excluded.embedding_status, embedding_model = excluded.embedding_model, embedding_updated_at = excluded.embedding_updated_at, - embedding_error = excluded.embedding_error + embedding_error = excluded.embedding_error, + generated_caption = excluded.generated_caption, + caption_model = excluded.caption_model, + caption_updated_at = excluded.caption_updated_at, + caption_error = excluded.caption_error RETURNING id", params![ img.folder_id, @@ -275,6 +307,10 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { img.embedding_model, img.embedding_updated_at, img.embedding_error, + img.generated_caption, + img.caption_model, + img.caption_updated_at, + img.caption_error, ], |row| row.get(0), )?; @@ -307,6 +343,51 @@ pub fn backfill_embedding_jobs(conn: &Connection) -> Result { Ok(inserted) } +pub fn repair_embedding_consistency(conn: &Connection) -> Result<(usize, usize)> { + let orphaned_vectors = vector::delete_orphaned_embeddings(conn)?; + let ready_ids = { + let mut stmt = conn.prepare( + "SELECT id + FROM images + WHERE embedding_status = 'ready'", + )?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + + let mut missing_vector_ids = Vec::new(); + for image_id in ready_ids { + if !vector::has_image_vector(conn, image_id)? { + missing_vector_ids.push(image_id); + } + } + + let tx = conn.unchecked_transaction()?; + for image_id in &missing_vector_ids { + tx.execute( + "INSERT INTO embedding_jobs (image_id, status, attempts, last_error, created_at, updated_at) + VALUES (?1, 'pending', 0, NULL, datetime('now'), datetime('now')) + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [image_id], + )?; + tx.execute( + "UPDATE images + SET embedding_status = 'pending', + embedding_error = NULL + WHERE id = ?1", + [image_id], + )?; + } + tx.commit()?; + + Ok((orphaned_vectors, missing_vector_ids.len())) +} + pub fn retry_failed_embedding_jobs(conn: &Connection, folder_id: i64) -> Result { // Only re-queue images that are actually embeddable right now. // Videos without a thumbnail would just fail again immediately, so skip them — @@ -348,6 +429,10 @@ pub fn reset_inflight_jobs(conn: &Connection) -> Result<()> { "UPDATE embedding_jobs SET status = 'pending' WHERE status = 'processing'", [], )?; + conn.execute( + "UPDATE caption_jobs SET status = 'pending' WHERE status = 'processing'", + [], + )?; Ok(()) } @@ -377,16 +462,102 @@ pub fn enqueue_metadata_job(conn: &Connection, image_id: i64) -> Result<()> { Ok(()) } -pub fn get_pending_embedding_jobs(conn: &Connection, limit: usize) -> Result> { - let mut stmt = conn.prepare( +pub fn enqueue_caption_job(conn: &Connection, image_id: i64) -> Result<()> { + conn.execute( + "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at) + VALUES (?1, 'pending', 0, NULL, datetime('now'), datetime('now')) + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [image_id], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE id = ?1", + [image_id], + )?; + Ok(()) +} + +pub fn requeue_caption_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()> { + for image_id in image_ids { + conn.execute( + "UPDATE caption_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'processing'", + [image_id], + )?; + } + Ok(()) +} + +pub fn enqueue_missing_caption_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result { + let inserted = conn.execute( + "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at) + SELECT id, 'pending', 0, NULL, datetime('now'), datetime('now') + FROM images + WHERE folder_id = ?1 + AND media_kind = 'image' + AND generated_caption IS NULL + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [folder_id], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE folder_id = ?1 + AND media_kind = 'image' + AND generated_caption IS NULL", + [folder_id], + )?; + Ok(inserted) +} + +pub fn enqueue_missing_caption_jobs(conn: &Connection) -> Result { + let inserted = conn.execute( + "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at) + SELECT id, 'pending', 0, NULL, datetime('now'), datetime('now') + FROM images + WHERE media_kind = 'image' + AND generated_caption IS NULL + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE media_kind = 'image' + AND generated_caption IS NULL", + [], + )?; + Ok(inserted) +} + +pub fn get_pending_embedding_jobs( + conn: &Connection, + excluded_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let sql = format!( "SELECT j.image_id, i.folder_id, i.path, i.thumbnail_path, i.media_kind, j.status, j.attempts, j.last_error, j.created_at, j.updated_at FROM embedding_jobs j JOIN images i ON i.id = j.image_id WHERE status = 'pending' + {} ORDER BY j.updated_at, j.image_id LIMIT ?1", - )?; + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; let rows = stmt.query_map([limit as i64], |row| { Ok(EmbeddingJob { image_id: row.get(0)?, @@ -404,9 +575,13 @@ pub fn get_pending_embedding_jobs(conn: &Connection, limit: usize) -> Result>>()?) } -pub fn claim_embedding_jobs(conn: &mut Connection, limit: usize) -> Result> { +pub fn claim_embedding_jobs( + conn: &mut Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { let tx = conn.transaction()?; - let candidates = get_pending_embedding_jobs(&tx, limit * 2)?; + let candidates = get_pending_embedding_jobs(&tx, paused_folder_ids, limit * 2)?; let mut claimed = Vec::with_capacity(limit); for job in candidates { @@ -430,6 +605,64 @@ pub fn claim_embedding_jobs(conn: &mut Connection, limit: usize) -> Result, + limit: usize, +) -> Result> { + let sql = format!( + "SELECT j.image_id, i.folder_id, i.path + FROM caption_jobs j + JOIN images i ON i.id = j.image_id + WHERE j.status = 'pending' + AND i.media_kind = 'image' + AND i.generated_caption IS NULL + {} + ORDER BY j.updated_at, j.image_id + LIMIT ?1", + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map([limit as i64], |row| { + Ok(CaptionJob { + image_id: row.get(0)?, + folder_id: row.get(1)?, + path: row.get(2)?, + }) + })?; + Ok(rows.collect::>>()?) +} + +pub fn claim_caption_jobs( + conn: &mut Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let tx = conn.transaction()?; + let candidates = get_pending_caption_jobs_excluding(&tx, paused_folder_ids, limit * 2)?; + let mut claimed = Vec::with_capacity(limit); + + for job in candidates { + let updated = tx.execute( + "UPDATE caption_jobs + SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'pending'", + [job.image_id], + )?; + + if updated == 1 { + claimed.push(job); + } + + if claimed.len() >= limit { + break; + } + } + + tx.commit()?; + Ok(claimed) +} + #[allow(dead_code)] pub fn mark_embedding_ready(conn: &Connection, image_id: i64, model: &str) -> Result<()> { conn.execute( @@ -489,6 +722,7 @@ pub fn delete_images_by_ids(conn: &Connection, image_ids: &[i64]) -> Result<()> let tx = conn.unchecked_transaction()?; for image_id in image_ids { vector::delete_embedding(&tx, *image_id)?; + vector::delete_caption_embedding(&tx, *image_id)?; tx.execute("DELETE FROM images WHERE id = ?1", [image_id])?; } tx.commit()?; @@ -539,6 +773,31 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result Result Result Result> { - let mut stmt = conn.prepare( +fn get_pending_thumbnail_jobs_excluding( + conn: &Connection, + excluded_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let sql = format!( "SELECT j.image_id, i.folder_id, i.path, i.media_kind FROM thumbnail_jobs j JOIN images i ON i.id = j.image_id WHERE j.status = 'pending' + {} ORDER BY j.updated_at, j.image_id LIMIT ?1", - )?; + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; let rows = stmt.query_map([limit as i64], |row| { Ok(ThumbnailJob { image_id: row.get(0)?, @@ -585,18 +854,20 @@ pub fn get_pending_thumbnail_jobs(conn: &Connection, limit: usize) -> Result, + paused_folder_ids: &std::collections::HashSet, fetch_limit: usize, claim_limit: usize, ) -> Result> { let tx = conn.transaction()?; - let candidates = get_pending_thumbnail_jobs(&tx, fetch_limit)?; + let excluded_folder_ids = active_folder_ids + .union(paused_folder_ids) + .copied() + .collect::>(); + let candidates = get_pending_thumbnail_jobs_excluding(&tx, &excluded_folder_ids, fetch_limit)?; let mut claimed = Vec::with_capacity(claim_limit); for job in candidates { - if active_folder_ids.contains(&job.folder_id) { - continue; - } - + debug_assert!(!excluded_folder_ids.contains(&job.folder_id)); let updated = tx.execute( "UPDATE thumbnail_jobs SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') @@ -617,15 +888,22 @@ pub fn claim_thumbnail_jobs( Ok(claimed) } -pub fn get_pending_metadata_jobs(conn: &Connection, limit: usize) -> Result> { - let mut stmt = conn.prepare( +fn get_pending_metadata_jobs_excluding( + conn: &Connection, + excluded_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let sql = format!( "SELECT j.image_id, i.folder_id, i.path FROM metadata_jobs j JOIN images i ON i.id = j.image_id WHERE j.status = 'pending' AND i.media_kind = 'video' + {} ORDER BY j.updated_at, j.image_id LIMIT ?1", - )?; + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; let rows = stmt.query_map([limit as i64], |row| { Ok(MetadataJob { image_id: row.get(0)?, @@ -639,18 +917,20 @@ pub fn get_pending_metadata_jobs(conn: &Connection, limit: usize) -> Result, + paused_folder_ids: &std::collections::HashSet, fetch_limit: usize, claim_limit: usize, ) -> Result> { let tx = conn.transaction()?; - let candidates = get_pending_metadata_jobs(&tx, fetch_limit)?; + let excluded_folder_ids = active_folder_ids + .union(paused_folder_ids) + .copied() + .collect::>(); + let candidates = get_pending_metadata_jobs_excluding(&tx, &excluded_folder_ids, fetch_limit)?; let mut claimed = Vec::with_capacity(claim_limit); for job in candidates { - if active_folder_ids.contains(&job.folder_id) { - continue; - } - + debug_assert!(!excluded_folder_ids.contains(&job.folder_id)); let updated = tx.execute( "UPDATE metadata_jobs SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') @@ -765,7 +1045,8 @@ pub fn update_image_details( conn.query_row( "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, - favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error + favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, + generated_caption, caption_model, caption_updated_at, caption_error FROM images WHERE id = ?1", [image_id], @@ -778,7 +1059,8 @@ pub fn get_image_by_id(conn: &Connection, image_id: i64) -> Result conn.query_row( "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, - favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error + favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, + generated_caption, caption_model, caption_updated_at, caption_error FROM images WHERE id = ?1", [image_id], @@ -790,7 +1072,11 @@ pub fn get_image_by_id(conn: &Connection, image_id: i64) -> Result pub fn get_images_by_ids(conn: &Connection, image_ids: &[i64]) -> Result> { let mut images = Vec::with_capacity(image_ids.len()); for image_id in image_ids { - images.push(get_image_by_id(conn, *image_id)?); + match get_image_by_id(conn, *image_id) { + Ok(image) => images.push(image), + Err(error) if is_query_returned_no_rows(&error) => {} + Err(error) => return Err(error), + } } Ok(images) } @@ -839,7 +1125,8 @@ pub fn get_images( let sql = format!( "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, - favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error + favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, + generated_caption, caption_model, caption_updated_at, caption_error FROM images WHERE (?1 IS NULL OR folder_id = ?1) AND (?2 IS NULL OR filename LIKE ?2) @@ -918,11 +1205,104 @@ pub fn get_failed_embedding_images( Ok(rows.collect::>>()?) } -pub fn delete_folder(conn: &Connection, folder_id: i64) -> Result<()> { - conn.execute("DELETE FROM folders WHERE id = ?1", params![folder_id])?; +pub fn update_generated_caption( + conn: &Connection, + image_id: i64, + caption: &str, + model: &str, +) -> Result { + conn.execute( + "UPDATE images + SET generated_caption = ?2, + caption_model = ?3, + caption_updated_at = datetime('now'), + caption_error = NULL + WHERE id = ?1", + params![image_id, caption, model], + )?; + conn.execute("DELETE FROM caption_jobs WHERE image_id = ?1", [image_id])?; + get_image_by_id(conn, image_id) +} + +pub fn mark_caption_failed(conn: &Connection, image_id: i64, error: &str) -> Result<()> { + conn.execute( + "UPDATE images + SET caption_error = ?2 + WHERE id = ?1", + params![image_id, error], + )?; + conn.execute( + "UPDATE caption_jobs + SET status = 'failed', last_error = ?2, updated_at = datetime('now') + WHERE image_id = ?1", + params![image_id, error], + )?; Ok(()) } +pub fn suggest_tags_from_caption( + conn: &Connection, + image_id: i64, + limit: usize, +) -> Result> { + let caption = conn.query_row( + "SELECT generated_caption FROM images WHERE id = ?1", + [image_id], + |row| row.get::<_, Option>(0), + )?; + + Ok(caption + .as_deref() + .map(|caption| derive_caption_tags(caption, limit)) + .unwrap_or_default()) +} + +pub fn delete_folder(conn: &Connection, folder_id: i64) -> Result<()> { + let image_ids = { + let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?; + let rows = stmt + .query_map([folder_id], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + + let tx = conn.unchecked_transaction()?; + for image_id in image_ids { + vector::delete_embedding(&tx, image_id)?; + vector::delete_caption_embedding(&tx, image_id)?; + } + tx.execute("DELETE FROM folders WHERE id = ?1", params![folder_id])?; + tx.commit()?; + Ok(()) +} + +fn derive_caption_tags(caption: &str, limit: usize) -> Vec { + const STOPWORDS: &[&str] = &[ + "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "in", "into", "is", "it", + "near", "of", "on", "or", "the", "to", "with", "without", "image", "photo", "picture", + "showing", "shows", "there", "this", "that", "over", "under", + ]; + + let stopwords = STOPWORDS + .iter() + .copied() + .collect::>(); + let mut tags = Vec::new(); + for word in caption + .split(|ch: char| !ch.is_alphanumeric()) + .map(|word| word.trim().to_lowercase()) + .filter(|word| word.len() >= 3 && !stopwords.contains(word.as_str())) + { + if !tags.contains(&word) { + tags.push(word); + } + if tags.len() >= limit { + break; + } + } + tags +} + fn map_image_row(row: &Row<'_>) -> rusqlite::Result { Ok(ImageRecord { id: row.get(0)?, @@ -948,6 +1328,10 @@ fn map_image_row(row: &Row<'_>) -> rusqlite::Result { embedding_model: row.get(20)?, embedding_updated_at: row.get(21)?, embedding_error: row.get(22)?, + generated_caption: row.get(23)?, + caption_model: row.get(24)?, + caption_updated_at: row.get(25)?, + caption_error: row.get(26)?, }) } @@ -1005,3 +1389,27 @@ fn ensure_column(conn: &Connection, table: &str, column: &str, definition: &str) )?; Ok(()) } + +fn is_query_returned_no_rows(error: &anyhow::Error) -> bool { + error + .downcast_ref::() + .is_some_and(|error| matches!(error, rusqlite::Error::QueryReturnedNoRows)) +} + +fn folder_exclusion_clause( + image_alias: &str, + excluded_folder_ids: &std::collections::HashSet, +) -> String { + if excluded_folder_ids.is_empty() { + return String::new(); + } + + let mut ids = excluded_folder_ids.iter().copied().collect::>(); + ids.sort_unstable(); + let id_list = ids + .into_iter() + .map(|id| id.to_string()) + .collect::>() + .join(","); + format!("AND {}.folder_id NOT IN ({})", image_alias, id_list) +} diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index cd7e8a2..33bcbcd 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -1,3 +1,4 @@ +use crate::captioner::{self, FlorenceCaptioner}; use crate::db::{self, DbPool, EmbeddingJob, FolderJobProgress, ImageRecord, IndexedMediaEntry}; use crate::embedder::{embedding_source_path, ClipImageEmbedder}; use crate::media::{probe_video_metadata, MediaTools}; @@ -9,7 +10,6 @@ use rayon::prelude::*; use serde::Serialize; use std::collections::{HashMap, HashSet}; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Mutex, OnceLock}; use std::time::{Duration, Instant}; use tauri::{AppHandle, Emitter}; @@ -22,29 +22,91 @@ const IMAGE_EXTENSIONS: &[&str] = &[ const VIDEO_EXTENSIONS: &[&str] = &["mp4", "mov", "m4v", "webm"]; const JOB_PROGRESS_EMIT_INTERVAL: Duration = Duration::from_millis(750); +const CAPTION_BATCH_SIZE: usize = 1; static LAST_JOB_PROGRESS_EMIT: OnceLock>> = OnceLock::new(); static ACTIVE_INDEXING_FOLDERS: OnceLock>> = OnceLock::new(); -static THUMBNAIL_WORKER_PAUSED: AtomicBool = AtomicBool::new(false); -static METADATA_WORKER_PAUSED: AtomicBool = AtomicBool::new(false); -static EMBEDDING_WORKER_PAUSED: AtomicBool = AtomicBool::new(false); +static PAUSED_WORKER_FOLDERS: OnceLock> = OnceLock::new(); -pub fn set_worker_paused(worker: &str, paused: bool) { - match worker { - "thumbnail" => THUMBNAIL_WORKER_PAUSED.store(paused, Ordering::Relaxed), - "metadata" => METADATA_WORKER_PAUSED.store(paused, Ordering::Relaxed), - "embedding" => EMBEDDING_WORKER_PAUSED.store(paused, Ordering::Relaxed), - _ => {} +#[derive(Default)] +struct PausedWorkerFolders { + thumbnail: HashSet, + metadata: HashSet, + embedding: HashSet, + caption: HashSet, +} + +#[derive(Clone, Copy)] +pub struct FolderWorkerPausedState { + pub thumbnail: bool, + pub metadata: bool, + pub embedding: bool, + pub caption: bool, +} + +pub fn set_worker_paused(worker: &str, folder_id: i64, paused: bool) { + if let Ok(mut paused_folders) = PAUSED_WORKER_FOLDERS + .get_or_init(|| Mutex::new(PausedWorkerFolders::default())) + .lock() + { + let folder_set = match worker { + "thumbnail" => Some(&mut paused_folders.thumbnail), + "metadata" => Some(&mut paused_folders.metadata), + "embedding" => Some(&mut paused_folders.embedding), + "caption" => Some(&mut paused_folders.caption), + _ => None, + }; + + if let Some(folder_set) = folder_set { + if paused { + folder_set.insert(folder_id); + } else { + folder_set.remove(&folder_id); + } + } } } -pub fn get_worker_paused_states() -> [bool; 3] { - [ - THUMBNAIL_WORKER_PAUSED.load(Ordering::Relaxed), - METADATA_WORKER_PAUSED.load(Ordering::Relaxed), - EMBEDDING_WORKER_PAUSED.load(Ordering::Relaxed), - ] +pub fn get_worker_paused_states(folder_ids: &[i64]) -> HashMap { + let Ok(paused_folders) = PAUSED_WORKER_FOLDERS + .get_or_init(|| Mutex::new(PausedWorkerFolders::default())) + .lock() + else { + return HashMap::new(); + }; + + folder_ids + .iter() + .copied() + .map(|folder_id| { + ( + folder_id, + FolderWorkerPausedState { + thumbnail: paused_folders.thumbnail.contains(&folder_id), + metadata: paused_folders.metadata.contains(&folder_id), + embedding: paused_folders.embedding.contains(&folder_id), + caption: paused_folders.caption.contains(&folder_id), + }, + ) + }) + .collect() +} + +fn paused_folder_ids(worker: &str) -> HashSet { + let Ok(paused_folders) = PAUSED_WORKER_FOLDERS + .get_or_init(|| Mutex::new(PausedWorkerFolders::default())) + .lock() + else { + return HashSet::new(); + }; + + match worker { + "thumbnail" => paused_folders.thumbnail.clone(), + "metadata" => paused_folders.metadata.clone(), + "embedding" => paused_folders.embedding.clone(), + _ => HashSet::new(), + } } static FOLDER_STORAGE_PROFILES: OnceLock>> = OnceLock::new(); @@ -95,10 +157,6 @@ pub fn start_thumbnail_worker( cache_dir: PathBuf, ) { std::thread::spawn(move || loop { - if THUMBNAIL_WORKER_PAUSED.load(Ordering::Relaxed) { - std::thread::sleep(std::time::Duration::from_millis(500)); - continue; - } if let Err(error) = process_thumbnail_batch(&app, &pool, &media_tools, &cache_dir) { eprintln!("Thumbnail worker error: {}", error); } @@ -108,10 +166,6 @@ pub fn start_thumbnail_worker( pub fn start_metadata_worker(app: AppHandle, pool: DbPool, media_tools: MediaTools) { std::thread::spawn(move || loop { - if METADATA_WORKER_PAUSED.load(Ordering::Relaxed) { - std::thread::sleep(std::time::Duration::from_millis(500)); - continue; - } if let Err(error) = process_metadata_batch(&app, &pool, &media_tools) { eprintln!("Metadata worker error: {}", error); } @@ -124,10 +178,6 @@ pub fn start_embedding_worker(app: AppHandle, pool: DbPool) { let mut embedder: Option = None; println!("Embedding worker started."); loop { - if EMBEDDING_WORKER_PAUSED.load(Ordering::Relaxed) { - std::thread::sleep(std::time::Duration::from_millis(500)); - continue; - } if let Err(error) = process_embedding_batch(&app, &pool, &mut embedder) { eprintln!("Embedding worker error: {}", error); } @@ -136,6 +186,20 @@ pub fn start_embedding_worker(app: AppHandle, pool: DbPool) { }); } +pub fn start_caption_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) { + std::thread::spawn(move || { + let mut captioner: Option = None; + println!("Caption worker started."); + loop { + if let Err(error) = process_caption_batch(&app, &pool, &app_data_dir, &mut captioner) { + eprintln!("Caption worker error: {}", error); + captioner = None; + } + std::thread::sleep(std::time::Duration::from_millis(750)); + } + }); +} + fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> { let existing_entries = { let conn = pool.get()?; @@ -197,7 +261,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) images: committed, }, ); - emit_folder_job_progress(&app, &pool, &[folder_id]); + emit_folder_job_progress(&app, &pool, &[folder_id], false); } processed += path_chunk.len(); @@ -245,7 +309,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) done: true, }, ); - emit_folder_job_progress(&app, &pool, &[folder_id]); + emit_folder_job_progress(&app, &pool, &[folder_id], true); Ok(()) } @@ -302,6 +366,10 @@ fn build_record( embedding_model: Some(vector::CLIP_MODEL_NAME.to_string()), embedding_updated_at: None, embedding_error: None, + generated_caption: None, + caption_model: None, + caption_updated_at: None, + caption_error: None, }) } @@ -335,11 +403,13 @@ fn process_thumbnail_batch( with_db_write_lock(|| { let mut conn = pool.get()?; let active_folders = active_indexing_folders(); + let paused_folders = paused_folder_ids("thumbnail"); let worker_batch_size = max_worker_batch_size(&active_folders); let worker_fetch_size = max_worker_fetch_size(&active_folders); db::claim_thumbnail_jobs( &mut conn, &active_folders, + &paused_folders, worker_fetch_size, worker_batch_size, ) @@ -428,7 +498,7 @@ fn process_thumbnail_batch( images: updated_images, }, ); - emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>()); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); } Ok(()) @@ -439,11 +509,13 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo with_db_write_lock(|| { let mut conn = pool.get()?; let active_folders = active_indexing_folders(); + let paused_folders = paused_folder_ids("metadata"); let worker_batch_size = max_worker_batch_size(&active_folders); let worker_fetch_size = max_worker_fetch_size(&active_folders); db::claim_metadata_jobs( &mut conn, &active_folders, + &paused_folders, worker_fetch_size, worker_batch_size, ) @@ -506,7 +578,7 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo images: updated_images, }, ); - emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>()); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); } Ok(()) @@ -518,14 +590,11 @@ fn process_embedding_batch( embedder: &mut Option, ) -> Result<()> { let batch_started_at = Instant::now(); - if embedder.is_none() { - *embedder = Some(ClipImageEmbedder::new()?); - } - let claim_started_at = Instant::now(); + let paused_folders = paused_folder_ids("embedding"); let jobs = with_db_write_lock(|| { let mut conn = pool.get()?; - db::claim_embedding_jobs(&mut conn, EMBEDDING_BATCH_SIZE) + db::claim_embedding_jobs(&mut conn, &paused_folders, EMBEDDING_BATCH_SIZE) })?; let claim_elapsed = claim_started_at.elapsed(); @@ -533,9 +602,18 @@ fn process_embedding_batch( return Ok(()); } + if embedder.is_none() { + *embedder = Some(ClipImageEmbedder::new()?); + } + println!("Embedding batch claimed: {} items", jobs.len()); let folder_ids = jobs.iter().map(|job| job.folder_id).collect::>(); - emit_folder_job_progress(app, pool, &folder_ids.iter().copied().collect::>()); + emit_folder_job_progress( + app, + pool, + &folder_ids.iter().copied().collect::>(), + false, + ); let embedder = embedder.as_ref().expect("embedder should be initialized"); let infer_started_at = Instant::now(); @@ -639,7 +717,7 @@ fn process_embedding_batch( images: updated_images, }, ); - emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>()); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); } let write_elapsed = write_started_at.elapsed(); @@ -652,6 +730,101 @@ fn process_embedding_batch( Ok(()) } +fn process_caption_batch( + app: &AppHandle, + pool: &DbPool, + app_data_dir: &Path, + captioner: &mut Option, +) -> Result<()> { + if !captioner::caption_model_status(app_data_dir).ready { + return Ok(()); + } + + let paused_folders = paused_folder_ids("caption"); + let jobs = with_db_write_lock(|| { + let mut conn = pool.get()?; + db::claim_caption_jobs(&mut conn, &paused_folders, CAPTION_BATCH_SIZE) + })?; + + if jobs.is_empty() { + return Ok(()); + } + + if captioner.is_none() { + match FlorenceCaptioner::new(app_data_dir) { + Ok(model) => *captioner = Some(model), + Err(error) => { + with_db_write_lock(|| { + let conn = pool.get()?; + db::requeue_caption_jobs( + &conn, + &jobs.iter().map(|job| job.image_id).collect::>(), + ) + })?; + return Err(error); + } + } + } + + let folder_ids = jobs.iter().map(|job| job.folder_id).collect::>(); + emit_folder_job_progress( + app, + pool, + &folder_ids.iter().copied().collect::>(), + false, + ); + + let captioner = captioner + .as_mut() + .expect("captioner should be initialized before caption batch processing"); + let caption_results = jobs + .iter() + .map(|job| (job.clone(), captioner.generate(Path::new(&job.path)))) + .collect::>(); + + let updated_images = with_db_write_lock(|| { + let mut conn = pool.get()?; + let tx = conn.transaction()?; + let mut updated_images = Vec::with_capacity(caption_results.len()); + + for (job, caption_result) in &caption_results { + match caption_result { + Ok(caption) => { + updated_images.push(db::update_generated_caption( + &tx, + job.image_id, + caption, + captioner::FLORENCE_CAPTION_MODEL_NAME, + )?); + } + Err(error) => { + db::mark_caption_failed(&tx, job.image_id, &error.to_string())?; + updated_images.push(db::get_image_by_id(&tx, job.image_id)?); + } + } + } + + tx.commit()?; + Ok(updated_images) + })?; + + if !updated_images.is_empty() { + let folder_ids = updated_images + .iter() + .map(|image| image.folder_id) + .collect::>(); + emit_media_updates( + app, + &MediaUpdateBatch { + images: updated_images, + }, + ); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); + } + + Ok(()) +} + fn active_indexing_folders() -> HashSet { ACTIVE_INDEXING_FOLDERS .get_or_init(|| Mutex::new(HashSet::new())) @@ -737,7 +910,7 @@ fn emit_media_updates(app: &AppHandle, batch: &MediaUpdateBatch) { let _ = app.emit("media-updated", batch); } -fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64]) { +fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64], force: bool) { let mut unique_folder_ids = folder_ids.iter().copied().collect::>(); unique_folder_ids.sort_unstable(); unique_folder_ids.dedup(); @@ -749,10 +922,11 @@ fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64]) Err(_) => return, }; unique_folder_ids.retain(|folder_id| { - let should_emit = tracker - .get(folder_id) - .map(|last_emit| now.duration_since(*last_emit) >= JOB_PROGRESS_EMIT_INTERVAL) - .unwrap_or(true); + let should_emit = force + || tracker + .get(folder_id) + .map(|last_emit| now.duration_since(*last_emit) >= JOB_PROGRESS_EMIT_INTERVAL) + .unwrap_or(true); if should_emit { tracker.insert(*folder_id, now); } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 3592fbf..bf4d4f6 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1,3 +1,4 @@ +mod captioner; mod commands; mod db; mod embedder; @@ -7,8 +8,8 @@ mod storage; mod thumbnail; mod vector; -use tauri::Manager; use crate::storage::StorageProfile; +use tauri::Manager; #[cfg_attr(mobile, tauri::mobile_entry_point)] pub fn run() { @@ -34,11 +35,19 @@ pub fn run() { let conn = pool.get().expect("Failed to get connection for migration"); db::migrate(&conn).expect("Failed to run migrations"); db::reset_inflight_jobs(&conn).expect("Failed to reset inflight jobs"); - let backfilled = db::backfill_embedding_jobs(&conn) - .expect("Failed to backfill embedding jobs"); + let backfilled = + db::backfill_embedding_jobs(&conn).expect("Failed to backfill embedding jobs"); if backfilled > 0 { println!("Backfilled {} embedding jobs.", backfilled); } + let (orphaned_vectors, missing_vectors) = db::repair_embedding_consistency(&conn) + .expect("Failed to repair embedding consistency"); + if orphaned_vectors > 0 || missing_vectors > 0 { + println!( + "Repaired embedding consistency: removed {} orphaned vectors, requeued {} missing vectors.", + orphaned_vectors, missing_vectors + ); + } } let thumb_dir = app_dir.join("thumbnails"); @@ -58,6 +67,7 @@ pub fn run() { } indexer::start_metadata_worker(app.handle().clone(), pool.clone(), media_tools.clone()); indexer::start_embedding_worker(app.handle().clone(), pool.clone()); + indexer::start_caption_worker(app.handle().clone(), pool.clone(), app_dir.clone()); app.manage(pool); app.manage(media_tools); @@ -73,8 +83,18 @@ pub fn run() { commands::reindex_folder, commands::update_image_details, commands::find_similar_images, + commands::debug_similar_images, commands::retry_failed_embeddings, commands::semantic_search_images, + commands::get_caption_model_status, + commands::prepare_caption_model, + commands::delete_caption_model, + commands::probe_caption_runtime, + commands::probe_caption_image, + commands::generate_caption_for_image, + commands::queue_caption_jobs, + commands::set_generated_caption, + commands::suggest_image_tags, commands::set_worker_paused, commands::get_worker_states, commands::get_tag_cloud, diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 61217d5..6b99c8b 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -1,5 +1,5 @@ use anyhow::{anyhow, Result}; -use rusqlite::{ffi::sqlite3_auto_extension, Connection}; +use rusqlite::{ffi::sqlite3_auto_extension, Connection, Error as SqliteError}; use sqlite_vec::sqlite3_vec_init; use std::sync::Once; @@ -19,8 +19,13 @@ pub fn migrate(conn: &Connection) -> Result<()> { "CREATE VIRTUAL TABLE IF NOT EXISTS image_vec USING vec0( image_id INTEGER PRIMARY KEY, embedding FLOAT[{}] distance_metric=cosine + ); + + CREATE VIRTUAL TABLE IF NOT EXISTS caption_vec USING vec0( + image_id INTEGER PRIMARY KEY, + embedding FLOAT[{}] distance_metric=cosine );", - CLIP_VECTOR_DIM + CLIP_VECTOR_DIM, CLIP_VECTOR_DIM ))?; Ok(()) } @@ -31,6 +36,12 @@ pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> { Ok(()) } +#[allow(dead_code)] +pub fn delete_caption_embedding(conn: &Connection, image_id: i64) -> Result<()> { + conn.execute("DELETE FROM caption_vec WHERE image_id = ?1", [image_id])?; + Ok(()) +} + #[allow(dead_code)] pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) -> Result<()> { if embedding.len() != CLIP_VECTOR_DIM { @@ -50,12 +61,35 @@ pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) -> Ok(()) } +#[allow(dead_code)] +pub fn upsert_caption_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) -> Result<()> { + if embedding.len() != CLIP_VECTOR_DIM { + return Err(anyhow!( + "expected {}-dimensional embedding, got {}", + CLIP_VECTOR_DIM, + embedding.len() + )); + } + + let packed = pack_f32(embedding); + conn.execute("DELETE FROM caption_vec WHERE image_id = ?1", [image_id])?; + conn.execute( + "INSERT INTO caption_vec (image_id, embedding) VALUES (?1, ?2)", + (&image_id, &packed), + )?; + Ok(()) +} + pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Result> { - let embedding: Vec = conn.query_row( + let embedding: Vec = match conn.query_row( "SELECT embedding FROM image_vec WHERE image_id = ?1", [image_id], |row| row.get(0), - )?; + ) { + Ok(embedding) => embedding, + Err(SqliteError::QueryReturnedNoRows) => return Ok(Vec::new()), + Err(error) => return Err(error.into()), + }; let mut stmt = conn.prepare( "SELECT image_id @@ -155,6 +189,115 @@ pub fn search_image_ids_by_embedding( Ok(ids) } +#[allow(dead_code)] +pub fn search_caption_ids_by_embedding( + conn: &Connection, + embedding: &[f32], + limit: usize, +) -> Result> { + if embedding.len() != CLIP_VECTOR_DIM { + return Err(anyhow!( + "expected {}-dimensional embedding, got {}", + CLIP_VECTOR_DIM, + embedding.len() + )); + } + + let packed = pack_f32(embedding); + let mut stmt = conn.prepare( + "SELECT image_id + FROM caption_vec + WHERE embedding MATCH vec_f32(?1) + AND k = ?2", + )?; + let rows = stmt.query_map((&packed, limit as i64), |row| row.get::<_, i64>(0))?; + + let mut ids = Vec::new(); + for row in rows { + ids.push(row?); + if ids.len() >= limit { + break; + } + } + Ok(ids) +} + +pub fn count_image_vectors(conn: &Connection) -> Result { + conn.query_row("SELECT COUNT(*) FROM image_vec", [], |row| row.get(0)) + .map_err(Into::into) +} + +#[allow(dead_code)] +pub fn count_caption_vectors(conn: &Connection) -> Result { + conn.query_row("SELECT COUNT(*) FROM caption_vec", [], |row| row.get(0)) + .map_err(Into::into) +} + +pub fn delete_orphaned_embeddings(conn: &Connection) -> Result { + let image_ids = { + let mut stmt = conn.prepare("SELECT id FROM images")?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + let vector_ids = { + let mut stmt = conn.prepare("SELECT image_id FROM image_vec")?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + let orphaned_ids = vector_ids + .into_iter() + .filter(|image_id| !image_ids.contains(image_id)) + .collect::>(); + + for image_id in &orphaned_ids { + delete_embedding(conn, *image_id)?; + } + + Ok(orphaned_ids.len()) +} + +#[allow(dead_code)] +pub fn delete_orphaned_caption_embeddings(conn: &Connection) -> Result { + let image_ids = { + let mut stmt = conn.prepare("SELECT id FROM images")?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + let vector_ids = { + let mut stmt = conn.prepare("SELECT image_id FROM caption_vec")?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + let orphaned_ids = vector_ids + .into_iter() + .filter(|image_id| !image_ids.contains(image_id)) + .collect::>(); + + for image_id in &orphaned_ids { + delete_caption_embedding(conn, *image_id)?; + } + + Ok(orphaned_ids.len()) +} + +pub fn has_image_vector(conn: &Connection, image_id: i64) -> Result { + conn.query_row( + "SELECT EXISTS(SELECT 1 FROM image_vec WHERE image_id = ?1)", + [image_id], + |row| row.get::<_, i64>(0), + ) + .map(|value| value != 0) + .map_err(Into::into) +} + #[allow(dead_code)] fn pack_f32(values: &[f32]) -> Vec { let mut out = Vec::with_capacity(values.len() * std::mem::size_of::()); diff --git a/src/App.tsx b/src/App.tsx index 9884233..4b6039b 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -7,17 +7,20 @@ import { Gallery } from "./components/Gallery"; import { Lightbox } from "./components/Lightbox"; import { TagCloud } from "./components/TagCloud"; import { TitleBar } from "./components/TitleBar"; +import { SettingsModal } from "./components/SettingsModal"; export default function App() { const loadFolders = useGalleryStore((state) => state.loadFolders); const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress); const loadImages = useGalleryStore((state) => state.loadImages); + const loadCaptionModelStatus = useGalleryStore((state) => state.loadCaptionModelStatus); const subscribeToProgress = useGalleryStore((state) => state.subscribeToProgress); const activeView = useGalleryStore((state) => state.activeView); useEffect(() => { loadFolders().then(() => { void loadBackgroundJobProgress(); + void loadCaptionModelStatus(); return loadImages(true); }); let unlisten: (() => void) | undefined; @@ -54,6 +57,7 @@ export default function App() { + ); } diff --git a/src/components/BackgroundTasks.tsx b/src/components/BackgroundTasks.tsx index 0433f63..6d4014c 100644 --- a/src/components/BackgroundTasks.tsx +++ b/src/components/BackgroundTasks.tsx @@ -2,12 +2,13 @@ import { useEffect, useMemo, useState } from "react"; import { invoke } from "@tauri-apps/api/core"; import { useGalleryStore } from "../store"; -type WorkerKey = "thumbnail" | "metadata" | "embedding"; +type WorkerKey = "thumbnail" | "metadata" | "embedding" | "caption"; const WORKER_FOR_STAGE: Record = { Thumbnails: "thumbnail", Metadata: "metadata", Embeddings: "embedding", + Captions: "caption", }; interface TaskStage { @@ -22,6 +23,7 @@ interface Task { name: string; stages: TaskStage[]; hasFailedEmbeddings: boolean; + hasFailedCaptions: boolean; pendingMediaWork: number; embeddingProcessed: number; embeddingTotal: number; @@ -35,6 +37,21 @@ interface FailedEmbeddingItem { error: string | null; } +interface FolderWorkerStates { + folder_id: number; + thumbnail_paused: boolean; + metadata_paused: boolean; + embedding_paused: boolean; + caption_paused: boolean; +} + +const DEFAULT_PAUSED_STATE: Record = { + thumbnail: false, + metadata: false, + embedding: false, + caption: false, +}; + export function BackgroundTasks() { const folders = useGalleryStore((state) => state.folders); const indexingProgress = useGalleryStore((state) => state.indexingProgress); @@ -42,24 +59,32 @@ export function BackgroundTasks() { const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings); const [expanded, setExpanded] = useState(false); const [dismissed, setDismissed] = useState>({}); - const [paused, setPaused] = useState>({ - thumbnail: false, - metadata: false, - embedding: false, - }); + const [paused, setPaused] = useState>>({}); const [failedItems, setFailedItems] = useState>({}); useEffect(() => { - invoke<{ thumbnail_paused: boolean; metadata_paused: boolean; embedding_paused: boolean }>( - "get_worker_states", - ).then((states) => { - setPaused({ - thumbnail: states.thumbnail_paused, - metadata: states.metadata_paused, - embedding: states.embedding_paused, - }); + const folderIds = folders.map((folder) => folder.id); + if (folderIds.length === 0) { + setPaused({}); + return; + } + + invoke("get_worker_states", { folderIds }).then((states) => { + setPaused( + Object.fromEntries( + states.map((state) => [ + state.folder_id, + { + thumbnail: state.thumbnail_paused, + metadata: state.metadata_paused, + embedding: state.embedding_paused, + caption: state.caption_paused, + }, + ]), + ), + ); }); - }, []); + }, [folders]); // Fetch failed embedding filenames whenever the expanded panel opens or failure counts change. const failedCounts = useMemo( @@ -83,10 +108,21 @@ export function BackgroundTasks() { } }, [expanded, failedCounts]); - const toggleWorker = (worker: WorkerKey) => { - const next = !paused[worker]; - setPaused((prev) => ({ ...prev, [worker]: next })); - void invoke("set_worker_paused", { worker, paused: next }); + const isWorkerPaused = (folderId: number, worker: WorkerKey) => { + return paused[folderId]?.[worker] ?? DEFAULT_PAUSED_STATE[worker]; + }; + + const toggleWorker = (folderId: number, worker: WorkerKey) => { + const next = !isWorkerPaused(folderId, worker); + setPaused((prev) => ({ + ...prev, + [folderId]: { + ...DEFAULT_PAUSED_STATE, + ...prev[folderId], + [worker]: next, + }, + })); + void invoke("set_worker_paused", { worker, folderId, paused: next }); }; const dismissTask = (id: number, snapshot: string) => { @@ -105,13 +141,16 @@ export function BackgroundTasks() { const embeddingPending = jobs?.embedding_pending ?? 0; const embeddingReady = jobs?.embedding_ready ?? 0; const embeddingFailed = jobs?.embedding_failed ?? 0; + const captionPending = jobs?.caption_pending ?? 0; + const captionFailed = jobs?.caption_failed ?? 0; - const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending; + const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending + captionPending; const embeddingProcessed = embeddingReady + embeddingFailed; const embeddingTotal = embeddingProcessed + embeddingPending; const hasFailedEmbeddings = embeddingFailed > 0; + const hasFailedCaptions = captionFailed > 0; - if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings) return null; + if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings && !hasFailedCaptions) return null; const stages: TaskStage[] = []; @@ -153,6 +192,15 @@ export function BackgroundTasks() { }); } + if (captionPending > 0) { + stages.push({ + label: "Captions", + detail: captionPending.toLocaleString(), + progress: null, + failed: false, + }); + } + if (hasFailedEmbeddings && pendingMediaWork === 0) { stages.push({ label: "Failed", @@ -162,13 +210,23 @@ export function BackgroundTasks() { }); } - const snapshot = `${pendingMediaWork}:${embeddingFailed}`; + if (hasFailedCaptions && pendingMediaWork === 0) { + stages.push({ + label: "Failed", + detail: `${captionFailed.toLocaleString()} captions`, + progress: null, + failed: true, + }); + } + + const snapshot = `${pendingMediaWork}:${embeddingFailed}:${captionFailed}`; return { id: folder.id, name: folder.name, stages, hasFailedEmbeddings, + hasFailedCaptions, pendingMediaWork, embeddingProcessed, embeddingTotal, @@ -184,7 +242,7 @@ export function BackgroundTasks() { const primary = tasks[0]; const extraCount = tasks.length - 1; - const hasFailed = tasks.some((t) => t.hasFailedEmbeddings && t.pendingMediaWork === 0); + const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedCaptions) && t.pendingMediaWork === 0); // Best progress bar value: use embedding progress if available (most informative), // otherwise fall back to scanning progress, otherwise indeterminate. @@ -214,7 +272,7 @@ export function BackgroundTasks() {
{primary.stages.map((stage) => { const workerKey = WORKER_FOR_STAGE[stage.label]; - const isPaused = workerKey ? paused[workerKey] : false; + const isPaused = workerKey ? isWorkerPaused(primary.id, workerKey) : false; return ( { e.stopPropagation(); toggleWorker(workerKey); }} + onClick={(e) => { e.stopPropagation(); toggleWorker(primary.id, workerKey); }} > {isPaused ? ( @@ -312,7 +370,7 @@ export function BackgroundTasks() { const taskEmbeddingStage = task.stages.find((s) => s.label === "Embeddings"); const taskScanningStage = task.stages.find((s) => s.label === "Scanning"); const taskBarProgress = taskEmbeddingStage?.progress ?? taskScanningStage?.progress ?? null; - const taskHasFailed = task.hasFailedEmbeddings && task.pendingMediaWork === 0; + const taskHasFailed = (task.hasFailedEmbeddings || task.hasFailedCaptions) && task.pendingMediaWork === 0; return (
@@ -322,7 +380,7 @@ export function BackgroundTasks() {
{task.stages.map((stage) => { const workerKey = WORKER_FOR_STAGE[stage.label]; - const isPaused = workerKey ? paused[workerKey] : false; + const isPaused = workerKey ? isWorkerPaused(task.id, workerKey) : false; return ( toggleWorker(workerKey)} + onClick={() => toggleWorker(task.id, workerKey)} > {isPaused ? ( diff --git a/src/components/Gallery.tsx b/src/components/Gallery.tsx index 95a5710..ee0e6ea 100644 --- a/src/components/Gallery.tsx +++ b/src/components/Gallery.tsx @@ -251,6 +251,10 @@ export function Gallery() { const zoomPreset = useGalleryStore((state) => state.zoomPreset); const search = useGalleryStore((state) => state.search); const searchMode = useGalleryStore((state) => state.searchMode); + const collectionTitle = useGalleryStore((state) => state.collectionTitle); + const imageLoadError = useGalleryStore((state) => state.imageLoadError); + const galleryScrollResetKey = useGalleryStore((state) => state.galleryScrollResetKey); + const isSimilarResults = collectionTitle === "Similar Images"; const parentRef = useRef(null); const [contextMenu, setContextMenu] = useState<{ x: number; y: number; image: ImageRecord } | null>(null); @@ -271,6 +275,10 @@ export function Gallery() { return () => element.removeEventListener("scroll", handleScroll); }, [handleScroll]); + useEffect(() => { + parentRef.current?.scrollTo({ top: 0, left: 0 }); + }, [galleryScrollResetKey]); + useEffect(() => { const close = (event: PointerEvent) => { if ((event.target as HTMLElement | null)?.closest("[data-gallery-context-menu]")) return; @@ -293,12 +301,16 @@ export function Gallery() {

- {searchMode === "semantic" && search.trim().length > 0 + {isSimilarResults + ? "Finding similar images" + : searchMode === "semantic" && search.trim().length > 0 ? `Searching for matches to "${search}"` : "Loading media"}

- {searchMode === "semantic" && search.trim().length > 0 + {isSimilarResults + ? "Comparing visual embeddings" + : searchMode === "semantic" && search.trim().length > 0 ? "Semantic search can take a little longer than filename search" : "Fetching results"}

@@ -316,12 +328,20 @@ export function Gallery() { d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-6-6h.01M6 20h12a2 2 0 002-2V6a2 2 0 00-2-2H6a2 2 0 00-2 2v12a2 2 0 002 2z" />

- {searchMode === "semantic" && search.trim().length > 0 + {imageLoadError + ? "Could not load results" + : isSimilarResults + ? "No similar images found" + : searchMode === "semantic" && search.trim().length > 0 ? "No semantic matches found" : "No media found"}

- {searchMode === "semantic" && search.trim().length > 0 + {imageLoadError + ? imageLoadError + : isSimilarResults + ? "This item may be visually isolated, or more embeddings may need to finish processing" + : searchMode === "semantic" && search.trim().length > 0 ? "Try a broader phrase, or wait for more embeddings to finish processing" : "Try adjusting your filters or add a new folder"}

diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index 1edfc37..5460b16 100644 --- a/src/components/Lightbox.tsx +++ b/src/components/Lightbox.tsx @@ -47,12 +47,27 @@ function embeddingLabel(status: string, model: string | null): string { export function Lightbox() { const selectedImage = useGalleryStore((state) => state.selectedImage); + const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); const closeImage = useGalleryStore((state) => state.closeImage); const images = useGalleryStore((state) => state.images); const openImage = useGalleryStore((state) => state.openImage); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); const updateImageDetails = useGalleryStore((state) => state.updateImageDetails); + const suggestImageTags = useGalleryStore((state) => state.suggestImageTags); + const captionModelStatus = useGalleryStore((state) => state.captionModelStatus); + const captionModelPreparing = useGalleryStore((state) => state.captionModelPreparing); + const captionModelError = useGalleryStore((state) => state.captionModelError); + const captionModelProgress = useGalleryStore((state) => state.captionModelProgress); + const aiCaptionsEnabled = useGalleryStore((state) => state.aiCaptionsEnabled); + const setAiCaptionsEnabled = useGalleryStore((state) => state.setAiCaptionsEnabled); + const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); + const prepareCaptionModel = useGalleryStore((state) => state.prepareCaptionModel); + const queueCaptionForImage = useGalleryStore((state) => state.queueCaptionForImage); + const queueCaptionJobs = useGalleryStore((state) => state.queueCaptionJobs); const [zoom, setZoom] = useState(1); + const [suggestedTags, setSuggestedTags] = useState([]); + const [captionQueueing, setCaptionQueueing] = useState(false); + const [captionQueueStatus, setCaptionQueueStatus] = useState(null); const imageViewportRef = useRef(null); const currentIndex = selectedImage ? images.findIndex((image) => image.id === selectedImage.id) : -1; @@ -68,8 +83,17 @@ export function Lightbox() { useEffect(() => { setZoom(1); + setSuggestedTags([]); + setCaptionQueueStatus(null); }, [selectedImage?.id]); + useEffect(() => { + if (!selectedImage?.generated_caption) return; + void suggestImageTags(selectedImage.id) + .then(setSuggestedTags) + .catch(() => setSuggestedTags([])); + }, [selectedImage?.id, selectedImage?.generated_caption, suggestImageTags]); + useEffect(() => { const viewport = imageViewportRef.current; if (!viewport || !selectedImage || selectedImage.media_kind !== "image") return; @@ -316,6 +340,119 @@ export function Lightbox() { ) : null}
+
+

AI Caption

+ {selectedImage.generated_caption ? ( + <> +

{selectedImage.generated_caption}

+ {suggestedTags.length > 0 ? ( +
+ {suggestedTags.map((tag) => ( + + {tag} + + ))} +
+ ) : null} + {selectedImage.caption_model ? ( +

{selectedImage.caption_model}

+ ) : null} + + ) : ( + <> +

+ {selectedImage.caption_error ?? "Not generated"} +

+ {captionModelStatus?.ready ? ( + <> +

+ {aiCaptionsEnabled ? "Florence-2 enabled locally" : "Florence-2 downloaded but disabled"} +

+
+ + {selectedImage.media_kind === "image" && aiCaptionsEnabled ? ( + + ) : null} + {aiCaptionsEnabled ? ( + + ) : null} +
+ {captionQueueStatus ? ( +

{captionQueueStatus}

+ ) : null} + + ) : ( + + )} + {captionModelProgress?.current_file ? ( +

{captionModelProgress.current_file}

+ ) : null} + {captionModelError ? ( +

{captionModelError}

+ ) : null} + {!captionModelStatus?.ready && !captionModelError ? ( +

Downloads Florence-2 on demand for offline captions.

+ ) : null} + + + )} +
+

Path

{selectedImage.path}

diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx new file mode 100644 index 0000000..3afb757 --- /dev/null +++ b/src/components/SettingsModal.tsx @@ -0,0 +1,354 @@ +import { useEffect, useState } from "react"; +import { useGalleryStore } from "../store"; + +type SettingsSection = "ai" | "library" | "display" | "storage"; + +const SECTIONS: { id: SettingsSection; label: string; detail: string }[] = [ + { id: "ai", label: "Local AI", detail: "Captions and suggestions" }, + { id: "library", label: "Library", detail: "Indexing and scanning" }, + { id: "display", label: "Display", detail: "Gallery preferences" }, + { id: "storage", label: "Storage", detail: "Cache and model files" }, +]; + +function ToggleSwitch({ + checked, + disabled, + onChange, + label, +}: { + checked: boolean; + disabled?: boolean; + onChange: (checked: boolean) => void; + label: string; +}) { + return ( + + ); +} + +function StatusPill({ children, tone }: { children: React.ReactNode; tone: "ready" | "muted" | "busy" }) { + const className = + tone === "ready" + ? "border-emerald-400/25 bg-emerald-500/10 text-emerald-300" + : tone === "busy" + ? "border-sky-400/25 bg-sky-500/10 text-sky-300" + : "border-white/10 bg-white/[0.04] text-gray-500"; + + return ( + + {children} + + ); +} + +function SettingsRow({ + title, + description, + children, +}: { + title: string; + description: string; + children: React.ReactNode; +}) { + return ( +
+
+

{title}

+

{description}

+
+
{children}
+
+ ); +} + +function SectionShell({ + eyebrow, + title, + children, +}: { + eyebrow: string; + title: string; + children: React.ReactNode; +}) { + return ( +
+

{eyebrow}

+

{title}

+
{children}
+
+ ); +} + +export function SettingsModal() { + const [activeSection, setActiveSection] = useState("ai"); + const [captionQueueStatus, setCaptionQueueStatus] = useState(null); + const [captionQueueing, setCaptionQueueing] = useState(false); + const settingsOpen = useGalleryStore((state) => state.settingsOpen); + const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); + const folders = useGalleryStore((state) => state.folders); + const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); + const captionModelStatus = useGalleryStore((state) => state.captionModelStatus); + const captionModelPreparing = useGalleryStore((state) => state.captionModelPreparing); + const captionModelProgress = useGalleryStore((state) => state.captionModelProgress); + const captionModelError = useGalleryStore((state) => state.captionModelError); + const captionRuntimeProbe = useGalleryStore((state) => state.captionRuntimeProbe); + const captionRuntimeChecking = useGalleryStore((state) => state.captionRuntimeChecking); + const aiCaptionsEnabled = useGalleryStore((state) => state.aiCaptionsEnabled); + const setAiCaptionsEnabled = useGalleryStore((state) => state.setAiCaptionsEnabled); + const prepareCaptionModel = useGalleryStore((state) => state.prepareCaptionModel); + const deleteCaptionModel = useGalleryStore((state) => state.deleteCaptionModel); + const probeCaptionRuntime = useGalleryStore((state) => state.probeCaptionRuntime); + const queueCaptionJobs = useGalleryStore((state) => state.queueCaptionJobs); + + useEffect(() => { + if (!settingsOpen) return; + const handleKeyDown = (event: KeyboardEvent) => { + if (event.key === "Escape") setSettingsOpen(false); + }; + window.addEventListener("keydown", handleKeyDown); + return () => window.removeEventListener("keydown", handleKeyDown); + }, [settingsOpen, setSettingsOpen]); + + if (!settingsOpen) return null; + + const modelReady = captionModelStatus?.ready ?? false; + const downloadLabel = captionModelProgress + ? `Downloading ${captionModelProgress.completed_files}/${captionModelProgress.total_files}` + : captionModelPreparing + ? "Preparing Florence-2..." + : modelReady + ? "Downloaded" + : "Download Florence-2"; + const downloadPercent = captionModelProgress + ? Math.round((captionModelProgress.completed_files / Math.max(captionModelProgress.total_files, 1)) * 100) + : 0; + const selectedFolder = folders.find((folder) => folder.id === selectedFolderId); + const captionScopeLabel = selectedFolder ? selectedFolder.name : "all libraries"; + + return ( +
setSettingsOpen(false)} + > +
event.stopPropagation()} + > + + +
+
+
+ {modelReady ? Florence-2 ready : Optional} + {captionModelPreparing ? Working : null} +
+ +
+ +
+ {activeSection === "ai" ? ( + + + + + + +
+ + {modelReady ? ( + <> + + + + ) : null} +
+ + + + + + +
+

Model location

+

+ {modelReady ? captionModelStatus?.local_dir : "Not downloaded"} +

+ {captionModelProgress?.current_file ? ( +

{captionModelProgress.current_file}

+ ) : null} + {captionModelError ? ( +

{captionModelError}

+ ) : null} + {captionQueueStatus ? ( +

{captionQueueStatus}

+ ) : null} + {captionRuntimeProbe ? ( +
+
+

Runtime check

+ Ready +
+

+ Tokenizer vocabulary: {captionRuntimeProbe.tokenizer_vocab_size.toLocaleString()} +

+
+ {captionRuntimeProbe.sessions.map((session) => ( +
+

{session.file}

+

+ {session.inputs.length} input{session.inputs.length === 1 ? "" : "s"} · {session.outputs.length} output{session.outputs.length === 1 ? "" : "s"} +

+
+ ))} +
+
+ ) : null} +
+ + ) : null} + + {activeSection === "library" ? ( + + + Managed per folder + + + Available + + + ) : null} + + {activeSection === "display" ? ( + + + Toolbar + + + Enabled + + + ) : null} + + {activeSection === "storage" ? ( + + + + + + ) : null} +
+
+
+
+ ); +} diff --git a/src/components/TitleBar.tsx b/src/components/TitleBar.tsx index 274367d..56a7f4e 100644 --- a/src/components/TitleBar.tsx +++ b/src/components/TitleBar.tsx @@ -1,5 +1,6 @@ import { useState, useEffect } from "react"; import { getCurrentWindow } from "@tauri-apps/api/window"; +import { useGalleryStore } from "../store"; // SVG icons for window controls function MinimizeIcon() { @@ -37,6 +38,7 @@ function CloseIcon() { export function TitleBar() { const [isMaximized, setIsMaximized] = useState(false); + const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); const appWindow = getCurrentWindow(); useEffect(() => { @@ -85,6 +87,17 @@ export function TitleBar() { className="flex items-stretch h-full" style={{ WebkitAppRegion: "no-drag" } as React.CSSProperties} > + + {/* Minimize */}
-
+

Rating

@@ -341,116 +351,102 @@ export function Lightbox() {
-

AI Caption

- {selectedImage.generated_caption ? ( - <> -

{selectedImage.generated_caption}

- {suggestedTags.length > 0 ? ( -
- {suggestedTags.map((tag) => ( - - {tag} - - ))} -
- ) : null} - {selectedImage.caption_model ? ( -

{selectedImage.caption_model}

- ) : null} - - ) : ( - <> -

- {selectedImage.caption_error ?? "Not generated"} -

- {captionModelStatus?.ready ? ( - <> -

- {aiCaptionsEnabled ? "Florence-2 enabled locally" : "Florence-2 downloaded but disabled"} -

-
- - {selectedImage.media_kind === "image" && aiCaptionsEnabled ? ( - - ) : null} - {aiCaptionsEnabled ? ( - - ) : null} -
- {captionQueueStatus ? ( -

{captionQueueStatus}

- ) : null} - - ) : ( - - )} - {captionModelProgress?.current_file ? ( -

{captionModelProgress.current_file}

- ) : null} - {captionModelError ? ( -

{captionModelError}

- ) : null} - {!captionModelStatus?.ready && !captionModelError ? ( -

Downloads Florence-2 on demand for offline captions.

+
+

Tags

+
+ {selectedImage.ai_rating ? ( + + {ratingPill(selectedImage.ai_rating).label} + ) : null} +
+
+ + {imageTags.length > 0 ? ( + <> +
+ {(tagsExpanded ? imageTags : imageTags.slice(0, 8)).map((t) => ( + + {t.tag} + + + ))} +
+ {imageTags.length > 8 && ( + + )} + ) : ( +

No tags yet

)} + +
{ + e.preventDefault(); + const raw = tagInput.trim(); + if (!raw || tagAdding) return; + setTagAdding(true); + void addUserTag(selectedImage.id, raw) + .then((newTag) => { + setImageTags((prev) => [...prev, newTag]); + setTagInput(""); + }) + .catch(() => undefined) + .finally(() => setTagAdding(false)); + }} + > + setTagInput(e.target.value)} + disabled={tagAdding} + /> + +
@@ -459,7 +455,7 @@ export function Lightbox() {
-
+
{currentIndex + 1} / {images.length}
diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx index 3afb757..cdf9ddd 100644 --- a/src/components/SettingsModal.tsx +++ b/src/components/SettingsModal.tsx @@ -1,48 +1,15 @@ import { useEffect, useState } from "react"; -import { useGalleryStore } from "../store"; +import { TaggerAcceleration, useGalleryStore } from "../store"; -type SettingsSection = "ai" | "library" | "display" | "storage"; +type SettingsSection = "tagging" | "library" | "display" | "storage"; const SECTIONS: { id: SettingsSection; label: string; detail: string }[] = [ - { id: "ai", label: "Local AI", detail: "Captions and suggestions" }, + { id: "tagging", label: "AI Tagging", detail: "WD tagger model" }, { id: "library", label: "Library", detail: "Indexing and scanning" }, { id: "display", label: "Display", detail: "Gallery preferences" }, { id: "storage", label: "Storage", detail: "Cache and model files" }, ]; -function ToggleSwitch({ - checked, - disabled, - onChange, - label, -}: { - checked: boolean; - disabled?: boolean; - onChange: (checked: boolean) => void; - label: string; -}) { - return ( - - ); -} function StatusPill({ children, tone }: { children: React.ReactNode; tone: "ready" | "muted" | "busy" }) { const className = @@ -59,6 +26,34 @@ function StatusPill({ children, tone }: { children: React.ReactNode; tone: "read ); } +function TaggerAccelerationButton({ + acceleration, + current, + onSelect, + children, +}: { + acceleration: TaggerAcceleration; + current: TaggerAcceleration; + onSelect: (acceleration: TaggerAcceleration) => void; + children: React.ReactNode; +}) { + const active = acceleration === current; + return ( + + ); +} + + function SettingsRow({ title, description, @@ -98,50 +93,52 @@ function SectionShell({ } export function SettingsModal() { - const [activeSection, setActiveSection] = useState("ai"); - const [captionQueueStatus, setCaptionQueueStatus] = useState(null); - const [captionQueueing, setCaptionQueueing] = useState(false); + const [activeSection, setActiveSection] = useState("tagging"); + const [taggerQueueStatus, setTaggerQueueStatus] = useState(null); + const [taggerQueueing, setTaggerQueueing] = useState(false); + const [taggerClearing, setTaggerClearing] = useState(false); + const [taggerAccelerationSaving, setTaggerAccelerationSaving] = useState(false); + const [taggerThresholdDraft, setTaggerThresholdDraft] = useState(null); + const [taggerThresholdSaving, setTaggerThresholdSaving] = useState(false); const settingsOpen = useGalleryStore((state) => state.settingsOpen); const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); - const folders = useGalleryStore((state) => state.folders); const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); - const captionModelStatus = useGalleryStore((state) => state.captionModelStatus); - const captionModelPreparing = useGalleryStore((state) => state.captionModelPreparing); - const captionModelProgress = useGalleryStore((state) => state.captionModelProgress); - const captionModelError = useGalleryStore((state) => state.captionModelError); - const captionRuntimeProbe = useGalleryStore((state) => state.captionRuntimeProbe); - const captionRuntimeChecking = useGalleryStore((state) => state.captionRuntimeChecking); - const aiCaptionsEnabled = useGalleryStore((state) => state.aiCaptionsEnabled); - const setAiCaptionsEnabled = useGalleryStore((state) => state.setAiCaptionsEnabled); - const prepareCaptionModel = useGalleryStore((state) => state.prepareCaptionModel); - const deleteCaptionModel = useGalleryStore((state) => state.deleteCaptionModel); - const probeCaptionRuntime = useGalleryStore((state) => state.probeCaptionRuntime); - const queueCaptionJobs = useGalleryStore((state) => state.queueCaptionJobs); + const folders = useGalleryStore((state) => state.folders); + const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus); + const taggerModelPreparing = useGalleryStore((state) => state.taggerModelPreparing); + const taggerModelProgress = useGalleryStore((state) => state.taggerModelProgress); + const taggerModelError = useGalleryStore((state) => state.taggerModelError); + const taggerAcceleration = useGalleryStore((state) => state.taggerAcceleration); + const taggerThreshold = useGalleryStore((state) => state.taggerThreshold); + const taggerRuntimeProbe = useGalleryStore((state) => state.taggerRuntimeProbe); + const taggerRuntimeChecking = useGalleryStore((state) => state.taggerRuntimeChecking); + const loadTaggerModelStatus = useGalleryStore((state) => state.loadTaggerModelStatus); + const prepareTaggerModel = useGalleryStore((state) => state.prepareTaggerModel); + const deleteTaggerModel = useGalleryStore((state) => state.deleteTaggerModel); + const loadTaggerAcceleration = useGalleryStore((state) => state.loadTaggerAcceleration); + const setTaggerAcceleration = useGalleryStore((state) => state.setTaggerAcceleration); + const loadTaggerThreshold = useGalleryStore((state) => state.loadTaggerThreshold); + const setTaggerThreshold = useGalleryStore((state) => state.setTaggerThreshold); + const probeTaggerRuntime = useGalleryStore((state) => state.probeTaggerRuntime); + const queueTaggingJobs = useGalleryStore((state) => state.queueTaggingJobs); + const clearTaggingJobs = useGalleryStore((state) => state.clearTaggingJobs); useEffect(() => { if (!settingsOpen) return; + void loadTaggerModelStatus(); + void loadTaggerAcceleration(); + void loadTaggerThreshold(); const handleKeyDown = (event: KeyboardEvent) => { if (event.key === "Escape") setSettingsOpen(false); }; window.addEventListener("keydown", handleKeyDown); return () => window.removeEventListener("keydown", handleKeyDown); - }, [settingsOpen, setSettingsOpen]); + }, [settingsOpen, loadTaggerModelStatus, loadTaggerAcceleration, loadTaggerThreshold, setSettingsOpen]); if (!settingsOpen) return null; - const modelReady = captionModelStatus?.ready ?? false; - const downloadLabel = captionModelProgress - ? `Downloading ${captionModelProgress.completed_files}/${captionModelProgress.total_files}` - : captionModelPreparing - ? "Preparing Florence-2..." - : modelReady - ? "Downloaded" - : "Download Florence-2"; - const downloadPercent = captionModelProgress - ? Math.round((captionModelProgress.completed_files / Math.max(captionModelProgress.total_files, 1)) * 100) - : 0; const selectedFolder = folders.find((folder) => folder.id === selectedFolderId); - const captionScopeLabel = selectedFolder ? selectedFolder.name : "all libraries"; + const scopeLabel = selectedFolder ? selectedFolder.name : "all libraries"; return (
-
-
- {modelReady ? Florence-2 ready : Optional} - {captionModelPreparing ? Working : null} -
+
- {activeSection === "ai" ? ( - - - - - - -
- - {modelReady ? ( - <> - - - - ) : null} -
- - - - - +
+ + {taggerReady ? ( + <> + + + + ) : null} +
+ -
-

Model location

-

- {modelReady ? captionModelStatus?.local_dir : "Not downloaded"} -

- {captionModelProgress?.current_file ? ( -

{captionModelProgress.current_file}

- ) : null} - {captionModelError ? ( -

{captionModelError}

- ) : null} - {captionQueueStatus ? ( -

{captionQueueStatus}

- ) : null} - {captionRuntimeProbe ? ( -
-
-

Runtime check

- Ready + +
+
+ { + setTaggerAccelerationSaving(true); + void setTaggerAcceleration(acceleration) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => setTaggerAccelerationSaving(false)); + }} + > + Auto + + { + setTaggerAccelerationSaving(true); + void setTaggerAcceleration(acceleration) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => setTaggerAccelerationSaving(false)); + }} + > + DirectML + + { + setTaggerAccelerationSaving(true); + void setTaggerAcceleration(acceleration) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => setTaggerAccelerationSaving(false)); + }} + > + CPU +
-

- Tokenizer vocabulary: {captionRuntimeProbe.tokenizer_vocab_size.toLocaleString()} +

+ {taggerAccelerationSaving ? "Saving..." : `Current: ${taggerAcceleration}`}

-
- {captionRuntimeProbe.sessions.map((session) => ( -
-

{session.file}

+
+ + + +
+
+ setTaggerThresholdDraft(e.target.value)} + onBlur={() => { + const val = parseFloat(thresholdDisplay); + if (!isNaN(val) && val >= 0.05 && val <= 0.99) { + setTaggerThresholdSaving(true); + void setTaggerThreshold(val) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => { + setTaggerThresholdDraft(null); + setTaggerThresholdSaving(false); + }); + } else { + setTaggerThresholdDraft(null); + } + }} + /> +
+

+ {taggerThresholdSaving ? "Saving..." : "Default: 0.35"} +

+
+ + + +
+ + +
+ + +
+

Model location

+

+ {taggerReady ? taggerModelStatus?.local_dir : "Not downloaded"} +

+ {taggerModelProgress?.current_file ? ( +

{taggerModelProgress.current_file}

+ ) : null} + {taggerModelError ? ( +

{taggerModelError}

+ ) : null} + {taggerQueueStatus ? ( +

{taggerQueueStatus}

+ ) : null} + {taggerRuntimeProbe ? ( +
+
+

Runtime check

+ Ready +
+

+ Tagger acceleration: {taggerRuntimeProbe.acceleration} +

+
+
+

{taggerRuntimeProbe.session.file}

- {session.inputs.length} input{session.inputs.length === 1 ? "" : "s"} · {session.outputs.length} output{session.outputs.length === 1 ? "" : "s"} + {taggerRuntimeProbe.session.inputs.length} input{taggerRuntimeProbe.session.inputs.length === 1 ? "" : "s"} · {taggerRuntimeProbe.session.outputs.join(", ")}

- ))} +
-
- ) : null} -
- - ) : null} + ) : null} +
+ + ); + })() : null} {activeSection === "library" ? ( @@ -335,11 +443,11 @@ export function SettingsModal() { {activeSection === "storage" ? ( - + diff --git a/src/store.ts b/src/store.ts index 5af3a44..32ae1fc 100644 --- a/src/store.ts +++ b/src/store.ts @@ -15,6 +15,10 @@ export type MediaKind = "image" | "video"; export type MediaFilter = "all" | MediaKind; export type ZoomPreset = "compact" | "comfortable" | "detail"; export type SearchMode = "filename" | "semantic"; +export type CaptionAcceleration = "auto" | "cpu" | "directml"; +export type CaptionDetail = "short" | "detailed" | "paragraph"; +export type TaggerAcceleration = "auto" | "cpu" | "directml"; +export type AiRating = "general" | "sensitive" | "questionable" | "explicit"; export interface ImageRecord { id: number; @@ -44,6 +48,35 @@ export interface ImageRecord { caption_model: string | null; caption_updated_at: string | null; caption_error: string | null; + ai_rating: AiRating | null; + ai_tagger_model: string | null; + ai_tagged_at: string | null; + ai_tagger_error: string | null; +} + +export interface ImageTag { + id: number; + image_id: number; + tag: string; + source: "user" | "ai"; + ai_model: string | null; + confidence: number | null; + created_at: string; +} + +export interface TaggerModelStatus { + model_id: string; + model_name: string; + local_dir: string; + ready: boolean; + missing_files: string[]; +} + +export interface TaggerModelProgress { + total_files: number; + completed_files: number; + current_file: string | null; + done: boolean; } export interface IndexProgress { @@ -64,6 +97,9 @@ export interface FolderJobProgress { caption_pending: number; caption_ready: number; caption_failed: number; + tagging_pending: number; + tagging_ready: number; + tagging_failed: number; } export interface MediaJobProgressEvent { @@ -110,6 +146,8 @@ export interface CaptionRuntimeSessionProbe { export interface CaptionRuntimeProbe { ready: boolean; + acceleration: CaptionAcceleration; + detail: CaptionDetail; tokenizer_vocab_size: number; sessions: CaptionRuntimeSessionProbe[]; } @@ -118,6 +156,19 @@ export interface CaptionVisionProbe { input_shape: number[]; output_shape: number[]; output_values: number; + acceleration: CaptionAcceleration; +} + +export interface TaggerRuntimeSessionProbe { + file: string; + inputs: string[]; + outputs: string[]; +} + +export interface TaggerRuntimeProbe { + ready: boolean; + acceleration: TaggerAcceleration; + session: TaggerRuntimeSessionProbe; } export type SortOrder = @@ -163,9 +214,20 @@ interface GalleryState { captionModelProgress: CaptionModelProgress | null; captionRuntimeProbe: CaptionRuntimeProbe | null; captionRuntimeChecking: boolean; + captionAcceleration: CaptionAcceleration; + captionDetail: CaptionDetail; aiCaptionsEnabled: boolean; settingsOpen: boolean; + taggerModelStatus: TaggerModelStatus | null; + taggerModelPreparing: boolean; + taggerModelError: string | null; + taggerModelProgress: TaggerModelProgress | null; + taggerAcceleration: TaggerAcceleration; + taggerThreshold: number; + taggerRuntimeProbe: TaggerRuntimeProbe | null; + taggerRuntimeChecking: boolean; + loadFolders: () => Promise; loadBackgroundJobProgress: () => Promise; addFolder: (path: string) => Promise; @@ -198,12 +260,33 @@ interface GalleryState { generateCaptionForImage: (imageId: number) => Promise; queueCaptionJobs: (folderId?: number | null) => Promise; queueCaptionForImage: (imageId: number) => Promise; + clearCaptionJobs: (folderId?: number | null) => Promise; + resetGeneratedCaptions: (folderId?: number | null) => Promise; + loadCaptionAcceleration: () => Promise; + setCaptionAcceleration: (acceleration: CaptionAcceleration) => Promise; + loadCaptionDetail: () => Promise; + setCaptionDetail: (detail: CaptionDetail) => Promise; setAiCaptionsEnabled: (enabled: boolean) => void; setSettingsOpen: (open: boolean) => void; retryFailedEmbeddings: (folderId: number) => Promise; updateImageDetails: (imageId: number, updates: { favorite?: boolean; rating?: number }) => Promise; setCacheDir: (dir: string) => void; subscribeToProgress: () => Promise; + + loadTaggerModelStatus: () => Promise; + prepareTaggerModel: () => Promise; + deleteTaggerModel: () => Promise; + loadTaggerAcceleration: () => Promise; + setTaggerAcceleration: (acceleration: TaggerAcceleration) => Promise; + loadTaggerThreshold: () => Promise; + setTaggerThreshold: (threshold: number) => Promise; + probeTaggerRuntime: () => Promise; + queueTaggingJobs: (folderId?: number | null) => Promise; + queueTaggingForImage: (imageId: number) => Promise; + clearTaggingJobs: (folderId?: number | null) => Promise; + getImageTags: (imageId: number) => Promise; + addUserTag: (imageId: number, tag: string) => Promise; + removeTag: (tagId: number) => Promise; } const PAGE_SIZE = 200; @@ -361,9 +444,20 @@ export const useGalleryStore = create((set, get) => ({ captionModelProgress: null, captionRuntimeProbe: null, captionRuntimeChecking: false, + captionAcceleration: "auto", + captionDetail: "paragraph", aiCaptionsEnabled: initialAiCaptionsEnabled(), settingsOpen: false, + taggerModelStatus: null, + taggerModelPreparing: false, + taggerModelError: null, + taggerModelProgress: null, + taggerAcceleration: "auto", + taggerThreshold: 0.35, + taggerRuntimeProbe: null, + taggerRuntimeChecking: false, + setCacheDir: (cacheDir) => set({ cacheDir }), loadFolders: async () => { @@ -567,10 +661,10 @@ export const useGalleryStore = create((set, get) => ({ similarSourceImageId: imageId, galleryScrollResetKey: reset ? state.galleryScrollResetKey + 1 : state.galleryScrollResetKey, })); - try { - const images = await invoke("find_similar_images", { - params: { image_id: imageId, limit: requestedLimit }, - }); + try { + const images = await invoke("find_similar_images", { + params: { image_id: imageId, folder_id: folderId ?? null, limit: requestedLimit }, + }); const hasMore = images.length >= requestedLimit; set({ images, @@ -616,6 +710,38 @@ export const useGalleryStore = create((set, get) => ({ } }, + loadCaptionAcceleration: async () => { + try { + const captionAcceleration = await invoke("get_caption_acceleration"); + set({ captionAcceleration }); + } catch (error) { + set({ captionModelError: String(error) }); + } + }, + + setCaptionAcceleration: async (acceleration) => { + const captionAcceleration = await invoke("set_caption_acceleration", { + params: { acceleration }, + }); + set({ captionAcceleration, captionRuntimeProbe: null }); + }, + + loadCaptionDetail: async () => { + try { + const captionDetail = await invoke("get_caption_detail"); + set({ captionDetail }); + } catch (error) { + set({ captionModelError: String(error) }); + } + }, + + setCaptionDetail: async (detail) => { + const captionDetail = await invoke("set_caption_detail", { + params: { detail }, + }); + set({ captionDetail, captionRuntimeProbe: null }); + }, + prepareCaptionModel: async () => { set({ captionModelPreparing: true, captionModelError: null, captionModelProgress: null }); try { @@ -683,6 +809,23 @@ export const useGalleryStore = create((set, get) => ({ return queued; }, + clearCaptionJobs: async (folderId = get().selectedFolderId) => { + const cleared = await invoke("clear_caption_jobs", { + params: { folder_id: folderId ?? null }, + }); + await get().loadBackgroundJobProgress(); + return cleared; + }, + + resetGeneratedCaptions: async (folderId = get().selectedFolderId) => { + const reset = await invoke("reset_generated_captions", { + params: { folder_id: folderId ?? null }, + }); + await get().loadBackgroundJobProgress(); + await get().loadImages(true); + return reset; + }, + setAiCaptionsEnabled: (aiCaptionsEnabled) => { window.localStorage.setItem(AI_CAPTIONS_ENABLED_KEY, String(aiCaptionsEnabled)); set({ aiCaptionsEnabled }); @@ -690,6 +833,119 @@ export const useGalleryStore = create((set, get) => ({ setSettingsOpen: (settingsOpen) => set({ settingsOpen }), + loadTaggerModelStatus: async () => { + try { + const taggerModelStatus = await invoke("get_tagger_model_status"); + set({ taggerModelStatus, taggerModelError: null }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + loadTaggerAcceleration: async () => { + try { + const taggerAcceleration = await invoke("get_tagger_acceleration"); + set({ taggerAcceleration }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + setTaggerAcceleration: async (acceleration) => { + const taggerAcceleration = await invoke("set_tagger_acceleration", { + params: { acceleration }, + }); + set({ taggerAcceleration, taggerRuntimeProbe: null }); + }, + + loadTaggerThreshold: async () => { + try { + const taggerThreshold = await invoke("get_tagger_threshold"); + set({ taggerThreshold }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + setTaggerThreshold: async (threshold) => { + const taggerThreshold = await invoke("set_tagger_threshold", { + params: { threshold }, + }); + set({ taggerThreshold }); + }, + + prepareTaggerModel: async () => { + set({ taggerModelPreparing: true, taggerModelError: null, taggerModelProgress: null }); + try { + const taggerModelStatus = await invoke("prepare_tagger_model"); + set({ taggerModelStatus, taggerModelPreparing: false, taggerModelError: null, taggerModelProgress: null }); + } catch (error) { + set({ taggerModelPreparing: false, taggerModelError: String(error), taggerModelProgress: null }); + } + }, + + deleteTaggerModel: async () => { + set({ taggerModelPreparing: true, taggerModelError: null, taggerModelProgress: null }); + try { + const taggerModelStatus = await invoke("delete_tagger_model"); + set({ taggerModelStatus, taggerModelPreparing: false, taggerModelError: null, taggerModelProgress: null, taggerRuntimeProbe: null }); + } catch (error) { + set({ taggerModelPreparing: false, taggerModelError: String(error), taggerModelProgress: null }); + } + }, + + probeTaggerRuntime: async () => { + set({ taggerRuntimeChecking: true, taggerModelError: null }); + try { + const taggerRuntimeProbe = await invoke("probe_tagger_runtime"); + set({ taggerRuntimeProbe, taggerRuntimeChecking: false, taggerModelError: null }); + } catch (error) { + set({ taggerRuntimeChecking: false, taggerModelError: String(error), taggerRuntimeProbe: null }); + } + }, + + queueTaggingJobs: async (folderId = get().selectedFolderId) => { + const queued = await invoke("queue_tagging_jobs", { + params: { folder_id: folderId ?? null, image_id: null }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + + queueTaggingForImage: async (imageId) => { + const queued = await invoke("queue_tagging_jobs", { + params: { folder_id: null, image_id: imageId }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + + clearTaggingJobs: async (folderId = get().selectedFolderId) => { + const cleared = await invoke("clear_tagging_jobs", { + params: { folder_id: folderId ?? null }, + }); + await get().loadBackgroundJobProgress(); + return cleared; + }, + + getImageTags: async (imageId) => { + return invoke("get_image_tags", { + params: { image_id: imageId }, + }); + }, + + addUserTag: async (imageId, tag) => { + return invoke("add_user_tag", { + params: { image_id: imageId, tag }, + }); + }, + + removeTag: async (tagId) => { + await invoke("remove_tag", { + params: { tag_id: tagId }, + }); + }, + retryFailedEmbeddings: async (folderId) => { await invoke("retry_failed_embeddings", { params: { folder_id: folderId } }); await get().loadBackgroundJobProgress(); @@ -752,6 +1008,13 @@ export const useGalleryStore = create((set, get) => ({ }); }); + const unlistenTaggerModelProgress = await listen("tagger-model-progress", (event) => { + set({ + taggerModelProgress: event.payload.done ? null : event.payload, + taggerModelPreparing: !event.payload.done, + }); + }); + const unlistenImages = await listen("indexed-images", (event) => { const batch = event.payload; @@ -817,6 +1080,7 @@ export const useGalleryStore = create((set, get) => ({ unlistenProgress(); unlistenMediaJobs(); unlistenCaptionModelProgress(); + unlistenTaggerModelProgress(); unlistenImages(); unlistenThumbnails(); }; -- 2.52.0 From f93a80bc87a041f4463bb787867bebb7b8ec5c5d Mon Sep 17 00:00:00 2001 From: LyAhn Date: Wed, 8 Apr 2026 17:15:44 +0100 Subject: [PATCH 03/25] fix: refresh lightbox tags live after AI tagging, hide AI tags button for video --- src/components/Lightbox.tsx | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index cb3af78..2af2a0b 100644 --- a/src/components/Lightbox.tsx +++ b/src/components/Lightbox.tsx @@ -102,7 +102,12 @@ export function Lightbox() { void getImageTags(selectedImage.id) .then(setImageTags) .catch(() => setImageTags([])); - }, [selectedImage?.id, getImageTags]); + }, [selectedImage?.id, selectedImage?.ai_tagged_at, getImageTags]); + + // Reset the queued state once the worker finishes so the button is usable again + useEffect(() => { + if (selectedImage?.ai_tagged_at) setTaggingQueued(false); + }, [selectedImage?.ai_tagged_at]); useEffect(() => { const viewport = imageViewportRef.current; @@ -359,17 +364,19 @@ export function Lightbox() { {ratingPill(selectedImage.ai_rating).label} ) : null} - + {selectedImage.media_kind === "image" ? ( + + ) : null}
-- 2.52.0 From b2826d1143ccc9e4f3ff35acef24a9bf117fe440 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Wed, 8 Apr 2026 20:04:37 +0100 Subject: [PATCH 04/25] feat: add WD tagger with CSV tag support and model download via ureq/zip Introduces tagger.rs for WD-based image tagging with CSV label support. Adds csv, ureq, and zip dependencies; updates gitignore to exclude Python/JSON files. --- .gitignore | 6 + src-tauri/Cargo.lock | 24 ++ src-tauri/Cargo.toml | 5 +- src-tauri/src/captioner.rs | 602 ++++++++++++++++++++++++++++------- src-tauri/src/commands.rs | 313 ++++++++++++++++++- src-tauri/src/db.rs | 510 +++++++++++++++++++++++++++++- src-tauri/src/indexer.rs | 169 +++++++++- src-tauri/src/tagger.rs | 625 +++++++++++++++++++++++++++++++++++++ src-tauri/src/vector.rs | 42 ++- 9 files changed, 2161 insertions(+), 135 deletions(-) create mode 100644 src-tauri/src/tagger.rs diff --git a/.gitignore b/.gitignore index 8b9cd9f..7e7224f 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,9 @@ dist-ssr # Local staging area /staging + +# Misc +*.py +*.json + +*.pyc diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index db7b209..0a27e7a 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -874,6 +874,27 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + [[package]] name = "ctor" version = "0.2.9" @@ -4009,6 +4030,7 @@ dependencies = [ "candle-nn", "candle-transformers", "chrono", + "csv", "fast_image_resize", "ffmpeg-sidecar", "hf-hub", @@ -4030,9 +4052,11 @@ dependencies = [ "tauri-plugin-opener", "tokenizers", "tokio", + "ureq", "uuid", "walkdir", "xxhash-rust", + "zip 4.6.1", ] [[package]] diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index ce47032..38356e2 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -44,4 +44,7 @@ candle-nn = { version = "0.10.2", features = ["cuda"] } candle-transformers = { version = "0.10.2", features = ["cuda"] } hf-hub = { version = "0.5.0", default-features = false, features = ["ureq", "native-tls"] } tokenizers = "0.22.1" -ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "download-binaries", "copy-dylibs", "load-dynamic", "api-24", "tls-native"] } +ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "download-binaries", "copy-dylibs", "load-dynamic", "directml", "api-24", "tls-native"] } +ureq = { version = "3.3.0", default-features = false, features = ["native-tls"] } +zip = { version = "4.6.1", default-features = false, features = ["deflate"] } +csv = "1" diff --git a/src-tauri/src/captioner.rs b/src-tauri/src/captioner.rs index 112aa7c..27bb5e5 100644 --- a/src-tauri/src/captioner.rs +++ b/src-tauri/src/captioner.rs @@ -1,18 +1,54 @@ use anyhow::Result; use hf_hub::{api::sync::Api, Repo, RepoType}; use image::{imageops::FilterType, ImageReader}; +use ort::ep; use ort::session::SessionInputValue; +use ort::session::SessionOutputs; use ort::session::{builder::GraphOptimizationLevel, Session}; use ort::value::{Shape, Tensor}; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use std::borrow::Cow; +use std::io::{Cursor, Read}; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::OnceLock; +use std::time::Instant; use tokenizers::Tokenizer; pub const FLORENCE_MODEL_ID: &str = "onnx-community/Florence-2-base-ft"; pub const FLORENCE_CAPTION_MODEL_NAME: &str = "florence-2-base-ft-onnx-q4"; +const ONNX_RUNTIME_NUGET_URL: &str = + "https://www.nuget.org/api/v2/package/Microsoft.ML.OnnxRuntime.DirectML/1.24.2"; +const DIRECTML_NUGET_URL: &str = + "https://www.nuget.org/api/v2/package/Microsoft.AI.DirectML/1.15.4"; +const ONNX_RUNTIME_DLL_FILE: &str = "onnxruntime/onnxruntime.dll"; +const ONNX_RUNTIME_PROVIDERS_DLL_FILE: &str = "onnxruntime/onnxruntime_providers_shared.dll"; +const DIRECTML_DLL_FILE: &str = "onnxruntime/DirectML.dll"; +const CAPTION_ACCELERATION_FILE: &str = "settings/caption_acceleration.txt"; +const CAPTION_DETAIL_FILE: &str = "settings/caption_detail.txt"; + +const ONNX_RUNTIME_FILES: &[(&str, &str, &str)] = &[ + ( + ONNX_RUNTIME_DLL_FILE, + ONNX_RUNTIME_NUGET_URL, + "runtimes/win-x64/native/onnxruntime.dll", + ), + ( + ONNX_RUNTIME_PROVIDERS_DLL_FILE, + ONNX_RUNTIME_NUGET_URL, + "runtimes/win-x64/native/onnxruntime_providers_shared.dll", + ), + ( + DIRECTML_DLL_FILE, + DIRECTML_NUGET_URL, + "bin/x64-win/DirectML.dll", + ), +]; const REQUIRED_FILES: &[&str] = &[ + ONNX_RUNTIME_DLL_FILE, + ONNX_RUNTIME_PROVIDERS_DLL_FILE, + DIRECTML_DLL_FILE, "config.json", "generation_config.json", "preprocessor_config.json", @@ -21,10 +57,81 @@ const REQUIRED_FILES: &[&str] = &[ "special_tokens_map.json", "onnx/vision_encoder_fp16.onnx", "onnx/encoder_model_q4.onnx", + "onnx/decoder_model_q4.onnx", "onnx/decoder_model_merged_q4.onnx", "onnx/embed_tokens_fp16.onnx", ]; +static ORT_RUNTIME_INIT: OnceLock> = OnceLock::new(); + +/// Set to `true` by `set_caption_acceleration` so the caption worker loop +/// knows to drop its cached `FlorenceCaptioner` and reload with the new EP. +pub static CAPTION_SESSION_DIRTY: AtomicBool = AtomicBool::new(false); + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum CaptionAcceleration { + Auto, + Cpu, + Directml, +} + +impl CaptionAcceleration { + fn as_str(self) -> &'static str { + match self { + Self::Auto => "auto", + Self::Cpu => "cpu", + Self::Directml => "directml", + } + } +} + +impl Default for CaptionAcceleration { + fn default() -> Self { + Self::Auto + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum CaptionDetail { + Short, + Detailed, + Paragraph, +} + +impl CaptionDetail { + fn as_str(self) -> &'static str { + match self { + Self::Short => "short", + Self::Detailed => "detailed", + Self::Paragraph => "paragraph", + } + } + + fn prompt(self) -> &'static str { + match self { + Self::Short => "What does the image describe?", + Self::Detailed => "Describe in detail what is shown in the image.", + Self::Paragraph => "Describe with a paragraph what is shown in the image.", + } + } + + fn max_new_tokens(self) -> usize { + match self { + Self::Short => 32, + Self::Detailed => 56, + Self::Paragraph => 96, + } + } +} + +impl Default for CaptionDetail { + fn default() -> Self { + Self::Paragraph + } +} + #[derive(Serialize)] pub struct CaptionModelStatus { pub model_id: &'static str, @@ -45,6 +152,8 @@ pub struct CaptionModelProgress { #[derive(Serialize)] pub struct CaptionRuntimeProbe { pub ready: bool, + pub acceleration: CaptionAcceleration, + pub detail: CaptionDetail, pub tokenizer_vocab_size: usize, pub sessions: Vec, } @@ -61,6 +170,7 @@ pub struct CaptionVisionProbe { pub input_shape: Vec, pub output_shape: Vec, pub output_values: usize, + pub acceleration: CaptionAcceleration, } #[derive(Clone)] @@ -71,9 +181,11 @@ struct TensorData { pub struct FlorenceCaptioner { tokenizer: Tokenizer, + caption_detail: CaptionDetail, vision_session: Session, embed_session: Session, encoder_session: Session, + decoder_prefill_session: Session, decoder_session: Session, } @@ -81,6 +193,54 @@ pub fn model_dir(app_data_dir: &Path) -> PathBuf { app_data_dir.join("models").join("florence-2-base-ft") } +pub fn caption_acceleration(app_data_dir: &Path) -> CaptionAcceleration { + let path = app_data_dir.join(CAPTION_ACCELERATION_FILE); + let Ok(value) = std::fs::read_to_string(path) else { + return CaptionAcceleration::default(); + }; + + match value.trim().to_ascii_lowercase().as_str() { + "cpu" => CaptionAcceleration::Cpu, + "directml" => CaptionAcceleration::Directml, + _ => CaptionAcceleration::Auto, + } +} + +pub fn set_caption_acceleration( + app_data_dir: &Path, + acceleration: CaptionAcceleration, +) -> Result { + let path = app_data_dir.join(CAPTION_ACCELERATION_FILE); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, acceleration.as_str())?; + CAPTION_SESSION_DIRTY.store(true, Ordering::Relaxed); + Ok(acceleration) +} + +pub fn caption_detail(app_data_dir: &Path) -> CaptionDetail { + let path = app_data_dir.join(CAPTION_DETAIL_FILE); + let Ok(value) = std::fs::read_to_string(path) else { + return CaptionDetail::default(); + }; + + match value.trim().to_ascii_lowercase().as_str() { + "short" => CaptionDetail::Short, + "paragraph" => CaptionDetail::Paragraph, + _ => CaptionDetail::Detailed, + } +} + +pub fn set_caption_detail(app_data_dir: &Path, detail: CaptionDetail) -> Result { + let path = app_data_dir.join(CAPTION_DETAIL_FILE); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, detail.as_str())?; + Ok(detail) +} + pub fn caption_model_status(app_data_dir: &Path) -> CaptionModelStatus { let local_dir = model_dir(app_data_dir); let missing_files = REQUIRED_FILES @@ -133,9 +293,20 @@ pub fn prepare_caption_model_with_progress( if let Some(parent) = destination.parent() { std::fs::create_dir_all(parent)?; } - let cached = repo.get(file)?; - std::fs::copy(cached, destination)?; - completed_files += 1; + if ONNX_RUNTIME_FILES + .iter() + .any(|(runtime_file, _, _)| runtime_file == file) + { + download_onnx_runtime_files(&local_dir)?; + completed_files = REQUIRED_FILES + .iter() + .filter(|file| local_dir.join(file).exists()) + .count(); + } else { + let cached = repo.get(file)?; + std::fs::copy(cached, destination)?; + completed_files += 1; + } emit_progress(CaptionModelProgress { total_files: REQUIRED_FILES.len(), completed_files, @@ -172,21 +343,40 @@ pub fn probe_caption_runtime(app_data_dir: &Path) -> Result } let local_dir = model_dir(app_data_dir); + ensure_onnx_runtime(&local_dir)?; let tokenizer = Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; - let sessions = [ - "onnx/vision_encoder_fp16.onnx", + let acceleration = caption_acceleration(app_data_dir); + + // For the vision encoder (the only session that runs on the GPU EP) we + // create an actual ORT session so we can verify which EP was actually + // loaded, rather than just echoing the settings file. + let vision_session = probe_vision_session( + &local_dir.join("onnx/vision_encoder_fp16.onnx"), + acceleration, + )?; + + // The remaining entries are DLLs and CPU-only text/decoder models — probe + // them with the lightweight file-size check as before. + let other_files = [ + "onnxruntime/onnxruntime.dll", + "onnxruntime/onnxruntime_providers_shared.dll", + "onnxruntime/DirectML.dll", "onnx/embed_tokens_fp16.onnx", "onnx/encoder_model_q4.onnx", + "onnx/decoder_model_q4.onnx", "onnx/decoder_model_merged_q4.onnx", - ] - .into_iter() - .map(|file| probe_session(file, &local_dir.join(file))) - .collect::>>()?; + ]; + let mut sessions = vec![vision_session]; + for file in other_files { + sessions.push(probe_session(file, &local_dir.join(file))?); + } Ok(CaptionRuntimeProbe { ready: true, + acceleration, + detail: caption_detail(app_data_dir), tokenizer_vocab_size: tokenizer.get_vocab_size(false), sessions, }) @@ -202,11 +392,17 @@ pub fn probe_caption_vision(app_data_dir: &Path, image_path: &Path) -> Result input @@ -220,6 +416,7 @@ pub fn probe_caption_vision(app_data_dir: &Path, image_path: &Path) -> Result Result Result { + let started_at = Instant::now(); let status = caption_model_status(app_data_dir); if !status.ready { anyhow::bail!( @@ -239,48 +437,92 @@ impl FlorenceCaptioner { } let local_dir = model_dir(app_data_dir); + ensure_onnx_runtime(&local_dir)?; let tokenizer = Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + let caption_detail = caption_detail(app_data_dir); - let vision_session = create_session(&local_dir.join("onnx/vision_encoder_fp16.onnx"))?; - let embed_session = create_session(&local_dir.join("onnx/embed_tokens_fp16.onnx"))?; - let encoder_session = create_session(&local_dir.join("onnx/encoder_model_q4.onnx"))?; - let decoder_session = create_session(&local_dir.join("onnx/decoder_model_merged_q4.onnx"))?; + let sessions_started_at = Instant::now(); + let acceleration = caption_acceleration(app_data_dir); + let vision_session = create_session( + &local_dir.join("onnx/vision_encoder_fp16.onnx"), + acceleration, + true, + )?; + let embed_session = create_session( + &local_dir.join("onnx/embed_tokens_fp16.onnx"), + acceleration, + false, + )?; + let encoder_session = create_session( + &local_dir.join("onnx/encoder_model_q4.onnx"), + acceleration, + false, + )?; + let decoder_prefill_session = create_session( + &local_dir.join("onnx/decoder_model_q4.onnx"), + acceleration, + false, + )?; + let decoder_session = create_session( + &local_dir.join("onnx/decoder_model_merged_q4.onnx"), + acceleration, + false, + )?; + println!( + "Florence sessions loaded in {:?} with {:?} acceleration (total init {:?})", + sessions_started_at.elapsed(), + acceleration, + started_at.elapsed() + ); Ok(Self { tokenizer, + caption_detail, vision_session, embed_session, encoder_session, + decoder_prefill_session, decoder_session, }) } pub fn generate(&mut self, image_path: &Path) -> Result { + let started_at = Instant::now(); + println!("Florence caption started: {}", image_path.display()); let image_features = run_vision_encoder(&mut self.vision_session, image_path)?; + println!("Florence vision encoder done in {:?}", started_at.elapsed()); let prompt_ids = self .tokenizer - .encode("What does the image describe?", false) + .encode(self.caption_detail.prompt(), false) .map_err(anyhow::Error::msg)? .get_ids() .iter() .map(|id| i64::from(*id)) .collect::>(); let prompt_embeds = run_token_embedder(&mut self.embed_session, &prompt_ids)?; - let encoder_embeds = concatenate_sequence_embeddings(&prompt_embeds, &image_features)?; + println!( + "Florence token embeddings done in {:?}", + started_at.elapsed() + ); + let encoder_embeds = concatenate_sequence_embeddings(&image_features, &prompt_embeds)?; let encoder_attention_mask = vec![1_i64; encoder_embeds.shape[1] as usize]; let encoder_hidden_states = run_encoder( &mut self.encoder_session, &encoder_embeds, &encoder_attention_mask, )?; + println!("Florence encoder done in {:?}", started_at.elapsed()); let generated_ids = run_decoder( + &mut self.decoder_prefill_session, &mut self.decoder_session, &mut self.embed_session, &encoder_hidden_states, &encoder_attention_mask, + self.caption_detail.max_new_tokens(), )?; + println!("Florence decoder done in {:?}", started_at.elapsed()); let generated_u32 = generated_ids .into_iter() @@ -330,6 +572,26 @@ fn probe_session(file: &'static str, path: &Path) -> Result ( + vec!["OrtGetApiBase".to_string()], + vec!["ORT DirectML 1.24.2".to_string()], + ), + "onnxruntime/onnxruntime_providers_shared.dll" => ( + vec!["providers".to_string()], + vec!["shared provider bridge".to_string()], + ), + "onnxruntime/DirectML.dll" => ( + vec!["DirectX 12".to_string()], + vec!["DirectML 1.15.4".to_string()], + ), + "onnx/decoder_model_q4.onnx" => ( + vec![ + "encoder_attention_mask".to_string(), + "encoder_hidden_states".to_string(), + "inputs_embeds".to_string(), + ], + vec!["logits".to_string(), "present_key_values".to_string()], + ), _ => (Vec::new(), Vec::new()), }; @@ -340,6 +602,111 @@ fn probe_session(file: &'static str, path: &Path) -> Result Result { + let metadata = std::fs::metadata(path)?; + if metadata.len() == 0 { + anyhow::bail!("{} is empty", path.display()); + } + + // Try to create the session with the requested acceleration. If DirectML + // was requested but fails we report what actually loaded. + let loaded_acceleration = match acceleration { + CaptionAcceleration::Cpu => { + create_session(path, CaptionAcceleration::Cpu, false)?; + CaptionAcceleration::Cpu + } + CaptionAcceleration::Auto => { + // `fail_silently` — session will fall back to CPU if DirectML + // is unavailable; we detect this by trying Directml explicitly. + let directml_ok = create_session(path, CaptionAcceleration::Directml, true).is_ok(); + if directml_ok { + CaptionAcceleration::Directml + } else { + create_session(path, CaptionAcceleration::Cpu, false)?; + CaptionAcceleration::Cpu + } + } + CaptionAcceleration::Directml => { + create_session(path, CaptionAcceleration::Directml, true)?; + CaptionAcceleration::Directml + } + }; + + Ok(CaptionRuntimeSessionProbe { + file: "onnx/vision_encoder_fp16.onnx", + inputs: vec!["pixel_values".to_string()], + outputs: vec![format!("image_features [EP: {:?}]", loaded_acceleration)], + }) +} + +pub fn ensure_onnx_runtime(local_dir: &Path) -> Result<()> { + let dll_path = local_dir.join(ONNX_RUNTIME_DLL_FILE); + ORT_RUNTIME_INIT + .get_or_init(|| { + if !dll_path.exists() { + return Err(format!( + "ONNX Runtime DLL is missing: {}", + dll_path.display() + )); + } + ort::environment::init_from(&dll_path) + .map_err(|error| error.to_string())? + .with_name("phokus-florence") + .commit(); + Ok(()) + }) + .clone() + .map_err(anyhow::Error::msg) +} + +fn download_onnx_runtime_files(local_dir: &Path) -> Result<()> { + if !cfg!(target_os = "windows") { + anyhow::bail!( + "Florence-2 ONNX Runtime download is currently configured for Windows builds" + ); + } + + for (destination_file, source_url, archive_path) in ONNX_RUNTIME_FILES { + let destination = local_dir.join(destination_file); + download_nuget_file(source_url, archive_path, &destination)?; + } + + Ok(()) +} + +fn download_nuget_file(source_url: &str, archive_path: &str, destination: &Path) -> Result<()> { + let mut response = ureq::get(source_url) + .call() + .map_err(|error| anyhow::anyhow!("{error}"))?; + let mut bytes = Vec::new(); + response + .body_mut() + .as_reader() + .read_to_end(&mut bytes) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?; + let mut dll = archive.by_name(archive_path)?; + if let Some(parent) = destination.parent() { + std::fs::create_dir_all(parent)?; + } + + let temp_destination = destination.with_extension("tmp"); + { + let mut file = std::fs::File::create(&temp_destination)?; + std::io::copy(&mut dll, &mut file)?; + } + std::fs::rename(temp_destination, destination)?; + Ok(()) +} + fn run_vision_encoder(session: &mut Session, image_path: &Path) -> Result { let pixels = preprocess_image(image_path)?; let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice())) @@ -387,24 +754,50 @@ fn run_encoder( } fn run_decoder( + decoder_prefill_session: &mut Session, decoder_session: &mut Session, embed_session: &mut Session, encoder_hidden_states: &TensorData, encoder_attention_mask: &[i64], + max_new_tokens: usize, ) -> Result> { const DECODER_LAYERS: usize = 6; - const DECODER_HEADS: usize = 12; - const HEAD_DIM: usize = 64; const DECODER_START_TOKEN_ID: i64 = 2; + const FORCED_BOS_TOKEN_ID: i64 = 0; const EOS_TOKEN_ID: i64 = 2; - const MAX_NEW_TOKENS: usize = 32; let mut generated = Vec::new(); - let mut next_input_id = DECODER_START_TOKEN_ID; - let mut past: Vec = Vec::new(); - let mut use_cache_branch = false; + let encoder_attention_mask_tensor = Tensor::from_array(( + [1usize, encoder_attention_mask.len()], + encoder_attention_mask.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let encoder_hidden_states_tensor = tensor_from_data(encoder_hidden_states)?; + let prefill_inputs_embeds = run_token_embedder(embed_session, &[DECODER_START_TOKEN_ID])?; + let prefill_inputs_embeds_tensor = tensor_from_data(&prefill_inputs_embeds)?; + let prefill_started_at = Instant::now(); + let prefill_outputs = decoder_prefill_session + .run(ort::inputs! { + "encoder_attention_mask" => encoder_attention_mask_tensor, + "encoder_hidden_states" => encoder_hidden_states_tensor, + "inputs_embeds" => prefill_inputs_embeds_tensor + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + println!( + "Florence decoder prefill done in {:?}", + prefill_started_at.elapsed() + ); + let _prefill_logits = tensor_data(&prefill_outputs[0])?; + let mut next_input_id = FORCED_BOS_TOKEN_ID; + let encoder_kv = collect_present_key_values(&prefill_outputs, DECODER_LAYERS)?; + let mut decoder_kv = encoder_kv.clone(); + + for _ in 0..max_new_tokens { + if next_input_id == EOS_TOKEN_ID { + break; + } + generated.push(next_input_id); - for _ in 0..MAX_NEW_TOKENS { let inputs_embeds = run_token_embedder(embed_session, &[next_input_id])?; let encoder_attention_mask_tensor = Tensor::from_array(( [1usize, encoder_attention_mask.len()], @@ -413,9 +806,8 @@ fn run_decoder( .map_err(|error| anyhow::anyhow!("{error}"))?; let encoder_hidden_states_tensor = tensor_from_data(encoder_hidden_states)?; let inputs_embeds_tensor = tensor_from_data(&inputs_embeds)?; - let use_cache_branch_tensor = - Tensor::from_array(([1usize], vec![use_cache_branch].into_boxed_slice())) - .map_err(|error| anyhow::anyhow!("{error}"))?; + let use_cache_branch_tensor = Tensor::from_array(([1usize], vec![true].into_boxed_slice())) + .map_err(|error| anyhow::anyhow!("{error}"))?; let mut inputs = ort::inputs! { "encoder_attention_mask" => encoder_attention_mask_tensor, @@ -424,92 +816,82 @@ fn run_decoder( "use_cache_branch" => use_cache_branch_tensor }; - if past.is_empty() { - for layer in 0..DECODER_LAYERS { - push_tensor_input( - &mut inputs, - format!("past_key_values.{layer}.decoder.key"), - TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), - )?; - push_tensor_input( - &mut inputs, - format!("past_key_values.{layer}.decoder.value"), - TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), - )?; - push_tensor_input( - &mut inputs, - format!("past_key_values.{layer}.encoder.key"), - TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), - )?; - push_tensor_input( - &mut inputs, - format!("past_key_values.{layer}.encoder.value"), - TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]), - )?; - } - } else { - for layer in 0..DECODER_LAYERS { - for cache_name in [ - "decoder.key", - "decoder.value", - "encoder.key", - "encoder.value", - ] { - let past_index = layer * 4 - + match cache_name { - "decoder.key" => 0, - "decoder.value" => 1, - "encoder.key" => 2, - "encoder.value" => 3, - _ => unreachable!(), - }; - push_tensor_input( - &mut inputs, - format!("past_key_values.{layer}.{cache_name}"), - past[past_index].clone(), - )?; - } - } + for layer in 0..DECODER_LAYERS { + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.decoder.key"), + decoder_kv[layer * 4].clone(), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.decoder.value"), + decoder_kv[layer * 4 + 1].clone(), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.encoder.key"), + encoder_kv[layer * 4 + 2].clone(), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.encoder.value"), + encoder_kv[layer * 4 + 3].clone(), + )?; } let outputs = decoder_session .run(inputs) .map_err(|error| anyhow::anyhow!("{error}"))?; let logits = tensor_data(&outputs["logits"])?; - let token_id = argmax_last_token(&logits)?; - if token_id == EOS_TOKEN_ID { - break; - } - generated.push(token_id); - next_input_id = token_id; - - past.clear(); - for layer in 0..DECODER_LAYERS { - for cache_name in [ - "decoder.key", - "decoder.value", - "encoder.key", - "encoder.value", - ] { - past.push(tensor_data( - &outputs[format!("present.{layer}.{cache_name}").as_str()], - )?); - } - } - use_cache_branch = true; + next_input_id = argmax_last_token(&logits)?; + decoder_kv = collect_present_key_values(&outputs, DECODER_LAYERS)?; } + println!("Florence decoder produced {} token(s)", generated.len()); Ok(generated) } -fn create_session(path: &Path) -> Result { +fn create_session( + path: &Path, + acceleration: CaptionAcceleration, + allow_directml: bool, +) -> Result { let builder = Session::builder().map_err(|error| anyhow::anyhow!("{error}"))?; let builder = builder .with_optimization_level(GraphOptimizationLevel::Level3) .map_err(|error| anyhow::anyhow!("{error}"))?; - let mut builder = builder + let use_directml = allow_directml + && matches!( + acceleration, + CaptionAcceleration::Auto | CaptionAcceleration::Directml + ); + let builder = builder + .with_memory_pattern(!use_directml) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let builder = builder + .with_parallel_execution(false) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let builder = builder .with_intra_threads(1) .map_err(|error| anyhow::anyhow!("{error}"))?; + let mut builder = match acceleration { + CaptionAcceleration::Cpu => builder, + CaptionAcceleration::Auto if use_directml => builder + .with_execution_providers([ep::DirectML::default().build().fail_silently()]) + .unwrap_or_else(|error| error.recover()), + CaptionAcceleration::Directml if use_directml => builder + .with_execution_providers([ep::DirectML::default().build().error_on_failure()]) + .map_err(|error| anyhow::anyhow!("{error}"))?, + CaptionAcceleration::Auto | CaptionAcceleration::Directml => { + // `allow_directml` is false for the 4 text/decoder sessions — + // they intentionally run on CPU regardless of the setting. + println!( + "Florence: using CPU for {} (DirectML disabled for this session type)", + path.display() + ); + builder + } + }; let session = builder .commit_from_file(path) .map_err(|error| anyhow::anyhow!("{error}"))?; @@ -537,6 +919,24 @@ fn concatenate_sequence_embeddings(text: &TensorData, image: &TensorData) -> Res }) } +fn collect_present_key_values( + outputs: &SessionOutputs<'_>, + layers: usize, +) -> Result> { + let expected = layers * 4; + if outputs.len() < expected + 1 { + anyhow::bail!( + "Decoder returned {} output(s), expected at least {}", + outputs.len(), + expected + 1 + ); + } + + (1..=expected) + .map(|index| tensor_data(&outputs[index])) + .collect::>>() +} + fn tensor_data(value: &ort::value::DynValue) -> Result { let (shape, values) = value .try_extract_tensor::() @@ -608,13 +1008,3 @@ fn preprocess_image(image_path: &Path) -> Result> { Ok(pixel_values) } - -impl TensorData { - fn zeros(shape: Vec) -> Self { - let values_len = shape.iter().map(|value| (*value).max(0) as usize).product(); - Self { - shape, - values: vec![0.0; values_len], - } - } -} diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index d28e7d5..6c2a47a 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -1,7 +1,11 @@ -use crate::captioner::{self, CaptionModelStatus, CaptionRuntimeProbe, CaptionVisionProbe}; -use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord}; +use crate::captioner::{ + self, CaptionAcceleration, CaptionDetail, CaptionModelStatus, CaptionRuntimeProbe, + CaptionVisionProbe, +}; +use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord, ImageTag}; use crate::embedder; use crate::indexer; +use crate::tagger::{self, TaggerAcceleration, TaggerModelStatus, TaggerRuntimeProbe}; use crate::vector; use serde::{Deserialize, Serialize}; use std::path::PathBuf; @@ -39,12 +43,14 @@ pub struct UpdateImageDetailsParams { #[derive(Deserialize)] pub struct FindSimilarImagesParams { pub image_id: i64, + pub folder_id: Option, pub limit: Option, } #[derive(Deserialize)] pub struct DebugSimilarImagesParams { pub image_id: i64, + pub folder_id: Option, pub limit: Option, } @@ -72,6 +78,26 @@ pub struct QueueCaptionJobsParams { pub image_id: Option, } +#[derive(Deserialize)] +pub struct ClearCaptionJobsParams { + pub folder_id: Option, +} + +#[derive(Deserialize)] +pub struct ResetCaptionsParams { + pub folder_id: Option, +} + +#[derive(Deserialize)] +pub struct SetCaptionAccelerationParams { + pub acceleration: CaptionAcceleration, +} + +#[derive(Deserialize)] +pub struct SetCaptionDetailParams { + pub detail: CaptionDetail, +} + #[derive(Deserialize)] pub struct ProbeCaptionImageParams { pub image_id: i64, @@ -235,8 +261,8 @@ pub async fn find_similar_images( db::repair_embedding_consistency(&conn).map_err(|e| e.to_string())?; return Ok(Vec::new()); } - let image_ids = - vector::find_similar_image_ids(&conn, params.image_id, limit).map_err(|e| e.to_string())?; + let image_ids = vector::find_similar_image_ids(&conn, params.image_id, limit, params.folder_id) + .map_err(|e| e.to_string())?; db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string()) } @@ -258,7 +284,8 @@ pub async fn debug_similar_images( let vector_count = vector::count_image_vectors(&conn).map_err(|e| e.to_string())?; let has_vector = vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())?; let similar_ids = if has_vector { - vector::find_similar_image_ids(&conn, params.image_id, limit).map_err(|e| e.to_string())? + vector::find_similar_image_ids(&conn, params.image_id, limit, params.folder_id) + .map_err(|e| e.to_string())? } else { Vec::new() }; @@ -332,6 +359,36 @@ pub async fn get_caption_model_status(app: AppHandle) -> Result Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(captioner::caption_acceleration(&app_dir)) +} + +#[tauri::command] +pub async fn set_caption_acceleration( + app: AppHandle, + params: SetCaptionAccelerationParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + captioner::set_caption_acceleration(&app_dir, params.acceleration).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn get_caption_detail(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(captioner::caption_detail(&app_dir)) +} + +#[tauri::command] +pub async fn set_caption_detail( + app: AppHandle, + params: SetCaptionDetailParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + captioner::set_caption_detail(&app_dir, params.detail).map_err(|e| e.to_string()) +} + #[tauri::command] pub async fn prepare_caption_model(app: AppHandle) -> Result { let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; @@ -441,6 +498,24 @@ pub async fn queue_caption_jobs( } } +#[tauri::command] +pub async fn clear_caption_jobs( + db: State<'_, DbState>, + params: ClearCaptionJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + db::clear_caption_jobs(&conn, params.folder_id).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn reset_generated_captions( + db: State<'_, DbState>, + params: ResetCaptionsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + db::reset_generated_captions(&conn, params.folder_id).map_err(|e| e.to_string()) +} + #[derive(Serialize, Deserialize)] pub struct TagCloudEntry { pub count: usize, @@ -670,11 +745,27 @@ pub struct FolderWorkerStates { pub metadata_paused: bool, pub embedding_paused: bool, pub caption_paused: bool, + pub tagging_paused: bool, } #[tauri::command] -pub async fn set_worker_paused(worker: String, folder_id: i64, paused: bool) -> Result<(), String> { +pub async fn set_worker_paused( + db: State<'_, DbState>, + worker: String, + folder_id: i64, + paused: bool, +) -> Result<(), String> { indexer::set_worker_paused(&worker, folder_id, paused); + if worker == "caption" && paused { + let conn = db.get().map_err(|e| e.to_string())?; + db::requeue_processing_caption_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } + if worker == "tagging" && paused { + let conn = db.get().map_err(|e| e.to_string())?; + db::requeue_processing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } Ok(()) } @@ -684,23 +775,213 @@ pub async fn get_worker_states(folder_ids: Vec) -> Result, + pub image_id: Option, +} + +#[derive(Deserialize)] +pub struct ClearTaggingJobsParams { + pub folder_id: Option, +} + +#[derive(Deserialize)] +pub struct GetImageTagsParams { + pub image_id: i64, +} + +#[derive(Deserialize)] +pub struct AddUserTagParams { + pub image_id: i64, + pub tag: String, +} + +#[derive(Deserialize)] +pub struct RemoveTagParams { + pub tag_id: i64, +} + +#[tauri::command] +pub async fn get_tagger_model_status(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_model_status(&app_dir)) +} + +#[tauri::command] +pub async fn get_tagger_acceleration(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_acceleration(&app_dir)) +} + +#[tauri::command] +pub async fn set_tagger_acceleration( + app: AppHandle, + params: SetTaggerAccelerationParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tagger::set_tagger_acceleration(&app_dir, params.acceleration).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn probe_tagger_runtime(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || tagger::probe_tagger_runtime(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn get_tagger_threshold(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_threshold(&app_dir)) +} + +#[tauri::command] +pub async fn set_tagger_threshold( + app: AppHandle, + params: SetTaggerThresholdParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tagger::set_tagger_threshold(&app_dir, params.threshold).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn prepare_tagger_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || { + let app = app.clone(); + tagger::prepare_tagger_model_with_progress(&app_dir, move |progress| { + let _ = app.emit("tagger-model-progress", progress); + }) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn delete_tagger_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || tagger::delete_tagger_model(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn queue_tagging_jobs( + app: AppHandle, + db: State<'_, DbState>, + params: QueueTaggingJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let (total, folder_ids) = match (params.folder_id, params.image_id) { + (_, Some(image_id)) => { + db::enqueue_tagging_job(&conn, image_id).map_err(|e| e.to_string())?; + // Look up just this image's folder_id rather than fetching all folders + let image = db::get_image_by_id(&conn, image_id).map_err(|e| e.to_string())?; + (1usize, vec![image.folder_id]) + } + (Some(folder_id), None) => { + let n = db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + (n, vec![folder_id]) + } + (None, None) => { + let folders = db::get_folders(&conn).map_err(|e| e.to_string())?; + let folder_ids: Vec = folders.iter().map(|f| f.id).collect(); + let mut total = 0usize; + for &folder_id in &folder_ids { + total += db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } + (total, folder_ids) + } + }; + drop(conn); + indexer::emit_folder_job_progress(&app, db.inner(), &folder_ids, true); + Ok(total) +} + +#[tauri::command] +pub async fn clear_tagging_jobs( + app: AppHandle, + db: State<'_, DbState>, + params: ClearTaggingJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let n = db::clear_tagging_jobs(&conn, params.folder_id).map_err(|e| e.to_string())?; + let folder_ids: Vec = match params.folder_id { + Some(id) => vec![id], + None => db::get_folders(&conn) + .map_err(|e| e.to_string())? + .into_iter() + .map(|f| f.id) + .collect(), + }; + drop(conn); + indexer::emit_folder_job_progress(&app, db.inner(), &folder_ids, true); + Ok(n) +} + +#[tauri::command] +pub async fn get_image_tags( + db: State<'_, DbState>, + params: GetImageTagsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_image_tags(&conn, params.image_id).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn add_user_tag( + db: State<'_, DbState>, + params: AddUserTagParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + db::add_user_tag(&conn, params.image_id, ¶ms.tag).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn remove_tag(db: State<'_, DbState>, params: RemoveTagParams) -> Result<(), String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::remove_tag(&conn, params.tag_id).map_err(|e| e.to_string()) +} diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index b576687..81bd162 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -61,6 +61,10 @@ pub struct ImageRecord { pub caption_model: Option, pub caption_updated_at: Option, pub caption_error: Option, + pub ai_rating: Option, + pub ai_tagger_model: Option, + pub ai_tagged_at: Option, + pub ai_tagger_error: Option, } #[allow(dead_code)] @@ -100,6 +104,24 @@ pub struct CaptionJob { pub path: String, } +#[derive(Debug, Clone)] +pub struct TaggingJob { + pub image_id: i64, + pub folder_id: i64, + pub path: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageTag { + pub id: i64, + pub image_id: i64, + pub tag: String, + pub source: String, + pub ai_model: Option, + pub confidence: Option, + pub created_at: String, +} + #[derive(Debug, Clone)] pub struct IndexedMediaEntry { pub id: i64, @@ -120,6 +142,9 @@ pub struct FolderJobProgress { pub caption_pending: i64, pub caption_ready: i64, pub caption_failed: i64, + pub tagging_pending: i64, + pub tagging_ready: i64, + pub tagging_failed: i64, } pub fn create_pool(db_path: &Path) -> Result { @@ -195,6 +220,26 @@ pub fn migrate(conn: &Connection) -> Result<()> { updated_at TEXT NOT NULL DEFAULT (datetime('now')) ); + CREATE TABLE IF NOT EXISTS tagging_jobs ( + image_id INTEGER PRIMARY KEY REFERENCES images(id) ON DELETE CASCADE, + status TEXT NOT NULL DEFAULT 'pending', + attempts INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) + ); + + CREATE TABLE IF NOT EXISTS image_tags ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + image_id INTEGER NOT NULL REFERENCES images(id) ON DELETE CASCADE, + tag TEXT NOT NULL, + source TEXT NOT NULL DEFAULT 'user', + ai_model TEXT, + confidence REAL, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE(image_id, tag) + ); + CREATE TABLE IF NOT EXISTS tag_cloud_cache ( folder_scope TEXT PRIMARY KEY, image_ids_hash INTEGER NOT NULL, @@ -208,6 +253,10 @@ pub fn migrate(conn: &Connection) -> Result<()> { CREATE INDEX IF NOT EXISTS idx_thumbnail_jobs_status ON thumbnail_jobs(status); CREATE INDEX IF NOT EXISTS idx_metadata_jobs_status ON metadata_jobs(status); CREATE INDEX IF NOT EXISTS idx_caption_jobs_status ON caption_jobs(status); + CREATE INDEX IF NOT EXISTS idx_tagging_jobs_status ON tagging_jobs(status); + CREATE INDEX IF NOT EXISTS idx_image_tags_image_id ON image_tags(image_id); + CREATE INDEX IF NOT EXISTS idx_image_tags_source ON image_tags(source); + CREATE INDEX IF NOT EXISTS idx_image_tags_tag ON image_tags(tag); ", )?; @@ -237,6 +286,10 @@ pub fn migrate(conn: &Connection) -> Result<()> { ensure_column(conn, "images", "caption_model", "TEXT")?; ensure_column(conn, "images", "caption_updated_at", "TEXT")?; ensure_column(conn, "images", "caption_error", "TEXT")?; + ensure_column(conn, "images", "ai_rating", "TEXT")?; + ensure_column(conn, "images", "ai_tagger_model", "TEXT")?; + ensure_column(conn, "images", "ai_tagged_at", "TEXT")?; + ensure_column(conn, "images", "ai_tagger_error", "TEXT")?; vector::migrate(conn)?; Ok(()) @@ -257,8 +310,8 @@ pub fn insert_folder(conn: &Connection, path: &str, name: &str) -> Result { pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { let id = conn.query_row( - "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, generated_caption, caption_model, caption_updated_at, caption_error) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26) + "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, generated_caption, caption_model, caption_updated_at, caption_error, ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26, ?27, ?28, ?29, ?30) ON CONFLICT(path) DO UPDATE SET folder_id = excluded.folder_id, filename = excluded.filename, @@ -311,6 +364,10 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { img.caption_model, img.caption_updated_at, img.caption_error, + img.ai_rating, + img.ai_tagger_model, + img.ai_tagged_at, + img.ai_tagger_error, ], |row| row.get(0), )?; @@ -433,6 +490,13 @@ pub fn reset_inflight_jobs(conn: &Connection) -> Result<()> { "UPDATE caption_jobs SET status = 'pending' WHERE status = 'processing'", [], )?; + conn.execute( + "UPDATE tagging_jobs SET status = 'pending' WHERE status = 'processing'", + [], + )?; + // Delete any rows that were cancelled before the previous shutdown so + // they don't silently linger in the DB across restarts. + conn.execute("DELETE FROM tagging_jobs WHERE status = 'cancelled'", [])?; Ok(()) } @@ -493,6 +557,110 @@ pub fn requeue_caption_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()> Ok(()) } +pub fn requeue_processing_caption_jobs_for_folder( + conn: &Connection, + folder_id: i64, +) -> Result { + conn.execute( + "UPDATE caption_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN ( + SELECT id FROM images WHERE folder_id = ?1 + )", + [folder_id], + ) + .map_err(Into::into) +} + +pub fn clear_caption_jobs(conn: &Connection, folder_id: Option) -> Result { + match folder_id { + Some(folder_id) => { + let deleted = conn.execute( + "DELETE FROM caption_jobs + WHERE image_id IN ( + SELECT id FROM images WHERE folder_id = ?1 + )", + [folder_id], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE folder_id = ?1 + AND generated_caption IS NULL", + [folder_id], + )?; + Ok(deleted) + } + None => { + let deleted = conn.execute("DELETE FROM caption_jobs", [])?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE generated_caption IS NULL", + [], + )?; + Ok(deleted) + } + } +} + +pub fn reset_generated_captions(conn: &Connection, folder_id: Option) -> Result { + let image_ids = match folder_id { + Some(folder_id) => { + let mut stmt = conn.prepare( + "SELECT id FROM images WHERE folder_id = ?1 AND generated_caption IS NOT NULL", + )?; + let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?; + rows.collect::>>()? + } + None => { + let mut stmt = + conn.prepare("SELECT id FROM images WHERE generated_caption IS NOT NULL")?; + let rows = stmt.query_map([], |row| row.get::<_, i64>(0))?; + rows.collect::>>()? + } + }; + + let tx = conn.unchecked_transaction()?; + for image_id in &image_ids { + vector::delete_caption_embedding(&tx, *image_id)?; + } + + match folder_id { + Some(folder_id) => { + tx.execute( + "UPDATE images + SET generated_caption = NULL, + caption_model = NULL, + caption_updated_at = NULL, + caption_error = NULL + WHERE folder_id = ?1", + [folder_id], + )?; + tx.execute( + "DELETE FROM caption_jobs + WHERE image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; + } + None => { + tx.execute( + "UPDATE images + SET generated_caption = NULL, + caption_model = NULL, + caption_updated_at = NULL, + caption_error = NULL", + [], + )?; + tx.execute("DELETE FROM caption_jobs", [])?; + } + } + + tx.commit()?; + Ok(image_ids.len()) +} + pub fn enqueue_missing_caption_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result { let inserted = conn.execute( "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at) @@ -791,9 +959,36 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result Result Result "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, - generated_caption, caption_model, caption_updated_at, caption_error + generated_caption, caption_model, caption_updated_at, caption_error, + ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error FROM images WHERE id = ?1", [image_id], @@ -1126,7 +1326,8 @@ pub fn get_images( "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, - generated_caption, caption_model, caption_updated_at, caption_error + generated_caption, caption_model, caption_updated_at, caption_error, + ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error FROM images WHERE (?1 IS NULL OR folder_id = ?1) AND (?2 IS NULL OR filename LIKE ?2) @@ -1240,6 +1441,299 @@ pub fn mark_caption_failed(conn: &Connection, image_id: i64, error: &str) -> Res Ok(()) } +// --------------------------------------------------------------------------- +// Tagging jobs +// --------------------------------------------------------------------------- + +pub fn enqueue_tagging_job(conn: &Connection, image_id: i64) -> Result<()> { + conn.execute( + "INSERT INTO tagging_jobs (image_id, status, attempts, last_error, created_at, updated_at) + VALUES (?1, 'pending', 0, NULL, datetime('now'), datetime('now')) + ON CONFLICT(image_id) DO UPDATE SET + status = CASE WHEN status != 'processing' THEN 'pending' ELSE status END, + last_error = CASE WHEN status != 'processing' THEN NULL ELSE last_error END, + updated_at = datetime('now')", + [image_id], + )?; + Ok(()) +} + +pub fn claim_tagging_jobs( + conn: &mut Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let tx = conn.transaction()?; + let candidates = get_pending_tagging_jobs_excluding(&tx, paused_folder_ids, limit * 2)?; + let mut claimed = Vec::new(); + for job in candidates { + let updated = tx.execute( + "UPDATE tagging_jobs + SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'pending'", + [job.image_id], + )?; + if updated > 0 { + claimed.push(job); + if claimed.len() >= limit { + break; + } + } + } + tx.commit()?; + Ok(claimed) +} + +fn get_pending_tagging_jobs_excluding( + conn: &Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let mut stmt = conn.prepare( + "SELECT j.image_id, i.folder_id, i.path + FROM tagging_jobs j + JOIN images i ON i.id = j.image_id + WHERE j.status = 'pending' + ORDER BY j.created_at ASC + LIMIT ?1", + )?; + let rows = stmt + .query_map([limit as i64], |row| { + Ok(TaggingJob { + image_id: row.get(0)?, + folder_id: row.get(1)?, + path: row.get(2)?, + }) + })? + .collect::>>()?; + Ok(rows + .into_iter() + .filter(|job| !paused_folder_ids.contains(&job.folder_id)) + .collect()) +} + +pub fn update_ai_tags( + conn: &Connection, + image_id: i64, + tags: &[(String, f64)], + rating: &str, + model: &str, +) -> Result<()> { + // NOTE: callers are responsible for wrapping this in a transaction. + // Do NOT open a nested transaction here — rusqlite/SQLite do not support + // nested transactions and will error with "cannot start a transaction + // within a transaction". + + // Remove previous AI tags for this image before inserting fresh ones + conn.execute( + "DELETE FROM image_tags WHERE image_id = ?1 AND source = 'ai'", + [image_id], + )?; + + let mut stmt = conn.prepare( + "INSERT INTO image_tags (image_id, tag, source, ai_model, confidence, created_at) + VALUES (?1, ?2, 'ai', ?3, ?4, datetime('now')) + ON CONFLICT(image_id, tag) DO UPDATE SET + source = 'ai', + ai_model = excluded.ai_model, + confidence = excluded.confidence", + )?; + for (tag, confidence) in tags { + stmt.execute(params![image_id, tag, model, confidence])?; + } + drop(stmt); + + conn.execute( + "UPDATE images + SET ai_rating = ?2, + ai_tagger_model = ?3, + ai_tagged_at = datetime('now'), + ai_tagger_error = NULL + WHERE id = ?1", + params![image_id, rating, model], + )?; + conn.execute("DELETE FROM tagging_jobs WHERE image_id = ?1", [image_id])?; + Ok(()) +} + +pub fn mark_tagging_failed(conn: &Connection, image_id: i64, error: &str) -> Result<()> { + conn.execute( + "UPDATE images SET ai_tagger_error = ?2 WHERE id = ?1", + params![image_id, error], + )?; + conn.execute( + "UPDATE tagging_jobs + SET status = 'failed', last_error = ?2, updated_at = datetime('now') + WHERE image_id = ?1", + params![image_id, error], + )?; + Ok(()) +} + +pub fn clear_tagging_jobs(conn: &Connection, folder_id: Option) -> Result { + // Rows currently being processed by the worker must not be deleted mid-flight + // — the worker holds a reference and will write results back after inference. + // Mark them 'cancelled' so the worker discards the result instead of saving it, + // then delete every non-processing row. On the next poll the worker will find + // no pending work and the queue will appear empty. + let deleted = match folder_id { + Some(fid) => { + conn.execute( + "UPDATE tagging_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [fid], + )?; + let n = conn.execute( + "DELETE FROM tagging_jobs + WHERE status != 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [fid], + )?; + conn.execute( + "UPDATE images + SET ai_tagger_error = NULL + WHERE folder_id = ?1 AND ai_tagged_at IS NULL", + [fid], + )?; + n + } + None => { + conn.execute( + "UPDATE tagging_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing'", + [], + )?; + let n = conn.execute("DELETE FROM tagging_jobs WHERE status != 'processing'", [])?; + conn.execute( + "UPDATE images SET ai_tagger_error = NULL WHERE ai_tagged_at IS NULL", + [], + )?; + n + } + }; + Ok(deleted) +} + +pub fn get_image_tags(conn: &Connection, image_id: i64) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, image_id, tag, source, ai_model, confidence, created_at + FROM image_tags + WHERE image_id = ?1 + ORDER BY source DESC, confidence DESC NULLS LAST, tag ASC", + )?; + let rows = stmt + .query_map([image_id], |row| { + Ok(ImageTag { + id: row.get(0)?, + image_id: row.get(1)?, + tag: row.get(2)?, + source: row.get(3)?, + ai_model: row.get(4)?, + confidence: row.get(5)?, + created_at: row.get(6)?, + }) + })? + .collect::>>()?; + Ok(rows) +} + +pub fn add_user_tag(conn: &Connection, image_id: i64, tag: &str) -> Result { + conn.execute( + "INSERT INTO image_tags (image_id, tag, source, ai_model, confidence, created_at) + VALUES (?1, ?2, 'user', NULL, NULL, datetime('now')) + ON CONFLICT(image_id, tag) DO NOTHING", + params![image_id, tag], + )?; + let row = conn.query_row( + "SELECT id, image_id, tag, source, ai_model, confidence, created_at + FROM image_tags WHERE image_id = ?1 AND tag = ?2", + params![image_id, tag], + |row| { + Ok(ImageTag { + id: row.get(0)?, + image_id: row.get(1)?, + tag: row.get(2)?, + source: row.get(3)?, + ai_model: row.get(4)?, + confidence: row.get(5)?, + created_at: row.get(6)?, + }) + }, + )?; + Ok(row) +} + +pub fn remove_tag(conn: &Connection, tag_id: i64) -> Result<()> { + conn.execute("DELETE FROM image_tags WHERE id = ?1", [tag_id])?; + Ok(()) +} + +pub fn enqueue_missing_tagging_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result { + let inserted = conn.execute( + "INSERT INTO tagging_jobs (image_id, status, attempts, last_error, created_at, updated_at) + SELECT id, 'pending', 0, NULL, datetime('now'), datetime('now') + FROM images + WHERE folder_id = ?1 + AND media_kind = 'image' + AND ai_tagged_at IS NULL + ON CONFLICT(image_id) DO UPDATE SET + -- Only reset to 'pending' if the row is not currently being + -- processed or cancelled; leave 'processing' rows untouched so + -- the worker can complete them and clean up normally. + status = CASE WHEN status NOT IN ('processing', 'cancelled') THEN 'pending' ELSE status END, + last_error = CASE WHEN status != 'processing' THEN NULL ELSE last_error END, + updated_at = datetime('now')", + [folder_id], + )?; + conn.execute( + "UPDATE images + SET ai_tagger_error = NULL + WHERE folder_id = ?1 + AND media_kind = 'image' + AND ai_tagged_at IS NULL", + [folder_id], + )?; + Ok(inserted) +} + +pub fn requeue_processing_tagging_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result<()> { + conn.execute( + "UPDATE tagging_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; + Ok(()) +} + +/// Returns `true` when the job row for `image_id` currently has status = 'cancelled'. +/// Used by the worker to discard inference results for jobs that were cancelled +/// while inference was running. +pub fn is_tagging_job_cancelled(conn: &Connection, image_id: i64) -> Result { + let count: i64 = conn.query_row( + "SELECT COUNT(*) FROM tagging_jobs WHERE image_id = ?1 AND status = 'cancelled'", + [image_id], + |row| row.get(0), + )?; + Ok(count > 0) +} + +pub fn requeue_tagging_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()> { + for image_id in image_ids { + conn.execute( + "UPDATE tagging_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'processing'", + [image_id], + )?; + } + Ok(()) +} + pub fn suggest_tags_from_caption( conn: &Connection, image_id: i64, @@ -1332,6 +1826,10 @@ fn map_image_row(row: &Row<'_>) -> rusqlite::Result { caption_model: row.get(24)?, caption_updated_at: row.get(25)?, caption_error: row.get(26)?, + ai_rating: row.get(27)?, + ai_tagger_model: row.get(28)?, + ai_tagged_at: row.get(29)?, + ai_tagger_error: row.get(30)?, }) } diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index 33bcbcd..f509d1f 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -3,6 +3,7 @@ use crate::db::{self, DbPool, EmbeddingJob, FolderJobProgress, ImageRecord, Inde use crate::embedder::{embedding_source_path, ClipImageEmbedder}; use crate::media::{probe_video_metadata, MediaTools}; use crate::storage::{detect_storage_profile, RuntimeAdaptiveProfile, StorageProfile}; +use crate::tagger::{self, WdTagger}; use crate::thumbnail; use crate::vector; use anyhow::Result; @@ -35,6 +36,7 @@ struct PausedWorkerFolders { metadata: HashSet, embedding: HashSet, caption: HashSet, + tagging: HashSet, } #[derive(Clone, Copy)] @@ -43,6 +45,7 @@ pub struct FolderWorkerPausedState { pub metadata: bool, pub embedding: bool, pub caption: bool, + pub tagging: bool, } pub fn set_worker_paused(worker: &str, folder_id: i64, paused: bool) { @@ -55,6 +58,7 @@ pub fn set_worker_paused(worker: &str, folder_id: i64, paused: bool) { "metadata" => Some(&mut paused_folders.metadata), "embedding" => Some(&mut paused_folders.embedding), "caption" => Some(&mut paused_folders.caption), + "tagging" => Some(&mut paused_folders.tagging), _ => None, }; @@ -87,6 +91,7 @@ pub fn get_worker_paused_states(folder_ids: &[i64]) -> HashMap HashSet { "thumbnail" => paused_folders.thumbnail.clone(), "metadata" => paused_folders.metadata.clone(), "embedding" => paused_folders.embedding.clone(), + "caption" => paused_folders.caption.clone(), + "tagging" => paused_folders.tagging.clone(), _ => HashSet::new(), } } @@ -191,6 +198,12 @@ pub fn start_caption_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) let mut captioner: Option = None; println!("Caption worker started."); loop { + // If the acceleration setting changed, drop the cached session so + // the next batch picks it up with the new execution provider. + if captioner::CAPTION_SESSION_DIRTY.swap(false, std::sync::atomic::Ordering::Relaxed) { + println!("Caption worker: acceleration setting changed — resetting session."); + captioner = None; + } if let Err(error) = process_caption_batch(&app, &pool, &app_data_dir, &mut captioner) { eprintln!("Caption worker error: {}", error); captioner = None; @@ -200,6 +213,28 @@ pub fn start_caption_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) }); } +pub fn start_tagging_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) { + std::thread::spawn(move || { + let mut tagger_instance: Option = None; + println!("Tagging worker started."); + loop { + // If the acceleration setting changed, drop the cached session so + // the next batch picks it up with the new execution provider. + if tagger::TAGGER_SESSION_DIRTY.swap(false, std::sync::atomic::Ordering::Relaxed) { + 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; + } + std::thread::sleep(std::time::Duration::from_millis(750)); + } + }); +} + fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> { let existing_entries = { let conn = pool.get()?; @@ -370,6 +405,10 @@ fn build_record( caption_model: None, caption_updated_at: None, caption_error: None, + ai_rating: None, + ai_tagger_model: None, + ai_tagged_at: None, + ai_tagger_error: None, }) } @@ -825,6 +864,134 @@ fn process_caption_batch( Ok(()) } +const TAGGING_BATCH_SIZE: usize = 1; + +fn process_tagging_batch( + app: &AppHandle, + pool: &DbPool, + app_data_dir: &Path, + tagger_instance: &mut Option, +) -> Result<()> { + if !tagger::tagger_model_status(app_data_dir).ready { + return Ok(()); + } + + let paused_folders = paused_folder_ids("tagging"); + let jobs = with_db_write_lock(|| { + let mut conn = pool.get()?; + db::claim_tagging_jobs(&mut conn, &paused_folders, TAGGING_BATCH_SIZE) + })?; + + if jobs.is_empty() { + return Ok(()); + } + + if tagger_instance.is_none() { + match WdTagger::new(app_data_dir) { + Ok(model) => *tagger_instance = Some(model), + Err(error) => { + with_db_write_lock(|| { + let conn = pool.get()?; + db::requeue_tagging_jobs( + &conn, + &jobs.iter().map(|job| job.image_id).collect::>(), + ) + })?; + return Err(error); + } + } + } + + let folder_ids = jobs.iter().map(|job| job.folder_id).collect::>(); + emit_folder_job_progress( + app, + pool, + &folder_ids.iter().copied().collect::>(), + false, + ); + + let tagger_ref = tagger_instance + .as_mut() + .expect("tagger should be initialized before tagging batch processing"); + + let tag_results = jobs + .iter() + .map(|job| { + ( + job.clone(), + tagger_ref.run(Path::new(&job.path), tagger::DEFAULT_MAX_TAGS), + ) + }) + .collect::>(); + + let updated_images = with_db_write_lock(|| { + let mut conn = pool.get()?; + let tx = conn.transaction()?; + let mut updated_images = Vec::with_capacity(tag_results.len()); + + for (job, tag_result) in &tag_results { + // If the job was cancelled while inference was running, discard + // the result and delete the row — don't save tags or mark failed. + if db::is_tagging_job_cancelled(&tx, job.image_id)? { + tx.execute( + "DELETE FROM tagging_jobs WHERE image_id = ?1", + [job.image_id], + )?; + continue; + } + match tag_result { + Ok(output) => { + let tag_pairs: Vec<(String, f64)> = output + .tags + .iter() + .map(|t| (t.tag.clone(), t.confidence as f64)) + .collect(); + db::update_ai_tags( + &tx, + job.image_id, + &tag_pairs, + &output.rating, + tagger::WD_TAGGER_MODEL_NAME, + )?; + } + Err(error) => { + db::mark_tagging_failed(&tx, job.image_id, &error.to_string())?; + } + } + updated_images.push(db::get_image_by_id(&tx, job.image_id)?); + } + + tx.commit()?; + Ok(updated_images) + }) + .or_else(|db_err| { + // The DB write failed. Try to requeue the claimed jobs so they aren't + // left stuck in 'processing' until the next app restart. + let image_ids: Vec = jobs.iter().map(|job| job.image_id).collect(); + let _ = with_db_write_lock(|| { + let conn = pool.get()?; + db::requeue_tagging_jobs(&conn, &image_ids) + }); + Err(db_err) + })?; + + if !updated_images.is_empty() { + let folder_ids = updated_images + .iter() + .map(|image| image.folder_id) + .collect::>(); + emit_media_updates( + app, + &MediaUpdateBatch { + images: updated_images, + }, + ); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); + } + + Ok(()) +} + fn active_indexing_folders() -> HashSet { ACTIVE_INDEXING_FOLDERS .get_or_init(|| Mutex::new(HashSet::new())) @@ -910,7 +1077,7 @@ fn emit_media_updates(app: &AppHandle, batch: &MediaUpdateBatch) { let _ = app.emit("media-updated", batch); } -fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64], force: bool) { +pub fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64], force: bool) { let mut unique_folder_ids = folder_ids.iter().copied().collect::>(); unique_folder_ids.sort_unstable(); unique_folder_ids.dedup(); diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs new file mode 100644 index 0000000..d0cd179 --- /dev/null +++ b/src-tauri/src/tagger.rs @@ -0,0 +1,625 @@ +use anyhow::Result; +use hf_hub::{api::sync::Api, Repo, RepoType}; +use image::{imageops::FilterType, DynamicImage, ImageReader}; +use ort::ep; +use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::Tensor; +use serde::{Deserialize, Serialize}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Instant; + +pub const WD_TAGGER_MODEL_ID: &str = "SmilingWolf/wd-swinv2-tagger-v3"; +pub const WD_TAGGER_MODEL_NAME: &str = "wd-swinv2-tagger-v3"; + +const TAGGER_ACCELERATION_FILE: &str = "settings/tagger_acceleration.txt"; +const TAGGER_THRESHOLD_FILE: &str = "settings/tagger_threshold.txt"; + +// Files required on disk before the tagger can run. The ONNX runtime DLLs +// are shared with the captioner and live in the same `onnxruntime/` directory. +const TAGGER_REQUIRED_FILES: &[&str] = &[ + "onnxruntime/onnxruntime.dll", + "onnxruntime/onnxruntime_providers_shared.dll", + "onnxruntime/DirectML.dll", + "model.onnx", + "selected_tags.csv", +]; + +// Tags in these Danbooru categories are kept in the output. +// Category 0 = general, category 4 = character. +// Category 9 = rating (explicit/questionable/sensitive/general) – used for +// `ai_rating` but NOT emitted as individual tags. +const GENERAL_CATEGORY: u8 = 0; +const CHARACTER_CATEGORY: u8 = 4; +const RATING_CATEGORY: u8 = 9; + +pub const DEFAULT_THRESHOLD: f32 = 0.35; +pub const DEFAULT_MAX_TAGS: usize = 30; + +/// 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); + +// --------------------------------------------------------------------------- +// Settings types +// --------------------------------------------------------------------------- + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum TaggerAcceleration { + Auto, + Cpu, + Directml, +} + +impl TaggerAcceleration { + fn as_str(self) -> &'static str { + match self { + Self::Auto => "auto", + Self::Cpu => "cpu", + Self::Directml => "directml", + } + } +} + +impl Default for TaggerAcceleration { + fn default() -> Self { + Self::Auto + } +} + +// --------------------------------------------------------------------------- +// Status / probe types exposed to the frontend +// --------------------------------------------------------------------------- + +#[derive(Serialize)] +pub struct TaggerModelStatus { + pub model_id: &'static str, + pub model_name: &'static str, + pub local_dir: String, + pub ready: bool, + pub missing_files: Vec, +} + +#[derive(Clone, Serialize)] +pub struct TaggerModelProgress { + pub total_files: usize, + pub completed_files: usize, + pub current_file: Option, + pub done: bool, +} + +// --------------------------------------------------------------------------- +// Runtime probe types exposed to the frontend +// --------------------------------------------------------------------------- + +#[derive(Serialize)] +pub struct TaggerRuntimeProbe { + pub ready: bool, + pub acceleration: TaggerAcceleration, + pub session: TaggerRuntimeSessionProbe, +} + +#[derive(Serialize)] +pub struct TaggerRuntimeSessionProbe { + pub file: &'static str, + pub inputs: Vec, + pub outputs: Vec, +} + +// --------------------------------------------------------------------------- +// Tag record returned to callers +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Serialize)] +pub struct TagResult { + pub tag: String, + pub confidence: f32, +} + +#[derive(Debug, Clone, Serialize)] +pub struct TaggerOutput { + pub tags: Vec, + /// Highest-scoring rating label: "general" | "sensitive" | "questionable" | "explicit" + pub rating: String, +} + +// --------------------------------------------------------------------------- +// Internal label table built from selected_tags.csv +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +struct TagEntry { + name: String, + category: u8, +} + +// --------------------------------------------------------------------------- +// Path helpers +// --------------------------------------------------------------------------- + +pub fn model_dir(app_data_dir: &Path) -> PathBuf { + app_data_dir.join("models").join("wd-swinv2-tagger-v3") +} + +// --------------------------------------------------------------------------- +// Settings persistence +// --------------------------------------------------------------------------- + +pub fn tagger_acceleration(app_data_dir: &Path) -> TaggerAcceleration { + let path = app_data_dir.join(TAGGER_ACCELERATION_FILE); + let Ok(value) = std::fs::read_to_string(path) else { + return TaggerAcceleration::default(); + }; + match value.trim().to_ascii_lowercase().as_str() { + "cpu" => TaggerAcceleration::Cpu, + "directml" => TaggerAcceleration::Directml, + _ => TaggerAcceleration::Auto, + } +} + +pub fn set_tagger_acceleration( + app_data_dir: &Path, + acceleration: TaggerAcceleration, +) -> Result { + let path = app_data_dir.join(TAGGER_ACCELERATION_FILE); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, acceleration.as_str())?; + TAGGER_SESSION_DIRTY.store(true, Ordering::Relaxed); + Ok(acceleration) +} + +pub fn tagger_threshold(app_data_dir: &Path) -> f32 { + let path = app_data_dir.join(TAGGER_THRESHOLD_FILE); + let Ok(value) = std::fs::read_to_string(path) else { + return DEFAULT_THRESHOLD; + }; + value + .trim() + .parse::() + .unwrap_or(DEFAULT_THRESHOLD) + .clamp(0.01, 1.0) +} + +pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result { + let clamped = threshold.clamp(0.01, 1.0); + let path = app_data_dir.join(TAGGER_THRESHOLD_FILE); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, clamped.to_string())?; + Ok(clamped) +} + +// --------------------------------------------------------------------------- +// Model status / download +// --------------------------------------------------------------------------- + +pub fn tagger_model_status(app_data_dir: &Path) -> TaggerModelStatus { + let local_dir = model_dir(app_data_dir); + // The ONNX runtime DLLs live in the caption model dir; reuse them. + let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft"); + let missing_files = TAGGER_REQUIRED_FILES + .iter() + .filter(|file| { + let path = if file.starts_with("onnxruntime/") { + caption_model_dir.join(file) + } else { + local_dir.join(file) + }; + !path.exists() + }) + .map(|file| (*file).to_string()) + .collect::>(); + + TaggerModelStatus { + model_id: WD_TAGGER_MODEL_ID, + model_name: WD_TAGGER_MODEL_NAME, + local_dir: local_dir.to_string_lossy().to_string(), + ready: missing_files.is_empty(), + missing_files, + } +} + +pub fn prepare_tagger_model_with_progress( + app_data_dir: &Path, + emit_progress: impl Fn(TaggerModelProgress), +) -> Result { + let local_dir = model_dir(app_data_dir); + std::fs::create_dir_all(&local_dir)?; + + // Only download the two tagger-specific files; ONNX runtime DLLs are + // already handled by the captioner download flow. + const DOWNLOAD_FILES: &[&str] = &["model.onnx", "selected_tags.csv"]; + + let api = Api::new()?; + let repo = api.repo(Repo::new(WD_TAGGER_MODEL_ID.to_string(), RepoType::Model)); + + let mut completed_files = DOWNLOAD_FILES + .iter() + .filter(|file| local_dir.join(file).exists()) + .count(); + + emit_progress(TaggerModelProgress { + total_files: DOWNLOAD_FILES.len(), + completed_files, + current_file: None, + done: completed_files == DOWNLOAD_FILES.len(), + }); + + for file in DOWNLOAD_FILES { + let destination = local_dir.join(file); + if destination.exists() { + continue; + } + emit_progress(TaggerModelProgress { + total_files: DOWNLOAD_FILES.len(), + completed_files, + current_file: Some((*file).to_string()), + done: false, + }); + let cached = repo.get(file)?; + std::fs::copy(cached, destination)?; + completed_files += 1; + emit_progress(TaggerModelProgress { + total_files: DOWNLOAD_FILES.len(), + completed_files, + current_file: Some((*file).to_string()), + done: completed_files == DOWNLOAD_FILES.len(), + }); + } + + emit_progress(TaggerModelProgress { + total_files: DOWNLOAD_FILES.len(), + completed_files, + current_file: None, + done: true, + }); + + Ok(tagger_model_status(app_data_dir)) +} + +pub fn delete_tagger_model(app_data_dir: &Path) -> Result { + let local_dir = model_dir(app_data_dir); + if local_dir.exists() { + std::fs::remove_dir_all(&local_dir)?; + } + Ok(tagger_model_status(app_data_dir)) +} + +pub fn probe_tagger_runtime(app_data_dir: &Path) -> Result { + let status = tagger_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "WD Tagger model is missing {} required file(s): {}", + status.missing_files.len(), + status.missing_files.join(", ") + ); + } + + let local_dir = model_dir(app_data_dir); + let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft"); + crate::captioner::ensure_onnx_runtime(&caption_model_dir)?; + + let acceleration = tagger_acceleration(app_data_dir); + let model_path = local_dir.join("model.onnx"); + + // Verify that the model file exists and has non-zero size before trying + // to create an ORT session (better error message on corruption). + let metadata = std::fs::metadata(&model_path)?; + if metadata.len() == 0 { + anyhow::bail!("model.onnx is empty"); + } + + // Actually create a session to verify the EP loads correctly. + let loaded_acceleration = match acceleration { + TaggerAcceleration::Cpu => { + create_tagger_session(&model_path, TaggerAcceleration::Cpu)?; + TaggerAcceleration::Cpu + } + TaggerAcceleration::Auto => { + // Try DirectML explicitly; if it fails the real session would have + // fallen back to CPU silently — report what would actually run. + let directml_ok = + create_tagger_session(&model_path, TaggerAcceleration::Directml).is_ok(); + if directml_ok { + TaggerAcceleration::Directml + } else { + create_tagger_session(&model_path, TaggerAcceleration::Cpu)?; + TaggerAcceleration::Cpu + } + } + TaggerAcceleration::Directml => { + create_tagger_session(&model_path, TaggerAcceleration::Directml)?; + TaggerAcceleration::Directml + } + }; + + Ok(TaggerRuntimeProbe { + ready: true, + acceleration: loaded_acceleration, + session: TaggerRuntimeSessionProbe { + file: "model.onnx", + inputs: vec!["pixel_values".to_string()], + outputs: vec![format!("output [EP: {:?}]", loaded_acceleration)], + }, + }) +} + +// --------------------------------------------------------------------------- +// Top-level inference entry point +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// Tagger implementation +// --------------------------------------------------------------------------- + +pub struct WdTagger { + session: Session, + labels: Vec, + threshold: f32, + input_size: usize, +} + +impl WdTagger { + pub fn new(app_data_dir: &Path) -> Result { + let started_at = Instant::now(); + + let status = tagger_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "WD tagger model is missing {} required file(s): {}", + status.missing_files.len(), + status.missing_files.join(", ") + ); + } + + let local_dir = model_dir(app_data_dir); + + // The ONNX runtime DLLs are shared with the captioner; use the + // captioner's shared ORT init lock to avoid double-initialisation. + let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft"); + crate::captioner::ensure_onnx_runtime(&caption_model_dir).map_err(|e| { + anyhow::anyhow!( + "ONNX Runtime not initialised — download the Florence-2 caption model first \ + to get the shared runtime DLLs. Original error: {e}" + ) + })?; + + let acceleration = tagger_acceleration(app_data_dir); + let threshold = tagger_threshold(app_data_dir); + let model_path = local_dir.join("model.onnx"); + let labels_path = local_dir.join("selected_tags.csv"); + + 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 = { + let inputs = session.inputs(); + let dim = 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 }) + } else { + None + } + }) + .unwrap_or(448); + dim + }; + + let labels = load_labels(&labels_path)?; + + println!( + "WD tagger loaded in {:?} ({} labels, input {}x{}, {:?} acceleration)", + started_at.elapsed(), + labels.len(), + input_size, + input_size, + acceleration, + ); + + Ok(Self { + session, + labels, + threshold, + input_size, + }) + } + + pub fn run(&mut self, image_path: &Path, max_tags: usize) -> Result { + let started_at = Instant::now(); + + let image_array = preprocess_image(image_path, self.input_size)?; + let batch_size = 1usize; + let input = Tensor::from_array(( + [batch_size, self.input_size, self.input_size, 3usize], + image_array.into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let input_name: String = self.session.inputs()[0].name().to_string(); + let outputs = self + .session + .run(ort::inputs! { input_name.as_str() => input }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let (_, probabilities) = outputs[0] + .try_extract_tensor::() + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let probs: &[f32] = &probabilities; + + if probs.len() != self.labels.len() { + anyhow::bail!( + "Model output length {} does not match label count {}", + probs.len(), + self.labels.len() + ); + } + + // Collect rating scores (category 9) - pick the argmax as the rating. + let rating = self + .labels + .iter() + .zip(probs.iter()) + .filter(|(entry, _)| entry.category == RATING_CATEGORY) + .max_by(|(_, a), (_, b)| a.total_cmp(b)) + .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 + .iter() + .zip(probs.iter()) + .filter(|(entry, prob)| { + (entry.category == GENERAL_CATEGORY || entry.category == CHARACTER_CATEGORY) + && **prob >= self.threshold + }) + .map(|(entry, prob)| TagResult { + tag: entry.name.clone(), + confidence: *prob, + }) + .collect(); + + tags.sort_by(|a, b| b.confidence.total_cmp(&a.confidence)); + tags.truncate(max_tags); + + println!( + "WD tagger: {} tags (threshold={}, rating={}) in {:?} for {}", + tags.len(), + self.threshold, + rating, + started_at.elapsed(), + image_path.display(), + ); + + Ok(TaggerOutput { tags, rating }) + } +} + +// --------------------------------------------------------------------------- +// Session creation – mirrors captioner's `create_session` +// --------------------------------------------------------------------------- + +fn create_tagger_session(path: &Path, acceleration: TaggerAcceleration) -> Result { + let builder = Session::builder().map_err(|error| anyhow::anyhow!("{error}"))?; + let builder = builder + .with_optimization_level(GraphOptimizationLevel::Level3) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let use_directml = matches!( + acceleration, + TaggerAcceleration::Auto | TaggerAcceleration::Directml + ); + + let builder = builder + .with_memory_pattern(!use_directml) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let builder = builder + .with_parallel_execution(false) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let builder = builder + .with_intra_threads(1) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut builder = match acceleration { + TaggerAcceleration::Cpu => { + println!("WD tagger: using CPU execution provider"); + builder + } + TaggerAcceleration::Auto => builder + .with_execution_providers([ep::DirectML::default().build().fail_silently()]) + .unwrap_or_else(|error| { + println!("WD tagger: DirectML unavailable, falling back to CPU"); + error.recover() + }), + TaggerAcceleration::Directml => builder + .with_execution_providers([ep::DirectML::default().build().error_on_failure()]) + .map_err(|error| anyhow::anyhow!("{error}"))?, + }; + + let session = builder + .commit_from_file(path) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + Ok(session) +} +// --------------------------------------------------------------------------- +// Label loading +// --------------------------------------------------------------------------- + +fn load_labels(path: &Path) -> Result> { + let mut reader = csv::Reader::from_path(path)?; + let mut entries = Vec::new(); + for result in reader.records() { + let record = result?; + // CSV columns: tag_id, name, category, count + let name = record + .get(1) + .ok_or_else(|| anyhow::anyhow!("Missing name column in selected_tags.csv"))? + .replace('_', " "); + let category: u8 = record + .get(2) + .ok_or_else(|| anyhow::anyhow!("Missing category column in selected_tags.csv"))? + .parse() + .unwrap_or(0); + entries.push(TagEntry { name, category }); + } + if entries.is_empty() { + anyhow::bail!("selected_tags.csv is empty or could not be parsed"); + } + Ok(entries) +} + +// --------------------------------------------------------------------------- +// Image preprocessing +// WD tagger expects: (1, H, W, 3) float32, raw [0,255] values, BGR channel +// order, padded to square with white (255,255,255). +// --------------------------------------------------------------------------- + +fn preprocess_image(image_path: &Path, target_size: usize) -> Result> { + let image = ImageReader::open(image_path)?.decode()?; + + // Composite any alpha channel onto a white background. + let image_rgba = image.to_rgba8(); + let (width, height) = image_rgba.dimensions(); + let mut canvas_rgba = + image::RgbaImage::from_pixel(width, height, image::Rgba([255, 255, 255, 255])); + image::imageops::overlay(&mut canvas_rgba, &image_rgba, 0, 0); + let image_rgb = DynamicImage::ImageRgba8(canvas_rgba).to_rgb8(); + + // Pad to square. + let max_dim = width.max(height); + let pad_left = (max_dim - width) / 2; + let pad_top = (max_dim - height) / 2; + let mut square = image::RgbImage::from_pixel(max_dim, max_dim, image::Rgb([255, 255, 255])); + image::imageops::overlay(&mut square, &image_rgb, pad_left as i64, pad_top as i64); + + // Resize to model input size. + let resized = image::imageops::resize( + &square, + target_size as u32, + target_size as u32, + FilterType::CatmullRom, + ); + + // Flatten to (H, W, 3) float32 in BGR order, values in [0, 255]. + let mut pixel_values = vec![0.0f32; target_size * target_size * 3]; + for (x, y, pixel) in resized.enumerate_pixels() { + let base = (y as usize * target_size + x as usize) * 3; + // BGR order + pixel_values[base] = f32::from(pixel[2]); // B + pixel_values[base + 1] = f32::from(pixel[1]); // G + pixel_values[base + 2] = f32::from(pixel[0]); // R + } + + Ok(pixel_values) +} diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 6b99c8b..3b5d942 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -80,7 +80,12 @@ pub fn upsert_caption_embedding(conn: &Connection, image_id: i64, embedding: &[f Ok(()) } -pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Result> { +pub fn find_similar_image_ids( + conn: &Connection, + image_id: i64, + limit: usize, + folder_id: Option, +) -> Result> { let embedding: Vec = match conn.query_row( "SELECT embedding FROM image_vec WHERE image_id = ?1", [image_id], @@ -91,19 +96,37 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Err(error) => return Err(error.into()), }; + let allowed_folder_ids = match folder_id { + Some(folder_id) => Some(image_ids_for_folder(conn, folder_id)?), + None => None, + }; + let search_limit = if allowed_folder_ids.is_some() { + count_image_vectors(conn)?.max(1) as usize + } else { + limit + 1 + }; let mut stmt = conn.prepare( "SELECT image_id FROM image_vec WHERE embedding MATCH vec_f32(?1) AND k = ?2", )?; - let rows = stmt.query_map((&embedding, (limit as i64) + 1), |row| row.get::<_, i64>(0))?; + let rows = stmt + .query_map((&embedding, search_limit as i64), |row| { + row.get::<_, i64>(0) + })? + .collect::>>()?; let mut ids = Vec::new(); for row in rows { - let candidate_id = row?; - if candidate_id != image_id { - ids.push(candidate_id); + if row != image_id { + if allowed_folder_ids + .as_ref() + .is_some_and(|folder_ids| !folder_ids.contains(&row)) + { + continue; + } + ids.push(row); } if ids.len() >= limit { break; @@ -112,6 +135,15 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Ok(ids) } +fn image_ids_for_folder( + conn: &Connection, + folder_id: i64, +) -> Result> { + let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?; + let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?; + Ok(rows.collect::>>()?) +} + /// Returns all stored image embeddings with their image IDs, optionally filtered to one folder. /// Each entry is `(image_id, normalized_f32_embedding)`. pub fn get_all_image_embeddings_with_ids( -- 2.52.0 From ff4a568b574a96b01f2824eb26c2a2bba42beff1 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 12 Apr 2026 12:18:47 +0100 Subject: [PATCH 05/25] feat: expand media exploration and tagging controls --- package.json | 10 +- pnpm-lock.yaml | 39 ++ src-tauri/.cargo/config.toml | 12 + src-tauri/Cargo.lock | 308 ++++++++++++ src-tauri/Cargo.toml | 50 ++ src-tauri/src/commands.rs | 413 ++++++++++++++- src-tauri/src/db.rs | 200 +++++++- src-tauri/src/hnsw_index.rs | 129 +++++ src-tauri/src/indexer.rs | 6 +- src-tauri/src/lib.rs | 10 + src-tauri/src/tagger.rs | 23 + src-tauri/src/vector.rs | 139 ++++- src-tauri/tauri.conf.json | 10 +- src/App.tsx | 8 + src/components/BackgroundTasks.tsx | 82 ++- src/components/DuplicateFinder.tsx | 274 ++++++++++ src/components/Gallery.tsx | 174 ++++--- src/components/Lightbox.tsx | 3 +- src/components/SettingsModal.tsx | 780 ++++++++++++++++------------- src/components/Sidebar.tsx | 17 + src/components/TagCloud.tsx | 512 ++++++++++++------- src/components/Toolbar.tsx | 288 ++++++++--- src/store.ts | 517 +++++++++++++++++-- 23 files changed, 3206 insertions(+), 798 deletions(-) create mode 100644 src-tauri/.cargo/config.toml create mode 100644 src-tauri/src/hnsw_index.rs create mode 100644 src/components/DuplicateFinder.tsx diff --git a/package.json b/package.json index c90f202..5d541c0 100644 --- a/package.json +++ b/package.json @@ -4,9 +4,11 @@ "version": "0.1.0", "type": "module", "scripts": { - "dev": "vite", - "build": "tsc && vite build", - "preview": "vite preview", + "build:app": "tauri build", + "build:vite": "tsc && vite build", + "dev:app": "tauri dev", + "dev:vite": "vite", +"preview": "vite preview", "tauri": "tauri" }, "dependencies": { @@ -15,6 +17,7 @@ "@tauri-apps/plugin-dialog": "^2.7.0", "@tauri-apps/plugin-fs": "^2.5.0", "@tauri-apps/plugin-opener": "^2", + "d3-force": "^3.0.0", "framer-motion": "^12.38.0", "react": "^19.1.0", "react-dom": "^19.1.0", @@ -23,6 +26,7 @@ "devDependencies": { "@tailwindcss/vite": "^4.2.2", "@tauri-apps/cli": "^2", + "@types/d3-force": "^3.0.10", "@types/react": "^19.1.8", "@types/react-dom": "^19.1.6", "@vitejs/plugin-react": "^4.6.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4fd524e..b2dd733 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -23,6 +23,9 @@ importers: '@tauri-apps/plugin-opener': specifier: ^2 version: 2.5.3 + d3-force: + specifier: ^3.0.0 + version: 3.0.0 framer-motion: specifier: ^12.38.0 version: 12.38.0(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -42,6 +45,9 @@ importers: '@tauri-apps/cli': specifier: ^2 version: 2.10.1 + '@types/d3-force': + specifier: ^3.0.10 + version: 3.0.10 '@types/react': specifier: ^19.1.8 version: 19.2.14 @@ -662,6 +668,9 @@ packages: '@types/babel__traverse@7.28.0': resolution: {integrity: sha512-8PvcXf70gTDZBgt9ptxJ8elBeBjcLOAcOtoO/mPJjtji1+CdGbHgm77om1GrsPxsiE+uXIpNSK64UYaIwQXd4Q==} + '@types/d3-force@3.0.10': + resolution: {integrity: sha512-ZYeSaCF3p73RdOKcjj+swRlZfnYpK1EbaDiYICEEp5Q6sUiqFaFQ9qgoshp5CzIyyb/yD09kD9o2zEltCexlgw==} + '@types/estree@1.0.8': resolution: {integrity: sha512-dWHzHa2WqEXI/O1E9OjrocMTKJl2mSrEolh1Iomrv6U+JuNwaHXsXx9bLu5gG7BUWFIN0skIQJQ/L1rIex4X6w==} @@ -698,6 +707,22 @@ packages: csstype@3.2.3: resolution: {integrity: sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==} + d3-dispatch@3.0.1: + resolution: {integrity: sha512-rzUyPU/S7rwUflMyLc1ETDeBj0NRuHKKAcvukozwhshr6g6c5d8zh4c2gQjY2bZ0dXeGLWc1PF174P2tVvKhfg==} + engines: {node: '>=12'} + + d3-force@3.0.0: + resolution: {integrity: sha512-zxV/SsA+U4yte8051P4ECydjD/S+qeYtnaIyAs9tgHCqfguma/aAQDjo85A9Z6EKhBirHRJHXIgJUlffT4wdLg==} + engines: {node: '>=12'} + + d3-quadtree@3.0.1: + resolution: {integrity: sha512-04xDrxQTDTCFwP5H6hRhsRcb9xxv2RzkcsygFzmkSIOJy3PeRJP7sNk3VRIbKXcog561P9oU0/rVH6vDROAgUw==} + engines: {node: '>=12'} + + d3-timer@3.0.1: + resolution: {integrity: sha512-ndfJ/JxxMd3nw31uyKoY2naivF+r29V+Lc0svZxe1JvvIRmi8hUsrMvdOwgS1o6uBHmiz91geQ0ylPP0aj1VUA==} + engines: {node: '>=12'} + debug@4.4.3: resolution: {integrity: sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==} engines: {node: '>=6.0'} @@ -1448,6 +1473,8 @@ snapshots: dependencies: '@babel/types': 7.29.0 + '@types/d3-force@3.0.10': {} + '@types/estree@1.0.8': {} '@types/react-dom@19.2.3(@types/react@19.2.14)': @@ -1486,6 +1513,18 @@ snapshots: csstype@3.2.3: {} + d3-dispatch@3.0.1: {} + + d3-force@3.0.0: + dependencies: + d3-dispatch: 3.0.1 + d3-quadtree: 3.0.1 + d3-timer: 3.0.1 + + d3-quadtree@3.0.1: {} + + d3-timer@3.0.1: {} + debug@4.4.3: dependencies: ms: 2.1.3 diff --git a/src-tauri/.cargo/config.toml b/src-tauri/.cargo/config.toml new file mode 100644 index 0000000..22422f2 --- /dev/null +++ b/src-tauri/.cargo/config.toml @@ -0,0 +1,12 @@ +[target.x86_64-pc-windows-msvc] +rustflags = [ + # Disable MSVC incremental linking — avoids .ilk file corruption + # and removes one source of link failure on restart. + "-C", "link-arg=/INCREMENTAL:NO", + # Skip PDB generation entirely in dev builds. + # The .pdb file is the most common reason the linker fails after + # an unclean shutdown: the previous phokus.exe process holds the + # file open, so link.exe cannot write a new one → LNK error. + # Rust backtraces still work without MSVC PDBs. + "-C", "link-arg=/DEBUG:NONE", +] diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 0a27e7a..195eddf 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -61,6 +61,73 @@ dependencies = [ "libc", ] +[[package]] +name = "anndists" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8396b473aa0bceed68fb32462505387ea39fa47c7029417e0a49f10592b036" +dependencies = [ + "anyhow", + "cfg-if", + "cpu-time", + "env_logger", + "lazy_static", + "log", + "num-traits", + "num_cpus", + "rayon", +] + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "anyhow" version = "1.0.102" @@ -266,6 +333,15 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + [[package]] name = "bit-set" version = "0.8.0" @@ -595,6 +671,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chacha20" version = "0.10.0" @@ -626,6 +708,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b" +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "combine" version = "4.6.7" @@ -757,6 +845,16 @@ dependencies = [ "libc", ] +[[package]] +name = "cpu-time" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9e393a7668fe1fad3075085b86c781883000b4ede868f43627b34a87c8b7ded" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -1352,12 +1450,35 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "env_filter" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef" +dependencies = [ + "log", + "regex", +] + [[package]] name = "env_home" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" +[[package]] +name = "env_logger" +version = "0.11.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a" +dependencies = [ + "anstream", + "anstyle", + "env_filter", + "jiff", + "log", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -2340,6 +2461,8 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ + "allocator-api2", + "equivalent", "foldhash 0.1.5", ] @@ -2416,6 +2539,31 @@ version = "1.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f" +[[package]] +name = "hnsw_rs" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a5258f079b97bf2e8311ff9579e903c899dcbac0d9a138d62e9a066778bd07" +dependencies = [ + "anndists", + "anyhow", + "bincode", + "cfg-if", + "cpu-time", + "env_logger", + "hashbrown 0.15.5", + "indexmap 2.13.1", + "lazy_static", + "log", + "mmap-rs", + "num-traits", + "num_cpus", + "parking_lot", + "rand 0.9.2", + "rayon", + "serde", +] + [[package]] name = "html5ever" version = "0.29.1" @@ -2819,6 +2967,12 @@ dependencies = [ "once_cell", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itertools" version = "0.14.0" @@ -2857,6 +3011,30 @@ dependencies = [ "system-deps", ] +[[package]] +name = "jiff" +version = "0.2.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359" +dependencies = [ + "jiff-static", + "log", + "portable-atomic", + "portable-atomic-util", + "serde_core", +] + +[[package]] +name = "jiff-static" +version = "0.2.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "jni" version = "0.21.1" @@ -2958,6 +3136,12 @@ dependencies = [ "selectors 0.24.0", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "leb128fmt" version = "0.1.0" @@ -3109,6 +3293,15 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" +[[package]] +name = "mach2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d640282b302c0bb0a2a8e0233ead9035e3bed871f0b7e81fe4a1ec829765db44" +dependencies = [ + "libc", +] + [[package]] name = "macro_rules_attribute" version = "0.2.2" @@ -3235,6 +3428,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "mmap-rs" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ecce9d566cb9234ae3db9e249c8b55665feaaf32b0859ff1e27e310d2beb3d8" +dependencies = [ + "bitflags 2.11.0", + "combine", + "libc", + "mach2", + "nix", + "sysctl", + "thiserror 2.0.18", + "widestring", + "windows 0.48.0", +] + [[package]] name = "monostate" version = "0.1.18" @@ -3356,6 +3566,18 @@ version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags 2.11.0", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nodrop" version = "0.1.14" @@ -3634,6 +3856,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "onig" version = "6.5.1" @@ -4034,8 +4262,10 @@ dependencies = [ "fast_image_resize", "ffmpeg-sidecar", "hf-hub", + "hnsw_rs", "image", "log", + "memmap2", "ort", "r2d2", "r2d2_sqlite", @@ -6575,6 +6805,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.23.0" @@ -6938,6 +7174,12 @@ dependencies = [ "winsafe", ] +[[package]] +name = "widestring" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" + [[package]] name = "winapi" version = "0.3.9" @@ -6984,6 +7226,15 @@ dependencies = [ "windows-version", ] +[[package]] +name = "windows" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f" +dependencies = [ + "windows-targets 0.48.5", +] + [[package]] name = "windows" version = "0.61.3" @@ -7236,6 +7487,21 @@ dependencies = [ "windows_x86_64_msvc 0.42.2", ] +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", +] + [[package]] name = "windows-targets" version = "0.52.6" @@ -7302,6 +7568,12 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" @@ -7320,6 +7592,12 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + [[package]] name = "windows_aarch64_msvc" version = "0.52.6" @@ -7338,6 +7616,12 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -7368,6 +7652,12 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + [[package]] name = "windows_i686_msvc" version = "0.52.6" @@ -7386,6 +7676,12 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + [[package]] name = "windows_x86_64_gnu" version = "0.52.6" @@ -7404,6 +7700,12 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" @@ -7422,6 +7724,12 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" + [[package]] name = "windows_x86_64_msvc" version = "0.52.6" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 38356e2..2b2919e 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -27,6 +27,7 @@ rusqlite = { version = "0.32", features = ["bundled"] } r2d2 = "0.8" r2d2_sqlite = "0.25" sqlite-vec = "=0.1.9" +hnsw_rs = "0.3.4" image = { version = "0.25", default-features = false, features = ["jpeg", "png", "webp", "gif", "bmp", "tiff"] } fast_image_resize = { version = "6.0.0", features = ["image"] } walkdir = "2" @@ -38,6 +39,7 @@ anyhow = "1" log = "0.4" ffmpeg-sidecar = "2.5.0" xxhash-rust = { version = "0.8", features = ["xxh3"] } +memmap2 = "0.9" sysinfo = "0.38.4" candle-core = { version = "0.10.2", features = ["cuda"] } candle-nn = { version = "0.10.2", features = ["cuda"] } @@ -48,3 +50,51 @@ ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "n ureq = { version = "3.3.0", default-features = false, features = ["native-tls"] } zip = { version = "4.6.1", default-features = false, features = ["deflate"] } csv = "1" + +# ── Dev-mode performance ──────────────────────────────────────────────────── +# opt-level=1 on the main crate keeps incremental compile short. +# Only the packages that are genuine hot-path bottlenecks in dev get opt-level=3 +# (using "*" caused cargo to recheck all dependency fingerprints too aggressively, +# which caused spurious full rebuilds and linker conflicts). +[profile.dev] +opt-level = 1 + +# ML inference — without opt these run 20-50× slower than release +[profile.dev.package.candle-core] +opt-level = 3 +[profile.dev.package.candle-nn] +opt-level = 3 +[profile.dev.package.candle-transformers] +opt-level = 3 + +# ONNX runtime (WD tagger) +[profile.dev.package.ort] +opt-level = 3 +[profile.dev.package.ort-sys] +opt-level = 3 + +# Image decode/resize workers +[profile.dev.package.image] +opt-level = 3 +[profile.dev.package.fast_image_resize] +opt-level = 3 + +# Parallel work scheduler +[profile.dev.package.rayon] +opt-level = 3 +[profile.dev.package.rayon-core] +opt-level = 3 + +# Tokenisation (embedding model pre-processing) +[profile.dev.package.tokenizers] +opt-level = 3 + +# Hashing (duplicate finder) +[profile.dev.package.xxhash-rust] +opt-level = 3 + +# SQLite (frequent db calls in workers) +[profile.dev.package.rusqlite] +opt-level = 3 +[profile.dev.package.libsqlite3-sys] +opt-level = 3 diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 6c2a47a..d23eeb4 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -2,8 +2,9 @@ use crate::captioner::{ self, CaptionAcceleration, CaptionDetail, CaptionModelStatus, CaptionRuntimeProbe, CaptionVisionProbe, }; -use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord, ImageTag}; +use crate::db::{self, DbPool, ExploreTagEntry, Folder, FolderJobProgress, ImageRecord, ImageTag}; use crate::embedder; +use crate::hnsw_index; use crate::indexer; use crate::tagger::{self, TaggerAcceleration, TaggerModelStatus, TaggerRuntimeProbe}; use crate::vector; @@ -21,12 +22,21 @@ pub struct ImagesPage { pub limit: i64, } +#[derive(Serialize)] +pub struct SimilarImagesPage { + pub images: Vec, + pub offset: usize, + pub limit: usize, + pub has_more: bool, +} + #[derive(Deserialize)] pub struct GetImagesParams { pub folder_id: Option, pub search: Option, pub media_kind: Option, pub favorites_only: Option, + pub rating_min: Option, pub embedding_failed_only: Option, pub sort: Option, pub offset: Option, @@ -44,7 +54,9 @@ pub struct UpdateImageDetailsParams { pub struct FindSimilarImagesParams { pub image_id: i64, pub folder_id: Option, + pub offset: Option, pub limit: Option, + pub threshold: Option, } #[derive(Deserialize)] @@ -114,9 +126,56 @@ pub struct SemanticSearchParams { pub folder_id: Option, pub media_kind: Option, pub favorites_only: Option, + pub rating_min: Option, pub limit: Option, } +#[derive(Deserialize)] +pub struct TagSearchParams { + pub query: String, + pub folder_id: Option, + pub media_kind: Option, + pub favorites_only: Option, + pub rating_min: Option, + pub limit: Option, +} + +#[derive(Deserialize)] +pub struct GetExploreTagsParams { + pub folder_id: Option, + pub limit: Option, +} + +#[derive(Deserialize)] +pub struct SearchTagsAutocompleteParams { + pub query: String, + pub folder_id: Option, + pub limit: Option, +} + +#[derive(Serialize, Deserialize)] +pub struct DuplicateGroup { + pub file_hash: String, + pub file_size: u64, + pub images: Vec, +} + +#[derive(Serialize)] +pub struct DuplicateScanCache { + pub groups: Vec, + pub scanned_at: i64, +} + +#[derive(Deserialize)] +pub struct DeleteImagesFromDiskParams { + pub image_ids: Vec, +} + +#[derive(Deserialize)] +pub struct GetImagesByIdsParams { + pub image_ids: Vec, +} + #[tauri::command] pub async fn add_folder( app: AppHandle, @@ -187,6 +246,7 @@ pub async fn get_images( let search = params.search.as_deref(); let media_kind = params.media_kind.as_deref(); let favorites_only = params.favorites_only.unwrap_or(false); + let rating_min = params.rating_min.unwrap_or(0); let embedding_failed_only = params.embedding_failed_only.unwrap_or(false); let total = db::count_images( @@ -195,6 +255,7 @@ pub async fn get_images( search, media_kind, favorites_only, + rating_min, embedding_failed_only, ) .map_err(|e| e.to_string())?; @@ -205,6 +266,7 @@ pub async fn get_images( search, media_kind, favorites_only, + rating_min, embedding_failed_only, sort, offset, @@ -254,16 +316,42 @@ pub async fn reindex_folder( pub async fn find_similar_images( db: State<'_, DbState>, params: FindSimilarImagesParams, -) -> Result, String> { +) -> Result { let conn = db.get().map_err(|e| e.to_string())?; let limit = params.limit.unwrap_or(32); + let offset = params.offset.unwrap_or(0); + let threshold = params.threshold.unwrap_or(0.24); if !vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())? { db::repair_embedding_consistency(&conn).map_err(|e| e.to_string())?; - return Ok(Vec::new()); + return Ok(SimilarImagesPage { + images: Vec::new(), + offset, + limit, + has_more: false, + }); } - let image_ids = vector::find_similar_image_ids(&conn, params.image_id, limit, params.folder_id) + let matches = hnsw_index::find_similar_image_matches( + &conn, + params.image_id, + params.folder_id, + threshold, + offset, + limit + 1, + ) .map_err(|e| e.to_string())?; - db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string()) + let has_more = matches.len() > limit; + let image_ids = matches + .into_iter() + .take(limit) + .map(|(image_id, _)| image_id) + .collect::>(); + let images = db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string())?; + Ok(SimilarImagesPage { + images, + offset, + limit, + has_more, + }) } #[derive(Serialize)] @@ -328,10 +416,31 @@ pub async fn semantic_search_images( if params.favorites_only.unwrap_or(false) { images.retain(|image| image.favorite); } + if let Some(rating_min) = params.rating_min { + images.retain(|image| image.rating >= rating_min); + } Ok(images) } +#[tauri::command] +pub async fn search_images_by_tag( + db: State<'_, DbState>, + params: TagSearchParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::search_images_by_tag( + &conn, + ¶ms.query, + params.folder_id, + params.media_kind.as_deref(), + params.favorites_only.unwrap_or(false), + params.rating_min.unwrap_or(0), + params.limit.unwrap_or(64), + ) + .map_err(|e| e.to_string()) +} + #[tauri::command] pub async fn set_generated_caption( db: State<'_, DbState>, @@ -521,6 +630,8 @@ pub struct TagCloudEntry { pub count: usize, pub representative_image_id: i64, pub thumbnail_path: Option, + #[serde(default)] + pub image_ids: Vec, } fn fnv_hash_ids(ids: &[i64]) -> u64 { @@ -569,7 +680,11 @@ pub async fn get_tag_cloud( .map_err(|e| e.to_string())? { if let Ok(entries) = serde_json::from_str::>(&json) { - return Ok(entries); + // Reject cache entries written before image_ids were tracked — they all + // have empty image_ids which causes "No media found" when a cluster is opened. + if entries.iter().all(|e| !e.image_ids.is_empty()) { + return Ok(entries); + } } } } @@ -597,6 +712,13 @@ pub async fn get_tag_cloud( } let centroid = ¢roids[ci]; + let cluster_ids = points + .iter() + .enumerate() + .filter(|(i, _)| assignments[*i] == ci) + .map(|(i, _)| ids[i]) + .collect::>(); + let best_id = points .iter() .enumerate() @@ -614,6 +736,7 @@ pub async fn get_tag_cloud( count, representative_image_id: best_id, thumbnail_path, + image_ids: cluster_ids, }); } @@ -625,6 +748,224 @@ pub async fn get_tag_cloud( Ok(entries) } +#[tauri::command] +pub async fn get_explore_tags( + db: State<'_, DbState>, + params: GetExploreTagsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_explore_tags(&conn, params.folder_id, params.limit.unwrap_or(48)).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn search_tags_autocomplete( + db: State<'_, DbState>, + params: SearchTagsAutocompleteParams, +) -> Result, String> { + if params.query.trim().is_empty() { + return Ok(vec![]); + } + let conn = db.get().map_err(|e| e.to_string())?; + db::search_tags_autocomplete(&conn, ¶ms.query, params.folder_id, params.limit.unwrap_or(10)) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn find_duplicates( + app: AppHandle, + db: State<'_, DbState>, + folder_id: Option, +) -> Result, String> { + let records = { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_all_image_paths(&conn, folder_id).map_err(|e| e.to_string())? + }; + + let total = records.len(); + let _ = app.emit("duplicate_scan_progress", (0usize, total)); + + // Two-phase detection. No full-file read at any point. + // + // Phase 1 — stat: + // Read only file metadata (size). Files with a unique size cannot be + // duplicates; discard them immediately. Zero file content read. + // + // Phase 2 — sample hash: + // For each size-matched candidate, read four evenly-spaced 16 KB windows + // (head, 33%, 66%, tail) and hash them together. Total I/O per file is + // capped at 64 KB regardless of how large the file is. + // + // For small files (≤ 64 KB) the windows cover the whole file, so the + // hash is exact. For large files, matching all four windows at different + // offsets is effectively impossible for natural photo/video content unless + // the files are genuinely identical — eliminating the need for a full read. + let app_hash = app.clone(); + let pairs: Vec<(u64, u64, i64)> = tokio::task::spawn_blocking(move || { + use memmap2::Mmap; + use rayon::prelude::*; + use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; + use xxhash_rust::xxh3::{xxh3_64, Xxh3}; + + const WINDOW: usize = 16 * 1024; // 16 KB per sample window + const N: usize = 4; // windows at 0%, 33%, 66%, 100% + const COVERED: usize = WINDOW * N; // 64 KB total; also full-coverage threshold + + // Hash four evenly-spaced 16 KB windows. For files ≤ 64 KB this is + // equivalent to hashing the entire file. + fn sample(mmap: &[u8]) -> u64 { + if mmap.len() <= COVERED { + return xxh3_64(mmap); + } + let mut h = Xxh3::new(); + for i in 0..N { + let pos = (mmap.len() - WINDOW) * i / (N - 1); + h.update(&mmap[pos..pos + WINDOW]); + } + h.digest() + } + + // ── Phase 1: stat ───────────────────────────────────────────────── + let sized: Vec<(i64, String, u64)> = records + .par_iter() + .filter_map(|r| { + let size = std::fs::metadata(&r.path).ok()?.len(); + if size == 0 { return None; } + Some((r.id, r.path.clone(), size)) + }) + .collect(); + + let mut by_size: HashMap> = HashMap::new(); + for (i, (_, _, size)) in sized.iter().enumerate() { + by_size.entry(*size).or_default().push(i); + } + let candidates: Vec = by_size + .into_values() + .filter(|g| g.len() > 1) + .flatten() + .collect(); + + if candidates.is_empty() { + return vec![]; + } + + // ── Phase 2: sample hash ────────────────────────────────────────── + let c_total = candidates.len(); + let _ = app_hash.emit("duplicate_scan_progress", (0usize, c_total)); + let counter = AtomicUsize::new(0); + + candidates + .par_iter() + .filter_map(|&idx| { + let (id, path, size) = &sized[idx]; + let file = std::fs::File::open(path).ok()?; + // SAFETY: read-only; no external truncation expected during scan. + let mmap = unsafe { Mmap::map(&file).ok()? }; + let hash = sample(&mmap); + let done = counter.fetch_add(1, Ordering::Relaxed) + 1; + if done % 100 == 0 || done == c_total { + let _ = app_hash.emit("duplicate_scan_progress", (done, c_total)); + } + Some((hash, *size, *id)) + }) + .collect() + }) + .await + .map_err(|e| e.to_string())?; + + let mut size_hash_map: std::collections::HashMap<(u64, u64), Vec> = std::collections::HashMap::new(); + for (hash, file_size, id) in pairs { + size_hash_map.entry((hash, file_size)).or_default().push(id); + } + + // Resolve image records for each duplicate group + let conn = db.get().map_err(|e| e.to_string())?; + let mut groups: Vec = size_hash_map + .into_iter() + .filter(|(_, ids)| ids.len() > 1) + .filter_map(|((hash, file_size), ids)| { + let images = db::get_images_by_ids(&conn, &ids).ok()?; + Some(DuplicateGroup { + file_hash: format!("{:016x}", hash), + file_size, + images, + }) + }) + .collect(); + + // Largest duplicates first — wastes the most space + groups.sort_by(|a, b| b.file_size.cmp(&a.file_size)); + + // Persist results so they survive restart — best-effort, ignore errors. + let folder_scope = match folder_id { + Some(id) => format!("folder:{}", id), + None => "all".to_string(), + }; + if let Ok(json) = serde_json::to_string(&groups) { + let _ = db::set_duplicate_scan_cache(&conn, &folder_scope, &json); + } + + Ok(groups) +} + +#[tauri::command] +pub async fn load_duplicate_scan_cache( + db: State<'_, DbState>, + folder_id: Option, +) -> Result, String> { + let folder_scope = match folder_id { + Some(id) => format!("folder:{}", id), + None => "all".to_string(), + }; + let conn = db.get().map_err(|e| e.to_string())?; + match db::get_duplicate_scan_cache(&conn, &folder_scope).map_err(|e| e.to_string())? { + Some((json, scanned_at)) => { + let groups: Vec = + serde_json::from_str(&json).map_err(|e| e.to_string())?; + Ok(Some(DuplicateScanCache { groups, scanned_at })) + } + None => Ok(None), + } +} + +#[tauri::command] +pub async fn delete_images_from_disk( + db: State<'_, DbState>, + params: DeleteImagesFromDiskParams, +) -> Result { + if params.image_ids.is_empty() { + return Ok(0); + } + let conn = db.get().map_err(|e| e.to_string())?; + // Collect paths before deleting DB rows + let records = db::get_all_image_paths(&conn, None).map_err(|e| e.to_string())?; + let id_set: std::collections::HashSet = params.image_ids.iter().copied().collect(); + let paths: Vec = records + .into_iter() + .filter(|r| id_set.contains(&r.id)) + .map(|r| r.path) + .collect(); + + db::delete_images_by_ids(&conn, ¶ms.image_ids).map_err(|e| e.to_string())?; + + let mut deleted = 0usize; + for path in &paths { + if std::fs::remove_file(path).is_ok() { + deleted += 1; + } + } + Ok(deleted) +} + +#[tauri::command] +pub async fn get_images_by_ids( + db: State<'_, DbState>, + params: GetImagesByIdsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_images_by_ids(&conn, ¶ms.image_ids).map_err(|e| e.to_string()) +} + // ── k-means with cosine similarity (all vectors assumed to be unit-normalized) ── fn dot(a: &[f32], b: &[f32]) -> f32 { @@ -811,15 +1152,22 @@ pub struct SetTaggerThresholdParams { pub threshold: f32, } +#[derive(Deserialize)] +pub struct SetTaggerBatchSizeParams { + pub batch_size: usize, +} + #[derive(Deserialize)] pub struct QueueTaggingJobsParams { pub folder_id: Option, + pub folder_ids: Option>, pub image_id: Option, } #[derive(Deserialize)] pub struct ClearTaggingJobsParams { pub folder_id: Option, + pub folder_ids: Option>, } #[derive(Deserialize)] @@ -883,6 +1231,21 @@ pub async fn set_tagger_threshold( tagger::set_tagger_threshold(&app_dir, params.threshold).map_err(|e| e.to_string()) } +#[tauri::command] +pub async fn get_tagger_batch_size(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_batch_size(&app_dir)) +} + +#[tauri::command] +pub async fn set_tagger_batch_size( + app: AppHandle, + params: SetTaggerBatchSizeParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tagger::set_tagger_batch_size(&app_dir, params.batch_size).map_err(|e| e.to_string()) +} + #[tauri::command] pub async fn prepare_tagger_model(app: AppHandle) -> Result { let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; @@ -913,19 +1276,28 @@ pub async fn queue_tagging_jobs( params: QueueTaggingJobsParams, ) -> Result { let conn = db.get().map_err(|e| e.to_string())?; - let (total, folder_ids) = match (params.folder_id, params.image_id) { - (_, Some(image_id)) => { + let requested_folder_ids = params.folder_ids.unwrap_or_default(); + let (total, folder_ids) = match (params.folder_id, params.image_id, requested_folder_ids.is_empty()) { + (_, Some(image_id), _) => { db::enqueue_tagging_job(&conn, image_id).map_err(|e| e.to_string())?; // Look up just this image's folder_id rather than fetching all folders let image = db::get_image_by_id(&conn, image_id).map_err(|e| e.to_string())?; (1usize, vec![image.folder_id]) } - (Some(folder_id), None) => { + (Some(folder_id), None, _) => { let n = db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id) .map_err(|e| e.to_string())?; (n, vec![folder_id]) } - (None, None) => { + (None, None, false) => { + let mut total = 0usize; + for &folder_id in &requested_folder_ids { + total += db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } + (total, requested_folder_ids) + } + (None, None, true) => { let folders = db::get_folders(&conn).map_err(|e| e.to_string())?; let folder_ids: Vec = folders.iter().map(|f| f.id).collect(); let mut total = 0usize; @@ -948,14 +1320,27 @@ pub async fn clear_tagging_jobs( params: ClearTaggingJobsParams, ) -> Result { let conn = db.get().map_err(|e| e.to_string())?; - let n = db::clear_tagging_jobs(&conn, params.folder_id).map_err(|e| e.to_string())?; - let folder_ids: Vec = match params.folder_id { - Some(id) => vec![id], - None => db::get_folders(&conn) + let requested_folder_ids = params.folder_ids.unwrap_or_default(); + let (n, folder_ids): (usize, Vec) = match (params.folder_id, requested_folder_ids.is_empty()) { + (Some(id), _) => ( + db::clear_tagging_jobs(&conn, Some(id)).map_err(|e| e.to_string())?, + vec![id], + ), + (None, false) => { + let mut total = 0usize; + for &folder_id in &requested_folder_ids { + total += db::clear_tagging_jobs(&conn, Some(folder_id)).map_err(|e| e.to_string())?; + } + (total, requested_folder_ids) + } + (None, true) => ( + db::clear_tagging_jobs(&conn, None).map_err(|e| e.to_string())?, + db::get_folders(&conn) .map_err(|e| e.to_string())? .into_iter() .map(|f| f.id) .collect(), + ), }; drop(conn); indexer::emit_folder_job_progress(&app, db.inner(), &folder_ids, true); diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 81bd162..a11ca3b 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -122,6 +122,14 @@ pub struct ImageTag { pub created_at: String, } +#[derive(Debug, Clone, Serialize)] +pub struct ExploreTagEntry { + pub tag: String, + pub count: i64, + pub representative_image_id: i64, + pub thumbnail_path: Option, +} + #[derive(Debug, Clone)] pub struct IndexedMediaEntry { pub id: i64, @@ -247,6 +255,12 @@ pub fn migrate(conn: &Connection) -> Result<()> { created_at TEXT NOT NULL DEFAULT (datetime('now')) ); + CREATE TABLE IF NOT EXISTS duplicate_scan_cache ( + folder_scope TEXT PRIMARY KEY, + scanned_at INTEGER NOT NULL, + groups_json TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_images_folder_id ON images(folder_id); CREATE INDEX IF NOT EXISTS idx_images_modified_at ON images(modified_at); CREATE INDEX IF NOT EXISTS idx_embedding_jobs_status ON embedding_jobs(status); @@ -1302,6 +1316,7 @@ pub fn get_images( search: Option<&str>, media_kind: Option<&str>, favorites_only: bool, + rating_min: i64, embedding_failed_only: bool, sort: &str, offset: i64, @@ -1314,6 +1329,8 @@ pub fn get_images( "date_desc" => "modified_at DESC NULLS LAST", "size_asc" => "file_size ASC", "size_desc" => "file_size DESC", + "rating_asc" => "rating ASC, modified_at DESC NULLS LAST", + "rating_desc" => "rating DESC, modified_at DESC NULLS LAST", "duration_asc" => "duration_ms ASC NULLS LAST", "duration_desc" => "duration_ms DESC NULLS LAST", _ => "modified_at DESC NULLS LAST", @@ -1333,9 +1350,10 @@ pub fn get_images( AND (?2 IS NULL OR filename LIKE ?2) AND (?3 IS NULL OR media_kind = ?3) AND (?4 = 0 OR favorite = 1) - AND (?5 = 0 OR embedding_status = 'failed') + AND rating >= ?5 + AND (?6 = 0 OR embedding_status = 'failed') ORDER BY {} - LIMIT ?6 OFFSET ?7", + LIMIT ?7 OFFSET ?8", order ); let mut stmt = conn.prepare(&sql)?; @@ -1345,6 +1363,7 @@ pub fn get_images( search_pattern, media_kind, favorites_flag, + rating_min, embedding_failed_flag, limit, offset @@ -1360,6 +1379,7 @@ pub fn count_images( search: Option<&str>, media_kind: Option<&str>, favorites_only: bool, + rating_min: i64, embedding_failed_only: bool, ) -> Result { let search_pattern = search.map(|value| format!("%{}%", value)); @@ -1372,12 +1392,14 @@ pub fn count_images( AND (?2 IS NULL OR filename LIKE ?2) AND (?3 IS NULL OR media_kind = ?3) AND (?4 = 0 OR favorite = 1) - AND (?5 = 0 OR embedding_status = 'failed')", + AND rating >= ?5 + AND (?6 = 0 OR embedding_status = 'failed')", params![ folder_id, search_pattern, media_kind, favorites_flag, + rating_min, embedding_failed_flag ], |row| row.get(0), @@ -1386,6 +1408,141 @@ pub fn count_images( Ok(count) } +pub fn search_images_by_tag( + conn: &Connection, + query: &str, + folder_id: Option, + media_kind: Option<&str>, + favorites_only: bool, + rating_min: i64, + limit: usize, +) -> Result> { + let normalized_query = query.trim().to_ascii_lowercase(); + if normalized_query.is_empty() { + return Ok(Vec::new()); + } + let favorites_flag = i64::from(favorites_only); + + let mut stmt = conn.prepare( + "SELECT DISTINCT i.id + FROM images i + JOIN image_tags t ON t.image_id = i.id + WHERE (?1 IS NULL OR i.folder_id = ?1) + AND (?2 IS NULL OR i.media_kind = ?2) + AND (?3 = 0 OR i.favorite = 1) + AND i.rating >= ?4 + AND LOWER(TRIM(t.tag)) = ?5 + ORDER BY i.rating DESC, i.modified_at DESC NULLS LAST, i.filename ASC + LIMIT ?6", + )?; + + let image_ids = stmt + .query_map( + params![ + folder_id, + media_kind, + favorites_flag, + rating_min, + normalized_query, + limit as i64 + ], + |row| row.get::<_, i64>(0), + )? + .collect::>>()?; + + get_images_by_ids(conn, &image_ids) +} + +pub fn search_tags_autocomplete( + conn: &Connection, + query: &str, + folder_id: Option, + limit: usize, +) -> Result> { + let pattern = format!("%{}%", query.to_lowercase()); + let mut stmt = conn.prepare( + "SELECT t.tag, COUNT(DISTINCT t.image_id) AS tag_count, MIN(t.image_id) AS representative_image_id + FROM image_tags t + JOIN images i ON i.id = t.image_id + WHERE (?1 IS NULL OR i.folder_id = ?1) + AND LOWER(t.tag) LIKE ?2 + GROUP BY t.tag + ORDER BY tag_count DESC, t.tag ASC + LIMIT ?3", + )?; + let rows = stmt + .query_map(params![folder_id, pattern, limit as i64], |row| { + Ok(ExploreTagEntry { + tag: row.get(0)?, + count: row.get(1)?, + representative_image_id: row.get(2)?, + thumbnail_path: None, // skip per-suggestion thumbnail for speed + }) + })? + .collect::>>()?; + Ok(rows) +} + +pub struct ImagePathRecord { + pub id: i64, + pub path: String, + pub thumbnail_path: Option, +} + +pub fn get_all_image_paths( + conn: &Connection, + folder_id: Option, +) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, path, thumbnail_path FROM images WHERE (?1 IS NULL OR folder_id = ?1) ORDER BY id", + )?; + let rows = stmt + .query_map(params![folder_id], |row| { + Ok(ImagePathRecord { + id: row.get(0)?, + path: row.get(1)?, + thumbnail_path: row.get(2)?, + }) + })? + .collect::>>()?; + Ok(rows) +} + +pub fn get_explore_tags( + conn: &Connection, + folder_id: Option, + limit: usize, +) -> Result> { + let mut stmt = conn.prepare( + "SELECT t.tag, COUNT(DISTINCT t.image_id) AS tag_count, MIN(t.image_id) AS representative_image_id + FROM image_tags t + JOIN images i ON i.id = t.image_id + WHERE (?1 IS NULL OR i.folder_id = ?1) + GROUP BY t.tag + HAVING COUNT(DISTINCT t.image_id) >= 2 + ORDER BY tag_count DESC, t.tag ASC + LIMIT ?2", + )?; + + let rows = stmt + .query_map(params![folder_id, limit as i64], |row| { + let representative_image_id = row.get::<_, i64>(2)?; + let thumbnail_path = get_image_by_id(conn, representative_image_id) + .ok() + .and_then(|image| image.thumbnail_path); + + Ok(ExploreTagEntry { + tag: row.get(0)?, + count: row.get(1)?, + representative_image_id, + thumbnail_path, + }) + })? + .collect::>>()?; + + Ok(rows) +} + pub fn get_failed_embedding_images( conn: &Connection, folder_id: i64, @@ -1871,6 +2028,43 @@ pub fn set_tag_cloud_cache( Ok(()) } +/// Returns (groups_json, scanned_at_unix) for the given folder scope, if present. +pub fn get_duplicate_scan_cache( + conn: &Connection, + folder_scope: &str, +) -> Result> { + match conn.query_row( + "SELECT groups_json, scanned_at FROM duplicate_scan_cache WHERE folder_scope = ?1", + params![folder_scope], + |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + ) { + Ok(row) => Ok(Some(row)), + Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(e.into()), + } +} + +/// Upserts the duplicate scan cache for the given scope. +pub fn set_duplicate_scan_cache( + conn: &Connection, + folder_scope: &str, + groups_json: &str, +) -> Result<()> { + let scanned_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + conn.execute( + "INSERT INTO duplicate_scan_cache (folder_scope, scanned_at, groups_json) + VALUES (?1, ?2, ?3) + ON CONFLICT(folder_scope) DO UPDATE SET + scanned_at = excluded.scanned_at, + groups_json = excluded.groups_json", + params![folder_scope, scanned_at, groups_json], + )?; + Ok(()) +} + fn ensure_column(conn: &Connection, table: &str, column: &str, definition: &str) -> Result<()> { let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", table))?; let mut rows = stmt.query([])?; diff --git a/src-tauri/src/hnsw_index.rs b/src-tauri/src/hnsw_index.rs new file mode 100644 index 0000000..b7ef6c7 --- /dev/null +++ b/src-tauri/src/hnsw_index.rs @@ -0,0 +1,129 @@ +use crate::vector; +use anyhow::Result; +use hnsw_rs::prelude::{DistCosine, Hnsw, Neighbour}; +use rusqlite::Connection; +use std::collections::HashMap; +use std::sync::{OnceLock, RwLock}; + +const HNSW_MAX_CONNECTIONS: usize = 24; +const HNSW_EF_CONSTRUCTION: usize = 300; +const HNSW_EF_SEARCH: usize = 96; + +struct CachedHnswIndex { + revision: String, + image_ids_by_external: Vec, + external_by_image_id: HashMap, + hnsw: Hnsw<'static, f32, DistCosine>, +} + +static IMAGE_HNSW_INDEX: OnceLock>> = OnceLock::new(); + +fn cache() -> &'static RwLock> { + IMAGE_HNSW_INDEX.get_or_init(|| RwLock::new(None)) +} + +fn build_index(conn: &Connection) -> Result { + let embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?; + let max_elements = embeddings.len().max(1); + let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).max(1); + let mut hnsw = Hnsw::::new( + HNSW_MAX_CONNECTIONS, + max_elements, + max_layer, + HNSW_EF_CONSTRUCTION, + DistCosine {}, + ); + + let image_ids_by_external = embeddings + .iter() + .map(|(image_id, _)| *image_id) + .collect::>(); + let external_by_image_id = image_ids_by_external + .iter() + .enumerate() + .map(|(external_id, image_id)| (*image_id, external_id)) + .collect::>(); + let data_with_id = embeddings + .iter() + .enumerate() + .map(|(external_id, (_, embedding))| (embedding, external_id)) + .collect::>(); + + hnsw.parallel_insert(&data_with_id); + hnsw.set_searching_mode(true); + + Ok(CachedHnswIndex { + revision: vector::get_embedding_revision(conn)?, + image_ids_by_external, + external_by_image_id, + hnsw, + }) +} + +fn ensure_index(conn: &Connection) -> Result<()> { + let revision = vector::get_embedding_revision(conn)?; + + { + let guard = cache().read().expect("hnsw cache poisoned"); + if guard.as_ref().map(|cached| cached.revision.as_str()) == Some(revision.as_str()) { + return Ok(()); + } + } + + let next = build_index(conn)?; + let mut guard = cache().write().expect("hnsw cache poisoned"); + *guard = Some(next); + Ok(()) +} + +pub fn find_similar_image_matches( + conn: &Connection, + image_id: i64, + folder_id: Option, + threshold: f32, + offset: usize, + limit: usize, +) -> Result> { + ensure_index(conn)?; + + let query_embedding = match vector::get_image_embedding(conn, image_id)? { + Some(embedding) => embedding, + None => return Ok(Vec::new()), + }; + + let guard = cache().read().expect("hnsw cache poisoned"); + let Some(cached) = guard.as_ref() else { + return Ok(Vec::new()); + }; + + let knbn = (offset + limit).max(limit).saturating_add(32); + let neighbours: Vec = if let Some(folder_id) = folder_id { + let mut allowed_ids = vector::get_all_image_embeddings_with_ids(conn, Some(folder_id))? + .into_iter() + .filter_map(|(allowed_image_id, _)| { + cached.external_by_image_id.get(&allowed_image_id).copied() + }) + .collect::>(); + allowed_ids.sort_unstable(); + cached + .hnsw + .search_filter(&query_embedding, knbn, HNSW_EF_SEARCH, Some(&allowed_ids)) + } else { + cached.hnsw.search(&query_embedding, knbn, HNSW_EF_SEARCH) + }; + + let matches = neighbours + .into_iter() + .filter_map(|neighbour| { + let image_id_match = cached.image_ids_by_external.get(neighbour.d_id).copied()?; + if image_id_match == image_id || neighbour.distance > threshold { + return None; + } + Some((image_id_match, neighbour.distance)) + }) + .skip(offset) + .take(limit) + .collect::>(); + + Ok(matches) +} diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index f509d1f..cffc605 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -331,6 +331,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) if !missing_ids.is_empty() { db::delete_images_by_ids(&conn, &missing_ids)?; } + let _ = db::backfill_embedding_jobs(&conn)?; db::update_folder_count(&conn, folder_id)?; } @@ -864,8 +865,6 @@ fn process_caption_batch( Ok(()) } -const TAGGING_BATCH_SIZE: usize = 1; - fn process_tagging_batch( app: &AppHandle, pool: &DbPool, @@ -877,9 +876,10 @@ fn process_tagging_batch( } let paused_folders = paused_folder_ids("tagging"); + let batch_size = crate::tagger::tagger_batch_size(app_data_dir); let jobs = with_db_write_lock(|| { let mut conn = pool.get()?; - db::claim_tagging_jobs(&mut conn, &paused_folders, TAGGING_BATCH_SIZE) + db::claim_tagging_jobs(&mut conn, &paused_folders, batch_size) })?; if jobs.is_empty() { diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 217eabf..39c9922 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2,6 +2,7 @@ mod captioner; mod commands; mod db; mod embedder; +mod hnsw_index; mod indexer; mod media; mod storage; @@ -89,6 +90,7 @@ pub fn run() { commands::debug_similar_images, commands::retry_failed_embeddings, commands::semantic_search_images, + commands::search_images_by_tag, commands::get_caption_model_status, commands::get_caption_acceleration, commands::set_caption_acceleration, @@ -107,6 +109,8 @@ pub fn run() { commands::set_worker_paused, commands::get_worker_states, commands::get_tag_cloud, + commands::get_explore_tags, + commands::get_images_by_ids, commands::get_failed_embedding_images, commands::get_tagger_model_status, commands::get_tagger_acceleration, @@ -114,6 +118,8 @@ pub fn run() { commands::probe_tagger_runtime, commands::get_tagger_threshold, commands::set_tagger_threshold, + commands::get_tagger_batch_size, + commands::set_tagger_batch_size, commands::prepare_tagger_model, commands::delete_tagger_model, commands::queue_tagging_jobs, @@ -121,6 +127,10 @@ pub fn run() { commands::get_image_tags, commands::add_user_tag, commands::remove_tag, + commands::search_tags_autocomplete, + commands::find_duplicates, + commands::load_duplicate_scan_cache, + commands::delete_images_from_disk, ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index d0cd179..defc066 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -14,6 +14,7 @@ pub const WD_TAGGER_MODEL_NAME: &str = "wd-swinv2-tagger-v3"; const TAGGER_ACCELERATION_FILE: &str = "settings/tagger_acceleration.txt"; const TAGGER_THRESHOLD_FILE: &str = "settings/tagger_threshold.txt"; +const TAGGER_BATCH_SIZE_FILE: &str = "settings/tagger_batch_size.txt"; // Files required on disk before the tagger can run. The ONNX runtime DLLs // are shared with the captioner and live in the same `onnxruntime/` directory. @@ -193,6 +194,28 @@ pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result Ok(clamped) } +pub fn tagger_batch_size(app_data_dir: &Path) -> usize { + let path = app_data_dir.join(TAGGER_BATCH_SIZE_FILE); + let Ok(value) = std::fs::read_to_string(path) else { + return 8; + }; + value + .trim() + .parse::() + .unwrap_or(8) + .clamp(1, 100) +} + +pub fn set_tagger_batch_size(app_data_dir: &Path, batch_size: usize) -> Result { + let clamped = batch_size.clamp(1, 100); + let path = app_data_dir.join(TAGGER_BATCH_SIZE_FILE); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, clamped.to_string())?; + Ok(clamped) +} + // --------------------------------------------------------------------------- // Model status / download // --------------------------------------------------------------------------- diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 3b5d942..0fa753c 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -96,15 +96,25 @@ pub fn find_similar_image_ids( Err(error) => return Err(error.into()), }; - let allowed_folder_ids = match folder_id { - Some(folder_id) => Some(image_ids_for_folder(conn, folder_id)?), - None => None, - }; - let search_limit = if allowed_folder_ids.is_some() { - count_image_vectors(conn)?.max(1) as usize - } else { - limit + 1 - }; + if let Some(folder_id) = folder_id { + // Brute-force cosine scan scoped to the folder — avoids the KNN k=4096 limit + // and returns exact nearest neighbours within the folder. + let mut stmt = conn.prepare( + "SELECT v.image_id + FROM image_vec v + JOIN images i ON i.id = v.image_id + WHERE i.folder_id = ?2 + AND v.image_id != ?3 + ORDER BY vec_distance_cosine(v.embedding, vec_f32(?1)) ASC + LIMIT ?4", + )?; + let rows = stmt.query_map((&embedding, folder_id, image_id, limit as i64), |row| { + row.get::<_, i64>(0) + })?; + return Ok(rows.collect::>>()?); + } + + // Global KNN search (no folder filter) — use the ANN index. let mut stmt = conn.prepare( "SELECT image_id FROM image_vec @@ -112,20 +122,12 @@ pub fn find_similar_image_ids( AND k = ?2", )?; let rows = stmt - .query_map((&embedding, search_limit as i64), |row| { - row.get::<_, i64>(0) - })? + .query_map((&embedding, (limit + 1) as i64), |row| row.get::<_, i64>(0))? .collect::>>()?; let mut ids = Vec::new(); for row in rows { if row != image_id { - if allowed_folder_ids - .as_ref() - .is_some_and(|folder_ids| !folder_ids.contains(&row)) - { - continue; - } ids.push(row); } if ids.len() >= limit { @@ -135,15 +137,102 @@ pub fn find_similar_image_ids( Ok(ids) } -fn image_ids_for_folder( - conn: &Connection, - folder_id: i64, -) -> Result> { - let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?; - let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?; - Ok(rows.collect::>>()?) +// pub fn find_similar_image_matches( +// conn: &Connection, +// image_id: i64, +// folder_id: Option, +// threshold: f32, +// offset: usize, +// limit: usize, +// ) -> Result> { +// let embedding: Vec = match conn.query_row( +// "SELECT embedding FROM image_vec WHERE image_id = ?1", +// [image_id], +// |row| row.get(0), +// ) { +// Ok(embedding) => embedding, +// Err(SqliteError::QueryReturnedNoRows) => return Ok(Vec::new()), +// Err(error) => return Err(error.into()), +// }; + +// let query = match folder_id { +// Some(_) => { +// "SELECT v.image_id, vec_distance_cosine(v.embedding, vec_f32(?1)) AS distance +// FROM image_vec v +// JOIN images i ON i.id = v.image_id +// WHERE i.folder_id = ?2 +// AND v.image_id != ?3 +// AND vec_distance_cosine(v.embedding, vec_f32(?1)) <= ?4 +// ORDER BY distance ASC +// LIMIT ?5 OFFSET ?6" +// } +// None => { +// "SELECT v.image_id, vec_distance_cosine(v.embedding, vec_f32(?1)) AS distance +// FROM image_vec v +// WHERE v.image_id != ?2 +// AND vec_distance_cosine(v.embedding, vec_f32(?1)) <= ?3 +// ORDER BY distance ASC +// LIMIT ?4 OFFSET ?5" +// } +// }; + +// let mut stmt = conn.prepare(query)?; +// match folder_id { +// Some(folder_id) => Ok(stmt +// .query_map( +// ( +// &embedding, +// folder_id, +// image_id, +// threshold, +// limit as i64, +// offset as i64, +// ), +// |row| Ok((row.get::<_, i64>(0)?, row.get::<_, f32>(1)?)), +// )? +// .collect::>>()?), +// None => Ok(stmt +// .query_map( +// (&embedding, image_id, threshold, limit as i64, offset as i64), +// |row| Ok((row.get::<_, i64>(0)?, row.get::<_, f32>(1)?)), +// )? +// .collect::>>()?), +// } +// } + +pub fn get_image_embedding(conn: &Connection, image_id: i64) -> Result>> { + let embedding: Result, rusqlite::Error> = conn.query_row( + "SELECT embedding FROM image_vec WHERE image_id = ?1", + [image_id], + |row| row.get(0), + ); + + match embedding { + Ok(bytes) => Ok(Some(unpack_f32(&bytes))), + Err(SqliteError::QueryReturnedNoRows) => Ok(None), + Err(error) => Err(error.into()), + } } +pub fn get_embedding_revision(conn: &Connection) -> Result { + let count: i64 = conn.query_row("SELECT COUNT(*) FROM image_vec", [], |row| row.get(0))?; + let max_updated_at: Option = conn.query_row( + "SELECT MAX(embedding_updated_at) FROM images WHERE embedding_status = 'ready'", + [], + |row| row.get(0), + )?; + Ok(format!("{}:{}", count, max_updated_at.unwrap_or_default())) +} + +// fn image_ids_for_folder( +// conn: &Connection, +// folder_id: i64, +// ) -> Result> { +// let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?; +// let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?; +// Ok(rows.collect::>>()?) +// } + /// Returns all stored image embeddings with their image IDs, optionally filtered to one folder. /// Each entry is `(image_id, normalized_f32_embedding)`. pub fn get_all_image_embeddings_with_ids( diff --git a/src-tauri/tauri.conf.json b/src-tauri/tauri.conf.json index ea720ac..51aa343 100644 --- a/src-tauri/tauri.conf.json +++ b/src-tauri/tauri.conf.json @@ -4,9 +4,9 @@ "version": "0.1.0", "identifier": "wtf.jezz.phokus", "build": { - "beforeDevCommand": "pnpm dev", + "beforeDevCommand": "pnpm dev:vite", "devUrl": "http://localhost:1420", - "beforeBuildCommand": "pnpm build", + "beforeBuildCommand": "pnpm build:vite", "frontendDist": "../dist" }, "app": { @@ -25,7 +25,9 @@ "csp": null, "assetProtocol": { "enable": true, - "scope": ["**"] + "scope": [ + "**" + ] } } }, @@ -40,4 +42,4 @@ "icons/icon.ico" ] } -} +} \ No newline at end of file diff --git a/src/App.tsx b/src/App.tsx index 4b6039b..2486ba9 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -6,6 +6,7 @@ import { Toolbar } from "./components/Toolbar"; import { Gallery } from "./components/Gallery"; import { Lightbox } from "./components/Lightbox"; import { TagCloud } from "./components/TagCloud"; +import { DuplicateFinder } from "./components/DuplicateFinder"; import { TitleBar } from "./components/TitleBar"; import { SettingsModal } from "./components/SettingsModal"; @@ -14,6 +15,7 @@ export default function App() { const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress); const loadImages = useGalleryStore((state) => state.loadImages); const loadCaptionModelStatus = useGalleryStore((state) => state.loadCaptionModelStatus); + const loadDuplicateScanCache = useGalleryStore((state) => state.loadDuplicateScanCache); const subscribeToProgress = useGalleryStore((state) => state.subscribeToProgress); const activeView = useGalleryStore((state) => state.activeView); @@ -21,6 +23,7 @@ export default function App() { loadFolders().then(() => { void loadBackgroundJobProgress(); void loadCaptionModelStatus(); + void loadDuplicateScanCache(); return loadImages(true); }); let unlisten: (() => void) | undefined; @@ -46,6 +49,11 @@ export default function App() { + ) : activeView === "duplicates" ? ( + <> + + + ) : ( <> diff --git a/src/components/BackgroundTasks.tsx b/src/components/BackgroundTasks.tsx index bc007a5..4515215 100644 --- a/src/components/BackgroundTasks.tsx +++ b/src/components/BackgroundTasks.tsx @@ -58,6 +58,8 @@ export function BackgroundTasks() { const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings); const clearTaggingJobs = useGalleryStore((state) => state.clearTaggingJobs); + const duplicateScanning = useGalleryStore((state) => state.duplicateScanning); + const duplicateScanProgress = useGalleryStore((state) => state.duplicateScanProgress); const [expanded, setExpanded] = useState(false); const [dismissed, setDismissed] = useState>({}); const [paused, setPaused] = useState>>({}); @@ -127,6 +129,7 @@ export function BackgroundTasks() { }; const dismissTask = (id: number, snapshot: string) => { + if (id < 0) return; // system tasks (duplicate scan) cannot be dismissed void clearTaggingJobs(id); setDismissed((prev) => ({ ...prev, [id]: snapshot })); setExpanded(false); @@ -244,10 +247,35 @@ export function BackgroundTasks() { .filter((t) => dismissed[t.id] !== t.snapshot); }, [folders, indexingProgress, mediaJobProgress, dismissed]); - if (tasks.length === 0) return null; + // Synthetic task for duplicate scanning — negative id so dismiss/retry are suppressed + const duplicateScanTask: Task | null = duplicateScanning ? { + id: -1, + name: "Duplicate Scan", + stages: [{ + label: "Hashing", + detail: duplicateScanProgress + ? `${duplicateScanProgress.scanned.toLocaleString()} / ${duplicateScanProgress.total.toLocaleString()}` + : "Starting…", + progress: duplicateScanProgress && duplicateScanProgress.total > 0 + ? (duplicateScanProgress.scanned / duplicateScanProgress.total) * 100 + : null, + failed: false, + }], + hasFailedEmbeddings: false, + hasFailedTagging: false, + pendingMediaWork: 1, + embeddingProcessed: 0, + embeddingTotal: 0, + currentFile: null, + snapshot: "", + } : null; - const primary = tasks[0]; - const extraCount = tasks.length - 1; + const allTasks = duplicateScanTask ? [duplicateScanTask, ...tasks] : tasks; + + if (allTasks.length === 0) return null; + + const primary = allTasks[0]; + const extraCount = allTasks.length - 1; const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedTagging) && t.pendingMediaWork === 0); // Best progress bar value: use embedding progress if available (most informative), @@ -348,8 +376,8 @@ export function BackgroundTasks() { )} - {/* Expand chevron (only when multiple folders) */} - {tasks.length > 1 && ( + {/* Expand chevron (only when multiple tasks) */} + {allTasks.length > 1 && ( )} - {/* Dismiss */} - + {/* Dismiss — hidden for system tasks like duplicate scan */} + {primary.id >= 0 && ( + + )}
{/* Expanded panel — one row per folder */} {expanded && (
- {tasks.map((task) => { + {allTasks.map((task) => { const taskEmbeddingStage = task.stages.find((s) => s.label === "Embeddings"); const taskTaggingStage = task.stages.find((s) => s.label === "Tags"); const taskScanningStage = task.stages.find((s) => s.label === "Scanning"); @@ -453,15 +483,17 @@ export function BackgroundTasks() { )} - + {task.id >= 0 && ( + + )}
{task.currentFile && ( diff --git a/src/components/DuplicateFinder.tsx b/src/components/DuplicateFinder.tsx new file mode 100644 index 0000000..d01f05d --- /dev/null +++ b/src/components/DuplicateFinder.tsx @@ -0,0 +1,274 @@ +import { useState } from "react"; +import { convertFileSrc } from "@tauri-apps/api/core"; +import { DuplicateGroup, useGalleryStore } from "../store"; + +function formatBytes(bytes: number): string { + if (bytes >= 1_073_741_824) return `${(bytes / 1_073_741_824).toFixed(1)} GB`; + if (bytes >= 1_048_576) return `${(bytes / 1_048_576).toFixed(1)} MB`; + if (bytes >= 1024) return `${(bytes / 1024).toFixed(0)} KB`; + return `${bytes} B`; +} + +function DuplicateGroupCard({ group }: { group: DuplicateGroup }) { + const selectedIds = useGalleryStore((state) => state.duplicateSelectedIds); + const toggleDuplicateSelected = useGalleryStore((state) => state.toggleDuplicateSelected); + const selectAllDuplicates = useGalleryStore((state) => state.selectAllDuplicates); + const groupSelectedCount = group.images.filter((img) => selectedIds.has(img.id)).length; + const noneSelected = groupSelectedCount === 0; + + // "Keep all but the first" — a common quick action + const handleKeepFirst = () => { + const toDelete = group.images.slice(1).map((img) => img.id); + // Clear any selection for this group first, then add the ones to delete + for (const img of group.images) { + if (selectedIds.has(img.id)) toggleDuplicateSelected(img.id); + } + selectAllDuplicates(toDelete); + }; + + return ( +
+ {/* Group header */} +
+
+ + {group.images.length} copies + + {formatBytes(group.file_size)} each + + {formatBytes(group.file_size * (group.images.length - 1))} wasted + +
+
+ {noneSelected ? ( + + ) : ( + + )} +
+
+ + {/* Image grid */} +
+ {group.images.map((image) => { + const isSelected = selectedIds.has(image.id); + const src = image.thumbnail_path ? convertFileSrc(image.thumbnail_path) : null; + return ( + + ); + })} +
+
+ ); +} + +function formatRelativeTime(unixSecs: number): string { + const diff = Math.floor(Date.now() / 1000) - unixSecs; + if (diff < 60) return "just now"; + if (diff < 3600) return `${Math.floor(diff / 60)}m ago`; + if (diff < 86400) return `${Math.floor(diff / 3600)}h ago`; + return `${Math.floor(diff / 86400)}d ago`; +} + +export function DuplicateFinder() { + const duplicateGroups = useGalleryStore((state) => state.duplicateGroups); + const duplicateScanning = useGalleryStore((state) => state.duplicateScanning); + const duplicateScanProgress = useGalleryStore((state) => state.duplicateScanProgress); + const duplicateSelectedIds = useGalleryStore((state) => state.duplicateSelectedIds); + const duplicateLastScanned = useGalleryStore((state) => state.duplicateLastScanned); + const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); + const scanDuplicates = useGalleryStore((state) => state.scanDuplicates); + const clearDuplicateSelection = useGalleryStore((state) => state.clearDuplicateSelection); + const selectKeepFirstAllGroups = useGalleryStore((state) => state.selectKeepFirstAllGroups); + const deleteSelectedDuplicates = useGalleryStore((state) => state.deleteSelectedDuplicates); + + const [deleting, setDeleting] = useState(false); + const [deleteResult, setDeleteResult] = useState(null); + + const selectedCount = duplicateSelectedIds.size; + const hasResults = duplicateGroups.length > 0; + const hasScanned = hasResults || duplicateLastScanned !== null || (!duplicateScanning && duplicateScanProgress !== null); + const totalWasted = duplicateGroups.reduce( + (sum, g) => sum + g.file_size * (g.images.length - 1), + 0, + ); + const totalDuplicateImages = duplicateGroups.reduce((sum, g) => sum + g.images.length - 1, 0); + + const handleDelete = async () => { + setDeleting(true); + setDeleteResult(null); + try { + const deleted = await deleteSelectedDuplicates(); + setDeleteResult(`Deleted ${deleted} file${deleted === 1 ? "" : "s"}.`); + } catch (e) { + setDeleteResult(String(e)); + } finally { + setDeleting(false); + } + }; + + const progressPercent = + duplicateScanProgress && duplicateScanProgress.total > 0 + ? Math.round((duplicateScanProgress.scanned / duplicateScanProgress.total) * 100) + : 0; + + return ( +
+ {/* Header */} +
+
+
+

Duplicate Finder

+

+ {duplicateScanning + ? duplicateScanProgress + ? `Scanning… ${duplicateScanProgress.scanned.toLocaleString()} / ${duplicateScanProgress.total.toLocaleString()}` + : "Starting scan…" + : hasResults + ? `${duplicateGroups.length} group${duplicateGroups.length === 1 ? "" : "s"} · ${formatBytes(totalWasted)} reclaimable` + : duplicateLastScanned !== null + ? "No duplicates found" + : "Scan your library for identical files"} +

+ {!duplicateScanning && duplicateLastScanned !== null && ( +

+ Last scanned {formatRelativeTime(duplicateLastScanned)} +

+ )} +
+
+ {/* Batch select — only shown when there are groups and nothing is selected yet */} + {hasResults && selectedCount === 0 && !deleting && ( + + )} + {selectedCount > 0 ? ( + <> + {selectedCount} marked for deletion + + + + ) : null} + +
+
+ + {/* Progress bar */} + {duplicateScanning && duplicateScanProgress ? ( +
+
+
+ ) : null} + + {deleteResult ? ( +

{deleteResult}

+ ) : null} +
+ + {/* Body */} + {duplicateScanning && !hasResults ? ( +
+
+ Hashing files… +
+ ) : !hasScanned ? ( +
+
+ + + +

+ Finds files with identical content regardless of filename or location. + Click Scan for duplicates to begin. +

+

+ Large libraries may take a minute — files are hashed from disk. +

+
+
+ ) : duplicateGroups.length === 0 ? ( +
+

No duplicate files found.

+
+ ) : ( +
+
+ {duplicateGroups.map((group) => ( + + ))} +
+
+ )} +
+ ); +} + diff --git a/src/components/Gallery.tsx b/src/components/Gallery.tsx index ee0e6ea..0f820d4 100644 --- a/src/components/Gallery.tsx +++ b/src/components/Gallery.tsx @@ -1,6 +1,6 @@ import { useEffect, useRef, useCallback, useState } from "react"; import { convertFileSrc } from "@tauri-apps/api/core"; -import { ImageRecord, tileSizeForZoom, useGalleryStore } from "../store"; +import { ImageRecord, parseSearchValue, tileSizeForZoom, useGalleryStore } from "../store"; const GAP = 6; @@ -30,6 +30,7 @@ function ContextMenu({ const openImage = useGalleryStore((state) => state.openImage); const updateImageDetails = useGalleryStore((state) => state.updateImageDetails); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); + const similarScope = useGalleryStore((state) => state.similarScope); const canFindSimilar = image.embedding_status === "ready"; return ( @@ -59,7 +60,7 @@ function ContextMenu({ }`} onClick={async () => { if (!canFindSimilar) return; - await loadSimilarImages(image.id); + await loadSimilarImages(image.id, similarScope === "current_folder" ? image.folder_id : null, true, image.folder_id); onClose(); }} disabled={!canFindSimilar} @@ -115,6 +116,7 @@ function ImageTile({ const [loaded, setLoaded] = useState(false); const [errored, setErrored] = useState(false); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); + const similarScope = useGalleryStore((state) => state.similarScope); const canFindSimilar = image.embedding_status === "ready"; const src = image.thumbnail_path @@ -181,6 +183,15 @@ function ImageTile({
)} + {image.rating > 0 && ( +
+ {Array.from({ length: image.rating }, (_, index) => ( + + + + ))} +
+ )} {image.media_kind === "video" && image.duration_ms && (
{formatDuration(image.duration_ms)} @@ -230,7 +241,7 @@ function ImageTile({ onClick={(event) => { event.stopPropagation(); if (!canFindSimilar) return; - void loadSimilarImages(image.id); + void loadSimilarImages(image.id, similarScope === "current_folder" ? image.folder_id : null, true, image.folder_id); }} disabled={!canFindSimilar} > @@ -250,11 +261,11 @@ export function Gallery() { const loadingImages = useGalleryStore((state) => state.loadingImages); const zoomPreset = useGalleryStore((state) => state.zoomPreset); const search = useGalleryStore((state) => state.search); - const searchMode = useGalleryStore((state) => state.searchMode); const collectionTitle = useGalleryStore((state) => state.collectionTitle); const imageLoadError = useGalleryStore((state) => state.imageLoadError); const galleryScrollResetKey = useGalleryStore((state) => state.galleryScrollResetKey); const isSimilarResults = collectionTitle === "Similar Images"; + const parsedSearch = parseSearchValue(search); const parentRef = useRef(null); const [contextMenu, setContextMenu] = useState<{ x: number; y: number; image: ImageRecord } | null>(null); @@ -262,6 +273,7 @@ export function Gallery() { const handleScroll = useCallback(() => { const element = parentRef.current; if (!element) return; + if (element.scrollTop < 24) return; const nearBottom = element.scrollTop + element.clientHeight >= element.scrollHeight - 600; if (nearBottom && !loadingImages && images.length < totalImages) { void loadMoreImages(); @@ -295,85 +307,87 @@ export function Gallery() { }; }, []); - if (images.length === 0 && loadingImages) { - return ( -
-
-
-

- {isSimilarResults - ? "Finding similar images" - : searchMode === "semantic" && search.trim().length > 0 - ? `Searching for matches to "${search}"` - : "Loading media"} -

-

- {isSimilarResults - ? "Comparing visual embeddings" - : searchMode === "semantic" && search.trim().length > 0 - ? "Semantic search can take a little longer than filename search" - : "Fetching results"} -

-
-
- ); - } - - if (images.length === 0 && !loadingImages) { - return ( -
-
- - - -

- {imageLoadError - ? "Could not load results" - : isSimilarResults - ? "No similar images found" - : searchMode === "semantic" && search.trim().length > 0 - ? "No semantic matches found" - : "No media found"} -

-

- {imageLoadError - ? imageLoadError - : isSimilarResults - ? "This item may be visually isolated, or more embeddings may need to finish processing" - : searchMode === "semantic" && search.trim().length > 0 - ? "Try a broader phrase, or wait for more embeddings to finish processing" - : "Try adjusting your filters or add a new folder"} -

-
-
- ); - } - return (
-
- {images.map((image) => ( - openImage(image)} - onContextMenu={(event) => { - event.preventDefault(); - setContextMenu({ x: event.clientX, y: event.clientY, image }); - }} - /> - ))} -
+ {images.length === 0 && loadingImages ? ( +
+
+
+

+ {isSimilarResults + ? "Finding similar images" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? `Searching for matches to "${parsedSearch.query}"` + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? `Searching tags for "${parsedSearch.query}"` + : "Loading media"} +

+

+ {isSimilarResults + ? "Comparing visual embeddings" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? "Semantic search can take a little longer than filename search" + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? "Matching against AI and user tags" + : "Fetching results"} +

+
+
+ ) : images.length === 0 && !loadingImages ? ( +
+
+ + + +

+ {imageLoadError + ? "Could not load results" + : isSimilarResults + ? "No similar images found" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? "No semantic matches found" + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? "No tag matches found" + : "No media found"} +

+

+ {imageLoadError + ? imageLoadError + : isSimilarResults + ? "This item may be visually isolated, or more embeddings may need to finish processing" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? "Try a broader phrase, or wait for more embeddings to finish processing" + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? "Try a shorter tag, or wait for more tagging jobs to finish" + : "Try adjusting your filters or add a new folder"} +

+
+
+ ) : ( +
+ {images.map((image) => ( + openImage(image)} + onContextMenu={(event) => { + event.preventDefault(); + setContextMenu({ x: event.clientX, y: event.clientY, image }); + }} + /> + ))} +
+ )} - {loadingImages ? ( + {images.length > 0 && loadingImages ? (
diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index 2af2a0b..b357d0f 100644 --- a/src/components/Lightbox.tsx +++ b/src/components/Lightbox.tsx @@ -64,6 +64,7 @@ export function Lightbox() { const images = useGalleryStore((state) => state.images); const openImage = useGalleryStore((state) => state.openImage); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); + const similarScope = useGalleryStore((state) => state.similarScope); const updateImageDetails = useGalleryStore((state) => state.updateImageDetails); const getImageTags = useGalleryStore((state) => state.getImageTags); const addUserTag = useGalleryStore((state) => state.addUserTag); @@ -246,7 +247,7 @@ export function Lightbox() { }`} onClick={() => { if (!canFindSimilar) return; - void loadSimilarImages(selectedImage.id); + void loadSimilarImages(selectedImage.id, similarScope === "current_folder" ? selectedImage.folder_id : null, true, selectedImage.folder_id); }} disabled={!canFindSimilar} > diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx index cdf9ddd..d328ac6 100644 --- a/src/components/SettingsModal.tsx +++ b/src/components/SettingsModal.tsx @@ -1,37 +1,87 @@ -import { useEffect, useState } from "react"; -import { TaggerAcceleration, useGalleryStore } from "../store"; +import { useEffect, useMemo, useState } from "react"; +import { TaggerAcceleration, TaggingQueueScope, useGalleryStore } from "../store"; -type SettingsSection = "tagging" | "library" | "display" | "storage"; +type SettingsSection = "workspace" | "workers"; const SECTIONS: { id: SettingsSection; label: string; detail: string }[] = [ - { id: "tagging", label: "AI Tagging", detail: "WD tagger model" }, - { id: "library", label: "Library", detail: "Indexing and scanning" }, - { id: "display", label: "Display", detail: "Gallery preferences" }, - { id: "storage", label: "Storage", detail: "Cache and model files" }, + { id: "workspace", label: "AI Workspace", detail: "Tagging models and queue targets" }, + { id: "workers", label: "Workers", detail: "Queue activity and background processing" }, ]; - function StatusPill({ children, tone }: { children: React.ReactNode; tone: "ready" | "muted" | "busy" }) { const className = tone === "ready" ? "border-emerald-400/25 bg-emerald-500/10 text-emerald-300" : tone === "busy" - ? "border-sky-400/25 bg-sky-500/10 text-sky-300" - : "border-white/10 bg-white/[0.04] text-gray-500"; + ? "border-sky-400/25 bg-sky-500/10 text-sky-300" + : "border-white/10 bg-white/[0.04] text-gray-500"; + return {children}; +} + +function SectionShell({ eyebrow, title, description, children }: { + eyebrow: string; + title: string; + description?: string; + children: React.ReactNode; +}) { return ( - - {children} - +
+

{eyebrow}

+

{title}

+ {description ?

{description}

: null} +
{children}
+
); } -function TaggerAccelerationButton({ - acceleration, - current, - onSelect, - children, -}: { +function SettingsCard({ title, description, children }: { title: string; description?: string; children: React.ReactNode }) { + return ( +
+
+

{title}

+ {description ?

{description}

: null} +
+
{children}
+
+ ); +} + +function SettingsRow({ title, description, children }: { title: string; description: string; children: React.ReactNode }) { + return ( +
+
+

{title}

+

{description}

+
+
{children}
+
+ ); +} + +function ScopeButton({ scope, current, onSelect, children }: { + scope: TaggingQueueScope; + current: TaggingQueueScope; + onSelect: (scope: TaggingQueueScope) => void; + children: React.ReactNode; +}) { + const active = scope === current; + return ( + + ); +} + +function TaggerAccelerationButton({ acceleration, current, onSelect, children }: { acceleration: TaggerAcceleration; current: TaggerAcceleration; onSelect: (acceleration: TaggerAcceleration) => void; @@ -53,63 +103,32 @@ function TaggerAccelerationButton({ ); } - -function SettingsRow({ - title, - description, - children, -}: { - title: string; - description: string; - children: React.ReactNode; -}) { - return ( -
-
-

{title}

-

{description}

-
-
{children}
-
- ); -} - -function SectionShell({ - eyebrow, - title, - children, -}: { - eyebrow: string; - title: string; - children: React.ReactNode; -}) { - return ( -
-

{eyebrow}

-

{title}

-
{children}
-
- ); -} - export function SettingsModal() { - const [activeSection, setActiveSection] = useState("tagging"); + const [activeSection, setActiveSection] = useState("workspace"); const [taggerQueueStatus, setTaggerQueueStatus] = useState(null); const [taggerQueueing, setTaggerQueueing] = useState(false); const [taggerClearing, setTaggerClearing] = useState(false); const [taggerAccelerationSaving, setTaggerAccelerationSaving] = useState(false); const [taggerThresholdDraft, setTaggerThresholdDraft] = useState(null); const [taggerThresholdSaving, setTaggerThresholdSaving] = useState(false); + const [taggerBatchSizeDraft, setTaggerBatchSizeDraft] = useState(null); + const [taggerBatchSizeSaving, setTaggerBatchSizeSaving] = useState(false); const settingsOpen = useGalleryStore((state) => state.settingsOpen); const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); - const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); const folders = useGalleryStore((state) => state.folders); + const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); + const taggingQueueScope = useGalleryStore((state) => state.taggingQueueScope); + const taggingQueueFolderIds = useGalleryStore((state) => state.taggingQueueFolderIds); + const setTaggingQueueScope = useGalleryStore((state) => state.setTaggingQueueScope); + const toggleTaggingQueueFolder = useGalleryStore((state) => state.toggleTaggingQueueFolder); + const setTaggingQueueFolderIds = useGalleryStore((state) => state.setTaggingQueueFolderIds); const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus); const taggerModelPreparing = useGalleryStore((state) => state.taggerModelPreparing); const taggerModelProgress = useGalleryStore((state) => state.taggerModelProgress); const taggerModelError = useGalleryStore((state) => state.taggerModelError); const taggerAcceleration = useGalleryStore((state) => state.taggerAcceleration); const taggerThreshold = useGalleryStore((state) => state.taggerThreshold); + const taggerBatchSize = useGalleryStore((state) => state.taggerBatchSize); const taggerRuntimeProbe = useGalleryStore((state) => state.taggerRuntimeProbe); const taggerRuntimeChecking = useGalleryStore((state) => state.taggerRuntimeChecking); const loadTaggerModelStatus = useGalleryStore((state) => state.loadTaggerModelStatus); @@ -119,49 +138,123 @@ export function SettingsModal() { const setTaggerAcceleration = useGalleryStore((state) => state.setTaggerAcceleration); const loadTaggerThreshold = useGalleryStore((state) => state.loadTaggerThreshold); const setTaggerThreshold = useGalleryStore((state) => state.setTaggerThreshold); + const loadTaggerBatchSize = useGalleryStore((state) => state.loadTaggerBatchSize); + const setTaggerBatchSize = useGalleryStore((state) => state.setTaggerBatchSize); const probeTaggerRuntime = useGalleryStore((state) => state.probeTaggerRuntime); const queueTaggingJobs = useGalleryStore((state) => state.queueTaggingJobs); + const queueTaggingJobsForFolders = useGalleryStore((state) => state.queueTaggingJobsForFolders); const clearTaggingJobs = useGalleryStore((state) => state.clearTaggingJobs); + const clearTaggingJobsForFolders = useGalleryStore((state) => state.clearTaggingJobsForFolders); useEffect(() => { if (!settingsOpen) return; void loadTaggerModelStatus(); void loadTaggerAcceleration(); void loadTaggerThreshold(); + void loadTaggerBatchSize(); + const handleKeyDown = (event: KeyboardEvent) => { if (event.key === "Escape") setSettingsOpen(false); }; window.addEventListener("keydown", handleKeyDown); return () => window.removeEventListener("keydown", handleKeyDown); - }, [settingsOpen, loadTaggerModelStatus, loadTaggerAcceleration, loadTaggerThreshold, setSettingsOpen]); + }, [settingsOpen, loadTaggerModelStatus, loadTaggerAcceleration, loadTaggerThreshold, loadTaggerBatchSize, setSettingsOpen]); + + const selectedFolders = useMemo( + () => folders.filter((folder) => taggingQueueFolderIds.includes(folder.id)), + [folders, taggingQueueFolderIds], + ); + + const totalQueuedJobs = useMemo( + () => Object.values(mediaJobProgress).reduce((sum, progress) => sum + (progress?.tagging_pending ?? 0), 0), + [mediaJobProgress], + ); if (!settingsOpen) return null; - const selectedFolder = folders.find((folder) => folder.id === selectedFolderId); - const scopeLabel = selectedFolder ? selectedFolder.name : "all libraries"; + const taggerReady = taggerModelStatus?.ready ?? false; + const queueScopeLabel = + taggingQueueScope === "all" + ? "all media" + : selectedFolders.length > 0 + ? `${selectedFolders.length} selected folder${selectedFolders.length === 1 ? "" : "s"}` + : "no folders selected"; + const thresholdDisplay = taggerThresholdDraft ?? String(taggerThreshold); + const batchSizeDisplay = taggerBatchSizeDraft ?? String(taggerBatchSize); + const taggerDownloadLabel = taggerModelProgress + ? `Downloading ${taggerModelProgress.completed_files}/${taggerModelProgress.total_files}` + : taggerModelPreparing + ? "Preparing WD Tagger..." + : taggerReady + ? "Installed" + : "Install model"; + const taggerDownloadPercent = taggerModelProgress + ? Math.round((taggerModelProgress.completed_files / Math.max(taggerModelProgress.total_files, 1)) * 100) + : 0; + + const runQueueAction = (action: "queue" | "clear") => { + const selectedIds = taggingQueueFolderIds; + const perform = + taggingQueueScope === "all" + ? action === "queue" + ? queueTaggingJobs(null) + : clearTaggingJobs(null) + : selectedIds.length > 0 + ? action === "queue" + ? queueTaggingJobsForFolders(selectedIds) + : clearTaggingJobsForFolders(selectedIds) + : Promise.resolve(0); + + if (action === "queue") { + setTaggerQueueing(true); + } else { + setTaggerClearing(true); + } + setTaggerQueueStatus(null); + + void perform + .then((count) => { + if (taggingQueueScope === "selected" && selectedIds.length === 0) { + setTaggerQueueStatus("Choose at least one folder before running tagging jobs."); + return; + } + setTaggerQueueStatus( + count === 0 + ? action === "queue" + ? "No missing tags found for the current target." + : "No queued tagging jobs to clear for the current target." + : action === "queue" + ? `Queued ${count.toLocaleString()} image${count === 1 ? "" : "s"} for tagging.` + : `Cleared ${count.toLocaleString()} queued tagging job${count === 1 ? "" : "s"}.`, + ); + }) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => { + if (action === "queue") { + setTaggerQueueing(false); + } else { + setTaggerClearing(false); + } + }); + }; return ( -
setSettingsOpen(false)} - > +
setSettingsOpen(false)}>
event.stopPropagation()} > -
diff --git a/src/components/Sidebar.tsx b/src/components/Sidebar.tsx index 4a4a33f..042b769 100644 --- a/src/components/Sidebar.tsx +++ b/src/components/Sidebar.tsx @@ -138,6 +138,23 @@ export function Sidebar() { Explore
+ +
setView("duplicates")} + > + + + + + Duplicates + +
{/* Section label */} diff --git a/src/components/TagCloud.tsx b/src/components/TagCloud.tsx index 7c6d3f6..43d2fcc 100644 --- a/src/components/TagCloud.tsx +++ b/src/components/TagCloud.tsx @@ -1,243 +1,375 @@ -import { useEffect } from "react"; +import { useEffect, useLayoutEffect, useMemo, useRef, useState } from "react"; import { motion } from "framer-motion"; import { convertFileSrc } from "@tauri-apps/api/core"; -import { useGalleryStore, TagCloudEntry } from "../store"; +import { ExploreTagEntry, TagCloudEntry, useGalleryStore } from "../store"; -// Accent glow colours for the hover ring — cycled by index -const GLOWS: string[] = [ - "rgba(59,130,246,0.5)", - "rgba(168,85,247,0.5)", - "rgba(16,185,129,0.5)", - "rgba(245,158,11,0.5)", - "rgba(236,72,153,0.5)", - "rgba(6,182,212,0.5)", - "rgba(249,115,22,0.5)", - "rgba(34,197,94,0.5)", +const ACCENTS = [ + "#60a5fa", + "#c084fc", + "#4ade80", + "#fbbf24", + "#f472b4", + "#2dd4bf", + "#fb923c", + "#a78bfa", + "#34d399", + "#f87171", ]; -function pseudoRandom(seed: number): number { - const x = Math.sin(seed + 1) * 10000; +const GOLDEN_ANGLE = Math.PI * (3 - Math.sqrt(5)); + +function seeded(n: number): number { + const x = Math.sin(n * 9301 + 49297) * 233280; return x - Math.floor(x); } -// Map cluster size to a tile size bucket (px) -function getTileSize(count: number, maxCount: number): number { - if (maxCount === 0) return 72; - const ratio = count / maxCount; - if (ratio > 0.75) return 160; - if (ratio > 0.45) return 128; - if (ratio > 0.22) return 104; - if (ratio > 0.08) return 88; - return 72; -} - -function TagButton({ - entry, - index, - maxCount, - onSearch, -}: { +interface PlacedNode { entry: TagCloudEntry; index: number; - maxCount: number; - onSearch: (imageId: number) => void; -}) { - const size = getTileSize(entry.count, maxCount); - const glow = GLOWS[index % GLOWS.length]; + x: number; + y: number; + w: number; + h: number; + accent: string; + driftX: number; + driftY: number; + driftDuration: number; + rotateSeed: number; +} - // Small random rotation for organic feel — larger tiles stay flatter - const maxRot = size >= 128 ? 0 : size >= 104 ? 3 : size >= 88 ? 6 : 10; - const rotation = (pseudoRandom(index * 7) - 0.5) * 2 * maxRot; +function buildCloud(entries: TagCloudEntry[], containerW: number, containerH: number): PlacedNode[] { + if (!entries.length || containerW <= 0 || containerH <= 0) return []; - const mt = Math.floor(pseudoRandom(index * 3) * 10) + 4; - const mr = Math.floor(pseudoRandom(index * 5) * 12) + 4; - const mb = Math.floor(pseudoRandom(index * 11) * 10) + 4; - const ml = Math.floor(pseudoRandom(index * 13) * 12) + 4; + const maxCount = Math.max(...entries.map((e) => e.count)); + const cx = containerW / 2; + const cy = containerH / 2; + // Spread ellipse shrinks slightly to leave room for card half-widths at the edges + const spreadX = containerW * 0.42; + const spreadY = containerH * 0.36; + const n = entries.length; - const src = entry.thumbnail_path ? convertFileSrc(entry.thumbnail_path) : null; + // 1. Build initial positions using phyllotaxis spiral + const nodes: PlacedNode[] = entries.map((entry, i) => { + const ratio = Math.max(entry.count / maxCount, 0.08); + // Cards scale from 110px to 230px wide; height is 3/4 of width + const w = 110 + Math.sqrt(ratio) * 120; + const h = w * 0.75; + const radialRatio = Math.sqrt((i + 0.5) / n); + const angle = i * GOLDEN_ANGLE; + + return { + entry, + index: i, + x: cx + Math.cos(angle) * radialRatio * spreadX, + y: cy + Math.sin(angle) * radialRatio * spreadY, + w, + h, + accent: ACCENTS[i % ACCENTS.length], + driftX: (seeded(i + 11) - 0.5) * 18, + driftY: (seeded(i + 17) - 0.5) * 14, + driftDuration: 8 + seeded(i + 23) * 7, + rotateSeed: (seeded(i + 31) - 0.5) * 4, + }; + }); + + // 2. Iterative overlap resolution — no physics, just push apart + const PAD = 24; + for (let iter = 0; iter < 80; iter++) { + for (let a = 0; a < nodes.length; a++) { + const na = nodes[a]; + for (let b = a + 1; b < nodes.length; b++) { + const nb = nodes[b]; + const dx = nb.x - na.x; + const dy = nb.y - na.y; + const overlapX = (na.w + nb.w) / 2 + PAD - Math.abs(dx); + const overlapY = (na.h + nb.h) / 2 + PAD - Math.abs(dy); + if (overlapX <= 0 || overlapY <= 0) continue; + // Push along the smaller overlap axis + if (overlapX < overlapY) { + const push = overlapX * 0.5 * (dx >= 0 ? 1 : -1); + nb.x += push; + na.x -= push; + } else { + const push = overlapY * 0.5 * (dy >= 0 ? 1 : -1); + nb.y += push; + na.y -= push; + } + } + // Pull gently back toward anchor to prevent runaway drift + na.x += (cx + Math.cos(na.index * GOLDEN_ANGLE) * Math.sqrt((na.index + 0.5) / n) * spreadX - na.x) * 0.05; + na.y += (cy + Math.sin(na.index * GOLDEN_ANGLE) * Math.sqrt((na.index + 0.5) / n) * spreadY - na.y) * 0.05; + } + } + + // 3. Clamp so cards never poke outside the container + return nodes.map((node) => ({ + ...node, + x: Math.min(Math.max(node.x, node.w / 2 + 16), containerW - node.w / 2 - 16), + y: Math.min(Math.max(node.y, node.h / 2 + 16), containerH - node.h / 2 - 16), + })); +} + +function CloudCard({ node, onOpen }: { node: PlacedNode; onOpen: (imageIds: number[]) => void }) { + const src = node.entry.thumbnail_path ? convertFileSrc(node.entry.thumbnail_path) : null; + const { w, h, accent } = node; return ( onSearch(entry.representative_image_id)} - title={`${entry.count} similar ${entry.count === 1 ? "photo" : "photos"}`} - style={{ - width: size, - height: size, - margin: `${mt}px ${mr}px ${mb}px ${ml}px`, - borderRadius: 12, - border: "2px solid rgba(255,255,255,0.08)", - background: "rgba(255,255,255,0.04)", - cursor: "pointer", - padding: 0, - overflow: "hidden", - position: "relative", - flexShrink: 0, - boxShadow: "none", - transition: "border-color 0.15s, box-shadow 0.15s", - }} - onMouseEnter={(e) => { - const el = e.currentTarget; - el.style.borderColor = glow; - el.style.boxShadow = `0 0 16px ${glow}, 0 0 32px ${glow.replace("0.5", "0.25")}`; - }} - onMouseLeave={(e) => { - const el = e.currentTarget; - el.style.borderColor = "rgba(255,255,255,0.08)"; - el.style.boxShadow = "none"; + opacity: { duration: 0.3, delay: Math.min(node.index * 0.035, 0.7) }, + scale: { duration: 0.3, delay: Math.min(node.index * 0.035, 0.7) }, + x: { duration: node.driftDuration, repeat: Infinity, ease: "easeInOut", delay: seeded(node.index + 41) * 3 }, + y: { duration: node.driftDuration + 1.6, repeat: Infinity, ease: "easeInOut", delay: seeded(node.index + 51) * 3 }, + rotate: { duration: node.driftDuration + 0.9, repeat: Infinity, ease: "easeInOut" }, }} + whileHover={{ scale: 1.06, rotate: 0, transition: { duration: 0.18 } }} + onClick={() => onOpen(node.entry.image_ids)} + title={`Open cluster — ${node.entry.count.toLocaleString()} images`} > {src ? ( ) : ( - // Fallback placeholder when no thumbnail exists yet -
+
)} - - {/* Count badge — bottom-right corner */} +
+ {/* Accent glow on hover */}
- {entry.count} + className="absolute inset-0 opacity-0 transition-opacity duration-300 group-hover:opacity-100" + style={{ background: `radial-gradient(ellipse at bottom, ${accent}25, transparent 70%)` }} + /> +
+
+
+
+

Cluster

+

{node.entry.count.toLocaleString()}

+
+ + Open + +
); } +// Actual tag cloud — word size driven by log-scaled frequency +function TagWord({ + entry, + index, + logMin, + logRange, + onSearch, +}: { + entry: ExploreTagEntry; + index: number; + logMin: number; + logRange: number; + onSearch: (tag: string) => void; +}) { + const ratio = logRange > 0 ? (Math.log(Math.max(entry.count, 1)) - logMin) / logRange : 0.5; + const fontSize = 11 + ratio * 28; // 11px – 39px + const accent = ACCENTS[index % ACCENTS.length]; + const tilt = (seeded(index + 5) - 0.5) * 7; + + return ( + onSearch(entry.tag)} + title={`${entry.tag} — ${entry.count.toLocaleString()} images`} + > + 0.55 ? accent : "rgba(255,255,255,0.82)" }} + > + {entry.tag} + + + {entry.count.toLocaleString()} + + + ); +} + +function Spinner() { + return ( + + ); +} + +// Separate component so its useLayoutEffect fires when the canvas is actually +// mounted — not at TagCloud mount time when the container may still be hidden +// behind a loading state. +function ClusterCloud({ + entries, + onOpen, +}: { + entries: TagCloudEntry[]; + onOpen: (imageIds: number[]) => void; +}) { + const canvasRef = useRef(null); + const [canvasSize, setCanvasSize] = useState({ w: 0, h: 0 }); + + useLayoutEffect(() => { + const el = canvasRef.current; + if (!el) return; + const update = () => { + const r = el.getBoundingClientRect(); + setCanvasSize({ w: r.width, h: r.height }); + }; + update(); + const ro = new ResizeObserver(update); + ro.observe(el); + return () => ro.disconnect(); + }, []); + + const nodes = useMemo( + () => buildCloud(entries, canvasSize.w, canvasSize.h), + [entries, canvasSize.w, canvasSize.h], + ); + + return ( +
+
+ {nodes.map((node) => ( + + ))} +
+ ); +} + export function TagCloud() { + const exploreMode = useGalleryStore((state) => state.exploreMode); + const setExploreMode = useGalleryStore((state) => state.setExploreMode); const tagCloudEntries = useGalleryStore((state) => state.tagCloudEntries); const tagCloudLoading = useGalleryStore((state) => state.tagCloudLoading); const loadTagCloud = useGalleryStore((state) => state.loadTagCloud); - const searchByTag = useGalleryStore((state) => state.searchByTag); + const exploreTagEntries = useGalleryStore((state) => state.exploreTagEntries); + const exploreTagLoading = useGalleryStore((state) => state.exploreTagLoading); + const loadExploreTags = useGalleryStore((state) => state.loadExploreTags); + const showVisualCluster = useGalleryStore((state) => state.showVisualCluster); + const searchForTag = useGalleryStore((state) => state.searchForTag); const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); useEffect(() => { - void loadTagCloud(); - }, [selectedFolderId]); + if (exploreMode === "visual") void loadTagCloud(); + else void loadExploreTags(); + }, [exploreMode, selectedFolderId, loadTagCloud, loadExploreTags]); - const maxCount = - tagCloudEntries.length > 0 - ? Math.max(...tagCloudEntries.map((e) => e.count)) - : 1; + const { logMin, logRange } = useMemo(() => { + if (!exploreTagEntries.length) return { logMin: 0, logRange: 1 }; + const logs = exploreTagEntries.map((e) => Math.log(Math.max(e.count, 1))); + const lo = Math.min(...logs); + const hi = Math.max(...logs); + return { logMin: lo, logRange: hi - lo || 1 }; + }, [exploreTagEntries]); + + const loading = exploreMode === "visual" ? tagCloudLoading : exploreTagLoading; + const hasEntries = exploreMode === "visual" ? tagCloudEntries.length > 0 : exploreTagEntries.length > 0; + const entryCount = exploreMode === "visual" ? tagCloudEntries.length : exploreTagEntries.length; return ( -
+
{/* Header */} - -

- Explore your library -

-

- Visual clusters from your photos — sized by how many match -

- - - {/* Loading */} - {tagCloudLoading && ( -
- - - - -

Clustering your library…

+
+
+
+

Explore

+

+ {loading + ? exploreMode === "visual" ? "Computing visual clusters…" : "Loading tags…" + : hasEntries + ? exploreMode === "visual" + ? `${entryCount} cluster${entryCount !== 1 ? "s" : ""} — click any to open` + : `${entryCount} tag${entryCount !== 1 ? "s" : ""} — click any to search` + : exploreMode === "visual" + ? "No clusters — images need embeddings first" + : "No tags — run the AI tagger or add tags manually"} +

+
+
+ + +
- )} +
- {/* Empty state */} - {!tagCloudLoading && tagCloudEntries.length === 0 && ( -
-

- No embeddings yet. Add a folder and wait for the embedding worker to - finish, then come back here. + {loading ? ( +

+ + {exploreMode === "visual" ? "Computing clusters…" : "Loading tags…"} +
+ ) : !hasEntries ? ( +
+

+ {exploreMode === "visual" + ? "No visual clusters yet. Images need embeddings before they can be grouped. Check indexing progress in the sidebar." + : "No tags yet. Run the AI tagger from Settings, or add tags manually in the image preview."}

- )} - - {/* Cluster grid */} - {!tagCloudLoading && tagCloudEntries.length > 0 && ( -
- {tagCloudEntries.map((entry, index) => ( - - ))} + ) : exploreMode === "visual" ? ( + + ) : ( + /* Tag cloud — words sized by log-scaled frequency, wrapped freely */ +
+
+ {exploreTagEntries.map((entry, index) => ( + + ))} +
)} - - {!tagCloudLoading && tagCloudEntries.length > 0 && ( -

- Grouped by visual similarity · CLIP ViT-B/32 -

- )}
); } diff --git a/src/components/Toolbar.tsx b/src/components/Toolbar.tsx index 88cedee..b0514f4 100644 --- a/src/components/Toolbar.tsx +++ b/src/components/Toolbar.tsx @@ -1,11 +1,14 @@ import { useEffect, useRef, useState } from "react"; -import { tileSizeForZoom, useGalleryStore, SortOrder, MediaFilter, SearchMode } from "../store"; +import { invoke } from "@tauri-apps/api/core"; +import { tileSizeForZoom, useGalleryStore, SortOrder, MediaFilter, SearchCommand, parseSearchValue, searchModeLabel, ExploreTagEntry } from "../store"; const BASE_SORT_OPTIONS: { value: SortOrder; label: string }[] = [ { value: "date_desc", label: "Newest first" }, { value: "date_asc", label: "Oldest first" }, { value: "name_asc", label: "Name A–Z" }, { value: "name_desc", label: "Name Z–A" }, + { value: "rating_desc", label: "Highest rated" }, + { value: "rating_asc", label: "Lowest rated" }, { value: "size_desc", label: "Largest first" }, { value: "size_asc", label: "Smallest first" }, ]; @@ -116,12 +119,27 @@ function FilterPill({ ); } +function commandPrefix(command: SearchCommand | null): string | null { + switch (command) { + case "semantic": + return "/s"; + case "tag": + return "/t"; + default: + return null; + } +} + +function composeSearchValue(command: SearchCommand | null, query: string): string { + const prefix = commandPrefix(command); + if (!prefix) return query; + return query.length > 0 ? `${prefix} ${query}` : prefix; +} + export function Toolbar() { const search = useGalleryStore((state) => state.search); const setSearch = useGalleryStore((state) => state.setSearch); const clearSearch = useGalleryStore((state) => state.clearSearch); - const searchMode = useGalleryStore((state) => state.searchMode); - const setSearchMode = useGalleryStore((state) => state.setSearchMode); const sort = useGalleryStore((state) => state.sort); const setSort = useGalleryStore((state) => state.setSort); const totalImages = useGalleryStore((state) => state.totalImages); @@ -133,20 +151,26 @@ export function Toolbar() { const setMediaFilter = useGalleryStore((state) => state.setMediaFilter); const favoritesOnly = useGalleryStore((state) => state.favoritesOnly); const setFavoritesOnly = useGalleryStore((state) => state.setFavoritesOnly); + const minimumRating = useGalleryStore((state) => state.minimumRating); + const setMinimumRating = useGalleryStore((state) => state.setMinimumRating); const failedEmbeddingsOnly = useGalleryStore((state) => state.failedEmbeddingsOnly); const setFailedEmbeddingsOnly = useGalleryStore((state) => state.setFailedEmbeddingsOnly); + const similarScope = useGalleryStore((state) => state.similarScope); + const setSimilarScope = useGalleryStore((state) => state.setSimilarScope); const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); const zoomPreset = useGalleryStore((state) => state.zoomPreset); const setZoomPreset = useGalleryStore((state) => state.setZoomPreset); const hasAnyFailedEmbeddings = Object.values(mediaJobProgress).some((p) => p.embedding_failed > 0); - const [searchValue, setSearchValue] = useState(search); + const [searchCommand, setSearchCommand] = useState(null); + const [searchQuery, setSearchQuery] = useState(search); + const [searchPanelOpen, setSearchPanelOpen] = useState(false); + const [tagSuggestions, setTagSuggestions] = useState([]); const debounceRef = useRef | null>(null); + const suggestDebounceRef = useRef | null>(null); const searchInputRef = useRef(null); - // Tracks whether the user has typed in the search box at least once. - // Prevents the debounce effect from dispatching setSearch on initial mount - // when searchValue === search (which would wipe a loadSimilarImages result). + const searchShellRef = useRef(null); const userHasTyped = useRef(false); const selectedFolder = folders.find((folder) => folder.id === selectedFolderId); @@ -154,11 +178,8 @@ export function Toolbar() { const tileSize = tileSizeForZoom(zoomPreset); const sortOptions = getSortOptions(mediaFilter); const hasActiveSearch = search.trim().length > 0; - - const searchModes: { value: SearchMode; label: string }[] = [ - { value: "filename", label: "Filename" }, - { value: "semantic", label: "Semantic" }, - ]; + const parsedSearch = parseSearchValue(composeSearchValue(searchCommand, searchQuery)); + const isSimilarResults = collectionTitle === "Similar Images"; // If current sort is video-only but we switched away from video filter, reset to date_desc useEffect(() => { @@ -170,35 +191,62 @@ export function Toolbar() { useEffect(() => { if (!userHasTyped.current) return; if (debounceRef.current) clearTimeout(debounceRef.current); - debounceRef.current = setTimeout(() => { setSearch(searchValue); }, 200); + debounceRef.current = setTimeout(() => { setSearch(composeSearchValue(searchCommand, searchQuery)); }, 200); return () => { if (debounceRef.current) clearTimeout(debounceRef.current); }; - }, [searchValue, setSearch]); + }, [searchCommand, searchQuery, setSearch]); useEffect(() => { - setSearchValue(search); + const parsed = parseSearchValue(search); + setSearchCommand(parsed.prefix && parsed.mode !== "filename" ? parsed.mode : null); + setSearchQuery(parsed.prefix ? parsed.query : search); }, [search]); + // Fetch tag suggestions when in tag mode + useEffect(() => { + if (searchCommand !== "tag") { + setTagSuggestions([]); + return; + } + if (suggestDebounceRef.current) clearTimeout(suggestDebounceRef.current); + suggestDebounceRef.current = setTimeout(async () => { + try { + const results = await invoke("search_tags_autocomplete", { + params: { query: searchQuery.trim(), folder_id: selectedFolderId ?? null, limit: 10 }, + }); + setTagSuggestions(results); + } catch { + setTagSuggestions([]); + } + }, 120); + return () => { if (suggestDebounceRef.current) clearTimeout(suggestDebounceRef.current); }; + }, [searchCommand, searchQuery, selectedFolderId]); + useEffect(() => { const handleKeyDown = (event: KeyboardEvent) => { - const isModeToggle = (event.ctrlKey || event.metaKey) && event.shiftKey && event.key.toLowerCase() === "s"; - if (isModeToggle) { - event.preventDefault(); - setSearchMode(searchMode === "semantic" ? "filename" : "semantic"); - searchInputRef.current?.focus(); - return; - } - const activeElement = document.activeElement; const searchFocused = activeElement === searchInputRef.current; if (event.key === "Escape" && (searchFocused || hasActiveSearch)) { event.preventDefault(); - setSearchValue(""); + setSearchCommand(null); + setSearchQuery(""); clearSearch(); } }; window.addEventListener("keydown", handleKeyDown); return () => window.removeEventListener("keydown", handleKeyDown); - }, [clearSearch, hasActiveSearch, searchMode, setSearchMode]); + }, [clearSearch, hasActiveSearch]); + + useEffect(() => { + const close = (event: PointerEvent) => { + if (searchShellRef.current?.contains(event.target as Node)) return; + setSearchPanelOpen(false); + }; + window.addEventListener("pointerdown", close); + return () => window.removeEventListener("pointerdown", close); + }, []); + + const showTagSuggestions = searchCommand === "tag" && searchPanelOpen; + const showCommandHints = !searchCommand && searchPanelOpen; return (
@@ -212,9 +260,9 @@ export function Toolbar() { ? `${loadedCount.toLocaleString()} / ${totalImages.toLocaleString()}` : totalImages.toLocaleString()} - {(hasActiveSearch || searchMode === "semantic") && ( + {hasActiveSearch && ( - {searchMode === "semantic" ? "Semantic Search" : "Filename Search"} + {searchModeLabel(parsedSearch.mode)} )}
@@ -222,55 +270,140 @@ export function Toolbar() {
{/* Search */} -
-
- {searchModes.map((mode) => ( - - ))} -
-
- - - - { - userHasTyped.current = true; - setSearchValue(event.target.value); - }} - placeholder={searchMode === "semantic" ? "Search by meaning..." : "Search filenames..."} - className="w-64 bg-transparent py-1.5 pl-8 pr-9 text-sm text-white placeholder:text-gray-600 focus:outline-none transition-colors" - /> - {searchValue.trim().length > 0 && ( - - )} + onKeyDown={(event) => { + if (event.key === "Backspace" && searchQuery.length === 0 && searchCommand !== null) { + event.preventDefault(); + setSearchCommand(null); + } + }} + onFocus={() => setSearchPanelOpen(true)} + placeholder="Search files, or use /s /t" + className={`w-64 bg-transparent py-1.5 pr-9 text-sm text-white placeholder:text-gray-600 focus:outline-none transition-colors ${searchCommand !== null ? "pl-16" : "pl-8"}`} + /> + {searchCommand !== null ? ( +
+ +
+ ) : null} + {searchQuery.trim().length > 0 || searchCommand !== null ? ( + + ) : null} +
+ + {/* Tag autocomplete suggestions */} + {showTagSuggestions && tagSuggestions.length > 0 ? ( +
+ {tagSuggestions.map((entry) => ( + + ))} +
+ ) : null} + + {/* Tag mode with no suggestions yet — show a brief hint */} + {showTagSuggestions && tagSuggestions.length === 0 && searchQuery.trim().length > 0 ? ( +
+

No matching tags

+
+ ) : null} + + {/* Semantic mode hint */} + {searchCommand === "semantic" && searchPanelOpen ? ( +
+

Search by meaning and visual concepts

+
+ ) : null} + + {/* Command hints — only shown when no command is active */} + {showCommandHints ? ( +
+ {( + [ + { command: "tag" as SearchCommand, prefix: "/t", label: "Tags", description: "Search AI and user tags" }, + { command: "semantic" as SearchCommand, prefix: "/s", label: "Semantic", description: "Search by meaning" }, + ] as const + ).map((option) => ( + + ))} +
+ ) : null}
{/* Sort */} @@ -302,10 +435,14 @@ export function Toolbar() { {/* Filter row */}
- { setMediaFilter("all"); setFavoritesOnly(false); setFailedEmbeddingsOnly(false); }} /> + { setMediaFilter("all"); setFavoritesOnly(false); setMinimumRating(0); setFailedEmbeddingsOnly(false); }} /> { setMediaFilter("image"); setFavoritesOnly(false); setFailedEmbeddingsOnly(false); }} /> { setMediaFilter("video"); setFavoritesOnly(false); setFailedEmbeddingsOnly(false); }} /> { setFavoritesOnly(!favoritesOnly); setFailedEmbeddingsOnly(false); }} /> + { setMinimumRating(minimumRating === 1 ? 0 : 1); setFailedEmbeddingsOnly(false); }} /> + { setMinimumRating(minimumRating === 4 ? 0 : 4); setFailedEmbeddingsOnly(false); }} /> + setSimilarScope("current_folder")} /> + setSimilarScope("all_media")} /> {hasAnyFailedEmbeddings ? ( setFailedEmbeddingsOnly(!failedEmbeddingsOnly)} /> ) : null} + {isSimilarResults ? Current similar scope: {similarScope === "current_folder" ? "current folder" : "all media"} : null}
); diff --git a/src/store.ts b/src/store.ts index 32ae1fc..ad02181 100644 --- a/src/store.ts +++ b/src/store.ts @@ -15,10 +15,14 @@ export type MediaKind = "image" | "video"; export type MediaFilter = "all" | MediaKind; export type ZoomPreset = "compact" | "comfortable" | "detail"; export type SearchMode = "filename" | "semantic"; +export type SearchCommand = "filename" | "semantic" | "tag"; export type CaptionAcceleration = "auto" | "cpu" | "directml"; export type CaptionDetail = "short" | "detailed" | "paragraph"; export type TaggerAcceleration = "auto" | "cpu" | "directml"; export type AiRating = "general" | "sensitive" | "questionable" | "explicit"; +export type TaggingQueueScope = "all" | "selected"; +export type SimilarScope = "all_media" | "current_folder"; +export type ExploreMode = "visual" | "tags"; export interface ImageRecord { id: number; @@ -115,12 +119,33 @@ export interface ThumbnailBatch { images: ImageRecord[]; } -export type ActiveView = "gallery" | "explore"; +export type ActiveView = "gallery" | "explore" | "duplicates"; export interface TagCloudEntry { count: number; representative_image_id: number; thumbnail_path: string | null; + image_ids: number[]; +} + +export interface ExploreTagEntry { + tag: string; + count: number; + representative_image_id: number; + thumbnail_path: string | null; +} + +export interface DuplicateGroup { + file_hash: string; + file_size: number; + images: ImageRecord[]; +} + +export interface SimilarImagesPage { + images: ImageRecord[]; + offset: number; + limit: number; + has_more: boolean; } export interface CaptionModelStatus { @@ -171,6 +196,12 @@ export interface TaggerRuntimeProbe { session: TaggerRuntimeSessionProbe; } +export interface ParsedSearch { + mode: SearchCommand; + query: string; + prefix: string | null; +} + export type SortOrder = | "date_desc" | "date_asc" @@ -178,6 +209,8 @@ export type SortOrder = | "name_desc" | "size_desc" | "size_asc" + | "rating_desc" + | "rating_asc" | "duration_desc" | "duration_asc"; @@ -194,17 +227,25 @@ interface GalleryState { sort: SortOrder; mediaFilter: MediaFilter; favoritesOnly: boolean; + minimumRating: number; failedEmbeddingsOnly: boolean; zoomPreset: ZoomPreset; selectedImage: ImageRecord | null; collectionTitle: string | null; similarSourceImageId: number | null; + similarSourceFolderId: number | null; similarHasMore: boolean; + similarScope: SimilarScope; + similarFolderId: number | null; galleryScrollResetKey: number; activeView: ActiveView; + exploreMode: ExploreMode; tagCloudEntries: TagCloudEntry[]; tagCloudLoading: boolean; tagCloudFolderId: number | null | undefined; // undefined = never loaded + exploreTagEntries: ExploreTagEntry[]; + exploreTagLoading: boolean; + exploreTagsFolderId: number | null | undefined; indexingProgress: Record; mediaJobProgress: Record; cacheDir: string; @@ -218,6 +259,8 @@ interface GalleryState { captionDetail: CaptionDetail; aiCaptionsEnabled: boolean; settingsOpen: boolean; + taggingQueueScope: TaggingQueueScope; + taggingQueueFolderIds: number[]; taggerModelStatus: TaggerModelStatus | null; taggerModelPreparing: boolean; @@ -225,9 +268,16 @@ interface GalleryState { taggerModelProgress: TaggerModelProgress | null; taggerAcceleration: TaggerAcceleration; taggerThreshold: number; + taggerBatchSize: number; taggerRuntimeProbe: TaggerRuntimeProbe | null; taggerRuntimeChecking: boolean; + duplicateGroups: DuplicateGroup[]; + duplicateScanning: boolean; + duplicateScanProgress: { scanned: number; total: number } | null; + duplicateSelectedIds: Set; + duplicateLastScanned: number | null; // Unix timestamp (seconds) + loadFolders: () => Promise; loadBackgroundJobProgress: () => Promise; addFolder: (path: string) => Promise; @@ -243,14 +293,19 @@ interface GalleryState { setSort: (sort: SortOrder) => void; setMediaFilter: (filter: MediaFilter) => void; setFavoritesOnly: (favoritesOnly: boolean) => void; + setMinimumRating: (minimumRating: number) => void; setFailedEmbeddingsOnly: (failedEmbeddingsOnly: boolean) => void; setZoomPreset: (zoomPreset: ZoomPreset) => void; openImage: (image: ImageRecord) => void; closeImage: () => void; setView: (view: ActiveView) => void; + setExploreMode: (mode: ExploreMode) => void; loadTagCloud: () => Promise; - searchByTag: (imageId: number) => void; - loadSimilarImages: (imageId: number, folderId?: number | null, reset?: boolean) => Promise; + loadExploreTags: () => Promise; + showVisualCluster: (imageIds: number[]) => Promise; + searchForTag: (tag: string) => void; + loadSimilarImages: (imageId: number, folderId?: number | null, reset?: boolean, sourceFolderId?: number | null) => Promise; + setSimilarScope: (scope: SimilarScope) => void; suggestImageTags: (imageId: number) => Promise; loadCaptionModelStatus: () => Promise; prepareCaptionModel: () => Promise; @@ -268,6 +323,9 @@ interface GalleryState { setCaptionDetail: (detail: CaptionDetail) => Promise; setAiCaptionsEnabled: (enabled: boolean) => void; setSettingsOpen: (open: boolean) => void; + setTaggingQueueScope: (scope: TaggingQueueScope) => void; + toggleTaggingQueueFolder: (folderId: number) => void; + setTaggingQueueFolderIds: (folderIds: number[]) => void; retryFailedEmbeddings: (folderId: number) => Promise; updateImageDetails: (imageId: number, updates: { favorite?: boolean; rating?: number }) => Promise; setCacheDir: (dir: string) => void; @@ -280,10 +338,21 @@ interface GalleryState { setTaggerAcceleration: (acceleration: TaggerAcceleration) => Promise; loadTaggerThreshold: () => Promise; setTaggerThreshold: (threshold: number) => Promise; + loadTaggerBatchSize: () => Promise; + setTaggerBatchSize: (batchSize: number) => Promise; probeTaggerRuntime: () => Promise; queueTaggingJobs: (folderId?: number | null) => Promise; + queueTaggingJobsForFolders: (folderIds: number[]) => Promise; queueTaggingForImage: (imageId: number) => Promise; clearTaggingJobs: (folderId?: number | null) => Promise; + clearTaggingJobsForFolders: (folderIds: number[]) => Promise; + loadDuplicateScanCache: (folderId?: number | null) => Promise; + scanDuplicates: (folderId?: number | null) => Promise; + toggleDuplicateSelected: (imageId: number) => void; + selectAllDuplicates: (imageIds: number[]) => void; + selectKeepFirstAllGroups: () => void; + clearDuplicateSelection: () => void; + deleteSelectedDuplicates: () => Promise; getImageTags: (imageId: number) => Promise; addUserTag: (imageId: number, tag: string) => Promise; removeTag: (tagId: number) => Promise; @@ -291,6 +360,12 @@ interface GalleryState { const PAGE_SIZE = 200; const AI_CAPTIONS_ENABLED_KEY = "phokus.aiCaptionsEnabled"; +const SIMILAR_DISTANCE_THRESHOLD = 0.24; + +let galleryRequestToken = 0; +let similarRequestToken = 0; +let tagCloudRequestToken = 0; +let exploreTagRequestToken = 0; function initialAiCaptionsEnabled(): boolean { if (typeof window === "undefined") return false; @@ -312,19 +387,71 @@ function matchesSearch(image: ImageRecord, search: string): boolean { return image.filename.toLowerCase().includes(search.toLowerCase()); } +function isDerivedCollectionTitle(collectionTitle: string | null): boolean { + return collectionTitle !== null; +} + +export function parseSearchValue(search: string): ParsedSearch { + if (!search.trim()) { + return { mode: "filename", query: "", prefix: null }; + } + + const slashPrefix = search.match(/^\/([a-z])(?:\s|$)/i); + if (slashPrefix) { + const rawPrefix = slashPrefix[1].toLowerCase(); + const query = search.length > 3 ? search.slice(3) : ""; + if (rawPrefix === "s") { + return { mode: "semantic", query, prefix: "/s" }; + } + if (rawPrefix === "t") { + return { mode: "tag", query, prefix: "/t" }; + } + return { mode: "filename", query, prefix: rawPrefix === "f" ? "/f" : null }; + } + + const trimmed = search.trim(); + const match = trimmed.match(/^([a-z]):\s*(.*)$/i); + if (!match) { + return { mode: "filename", query: trimmed, prefix: null }; + } + + const rawPrefix = match[1].toLowerCase(); + const query = match[2].trim(); + if (rawPrefix === "s") { + return { mode: "semantic", query, prefix: "s:" }; + } + if (rawPrefix === "t") { + return { mode: "tag", query, prefix: "t:" }; + } + return { mode: "filename", query, prefix: rawPrefix === "f" ? "f:" : null }; +} + +export function searchModeLabel(mode: SearchCommand): string { + switch (mode) { + case "semantic": + return "Semantic Search"; + case "tag": + return "Tag Search"; + default: + return "Filename Search"; + } +} + function matchesFilters( image: ImageRecord, selectedFolderId: number | null, mediaFilter: MediaFilter, favoritesOnly: boolean, + minimumRating: number, failedEmbeddingsOnly: boolean, search: string, ): boolean { const matchesFolder = selectedFolderId === null || image.folder_id === selectedFolderId; const matchesMedia = mediaFilter === "all" || image.media_kind === mediaFilter; const matchesFavorite = !favoritesOnly || image.favorite; + const matchesRating = image.rating >= minimumRating; const matchesFailedEmbedding = !failedEmbeddingsOnly || image.embedding_status === "failed"; - return matchesFolder && matchesMedia && matchesFavorite && matchesFailedEmbedding && matchesSearch(image, search); + return matchesFolder && matchesMedia && matchesFavorite && matchesRating && matchesFailedEmbedding && matchesSearch(image, search); } function compareNullableNumber(a: number | null, b: number | null): number { @@ -349,6 +476,10 @@ function compareImages(a: ImageRecord, b: ImageRecord, sort: SortOrder): number return compareNullableNumber(a.file_size, b.file_size); case "size_desc": return compareNullableNumber(b.file_size, a.file_size); + case "rating_asc": + return compareNullableNumber(a.rating, b.rating); + case "rating_desc": + return compareNullableNumber(b.rating, a.rating); case "duration_asc": return compareNullableNumber(a.duration_ms, b.duration_ms); case "duration_desc": @@ -424,17 +555,25 @@ export const useGalleryStore = create((set, get) => ({ sort: "date_desc", mediaFilter: "all", favoritesOnly: false, + minimumRating: 0, failedEmbeddingsOnly: false, zoomPreset: "comfortable", selectedImage: null, collectionTitle: null, similarSourceImageId: null, + similarSourceFolderId: null, similarHasMore: false, + similarScope: "all_media", + similarFolderId: null, galleryScrollResetKey: 0, activeView: "gallery", + exploreMode: "visual", tagCloudEntries: [], tagCloudLoading: false, tagCloudFolderId: undefined, + exploreTagEntries: [], + exploreTagLoading: false, + exploreTagsFolderId: undefined, indexingProgress: {}, mediaJobProgress: {}, cacheDir: "", @@ -448,6 +587,8 @@ export const useGalleryStore = create((set, get) => ({ captionDetail: "paragraph", aiCaptionsEnabled: initialAiCaptionsEnabled(), settingsOpen: false, + taggingQueueScope: "all", + taggingQueueFolderIds: [], taggerModelStatus: null, taggerModelPreparing: false, @@ -455,14 +596,33 @@ export const useGalleryStore = create((set, get) => ({ taggerModelProgress: null, taggerAcceleration: "auto", taggerThreshold: 0.35, + taggerBatchSize: 8, taggerRuntimeProbe: null, taggerRuntimeChecking: false, + duplicateGroups: [], + duplicateScanning: false, + duplicateScanProgress: null, + duplicateSelectedIds: new Set(), + duplicateLastScanned: null, + setCacheDir: (cacheDir) => set({ cacheDir }), loadFolders: async () => { const folders = await invoke("get_folders"); - set({ folders }); + set((state) => { + const folderIds = new Set(folders.map((folder) => folder.id)); + const nextSelected = state.taggingQueueFolderIds.filter((folderId) => folderIds.has(folderId)); + return { + folders, + taggingQueueFolderIds: + nextSelected.length > 0 + ? nextSelected + : state.taggingQueueScope === "selected" && folders.length > 0 + ? [folders[0].id] + : nextSelected, + }; + }); }, loadBackgroundJobProgress: async () => { @@ -507,29 +667,64 @@ export const useGalleryStore = create((set, get) => ({ }, loadImages: async (reset = false) => { - const { selectedFolderId, search, searchMode, sort, loadedCount, mediaFilter, favoritesOnly, failedEmbeddingsOnly } = get(); + const { selectedFolderId, search, sort, loadedCount, mediaFilter, favoritesOnly, minimumRating, failedEmbeddingsOnly } = get(); + const parsedSearch = parseSearchValue(search); + const requestToken = ++galleryRequestToken; set({ loadingImages: true, imageLoadError: null }); try { - if (searchMode === "semantic" && search.trim()) { + if (parsedSearch.mode === "semantic" && parsedSearch.query) { const images = await invoke("semantic_search_images", { params: { - query: search, + query: parsedSearch.query, folder_id: selectedFolderId, media_kind: mediaFilter === "all" ? null : mediaFilter, favorites_only: favoritesOnly, + rating_min: minimumRating > 0 ? minimumRating : null, limit: PAGE_SIZE, }, }); + if (requestToken !== galleryRequestToken) return; set({ images, totalImages: images.length, loadedCount: images.length, loadingImages: false, - collectionTitle: `Semantic search: ${search}`, + collectionTitle: `Semantic search: ${parsedSearch.query}`, + selectedFolderId, similarSourceImageId: null, + similarSourceFolderId: null, similarHasMore: false, + similarFolderId: null, + }); + return; + } + + if (parsedSearch.mode === "tag" && parsedSearch.query) { + const images = await invoke("search_images_by_tag", { + params: { + query: parsedSearch.query, + folder_id: selectedFolderId, + media_kind: mediaFilter === "all" ? null : mediaFilter, + favorites_only: favoritesOnly, + rating_min: minimumRating > 0 ? minimumRating : null, + limit: PAGE_SIZE, + }, + }); + + if (requestToken !== galleryRequestToken) return; + set({ + images, + totalImages: images.length, + loadedCount: images.length, + loadingImages: false, + collectionTitle: `Tag search: ${parsedSearch.query}`, + selectedFolderId, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarFolderId: null, }); return; } @@ -543,9 +738,10 @@ export const useGalleryStore = create((set, get) => ({ }>("get_images", { params: { folder_id: selectedFolderId, - search: search || null, + search: parsedSearch.query || null, media_kind: mediaFilter === "all" ? null : mediaFilter, favorites_only: favoritesOnly, + rating_min: minimumRating > 0 ? minimumRating : null, embedding_failed_only: failedEmbeddingsOnly, sort, offset, @@ -553,6 +749,7 @@ export const useGalleryStore = create((set, get) => ({ }, }); + if (requestToken !== galleryRequestToken) return; set((state) => ({ images: reset ? result.images : [...state.images, ...result.images], totalImages: result.total, @@ -560,20 +757,24 @@ export const useGalleryStore = create((set, get) => ({ loadingImages: false, collectionTitle: reset ? null : state.collectionTitle, similarSourceImageId: null, + similarSourceFolderId: null, similarHasMore: false, + similarFolderId: null, })); } catch (error) { + if (requestToken !== galleryRequestToken) return; console.error("Failed to load media:", error); set({ loadingImages: false, imageLoadError: String(error) }); } }, loadMoreImages: async () => { - const { loadedCount, totalImages, loadingImages, collectionTitle, similarSourceImageId, similarHasMore, selectedFolderId } = get(); + const { loadedCount, totalImages, loadingImages, collectionTitle, similarSourceImageId, similarHasMore, similarFolderId } = get(); if (loadingImages || loadedCount >= totalImages) return; + if (collectionTitle === "Explore Cluster") return; if (collectionTitle === "Similar Images" && similarSourceImageId !== null) { if (!similarHasMore) return; - await get().loadSimilarImages(similarSourceImageId, selectedFolderId, false); + await get().loadSimilarImages(similarSourceImageId, similarFolderId, false); return; } await get().loadImages(false); @@ -614,6 +815,11 @@ export const useGalleryStore = create((set, get) => ({ void get().loadImages(true); }, + setMinimumRating: (minimumRating) => { + set({ minimumRating, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); + void get().loadImages(true); + }, + setFailedEmbeddingsOnly: (failedEmbeddingsOnly) => { set({ failedEmbeddingsOnly, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); @@ -626,59 +832,148 @@ export const useGalleryStore = create((set, get) => ({ setView: (activeView) => set({ activeView }), + setExploreMode: (exploreMode) => set({ exploreMode }), + loadTagCloud: async () => { const { selectedFolderId, tagCloudFolderId, tagCloudLoading } = get(); // Skip if already loaded for this folder and not currently loading if (!tagCloudLoading && tagCloudFolderId !== undefined && tagCloudFolderId === selectedFolderId) { return; } + const requestToken = ++tagCloudRequestToken; set({ tagCloudLoading: true, tagCloudFolderId: selectedFolderId }); try { const entries = await invoke("get_tag_cloud", { folderId: selectedFolderId, }); + if (requestToken !== tagCloudRequestToken) return; set({ tagCloudEntries: entries, tagCloudLoading: false }); } catch (error) { + if (requestToken !== tagCloudRequestToken) return; console.error("Failed to load tag cloud:", error); set({ tagCloudLoading: false }); } }, - searchByTag: (imageId) => { - const { selectedFolderId } = get(); - set((state) => ({ activeView: "gallery", images: [], loadedCount: 0, loadingImages: true, collectionTitle: "Similar Images", imageLoadError: null, galleryScrollResetKey: state.galleryScrollResetKey + 1 })); - void get().loadSimilarImages(imageId, selectedFolderId); + loadExploreTags: async () => { + const { selectedFolderId, exploreTagsFolderId, exploreTagLoading } = get(); + if (!exploreTagLoading && exploreTagsFolderId !== undefined && exploreTagsFolderId === selectedFolderId) { + return; + } + const requestToken = ++exploreTagRequestToken; + set({ exploreTagLoading: true, exploreTagsFolderId: selectedFolderId }); + try { + const entries = await invoke("get_explore_tags", { + params: { folder_id: selectedFolderId, limit: 48 }, + }); + if (requestToken !== exploreTagRequestToken) return; + set({ exploreTagEntries: entries, exploreTagLoading: false }); + } catch (error) { + if (requestToken !== exploreTagRequestToken) return; + console.error("Failed to load explore tags:", error); + set({ exploreTagLoading: false }); + } }, - loadSimilarImages: async (imageId, folderId = get().selectedFolderId, reset = true) => { - const requestedLimit = reset ? PAGE_SIZE : get().loadedCount + PAGE_SIZE; + showVisualCluster: async (imageIds) => { + const requestToken = ++similarRequestToken; set((state) => ({ - images: reset ? [] : get().images, - loadedCount: reset ? 0 : get().loadedCount, + activeView: "gallery", + search: "", + images: [], + totalImages: imageIds.length, + loadedCount: 0, + loadingImages: true, + collectionTitle: "Explore Cluster", + imageLoadError: null, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarFolderId: null, + galleryScrollResetKey: state.galleryScrollResetKey + 1, + })); + + try { + const images = await invoke("get_images_by_ids", { + params: { image_ids: imageIds }, + }); + if (requestToken !== similarRequestToken) return; + set({ + images, + totalImages: images.length, + loadedCount: images.length, + loadingImages: false, + imageLoadError: null, + collectionTitle: "Explore Cluster", + }); + } catch (error) { + if (requestToken !== similarRequestToken) return; + set({ + images: [], + totalImages: 0, + loadedCount: 0, + loadingImages: false, + imageLoadError: String(error), + collectionTitle: "Explore Cluster", + }); + } + }, + + searchForTag: (tag) => { + set({ activeView: "gallery", search: `/t ${tag}`, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, similarFolderId: null, imageLoadError: null }); + void get().loadImages(true); + }, + + loadSimilarImages: async (imageId, folderId = get().selectedFolderId, reset = true, sourceFolderId = folderId ?? null) => { + const requestToken = ++similarRequestToken; + const offset = reset ? 0 : get().loadedCount; + const similarScope = folderId === null ? "all_media" : "current_folder"; + set((state) => ({ + images: reset ? [] : state.images, + loadedCount: reset ? 0 : state.loadedCount, loadingImages: true, collectionTitle: "Similar Images", imageLoadError: null, similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarFolderId: folderId ?? null, + similarScope, galleryScrollResetKey: reset ? state.galleryScrollResetKey + 1 : state.galleryScrollResetKey, })); - try { - const images = await invoke("find_similar_images", { - params: { image_id: imageId, folder_id: folderId ?? null, limit: requestedLimit }, - }); - const hasMore = images.length >= requestedLimit; - set({ - images, - totalImages: hasMore ? images.length + PAGE_SIZE : images.length, - loadedCount: images.length, - loadingImages: false, - imageLoadError: null, - collectionTitle: "Similar Images", - similarSourceImageId: imageId, - similarHasMore: hasMore, - selectedFolderId: folderId ?? null, - selectedImage: reset ? null : get().selectedImage, + + try { + const result = await invoke("find_similar_images", { + params: { + image_id: imageId, + folder_id: folderId ?? null, + offset, + limit: PAGE_SIZE, + threshold: SIMILAR_DISTANCE_THRESHOLD, + }, + }); + + if (requestToken !== similarRequestToken) return; + + set((state) => { + const nextImages = reset ? result.images : [...state.images, ...result.images]; + const nextLoadedCount = nextImages.length; + return { + images: nextImages, + totalImages: result.has_more ? nextLoadedCount + 1 : nextLoadedCount, + loadedCount: nextLoadedCount, + loadingImages: false, + imageLoadError: null, + collectionTitle: "Similar Images", + similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarHasMore: result.has_more, + similarFolderId: folderId ?? null, + similarScope, + selectedImage: reset ? null : state.selectedImage, + }; }); } catch (error) { + if (requestToken !== similarRequestToken) return; console.error("Failed to load similar images:", error); set({ images: [], @@ -688,13 +983,23 @@ export const useGalleryStore = create((set, get) => ({ imageLoadError: String(error), collectionTitle: "Similar Images", similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, similarHasMore: false, - selectedFolderId: folderId ?? null, + similarFolderId: folderId ?? null, + similarScope, selectedImage: null, }); } }, + setSimilarScope: (similarScope) => { + set({ similarScope }); + const { similarSourceImageId, similarSourceFolderId, selectedFolderId } = get(); + if (similarSourceImageId === null) return; + const folderId = similarScope === "current_folder" ? (similarSourceFolderId ?? selectedFolderId) : null; + void get().loadSimilarImages(similarSourceImageId, folderId, true, similarSourceFolderId); + }, + suggestImageTags: async (imageId) => { return invoke("suggest_image_tags", { params: { image_id: imageId, limit: 2 }, @@ -833,6 +1138,27 @@ export const useGalleryStore = create((set, get) => ({ setSettingsOpen: (settingsOpen) => set({ settingsOpen }), + setTaggingQueueScope: (taggingQueueScope) => { + set((state) => ({ + taggingQueueScope, + taggingQueueFolderIds: + taggingQueueScope === "selected" && state.taggingQueueFolderIds.length === 0 && state.folders.length > 0 + ? [state.folders[0].id] + : state.taggingQueueFolderIds, + })); + }, + + toggleTaggingQueueFolder: (folderId) => { + set((state) => { + const next = state.taggingQueueFolderIds.includes(folderId) + ? state.taggingQueueFolderIds.filter((id) => id !== folderId) + : [...state.taggingQueueFolderIds, folderId].sort((a, b) => a - b); + return { taggingQueueFolderIds: next }; + }); + }, + + setTaggingQueueFolderIds: (taggingQueueFolderIds) => set({ taggingQueueFolderIds }), + loadTaggerModelStatus: async () => { try { const taggerModelStatus = await invoke("get_tagger_model_status"); @@ -874,6 +1200,22 @@ export const useGalleryStore = create((set, get) => ({ set({ taggerThreshold }); }, + loadTaggerBatchSize: async () => { + try { + const taggerBatchSize = await invoke("get_tagger_batch_size"); + set({ taggerBatchSize }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + setTaggerBatchSize: async (batchSize) => { + const taggerBatchSize = await invoke("set_tagger_batch_size", { + params: { batch_size: batchSize }, + }); + set({ taggerBatchSize }); + }, + prepareTaggerModel: async () => { set({ taggerModelPreparing: true, taggerModelError: null, taggerModelProgress: null }); try { @@ -912,6 +1254,14 @@ export const useGalleryStore = create((set, get) => ({ return queued; }, + queueTaggingJobsForFolders: async (folderIds) => { + const queued = await invoke("queue_tagging_jobs", { + params: { folder_id: null, folder_ids: folderIds, image_id: null }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + queueTaggingForImage: async (imageId) => { const queued = await invoke("queue_tagging_jobs", { params: { folder_id: null, image_id: imageId }, @@ -928,6 +1278,14 @@ export const useGalleryStore = create((set, get) => ({ return cleared; }, + clearTaggingJobsForFolders: async (folderIds) => { + const cleared = await invoke("clear_tagging_jobs", { + params: { folder_id: null, folder_ids: folderIds }, + }); + await get().loadBackgroundJobProgress(); + return cleared; + }, + getImageTags: async (imageId) => { return invoke("get_image_tags", { params: { image_id: imageId }, @@ -946,6 +1304,73 @@ export const useGalleryStore = create((set, get) => ({ }); }, + loadDuplicateScanCache: async (folderId = null) => { + interface CacheResult { groups: DuplicateGroup[]; scanned_at: number } + const cached = await invoke("load_duplicate_scan_cache", { folderId: folderId ?? null }); + if (cached) { + set({ duplicateGroups: cached.groups, duplicateLastScanned: cached.scanned_at }); + } + }, + + scanDuplicates: async (folderId = null) => { + const { listen } = await import("@tauri-apps/api/event"); + set({ duplicateScanning: true, duplicateGroups: [], duplicateScanProgress: null, duplicateSelectedIds: new Set() }); + const unlisten = await listen<[number, number]>("duplicate_scan_progress", (event) => { + const [scanned, total] = event.payload; + set({ duplicateScanProgress: { scanned, total } }); + }); + try { + const groups = await invoke("find_duplicates", { folderId: folderId ?? null }); + set({ duplicateGroups: groups, duplicateLastScanned: Math.floor(Date.now() / 1000) }); + } finally { + unlisten(); + set({ duplicateScanning: false }); + } + }, + + toggleDuplicateSelected: (imageId) => { + set((state) => { + const next = new Set(state.duplicateSelectedIds); + if (next.has(imageId)) next.delete(imageId); + else next.add(imageId); + return { duplicateSelectedIds: next }; + }); + }, + + selectAllDuplicates: (imageIds) => { + set((state) => { + const next = new Set(state.duplicateSelectedIds); + for (const id of imageIds) next.add(id); + return { duplicateSelectedIds: next }; + }); + }, + + selectKeepFirstAllGroups: () => { + const { duplicateGroups } = get(); + const toMark = new Set(); + for (const group of duplicateGroups) { + for (const img of group.images.slice(1)) toMark.add(img.id); + } + set({ duplicateSelectedIds: toMark }); + }, + + clearDuplicateSelection: () => set({ duplicateSelectedIds: new Set() }), + + deleteSelectedDuplicates: async () => { + const { duplicateSelectedIds } = get(); + const ids = Array.from(duplicateSelectedIds); + if (ids.length === 0) return 0; + const deleted = await invoke("delete_images_from_disk", { params: { image_ids: ids } }); + // Remove deleted images from groups and drop now-trivial groups + set((state) => ({ + duplicateSelectedIds: new Set(), + duplicateGroups: state.duplicateGroups + .map((g) => ({ ...g, images: g.images.filter((img) => !duplicateSelectedIds.has(img.id)) })) + .filter((g) => g.images.length > 1), + })); + return deleted; + }, + retryFailedEmbeddings: async (folderId) => { await invoke("retry_failed_embeddings", { params: { folder_id: folderId } }); await get().loadBackgroundJobProgress(); @@ -979,7 +1404,9 @@ export const useGalleryStore = create((set, get) => ({ if (progress.done) { void get().loadFolders(); void get().loadBackgroundJobProgress(); - void get().loadImages(true); + if (get().activeView !== "explore" && !isDerivedCollectionTitle(get().collectionTitle)) { + void get().loadImages(true); + } setTimeout(() => { set((state) => { @@ -1019,12 +1446,17 @@ export const useGalleryStore = create((set, get) => ({ const batch = event.payload; set((state) => { + if (isDerivedCollectionTitle(state.collectionTitle) || state.activeView === "explore") { + return state; + } + const visibleImages = batch.images.filter((image) => matchesFilters( image, state.selectedFolderId, state.mediaFilter, state.favoritesOnly, + state.minimumRating, state.failedEmbeddingsOnly, state.search, ), @@ -1049,12 +1481,21 @@ export const useGalleryStore = create((set, get) => ({ const batch = event.payload; set((state) => { + if (isDerivedCollectionTitle(state.collectionTitle) || state.activeView === "explore") { + const selectedImage = + state.selectedImage && batch.images.some((image) => image.id === state.selectedImage?.id) + ? batch.images.find((image) => image.id === state.selectedImage?.id) ?? state.selectedImage + : state.selectedImage; + return { selectedImage }; + } + const visibleImages = batch.images.filter((image) => matchesFilters( image, state.selectedFolderId, state.mediaFilter, state.favoritesOnly, + state.minimumRating, state.failedEmbeddingsOnly, state.search, ), -- 2.52.0 From df1749780841b218660c7f2e85a2e31f18615ad9 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Tue, 2 Jun 2026 19:39:31 +0100 Subject: [PATCH 06/25] feat: add region-based similarity search Adds a "Search within image" button to the lightbox that lets the user draw a crop region on an image and find visually similar results using that crop's embedding. The crop is embedded in-memory without a temp file via a new embed_image_crop method on ClipImageEmbedder. Introduces a folder-scoped cosine search (search_image_ids_by_embedding_in_folder) to support the current-folder scope option, and wires up the new find_similar_by_region Tauri command end-to-end from Rust through to the Zustand store and Lightbox UI. --- src-tauri/src/commands.rs | 73 ++++++++++ src-tauri/src/embedder.rs | 45 ++++++ src-tauri/src/lib.rs | 1 + src-tauri/src/vector.rs | 34 +++++ src/components/Lightbox.tsx | 281 +++++++++++++++++++++++++++++++++--- src/store.ts | 67 +++++++++ 6 files changed, 477 insertions(+), 24 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index d23eeb4..3c1a8ee 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -59,6 +59,19 @@ pub struct FindSimilarImagesParams { pub threshold: Option, } +#[derive(Deserialize)] +pub struct FindSimilarByRegionParams { + pub image_id: i64, + /// Normalized crop rect (0.0–1.0). + pub crop_x: f32, + pub crop_y: f32, + pub crop_w: f32, + pub crop_h: f32, + pub folder_id: Option, + pub offset: Option, + pub limit: Option, +} + #[derive(Deserialize)] pub struct DebugSimilarImagesParams { pub image_id: i64, @@ -354,6 +367,66 @@ pub async fn find_similar_images( }) } +#[tauri::command] +pub async fn find_similar_by_region( + db: State<'_, DbState>, + params: FindSimilarByRegionParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let limit = params.limit.unwrap_or(32); + let offset = params.offset.unwrap_or(0); + + // Look up the source image path + let image = db::get_image_by_id(&conn, params.image_id).map_err(|e| e.to_string())?; + let image_path = std::path::Path::new(&image.path); + + // Embed the cropped region in-memory (no temp file needed) + let embedder = embedder::ClipImageEmbedder::new().map_err(|e| e.to_string())?; + let embedding = embedder + .embed_image_crop( + image_path, + params.crop_x, + params.crop_y, + params.crop_w, + params.crop_h, + ) + .map_err(|e| e.to_string())?; + + // Search for similar images using the crop embedding + let image_ids = match params.folder_id { + Some(folder_id) => vector::search_image_ids_by_embedding_in_folder( + &conn, + &embedding, + folder_id, + Some(params.image_id), + offset + limit + 1, + ) + .map_err(|e| e.to_string())?, + None => { + let mut ids = vector::search_image_ids_by_embedding(&conn, &embedding, offset + limit + 1) + .map_err(|e| e.to_string())?; + // Exclude the source image from global results + ids.retain(|&id| id != params.image_id); + ids + } + }; + + let has_more = image_ids.len() > offset + limit; + let page_ids = image_ids + .into_iter() + .skip(offset) + .take(limit) + .collect::>(); + + let images = db::get_images_by_ids(&conn, &page_ids).map_err(|e| e.to_string())?; + Ok(SimilarImagesPage { + images, + offset, + limit, + has_more, + }) +} + #[derive(Serialize)] pub struct SimilarImagesDebug { pub image_id: i64, diff --git a/src-tauri/src/embedder.rs b/src-tauri/src/embedder.rs index 532417d..3a977c7 100644 --- a/src-tauri/src/embedder.rs +++ b/src-tauri/src/embedder.rs @@ -74,6 +74,51 @@ impl ClipImageEmbedder { Ok(self.embed_images(&[path.to_path_buf()])?.remove(0)) } + /// Embed a cropped region of an image without writing a temp file to disk. + /// `crop_x`, `crop_y`, `crop_w`, `crop_h` are normalized 0.0–1.0 coordinates. + pub fn embed_image_crop( + &self, + path: &Path, + crop_x: f32, + crop_y: f32, + crop_w: f32, + crop_h: f32, + ) -> Result> { + let img = image::ImageReader::open(path)? + .with_guessed_format()? + .decode()?; + + let img_w = img.width() as f32; + let img_h = img.height() as f32; + + let x = ((crop_x * img_w) as u32).min(img.width().saturating_sub(1)); + let y = ((crop_y * img_h) as u32).min(img.height().saturating_sub(1)); + let w = ((crop_w * img_w) as u32).max(1).min(img.width() - x); + let h = ((crop_h * img_h) as u32).max(1).min(img.height() - y); + + let cropped = img.crop_imm(x, y, w, h); + let resized = cropped.resize_to_fill( + self.image_size as u32, + self.image_size as u32, + image::imageops::FilterType::Triangle, + ); + + let raw = resized.to_rgb8().into_raw(); + let tensor = candle_core::Tensor::from_vec( + raw, + (self.image_size, self.image_size, 3), + &candle_core::Device::Cpu, + )? + .permute((2, 0, 1))? + .to_dtype(candle_core::DType::F32)? + .affine(2.0 / 255.0, -1.0)?; + + let batch = tensor.unsqueeze(0)?.to_device(&self.device)?; + let features = self.model.get_image_features(&batch)?; + let normalized = candle_transformers::models::clip::div_l2_norm(&features)?; + Ok(normalized.get(0)?.flatten_all()?.to_vec1::()?) + } + pub fn embed_images(&self, paths: &[PathBuf]) -> Result>> { let images = load_images(paths, self.image_size)?.to_device(&self.device)?; let features = self.model.get_image_features(&images)?; diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 39c9922..0ef2241 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -87,6 +87,7 @@ pub fn run() { commands::reindex_folder, commands::update_image_details, commands::find_similar_images, + commands::find_similar_by_region, commands::debug_similar_images, commands::retry_failed_embeddings, commands::semantic_search_images, diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 0fa753c..644ece8 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -310,6 +310,40 @@ pub fn search_image_ids_by_embedding( Ok(ids) } +/// Brute-force cosine search scoped to a single folder, ordered by ascending distance. +/// Used for region-based similarity search where we want folder-scoped results. +pub fn search_image_ids_by_embedding_in_folder( + conn: &Connection, + embedding: &[f32], + folder_id: i64, + exclude_image_id: Option, + limit: usize, +) -> Result> { + if embedding.len() != CLIP_VECTOR_DIM { + return Err(anyhow!( + "expected {}-dimensional embedding, got {}", + CLIP_VECTOR_DIM, + embedding.len() + )); + } + + let packed = pack_f32(embedding); + let exclude_id = exclude_image_id.unwrap_or(-1); + let mut stmt = conn.prepare( + "SELECT v.image_id + FROM image_vec v + JOIN images i ON i.id = v.image_id + WHERE i.folder_id = ?2 + AND v.image_id != ?3 + ORDER BY vec_distance_cosine(v.embedding, vec_f32(?1)) ASC + LIMIT ?4", + )?; + let rows = stmt.query_map((&packed, folder_id, exclude_id, limit as i64), |row| { + row.get::<_, i64>(0) + })?; + Ok(rows.collect::>>()?) +} + #[allow(dead_code)] pub fn search_caption_ids_by_embedding( conn: &Connection, diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index b357d0f..0ee42f5 100644 --- a/src/components/Lightbox.tsx +++ b/src/components/Lightbox.tsx @@ -58,12 +58,69 @@ function ratingPill(rating: AiRating): { label: string; className: string } { } } +interface DragRect { + startX: number; + startY: number; + endX: number; + endY: number; +} + +/** Compute a CSS-pixel rect (relative to the viewport container) from a DragRect. */ +function normaliseRect(r: DragRect): { left: number; top: number; width: number; height: number } { + return { + left: Math.min(r.startX, r.endX), + top: Math.min(r.startY, r.endY), + width: Math.abs(r.endX - r.startX), + height: Math.abs(r.endY - r.startY), + }; +} + +/** Convert a CSS-pixel drag rect (relative to viewport container) to normalised 0–1 crop coords + * relative to the actual rendered element bounds. */ +function rectToNormalisedCrop( + rect: DragRect, + imgEl: HTMLImageElement, +): { x: number; y: number; w: number; h: number } | null { + const imgBounds = imgEl.getBoundingClientRect(); + if (imgBounds.width === 0 || imgBounds.height === 0) return null; + + // rect coords are already in viewport space (client coords) + const rawX = Math.min(rect.startX, rect.endX); + const rawY = Math.min(rect.startY, rect.endY); + const rawW = Math.abs(rect.endX - rect.startX); + const rawH = Math.abs(rect.endY - rect.startY); + + // Clamp to image bounds + const clampedX = Math.max(rawX, imgBounds.left); + const clampedY = Math.max(rawY, imgBounds.top); + const clampedRight = Math.min(rawX + rawW, imgBounds.right); + const clampedBottom = Math.min(rawY + rawH, imgBounds.bottom); + + const croppedW = clampedRight - clampedX; + const croppedH = clampedBottom - clampedY; + + if (croppedW <= 0 || croppedH <= 0) return null; + + // Normalize by the CSS transform scale — getBoundingClientRect already returns + // the scaled (on-screen) size, so we normalize directly against that. + return { + x: (clampedX - imgBounds.left) / imgBounds.width, + y: (clampedY - imgBounds.top) / imgBounds.height, + w: croppedW / imgBounds.width, + h: croppedH / imgBounds.height, + }; +} + +/** Minimum selection size as a fraction of the viewport container dimension. */ +const MIN_SELECTION_FRACTION = 0.02; + export function Lightbox() { const selectedImage = useGalleryStore((state) => state.selectedImage); const closeImage = useGalleryStore((state) => state.closeImage); const images = useGalleryStore((state) => state.images); const openImage = useGalleryStore((state) => state.openImage); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); + const loadSimilarByRegion = useGalleryStore((state) => state.loadSimilarByRegion); const similarScope = useGalleryStore((state) => state.similarScope); const updateImageDetails = useGalleryStore((state) => state.updateImageDetails); const getImageTags = useGalleryStore((state) => state.getImageTags); @@ -71,16 +128,26 @@ export function Lightbox() { const removeTag = useGalleryStore((state) => state.removeTag); const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus); const queueTaggingForImage = useGalleryStore((state) => state.queueTaggingForImage); + const [zoom, setZoom] = useState(1); const [imageTags, setImageTags] = useState([]); const [tagInput, setTagInput] = useState(""); const [tagAdding, setTagAdding] = useState(false); const [tagsExpanded, setTagsExpanded] = useState(false); const [taggingQueued, setTaggingQueued] = useState(false); + + // Region selection state + const [regionSelectMode, setRegionSelectMode] = useState(false); + const [isDragging, setIsDragging] = useState(false); + const [dragRect, setDragRect] = useState(null); + const [regionSearching, setRegionSearching] = useState(false); + const imageViewportRef = useRef(null); + const imgRef = useRef(null); const currentIndex = selectedImage ? images.findIndex((image) => image.id === selectedImage.id) : -1; const canFindSimilar = selectedImage?.embedding_status === "ready"; + const canSearchRegion = canFindSimilar && selectedImage?.media_kind === "image"; const goPrev = useCallback(() => { if (currentIndex > 0) openImage(images[currentIndex - 1]); @@ -90,13 +157,21 @@ export function Lightbox() { if (currentIndex >= 0 && currentIndex < images.length - 1) openImage(images[currentIndex + 1]); }, [currentIndex, images, openImage]); + const exitRegionMode = useCallback(() => { + setRegionSelectMode(false); + setIsDragging(false); + setDragRect(null); + }, []); + useEffect(() => { setZoom(1); setImageTags([]); setTagInput(""); setTagsExpanded(false); setTaggingQueued(false); - }, [selectedImage?.id]); + exitRegionMode(); + setRegionSearching(false); + }, [selectedImage?.id, exitRegionMode]); useEffect(() => { if (!selectedImage) return; @@ -115,6 +190,7 @@ export function Lightbox() { if (!viewport || !selectedImage || selectedImage.media_kind !== "image") return; const handleWheel = (event: WheelEvent) => { + if (regionSelectMode) return; // don't zoom during selection if (!event.ctrlKey && Math.abs(event.deltaY) < Math.abs(event.deltaX)) return; event.preventDefault(); setZoom((value) => { @@ -125,12 +201,19 @@ export function Lightbox() { viewport.addEventListener("wheel", handleWheel, { passive: false }); return () => viewport.removeEventListener("wheel", handleWheel); - }, [selectedImage]); + }, [selectedImage, regionSelectMode]); useEffect(() => { const handler = (event: KeyboardEvent) => { if (!selectedImage) return; - if (event.key === "Escape") closeImage(); + if (event.key === "Escape") { + if (regionSelectMode) { + exitRegionMode(); + } else { + closeImage(); + } + } + if (regionSelectMode) return; // block nav keys during selection if (event.key === "ArrowLeft") goPrev(); if (event.key === "ArrowRight") goNext(); if (event.key === "+" || event.key === "=") setZoom((value) => Math.min(3, value + 0.25)); @@ -139,7 +222,75 @@ export function Lightbox() { window.addEventListener("keydown", handler); return () => window.removeEventListener("keydown", handler); - }, [selectedImage, closeImage, goPrev, goNext]); + }, [selectedImage, closeImage, goPrev, goNext, regionSelectMode, exitRegionMode]); + + // ── Region selection pointer handlers ─────────────────────────────────────── + const handleRegionPointerDown = useCallback( + (event: React.PointerEvent) => { + if (!regionSelectMode) return; + event.preventDefault(); + event.currentTarget.setPointerCapture(event.pointerId); + setIsDragging(true); + setDragRect({ + startX: event.clientX, + startY: event.clientY, + endX: event.clientX, + endY: event.clientY, + }); + }, + [regionSelectMode], + ); + + const handleRegionPointerMove = useCallback( + (event: React.PointerEvent) => { + if (!isDragging) return; + setDragRect((prev) => + prev ? { ...prev, endX: event.clientX, endY: event.clientY } : null, + ); + }, + [isDragging], + ); + + const handleRegionPointerUp = useCallback( + (event: React.PointerEvent) => { + if (!isDragging || !dragRect || !selectedImage || !imgRef.current) { + setIsDragging(false); + return; + } + event.currentTarget.releasePointerCapture(event.pointerId); + + const finalRect: DragRect = { ...dragRect, endX: event.clientX, endY: event.clientY }; + const crop = rectToNormalisedCrop(finalRect, imgRef.current); + + setIsDragging(false); + setDragRect(null); + + // Ignore tiny accidental clicks + const containerBounds = imageViewportRef.current?.getBoundingClientRect(); + const containerSize = containerBounds + ? Math.min(containerBounds.width, containerBounds.height) + : 500; + const selW = Math.abs(finalRect.endX - finalRect.startX); + const selH = Math.abs(finalRect.endY - finalRect.startY); + if (!crop || selW < containerSize * MIN_SELECTION_FRACTION || selH < containerSize * MIN_SELECTION_FRACTION) { + exitRegionMode(); + return; + } + + exitRegionMode(); + setRegionSearching(true); + + const folderId = + similarScope === "current_folder" ? selectedImage.folder_id : null; + void loadSimilarByRegion(selectedImage.id, crop, folderId, selectedImage.folder_id) + .finally(() => setRegionSearching(false)); + }, + [isDragging, dragRect, selectedImage, similarScope, loadSimilarByRegion, exitRegionMode], + ); + + // Build the CSS rect for the selection overlay (viewport-relative) + const selectionOverlay = + isDragging && dragRect ? normaliseRect(dragRect) : null; return ( @@ -151,11 +302,11 @@ export function Lightbox() { animate={{ opacity: 1 }} exit={{ opacity: 0 }} transition={{ duration: 0.15 }} - onClick={closeImage} + onClick={regionSelectMode ? undefined : closeImage} > - {Math.round(zoom * 100)}% - + {!regionSelectMode && ( +
+
+ + {Math.round(zoom * 100)}% + +
-
+ )} )} @@ -261,6 +447,53 @@ export function Lightbox() {
+ {/* Search region button row */} + {canSearchRegion && ( +
+ +
+ )} +

Rating

@@ -472,7 +705,7 @@ export function Lightbox() { - -
+ + +
+ ) : ( +
+ + +
+ )}
); } diff --git a/src/store.ts b/src/store.ts index 3fc221c..1aba18c 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1360,15 +1360,20 @@ export const useGalleryStore = create((set, get) => ({ }, addUserTag: async (imageId, tag) => { - return invoke("add_user_tag", { + const result = await invoke("add_user_tag", { params: { image_id: imageId, tag }, }); + // Invalidate explore tags cache so new tag appears immediately + set({ exploreTagsFolderId: undefined }); + return result; }, removeTag: async (tagId) => { await invoke("remove_tag", { params: { tag_id: tagId }, }); + // Invalidate explore tags cache so removed tag disappears immediately + set({ exploreTagsFolderId: undefined }); }, loadDuplicateScanCache: async (folderId = null) => { -- 2.52.0 From c1070649fa67872b838fc66be52e6863e723baca Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sat, 6 Jun 2026 19:52:24 +0100 Subject: [PATCH 11/25] feat(notifications): report completed background tasks Register the Tauri notification plugin and request notification permission during startup. Send completion notifications for folder scans, embeddings, AI tagging, and duplicate scans, including failure counts where available. --- package.json | 3 +- pnpm-lock.yaml | 10 +++++ src-tauri/Cargo.lock | 69 ++++++++++++++++++++++++++++- src-tauri/Cargo.toml | 1 + src-tauri/capabilities/default.json | 1 + src-tauri/src/lib.rs | 1 + src/App.tsx | 2 + src/notifications.ts | 31 +++++++++++++ src/store.ts | 50 +++++++++++++++++++++ 9 files changed, 166 insertions(+), 2 deletions(-) create mode 100644 src/notifications.ts diff --git a/package.json b/package.json index 5d541c0..1836830 100644 --- a/package.json +++ b/package.json @@ -8,7 +8,7 @@ "build:vite": "tsc && vite build", "dev:app": "tauri dev", "dev:vite": "vite", -"preview": "vite preview", + "preview": "vite preview", "tauri": "tauri" }, "dependencies": { @@ -16,6 +16,7 @@ "@tauri-apps/api": "^2", "@tauri-apps/plugin-dialog": "^2.7.0", "@tauri-apps/plugin-fs": "^2.5.0", + "@tauri-apps/plugin-notification": "^2.3.3", "@tauri-apps/plugin-opener": "^2", "d3-force": "^3.0.0", "framer-motion": "^12.38.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b2dd733..ec0add6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -20,6 +20,9 @@ importers: '@tauri-apps/plugin-fs': specifier: ^2.5.0 version: 2.5.0 + '@tauri-apps/plugin-notification': + specifier: ^2.3.3 + version: 2.3.3 '@tauri-apps/plugin-opener': specifier: ^2 version: 2.5.3 @@ -653,6 +656,9 @@ packages: '@tauri-apps/plugin-fs@2.5.0': resolution: {integrity: sha512-c83kbz61AK+rKjhS+je9+stIO27nXj7p9cqeg36TwkIUtxpCFTttlHHtqon6h6FN54cXjyAjlMPOJcW3mwE5XQ==} + '@tauri-apps/plugin-notification@2.3.3': + resolution: {integrity: sha512-Zw+ZH18RJb41G4NrfHgIuofJiymusqN+q8fGUIIV7vyCH+5sSn5coqRv/MWB9qETsUs97vmU045q7OyseCV3Qg==} + '@tauri-apps/plugin-opener@2.5.3': resolution: {integrity: sha512-CCcUltXMOfUEArbf3db3kCE7Ggy1ExBEBl51Ko2ODJ6GDYHRp1nSNlQm5uNCFY5k7/ufaK5Ib3Du/Zir19IYQQ==} @@ -1448,6 +1454,10 @@ snapshots: dependencies: '@tauri-apps/api': 2.10.1 + '@tauri-apps/plugin-notification@2.3.3': + dependencies: + '@tauri-apps/api': 2.10.1 + '@tauri-apps/plugin-opener@2.5.3': dependencies: '@tauri-apps/api': 2.10.1 diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 195eddf..edd9078 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -3293,6 +3293,18 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" +[[package]] +name = "mac-notification-sys" +version = "0.6.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50efa634682b3fc5a1ab6f3dd5b2bce7b848011fc485b53b063dc68f2f74feae" +dependencies = [ + "cc", + "objc2", + "objc2-foundation", + "time", +] + [[package]] name = "mach2" version = "0.4.3" @@ -3594,6 +3606,20 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "notify-rust" +version = "4.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50ff2e74231b72c832d82982193b417f230945be6bdb5575b251d941d31adb00" +dependencies = [ + "futures-lite", + "log", + "mac-notification-sys", + "serde", + "tauri-winrt-notification", + "zbus", +] + [[package]] name = "ntapi" version = "0.4.3" @@ -4279,6 +4305,7 @@ dependencies = [ "tauri-build", "tauri-plugin-dialog", "tauri-plugin-fs", + "tauri-plugin-notification", "tauri-plugin-opener", "tokenizers", "tokio", @@ -4326,7 +4353,7 @@ checksum = "740ebea15c5d1428f910cd1a5f52cebf8d25006245ed8ade92702f4943d91e07" dependencies = [ "base64 0.22.1", "indexmap 2.13.1", - "quick-xml", + "quick-xml 0.38.4", "serde", "time", ] @@ -4543,6 +4570,15 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3" +[[package]] +name = "quick-xml" +version = "0.37.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "331e97a1af0bf59823e6eadffe373d7b27f485be8748f71471c662c1f269b7fb" +dependencies = [ + "memchr", +] + [[package]] name = "quick-xml" version = "0.38.4" @@ -6016,6 +6052,25 @@ dependencies = [ "url", ] +[[package]] +name = "tauri-plugin-notification" +version = "2.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01fc2c5ff41105bd1f7242d8201fdf3efd70749b82fa013a17f2126357d194cc" +dependencies = [ + "log", + "notify-rust", + "rand 0.9.2", + "serde", + "serde_json", + "serde_repr", + "tauri", + "tauri-plugin", + "thiserror 2.0.18", + "time", + "url", +] + [[package]] name = "tauri-plugin-opener" version = "2.5.3" @@ -6138,6 +6193,18 @@ dependencies = [ "toml 0.9.12+spec-1.1.0", ] +[[package]] +name = "tauri-winrt-notification" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b1e66e07de489fe43a46678dd0b8df65e0c973909df1b60ba33874e297ba9b9" +dependencies = [ + "quick-xml 0.37.5", + "thiserror 2.0.18", + "windows 0.61.3", + "windows-version", +] + [[package]] name = "tempfile" version = "3.27.0" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 2b2919e..2bd2dde 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -50,6 +50,7 @@ ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "n ureq = { version = "3.3.0", default-features = false, features = ["native-tls"] } zip = { version = "4.6.1", default-features = false, features = ["deflate"] } csv = "1" +tauri-plugin-notification = "2" # ── Dev-mode performance ──────────────────────────────────────────────────── # opt-level=1 on the main crate keeps incremental compile short. diff --git a/src-tauri/capabilities/default.json b/src-tauri/capabilities/default.json index c3a561a..a526b2c 100644 --- a/src-tauri/capabilities/default.json +++ b/src-tauri/capabilities/default.json @@ -13,6 +13,7 @@ "fs:allow-read-dir", "fs:read-files", "fs:read-dirs", + "notification:default", "core:window:allow-minimize", "core:window:allow-close", "core:window:allow-toggle-maximize", diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 0ef2241..8c820db 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -19,6 +19,7 @@ pub fn run() { .plugin(tauri_plugin_opener::init()) .plugin(tauri_plugin_dialog::init()) .plugin(tauri_plugin_fs::init()) + .plugin(tauri_plugin_notification::init()) .setup(|app| { let app_dir = app .path() diff --git a/src/App.tsx b/src/App.tsx index 2486ba9..ce2a339 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -9,6 +9,7 @@ import { TagCloud } from "./components/TagCloud"; import { DuplicateFinder } from "./components/DuplicateFinder"; import { TitleBar } from "./components/TitleBar"; import { SettingsModal } from "./components/SettingsModal"; +import { initializeNotifications } from "./notifications"; export default function App() { const loadFolders = useGalleryStore((state) => state.loadFolders); @@ -20,6 +21,7 @@ export default function App() { const activeView = useGalleryStore((state) => state.activeView); useEffect(() => { + void initializeNotifications(); loadFolders().then(() => { void loadBackgroundJobProgress(); void loadCaptionModelStatus(); diff --git a/src/notifications.ts b/src/notifications.ts new file mode 100644 index 0000000..ccfcac8 --- /dev/null +++ b/src/notifications.ts @@ -0,0 +1,31 @@ +import { + isPermissionGranted, + requestPermission, + sendNotification, +} from "@tauri-apps/plugin-notification"; + +let permissionPromise: Promise | null = null; + +export function initializeNotifications(): Promise { + permissionPromise ??= (async () => { + try { + if (await isPermissionGranted()) return true; + return (await requestPermission()) === "granted"; + } catch (error) { + console.warn("Windows notifications are unavailable:", error); + return false; + } + })(); + + return permissionPromise; +} + +export async function notifyTaskComplete(title: string, body: string): Promise { + if (!(await initializeNotifications())) return; + + try { + sendNotification({ title, body }); + } catch (error) { + console.warn("Could not send task completion notification:", error); + } +} diff --git a/src/store.ts b/src/store.ts index 1aba18c..ac37299 100644 --- a/src/store.ts +++ b/src/store.ts @@ -2,6 +2,7 @@ import { create } from "zustand"; import { invoke } from "@tauri-apps/api/core"; import { listen, UnlistenFn } from "@tauri-apps/api/event"; import { appDataDir, join } from "@tauri-apps/api/path"; +import { notifyTaskComplete } from "./notifications"; export interface Folder { id: number; @@ -1394,6 +1395,10 @@ export const useGalleryStore = create((set, get) => ({ try { const groups = await invoke("find_duplicates", { folderId: folderId ?? null }); set({ duplicateGroups: groups, duplicateLastScanned: Math.floor(Date.now() / 1000) }); + void notifyTaskComplete( + "Duplicate scan complete", + groups.length === 1 ? "Found 1 duplicate group." : `Found ${groups.length.toLocaleString()} duplicate groups.`, + ); } finally { unlisten(); set({ duplicateScanning: false }); @@ -1466,6 +1471,7 @@ export const useGalleryStore = create((set, get) => ({ subscribeToProgress: async () => { const unlistenProgress = await listen("index-progress", (event) => { const progress = event.payload; + const previous = get().indexingProgress[progress.folder_id]; set((state) => ({ indexingProgress: { ...state.indexingProgress, @@ -1474,6 +1480,18 @@ export const useGalleryStore = create((set, get) => ({ })); if (progress.done) { + if ( + previous && + !previous.done && + progress.total > 0 && + progress.indexed >= progress.total + ) { + const folderName = get().folders.find((folder) => folder.id === progress.folder_id)?.name; + void notifyTaskComplete( + "Folder scan complete", + folderName ? `${folderName} has finished scanning.` : "A folder has finished scanning.", + ); + } void get().loadFolders(); void get().loadBackgroundJobProgress(); if (get().activeView !== "explore" && !isDerivedCollectionTitle(get().collectionTitle)) { @@ -1491,6 +1509,38 @@ export const useGalleryStore = create((set, get) => ({ }); const unlistenMediaJobs = await listen("media-job-progress", (event) => { + const previousProgress = get().mediaJobProgress; + + for (const progress of event.payload.progress) { + const previous = previousProgress[progress.folder_id]; + if (!previous) continue; + + const folderName = + get().folders.find((folder) => folder.id === progress.folder_id)?.name ?? "Folder"; + + if (previous.embedding_pending > 0 && progress.embedding_pending === 0) { + const failureDetail = + progress.embedding_failed > 0 + ? ` ${progress.embedding_failed.toLocaleString()} failed.` + : ""; + void notifyTaskComplete( + "Embeddings complete", + `${folderName} finished generating embeddings.${failureDetail}`, + ); + } + + if (previous.tagging_pending > 0 && progress.tagging_pending === 0) { + const failureDetail = + progress.tagging_failed > 0 + ? ` ${progress.tagging_failed.toLocaleString()} failed.` + : ""; + void notifyTaskComplete( + "AI tagging complete", + `${folderName} finished generating tags.${failureDetail}`, + ); + } + } + set((state) => { const next = { ...state.mediaJobProgress }; for (const progress of event.payload.progress) { -- 2.52.0 From f82002129abffa53fd91abfd7018c98f9d6c91dc Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sat, 6 Jun 2026 21:50:53 +0100 Subject: [PATCH 12/25] fix: address four correctness bugs in duplicate detection and tagging - Duplicate scanner now runs a full-file hash pass (Phase 3) for large files (> 64 KB) after the sample-hash shortlist, preventing false positives from reaching the destructive deletion workflow - delete_images_from_disk now attempts filesystem removal before touching the DB; only successfully deleted files have their rows removed, keeping locked/permission-denied files visible in the library - AI tag upsert guards user-sourced rows with WHERE source != 'user' so a tagger re-run cannot silently overwrite manually added tags - clear_tagging_jobs excludes 'cancelled' rows from the immediate DELETE so the worker can observe cancellation via is_tagging_job_cancelled before the rows are cleaned up at next startup --- src-tauri/src/commands.rs | 109 +++++++++++++++++++++++++++++--------- src-tauri/src/db.rs | 7 +-- 2 files changed, 87 insertions(+), 29 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 3c1a8ee..f400236 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -857,7 +857,7 @@ pub async fn find_duplicates( let total = records.len(); let _ = app.emit("duplicate_scan_progress", (0usize, total)); - // Two-phase detection. No full-file read at any point. + // Three-phase detection. // // Phase 1 — stat: // Read only file metadata (size). Files with a unique size cannot be @@ -865,15 +865,14 @@ pub async fn find_duplicates( // // Phase 2 — sample hash: // For each size-matched candidate, read four evenly-spaced 16 KB windows - // (head, 33%, 66%, tail) and hash them together. Total I/O per file is - // capped at 64 KB regardless of how large the file is. + // and hash them. For small files (≤ 64 KB) the windows cover the whole + // file so the hash is exact. For large files this is a fast shortlist. // - // For small files (≤ 64 KB) the windows cover the whole file, so the - // hash is exact. For large files, matching all four windows at different - // offsets is effectively impossible for natural photo/video content unless - // the files are genuinely identical — eliminating the need for a full read. + // Phase 3 — full hash (large files only): + // For sample-hash groups that contain files > 64 KB, compute a full-file + // hash to confirm the match before presenting them as deletable duplicates. let app_hash = app.clone(); - let pairs: Vec<(u64, u64, i64)> = tokio::task::spawn_blocking(move || { + let pairs: Vec<(u64, u64, i64, String)> = tokio::task::spawn_blocking(move || { use memmap2::Mmap; use rayon::prelude::*; use std::collections::HashMap; @@ -939,23 +938,82 @@ pub async fn find_duplicates( if done % 100 == 0 || done == c_total { let _ = app_hash.emit("duplicate_scan_progress", (done, c_total)); } - Some((hash, *size, *id)) + Some((hash, *size, *id, path.clone())) }) .collect() }) .await .map_err(|e| e.to_string())?; - let mut size_hash_map: std::collections::HashMap<(u64, u64), Vec> = std::collections::HashMap::new(); - for (hash, file_size, id) in pairs { - size_hash_map.entry((hash, file_size)).or_default().push(id); + // Group by (sample_hash, file_size). + let mut size_hash_map: std::collections::HashMap<(u64, u64), Vec<(i64, String)>> = + std::collections::HashMap::new(); + for (hash, file_size, id, path) in pairs { + size_hash_map.entry((hash, file_size)).or_default().push((id, path)); } + // Phase 3: for large-file groups (> 64 KB) the sample hash is not exact; + // re-hash the full file contents to avoid false positives before deletion. + const COVERED: u64 = (16 * 1024 * 4) as u64; + let confirmed: std::collections::HashMap<(u64, u64), Vec> = + tokio::task::spawn_blocking(move || { + use memmap2::Mmap; + use rayon::prelude::*; + use xxhash_rust::xxh3::xxh3_64; + + size_hash_map + .into_iter() + .filter(|(_, entries)| entries.len() > 1) + .filter_map(|((sample_hash, file_size), entries)| { + if file_size <= COVERED { + // Sample was already a full hash — no re-read needed. + let ids = entries.into_iter().map(|(id, _)| id).collect(); + return Some(((sample_hash, file_size), ids)); + } + // Full-file hash pass to confirm. + let full_hashes: Vec<(i64, Option)> = entries + .par_iter() + .map(|(id, path)| { + let hash = std::fs::File::open(path) + .ok() + .and_then(|f| unsafe { Mmap::map(&f).ok() }) + .map(|mmap| xxh3_64(&mmap)); + (*id, hash) + }) + .collect(); + // Group by full hash; keep only groups where every file hashed + // successfully and at least two share the same hash. + let mut by_full: std::collections::HashMap> = + std::collections::HashMap::new(); + for (id, maybe_hash) in full_hashes { + if let Some(h) = maybe_hash { + by_full.entry(h).or_default().push(id); + } + } + // Return the largest confirmed group (same sample_hash/size bucket). + // In practice a single full hash will dominate; emit each sub-group. + // We flatten into one entry using the sample_hash key — callers only + // need confirmed IDs, not the exact full hash. + let confirmed_ids: Vec = by_full + .into_values() + .filter(|g| g.len() > 1) + .flatten() + .collect(); + if confirmed_ids.is_empty() { + None + } else { + Some(((sample_hash, file_size), confirmed_ids)) + } + }) + .collect() + }) + .await + .map_err(|e| e.to_string())?; + // Resolve image records for each duplicate group let conn = db.get().map_err(|e| e.to_string())?; - let mut groups: Vec = size_hash_map + let mut groups: Vec = confirmed .into_iter() - .filter(|(_, ids)| ids.len() > 1) .filter_map(|((hash, file_size), ids)| { let images = db::get_images_by_ids(&conn, &ids).ok()?; Some(DuplicateGroup { @@ -1010,23 +1068,22 @@ pub async fn delete_images_from_disk( return Ok(0); } let conn = db.get().map_err(|e| e.to_string())?; - // Collect paths before deleting DB rows let records = db::get_all_image_paths(&conn, None).map_err(|e| e.to_string())?; let id_set: std::collections::HashSet = params.image_ids.iter().copied().collect(); - let paths: Vec = records - .into_iter() - .filter(|r| id_set.contains(&r.id)) - .map(|r| r.path) - .collect(); - db::delete_images_by_ids(&conn, ¶ms.image_ids).map_err(|e| e.to_string())?; - - let mut deleted = 0usize; - for path in &paths { - if std::fs::remove_file(path).is_ok() { - deleted += 1; + // Attempt filesystem removal first; only delete DB rows for files that + // were successfully removed. This prevents entries from disappearing from + // the library when a file is locked or permission-denied. + let mut succeeded_ids: Vec = Vec::new(); + for r in records.into_iter().filter(|r| id_set.contains(&r.id)) { + if std::fs::remove_file(&r.path).is_ok() { + succeeded_ids.push(r.id); } } + let deleted = succeeded_ids.len(); + if !succeeded_ids.is_empty() { + db::delete_images_by_ids(&conn, &succeeded_ids).map_err(|e| e.to_string())?; + } Ok(deleted) } diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 5df7fc2..ecb663e 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -1697,7 +1697,8 @@ pub fn update_ai_tags( ON CONFLICT(image_id, tag) DO UPDATE SET source = 'ai', ai_model = excluded.ai_model, - confidence = excluded.confidence", + confidence = excluded.confidence + WHERE source != 'user'", )?; for (tag, confidence) in tags { stmt.execute(params![image_id, tag, model, confidence])?; @@ -1748,7 +1749,7 @@ pub fn clear_tagging_jobs(conn: &Connection, folder_id: Option) -> Result) -> Result Date: Sat, 6 Jun 2026 22:00:20 +0100 Subject: [PATCH 13/25] fix: four more correctness bugs from PR review - Duplicate scanner now emits one DuplicateGroup per distinct full hash within a sample-hash bucket, preventing disjoint sets (A,A vs B,B) from being merged and presented as a single deletable group - Tagger model download now installs the shared ONNX runtime DLLs via ensure_onnx_runtime so the tagger is ready on a clean install without requiring the caption model to have been downloaded first - Caption queue clear now follows the tagging worker pattern: marks in-flight rows as cancelled and deletes only non-processing rows; the caption worker skips writing results for cancelled jobs; startup cleanup removes any leftover cancelled caption rows - Region search pagination now stores the crop rect in state and uses find_similar_by_region for subsequent pages instead of falling through to loadImages which appended unrelated folder results --- src-tauri/src/commands.rs | 40 ++++++++++++++++----------------------- src-tauri/src/db.rs | 36 +++++++++++++++++++++++++++++++---- src-tauri/src/indexer.rs | 7 +++++++ src-tauri/src/tagger.rs | 11 +++++++++-- src/store.ts | 28 ++++++++++++++++++++++++++- 5 files changed, 91 insertions(+), 31 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index f400236..4936847 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -953,9 +953,11 @@ pub async fn find_duplicates( } // Phase 3: for large-file groups (> 64 KB) the sample hash is not exact; - // re-hash the full file contents to avoid false positives before deletion. + // re-hash the full file contents to confirm matches. Each distinct full + // hash becomes its own DuplicateGroup so disjoint sets (A,A vs B,B that + // happened to share size and sample hash) are never merged. const COVERED: u64 = (16 * 1024 * 4) as u64; - let confirmed: std::collections::HashMap<(u64, u64), Vec> = + let confirmed: Vec<(u64, u64, Vec)> = tokio::task::spawn_blocking(move || { use memmap2::Mmap; use rayon::prelude::*; @@ -964,13 +966,13 @@ pub async fn find_duplicates( size_hash_map .into_iter() .filter(|(_, entries)| entries.len() > 1) - .filter_map(|((sample_hash, file_size), entries)| { + .flat_map(|((sample_hash, file_size), entries)| { if file_size <= COVERED { - // Sample was already a full hash — no re-read needed. - let ids = entries.into_iter().map(|(id, _)| id).collect(); - return Some(((sample_hash, file_size), ids)); + // Sample was already a full hash — one group, no re-read needed. + return vec![(sample_hash, file_size, entries.into_iter().map(|(id, _)| id).collect())]; } - // Full-file hash pass to confirm. + // Full-file hash pass. Group by full hash so that two sets of + // files that collide only in the sample produce separate groups. let full_hashes: Vec<(i64, Option)> = entries .par_iter() .map(|(id, path)| { @@ -981,8 +983,6 @@ pub async fn find_duplicates( (*id, hash) }) .collect(); - // Group by full hash; keep only groups where every file hashed - // successfully and at least two share the same hash. let mut by_full: std::collections::HashMap> = std::collections::HashMap::new(); for (id, maybe_hash) in full_hashes { @@ -990,20 +990,12 @@ pub async fn find_duplicates( by_full.entry(h).or_default().push(id); } } - // Return the largest confirmed group (same sample_hash/size bucket). - // In practice a single full hash will dominate; emit each sub-group. - // We flatten into one entry using the sample_hash key — callers only - // need confirmed IDs, not the exact full hash. - let confirmed_ids: Vec = by_full - .into_values() - .filter(|g| g.len() > 1) - .flatten() - .collect(); - if confirmed_ids.is_empty() { - None - } else { - Some(((sample_hash, file_size), confirmed_ids)) - } + // Emit one entry per distinct full hash that has ≥ 2 members. + by_full + .into_iter() + .filter(|(_, ids)| ids.len() > 1) + .map(|(full_hash, ids)| (full_hash, file_size, ids)) + .collect() }) .collect() }) @@ -1014,7 +1006,7 @@ pub async fn find_duplicates( let conn = db.get().map_err(|e| e.to_string())?; let mut groups: Vec = confirmed .into_iter() - .filter_map(|((hash, file_size), ids)| { + .filter_map(|(hash, file_size, ids)| { let images = db::get_images_by_ids(&conn, &ids).ok()?; Some(DuplicateGroup { file_hash: format!("{:016x}", hash), diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index ecb663e..22f009f 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -511,6 +511,7 @@ pub fn reset_inflight_jobs(conn: &Connection) -> Result<()> { // Delete any rows that were cancelled before the previous shutdown so // they don't silently linger in the DB across restarts. conn.execute("DELETE FROM tagging_jobs WHERE status = 'cancelled'", [])?; + conn.execute("DELETE FROM caption_jobs WHERE status = 'cancelled'", [])?; Ok(()) } @@ -588,13 +589,22 @@ pub fn requeue_processing_caption_jobs_for_folder( } pub fn clear_caption_jobs(conn: &Connection, folder_id: Option) -> Result { + // Mirror the tagging worker pattern: mark in-flight rows as cancelled so + // the worker can detect cancellation before writing results, then delete + // every non-processing row immediately (pending, failed, etc.). match folder_id { Some(folder_id) => { + conn.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; let deleted = conn.execute( "DELETE FROM caption_jobs - WHERE image_id IN ( - SELECT id FROM images WHERE folder_id = ?1 - )", + WHERE status NOT IN ('processing', 'cancelled') + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", [folder_id], )?; conn.execute( @@ -607,7 +617,16 @@ pub fn clear_caption_jobs(conn: &Connection, folder_id: Option) -> Result { - let deleted = conn.execute("DELETE FROM caption_jobs", [])?; + conn.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing'", + [], + )?; + let deleted = conn.execute( + "DELETE FROM caption_jobs WHERE status NOT IN ('processing', 'cancelled')", + [], + )?; conn.execute( "UPDATE images SET caption_error = NULL @@ -619,6 +638,15 @@ pub fn clear_caption_jobs(conn: &Connection, folder_id: Option) -> Result Result { + let count: i64 = conn.query_row( + "SELECT COUNT(*) FROM caption_jobs WHERE image_id = ?1 AND status = 'cancelled'", + [image_id], + |row| row.get(0), + )?; + Ok(count > 0) +} + pub fn reset_generated_captions(conn: &Connection, folder_id: Option) -> Result { let image_ids = match folder_id { Some(folder_id) => { diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index 097fc8f..f619dcf 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -839,6 +839,13 @@ fn process_caption_batch( let mut updated_images = Vec::with_capacity(caption_results.len()); for (job, caption_result) in &caption_results { + if db::is_caption_job_cancelled(&tx, job.image_id)? { + tx.execute( + "DELETE FROM caption_jobs WHERE image_id = ?1", + [job.image_id], + )?; + continue; + } match caption_result { Ok(caption) => { updated_images.push(db::update_generated_caption( diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index defc066..48522bd 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -253,8 +253,15 @@ pub fn prepare_tagger_model_with_progress( let local_dir = model_dir(app_data_dir); std::fs::create_dir_all(&local_dir)?; - // Only download the two tagger-specific files; ONNX runtime DLLs are - // already handled by the captioner download flow. + // Ensure the shared ONNX runtime DLLs are present. The tagger shares + // them with the captioner; install them here so the tagger is fully + // functional even on a clean install where the caption model has never + // been downloaded. + let caption_model_dir = crate::captioner::model_dir(app_data_dir); + std::fs::create_dir_all(&caption_model_dir)?; + crate::captioner::ensure_onnx_runtime(&caption_model_dir)?; + + // Download the two tagger-specific files. const DOWNLOAD_FILES: &[&str] = &["model.onnx", "selected_tags.csv"]; let api = Api::new()?; diff --git a/src/store.ts b/src/store.ts index ac37299..dd82273 100644 --- a/src/store.ts +++ b/src/store.ts @@ -238,6 +238,7 @@ interface GalleryState { similarHasMore: boolean; similarScope: SimilarScope; similarFolderId: number | null; + similarCrop: { x: number; y: number; w: number; h: number } | null; galleryScrollResetKey: number; activeView: ActiveView; exploreMode: ExploreMode; @@ -567,6 +568,7 @@ export const useGalleryStore = create((set, get) => ({ similarHasMore: false, similarScope: "all_media", similarFolderId: null, + similarCrop: null, galleryScrollResetKey: 0, activeView: "gallery", exploreMode: "visual", @@ -771,7 +773,7 @@ export const useGalleryStore = create((set, get) => ({ }, loadMoreImages: async () => { - const { loadedCount, totalImages, loadingImages, collectionTitle, similarSourceImageId, similarHasMore, similarFolderId } = get(); + const { loadedCount, totalImages, loadingImages, collectionTitle, similarSourceImageId, similarHasMore, similarFolderId, similarCrop } = get(); if (loadingImages || loadedCount >= totalImages) return; if (collectionTitle === "Explore Cluster") return; if (collectionTitle === "Similar Images" && similarSourceImageId !== null) { @@ -779,6 +781,28 @@ export const useGalleryStore = create((set, get) => ({ await get().loadSimilarImages(similarSourceImageId, similarFolderId, false); return; } + if (collectionTitle === "Region Search Results" && similarSourceImageId !== null && similarCrop !== null) { + if (!similarHasMore) return; + const result = await invoke("find_similar_by_region", { + params: { + image_id: similarSourceImageId, + crop_x: similarCrop.x, + crop_y: similarCrop.y, + crop_w: similarCrop.w, + crop_h: similarCrop.h, + folder_id: similarFolderId, + offset: loadedCount, + limit: PAGE_SIZE, + }, + }); + set((state) => ({ + images: [...state.images, ...result.images], + loadedCount: state.loadedCount + result.images.length, + totalImages: result.has_more ? state.loadedCount + result.images.length + 1 : state.loadedCount + result.images.length, + similarHasMore: result.has_more, + })); + return; + } await get().loadImages(false); }, @@ -1006,6 +1030,7 @@ export const useGalleryStore = create((set, get) => ({ similarSourceImageId: imageId, similarSourceFolderId: sourceFolderId, similarFolderId: folderId ?? null, + similarCrop: crop, similarScope, galleryScrollResetKey: state.galleryScrollResetKey + 1, selectedImage: null, @@ -1038,6 +1063,7 @@ export const useGalleryStore = create((set, get) => ({ similarSourceFolderId: sourceFolderId, similarHasMore: result.has_more, similarFolderId: folderId ?? null, + similarCrop: crop, similarScope, }); } catch (error) { -- 2.52.0 From 9e08d9cce36eed7c41741a97adb225cebba7f7b8 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sat, 6 Jun 2026 22:09:54 +0100 Subject: [PATCH 14/25] fix: three more correctness bugs from PR review - upsert_image now resets ai_tagged_at, ai_rating, ai_tagger_model, and ai_tagger_error to NULL and deletes AI-sourced image_tags when a file is reindexed after content changes, preventing stale tags/ratings from the previous version being shown and blocking re-tagging - get_pending_tagging_jobs_excluding now applies the paused-folder exclusion in SQL before LIMIT using folder_exclusion_clause, matching the pattern used by caption and embedding job claimers; previously a large paused-folder backlog could starve all unpaused folders - Merged similarRequestToken into galleryRequestToken so all gallery-producing requests share one generation counter; a late similarity response can no longer overwrite a subsequently loaded folder, and vice versa --- src-tauri/src/db.rs | 24 +++++++++++++++++------- src/store.ts | 22 ++++++++++++---------- 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 22f009f..15134a2 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -349,7 +349,11 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { generated_caption = excluded.generated_caption, caption_model = excluded.caption_model, caption_updated_at = excluded.caption_updated_at, - caption_error = excluded.caption_error + caption_error = excluded.caption_error, + ai_rating = NULL, + ai_tagger_model = NULL, + ai_tagged_at = NULL, + ai_tagger_error = NULL RETURNING id", params![ img.folder_id, @@ -385,6 +389,12 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { ], |row| row.get(0), )?; + // Remove stale AI tags when content changes so they aren't shown for the + // new file. User-sourced tags are preserved. + conn.execute( + "DELETE FROM image_tags WHERE image_id = ?1 AND source = 'ai'", + [id], + )?; Ok(id) } @@ -1678,14 +1688,17 @@ fn get_pending_tagging_jobs_excluding( paused_folder_ids: &std::collections::HashSet, limit: usize, ) -> Result> { - let mut stmt = conn.prepare( + let sql = format!( "SELECT j.image_id, i.folder_id, i.path FROM tagging_jobs j JOIN images i ON i.id = j.image_id WHERE j.status = 'pending' + {} ORDER BY j.created_at ASC LIMIT ?1", - )?; + folder_exclusion_clause("i", paused_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; let rows = stmt .query_map([limit as i64], |row| { Ok(TaggingJob { @@ -1695,10 +1708,7 @@ fn get_pending_tagging_jobs_excluding( }) })? .collect::>>()?; - Ok(rows - .into_iter() - .filter(|job| !paused_folder_ids.contains(&job.folder_id)) - .collect()) + Ok(rows) } pub fn update_ai_tags( diff --git a/src/store.ts b/src/store.ts index dd82273..9730662 100644 --- a/src/store.ts +++ b/src/store.ts @@ -365,8 +365,10 @@ const PAGE_SIZE = 200; const AI_CAPTIONS_ENABLED_KEY = "phokus.aiCaptionsEnabled"; const SIMILAR_DISTANCE_THRESHOLD = 0.24; +// Single token shared by all gallery-producing requests (folder loads, searches, +// similarity, region search). Any new request increments it so a stale response +// from a previous collection type cannot overwrite newer results. let galleryRequestToken = 0; -let similarRequestToken = 0; let tagCloudRequestToken = 0; let exploreTagRequestToken = 0; @@ -902,7 +904,7 @@ export const useGalleryStore = create((set, get) => ({ }, showVisualCluster: async (imageIds) => { - const requestToken = ++similarRequestToken; + const requestToken = ++galleryRequestToken; set((state) => ({ activeView: "gallery", search: "", @@ -923,7 +925,7 @@ export const useGalleryStore = create((set, get) => ({ const images = await invoke("get_images_by_ids", { params: { image_ids: imageIds }, }); - if (requestToken !== similarRequestToken) return; + if (requestToken !== galleryRequestToken) return; set({ images, totalImages: images.length, @@ -933,7 +935,7 @@ export const useGalleryStore = create((set, get) => ({ collectionTitle: "Explore Cluster", }); } catch (error) { - if (requestToken !== similarRequestToken) return; + if (requestToken !== galleryRequestToken) return; set({ images: [], totalImages: 0, @@ -951,7 +953,7 @@ export const useGalleryStore = create((set, get) => ({ }, loadSimilarImages: async (imageId, folderId = get().selectedFolderId, reset = true, sourceFolderId = folderId ?? null) => { - const requestToken = ++similarRequestToken; + const requestToken = ++galleryRequestToken; const offset = reset ? 0 : get().loadedCount; const similarScope = folderId === null ? "all_media" : "current_folder"; set((state) => ({ @@ -978,7 +980,7 @@ export const useGalleryStore = create((set, get) => ({ }, }); - if (requestToken !== similarRequestToken) return; + if (requestToken !== galleryRequestToken) return; set((state) => { const nextImages = reset ? result.images : [...state.images, ...result.images]; @@ -999,7 +1001,7 @@ export const useGalleryStore = create((set, get) => ({ }; }); } catch (error) { - if (requestToken !== similarRequestToken) return; + if (requestToken !== galleryRequestToken) return; console.error("Failed to load similar images:", error); set({ images: [], @@ -1019,7 +1021,7 @@ export const useGalleryStore = create((set, get) => ({ }, loadSimilarByRegion: async (imageId, crop, folderId = get().selectedFolderId, sourceFolderId = folderId ?? null) => { - const requestToken = ++similarRequestToken; + const requestToken = ++galleryRequestToken; const similarScope = folderId === null ? "all_media" : "current_folder"; set((state) => ({ images: [], @@ -1050,7 +1052,7 @@ export const useGalleryStore = create((set, get) => ({ }, }); - if (requestToken !== similarRequestToken) return; + if (requestToken !== galleryRequestToken) return; set({ images: result.images, @@ -1067,7 +1069,7 @@ export const useGalleryStore = create((set, get) => ({ similarScope, }); } catch (error) { - if (requestToken !== similarRequestToken) return; + if (requestToken !== galleryRequestToken) return; console.error("Failed to load region search results:", error); set({ images: [], -- 2.52.0 From 7871d52d39c9fe40026f597c6f2838202c4f9ffb Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 09:00:32 +0100 Subject: [PATCH 15/25] fix: five PR review bugs and folder-recovery UX for missing/renamed folders BackgroundTasks: dismiss no longer calls clearTaggingJobs; dismissal only updates local dismissed state DuplicateFinder: entering the view clears stale groups when the stored scan scope differs from selectedFolderId; duplicateScanFolderId is recorded on each scan and cache load so scope is always tracked Tagger threshold: set_tagger_threshold now sets TAGGER_SESSION_DIRTY so the worker rebuilds its WdTagger instance and applies the new threshold on the next batch Region search global scope: fetch offset+limit+2 candidates before removing the source image so has_more is accurate across pages Region search pagination: loadMoreImages now sets loadingImages and checks galleryRequestToken before appending results, preventing duplicate pages and stale responses corrupting a new collection Folder recovery: when index_folder detects a missing path it records the error in a new folders.scan_error column (ensure_column migration) and emits done without deleting images, preserving them for recovery. A new update_folder_path command rewrites image paths in the DB before reindexing so thumbnails and embeddings are not needlessly regenerated for unchanged files. The sidebar shows an amber warning icon on affected folders and a recovery banner with Locate Folder and Remove actions --- src-tauri/src/commands.rs | 36 ++++++++++++++- src-tauri/src/db.rs | 39 +++++++++++++++- src-tauri/src/indexer.rs | 35 +++++++++++++-- src-tauri/src/lib.rs | 1 + src-tauri/src/tagger.rs | 1 + src/components/BackgroundTasks.tsx | 2 - src/components/Sidebar.tsx | 52 ++++++++++++++++++--- src/store.ts | 72 +++++++++++++++++++++--------- 8 files changed, 203 insertions(+), 35 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 4936847..434cfd4 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -325,6 +325,37 @@ pub async fn reindex_folder( Ok(()) } +#[tauri::command] +pub async fn update_folder_path( + app: AppHandle, + db: State<'_, DbState>, + folder_id: i64, + new_path: String, +) -> Result<(), String> { + let new_path_buf = PathBuf::from(&new_path); + if !new_path_buf.is_dir() { + return Err(format!("Path is not a valid directory: {}", new_path)); + } + let new_name = new_path_buf + .file_name() + .map(|n| n.to_string_lossy().to_string()) + .unwrap_or_else(|| new_path.clone()); + { + let conn = db.get().map_err(|e| e.to_string())?; + // Fetch the old path before updating so image paths can be rewritten. + let old_path = db::get_folders(&conn) + .map_err(|e| e.to_string())? + .into_iter() + .find(|f| f.id == folder_id) + .map(|f| f.path) + .ok_or("Folder not found")?; + db::update_folder_path(&conn, folder_id, &old_path, &new_path, &new_name) + .map_err(|e| e.to_string())?; + } + indexer::index_folder(app, db.inner().clone(), folder_id, new_path_buf); + Ok(()) +} + #[tauri::command] pub async fn find_similar_images( db: State<'_, DbState>, @@ -403,9 +434,10 @@ pub async fn find_similar_by_region( ) .map_err(|e| e.to_string())?, None => { - let mut ids = vector::search_image_ids_by_embedding(&conn, &embedding, offset + limit + 1) + // Fetch one extra candidate to compensate for the source image that + // will be removed, so has_more is accurate and results span multiple pages. + let mut ids = vector::search_image_ids_by_embedding(&conn, &embedding, offset + limit + 2) .map_err(|e| e.to_string())?; - // Exclude the source image from global results ids.retain(|&id| id != params.image_id); ids } diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 15134a2..549e59d 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -30,6 +30,7 @@ pub struct Folder { pub name: String, pub image_count: i64, pub indexed_at: Option, + pub scan_error: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -172,7 +173,8 @@ pub fn migrate(conn: &Connection) -> Result<()> { path TEXT NOT NULL UNIQUE, name TEXT NOT NULL, image_count INTEGER NOT NULL DEFAULT 0, - indexed_at TEXT + indexed_at TEXT, + scan_error TEXT ); CREATE TABLE IF NOT EXISTS images ( @@ -304,6 +306,7 @@ pub fn migrate(conn: &Connection) -> Result<()> { ensure_column(conn, "images", "ai_tagger_model", "TEXT")?; ensure_column(conn, "images", "ai_tagged_at", "TEXT")?; ensure_column(conn, "images", "ai_tagger_error", "TEXT")?; + ensure_column(conn, "folders", "scan_error", "TEXT")?; vector::migrate(conn)?; Ok(()) @@ -1339,7 +1342,7 @@ pub fn get_images_by_ids(conn: &Connection, image_ids: &[i64]) -> Result Result> { let mut stmt = - conn.prepare("SELECT id, path, name, image_count, indexed_at FROM folders ORDER BY name")?; + conn.prepare("SELECT id, path, name, image_count, indexed_at, scan_error FROM folders ORDER BY name")?; let rows = stmt.query_map([], |row| { Ok(Folder { id: row.get(0)?, @@ -1347,11 +1350,43 @@ pub fn get_folders(conn: &Connection) -> Result> { name: row.get(2)?, image_count: row.get(3)?, indexed_at: row.get(4)?, + scan_error: row.get(5)?, }) })?; Ok(rows.collect::>>()?) } +pub fn update_folder_path(conn: &Connection, folder_id: i64, old_path: &str, new_path: &str, new_name: &str) -> Result<()> { + conn.execute( + "UPDATE folders SET path = ?2, name = ?3, scan_error = NULL WHERE id = ?1", + params![folder_id, new_path, new_name], + )?; + // Rewrite image paths so the indexer can match them by path and skip + // re-generating thumbnails and embeddings for unchanged files. + // SQLite's replace() does a literal prefix substitution on each path. + conn.execute( + "UPDATE images SET path = replace(path, ?1, ?2) WHERE folder_id = ?3", + params![old_path, new_path, folder_id], + )?; + Ok(()) +} + +pub fn set_folder_scan_error(conn: &Connection, folder_id: i64, error: &str) -> Result<()> { + conn.execute( + "UPDATE folders SET scan_error = ?2 WHERE id = ?1", + params![folder_id, error], + )?; + Ok(()) +} + +pub fn clear_folder_scan_error(conn: &Connection, folder_id: i64) -> Result<()> { + conn.execute( + "UPDATE folders SET scan_error = NULL WHERE id = ?1", + [folder_id], + )?; + Ok(()) +} + pub fn get_images( conn: &Connection, folder_id: Option, diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index f619dcf..a5d062d 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -147,11 +147,39 @@ pub struct MediaJobProgressEvent { pub fn index_folder(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) { std::thread::spawn(move || { + set_folder_indexing_state(folder_id, true); + + // If the folder path no longer exists on disk, record the error and + // emit done. Images are intentionally kept in the DB so the user can + // choose to relocate the folder or remove it explicitly — they should + // not be silently destroyed. + if !folder_path.is_dir() { + let error_msg = format!("Folder not found: {}", folder_path.display()); + eprintln!("Indexing error for folder {}: {}", folder_id, error_msg); + if let Ok(conn) = pool.get() { + let _ = db::set_folder_scan_error(&conn, folder_id, &error_msg); + } + emit_progress( + &app, + &IndexProgress { + folder_id, + total: 0, + indexed: 0, + current_file: String::new(), + done: true, + }, + ); + set_folder_indexing_state(folder_id, false); + return; + } + let storage_profile = detect_storage_profile(&folder_path); set_folder_storage_profile(folder_id, RuntimeAdaptiveProfile::new(storage_profile)); - set_folder_indexing_state(folder_id, true); - if let Err(error) = do_index(app.clone(), pool, folder_id, folder_path) { + if let Err(error) = do_index(app.clone(), &pool, folder_id, folder_path) { eprintln!("Indexing error for folder {}: {}", folder_id, error); + if let Ok(conn) = pool.get() { + let _ = db::set_folder_scan_error(&conn, folder_id, &error.to_string()); + } // Always emit done so the frontend reloads and recovers from partial state. emit_progress( &app, @@ -246,7 +274,7 @@ pub fn start_tagging_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) }); } -fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> { +fn do_index(app: AppHandle, pool: &DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> { let existing_entries = { let conn = pool.get()?; db::get_folder_media_index(&conn, folder_id)? @@ -344,6 +372,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) } let _ = db::backfill_embedding_jobs(&conn)?; db::update_folder_count(&conn, folder_id)?; + let _ = db::clear_folder_scan_error(&conn, folder_id); } emit_progress( diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 8c820db..26522b5 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -133,6 +133,7 @@ pub fn run() { commands::find_duplicates, commands::load_duplicate_scan_cache, commands::delete_images_from_disk, + commands::update_folder_path, ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index 48522bd..94a284e 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -191,6 +191,7 @@ pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result std::fs::create_dir_all(parent)?; } std::fs::write(path, clamped.to_string())?; + TAGGER_SESSION_DIRTY.store(true, Ordering::Relaxed); Ok(clamped) } diff --git a/src/components/BackgroundTasks.tsx b/src/components/BackgroundTasks.tsx index 4515215..62ead7b 100644 --- a/src/components/BackgroundTasks.tsx +++ b/src/components/BackgroundTasks.tsx @@ -57,7 +57,6 @@ export function BackgroundTasks() { const indexingProgress = useGalleryStore((state) => state.indexingProgress); const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings); - const clearTaggingJobs = useGalleryStore((state) => state.clearTaggingJobs); const duplicateScanning = useGalleryStore((state) => state.duplicateScanning); const duplicateScanProgress = useGalleryStore((state) => state.duplicateScanProgress); const [expanded, setExpanded] = useState(false); @@ -130,7 +129,6 @@ export function BackgroundTasks() { const dismissTask = (id: number, snapshot: string) => { if (id < 0) return; // system tasks (duplicate scan) cannot be dismissed - void clearTaggingJobs(id); setDismissed((prev) => ({ ...prev, [id]: snapshot })); setExpanded(false); }; diff --git a/src/components/Sidebar.tsx b/src/components/Sidebar.tsx index 1aad9d5..84c039f 100644 --- a/src/components/Sidebar.tsx +++ b/src/components/Sidebar.tsx @@ -11,9 +11,18 @@ function FolderItem({ selected: boolean; progress: IndexProgress | undefined; }) { - const { selectFolder, removeFolder, reindexFolder } = useGalleryStore(); + const { selectFolder, removeFolder, reindexFolder, updateFolderPath } = useGalleryStore(); const isIndexing = progress && !progress.done; const [confirmingRemoval, setConfirmingRemoval] = useState(false); + const isMissing = !!folder.scan_error && !isIndexing; + + const handleLocateFolder = async (e: React.MouseEvent) => { + e.stopPropagation(); + const selected = await open({ directory: true, multiple: false, title: `Locate "${folder.name}"` }); + if (selected && typeof selected === "string") { + await updateFolderPath(folder.id, selected); + } + }; useEffect(() => { if (!confirmingRemoval) return; @@ -31,6 +40,7 @@ function FolderItem({ }; return ( + <>
selectFolder(folder.id)} > - - - + {folder.scan_error ? ( + + + + + + ) : ( + + + + )}
@@ -112,6 +131,29 @@ function FolderItem({
)}
+ {isMissing && ( +
+

Folder not found

+

+ This folder may have been moved or renamed. Locate it to resume, or remove it from the app. +

+
+ + +
+
+ )} + ); } diff --git a/src/store.ts b/src/store.ts index 9730662..1ebb358 100644 --- a/src/store.ts +++ b/src/store.ts @@ -10,6 +10,7 @@ export interface Folder { name: string; image_count: number; indexed_at: string | null; + scan_error: string | null; } export type MediaKind = "image" | "video"; @@ -279,12 +280,14 @@ interface GalleryState { duplicateScanProgress: { scanned: number; total: number } | null; duplicateSelectedIds: Set; duplicateLastScanned: number | null; // Unix timestamp (seconds) + duplicateScanFolderId: number | null | undefined; // undefined = never scanned loadFolders: () => Promise; loadBackgroundJobProgress: () => Promise; addFolder: (path: string) => Promise; removeFolder: (folderId: number) => Promise; reindexFolder: (folderId: number) => Promise; + updateFolderPath: (folderId: number, newPath: string) => Promise; selectFolder: (folderId: number | null) => void; loadImages: (reset?: boolean) => Promise; loadMoreImages: () => Promise; @@ -611,6 +614,7 @@ export const useGalleryStore = create((set, get) => ({ duplicateScanProgress: null, duplicateSelectedIds: new Set(), duplicateLastScanned: null, + duplicateScanFolderId: undefined, setCacheDir: (cacheDir) => set({ cacheDir }), @@ -667,6 +671,13 @@ export const useGalleryStore = create((set, get) => ({ await loadBackgroundJobProgress(); }, + updateFolderPath: async (folderId, newPath) => { + const { loadFolders, loadBackgroundJobProgress } = get(); + await invoke("update_folder_path", { folderId, newPath }); + await loadFolders(); + await loadBackgroundJobProgress(); + }, + selectFolder: (folderId) => { set({ selectedFolderId: folderId, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, activeView: "gallery", failedEmbeddingsOnly: false, imageLoadError: null }); void get().loadImages(true); @@ -785,24 +796,33 @@ export const useGalleryStore = create((set, get) => ({ } if (collectionTitle === "Region Search Results" && similarSourceImageId !== null && similarCrop !== null) { if (!similarHasMore) return; - const result = await invoke("find_similar_by_region", { - params: { - image_id: similarSourceImageId, - crop_x: similarCrop.x, - crop_y: similarCrop.y, - crop_w: similarCrop.w, - crop_h: similarCrop.h, - folder_id: similarFolderId, - offset: loadedCount, - limit: PAGE_SIZE, - }, - }); - set((state) => ({ - images: [...state.images, ...result.images], - loadedCount: state.loadedCount + result.images.length, - totalImages: result.has_more ? state.loadedCount + result.images.length + 1 : state.loadedCount + result.images.length, - similarHasMore: result.has_more, - })); + const requestToken = ++galleryRequestToken; + set({ loadingImages: true }); + try { + const result = await invoke("find_similar_by_region", { + params: { + image_id: similarSourceImageId, + crop_x: similarCrop.x, + crop_y: similarCrop.y, + crop_w: similarCrop.w, + crop_h: similarCrop.h, + folder_id: similarFolderId, + offset: loadedCount, + limit: PAGE_SIZE, + }, + }); + if (requestToken !== galleryRequestToken) return; + set((state) => ({ + images: [...state.images, ...result.images], + loadedCount: state.loadedCount + result.images.length, + totalImages: result.has_more ? state.loadedCount + result.images.length + 1 : state.loadedCount + result.images.length, + similarHasMore: result.has_more, + loadingImages: false, + })); + } catch { + if (requestToken !== galleryRequestToken) return; + set({ loadingImages: false }); + } return; } await get().loadImages(false); @@ -858,7 +878,17 @@ export const useGalleryStore = create((set, get) => ({ openImage: (image) => set({ selectedImage: image }), closeImage: () => set({ selectedImage: null }), - setView: (activeView) => set({ activeView }), + setView: (activeView) => { + if (activeView === "duplicates") { + const { selectedFolderId, duplicateScanFolderId } = get(); + if (duplicateScanFolderId !== selectedFolderId) { + set({ activeView, duplicateGroups: [], duplicateLastScanned: null, duplicateScanFolderId: undefined }); + void get().loadDuplicateScanCache(selectedFolderId); + return; + } + } + set({ activeView }); + }, setExploreMode: (exploreMode) => set({ exploreMode }), @@ -1409,7 +1439,7 @@ export const useGalleryStore = create((set, get) => ({ interface CacheResult { groups: DuplicateGroup[]; scanned_at: number } const cached = await invoke("load_duplicate_scan_cache", { folderId: folderId ?? null }); if (cached) { - set({ duplicateGroups: cached.groups, duplicateLastScanned: cached.scanned_at }); + set({ duplicateGroups: cached.groups, duplicateLastScanned: cached.scanned_at, duplicateScanFolderId: folderId }); } }, @@ -1422,7 +1452,7 @@ export const useGalleryStore = create((set, get) => ({ }); try { const groups = await invoke("find_duplicates", { folderId: folderId ?? null }); - set({ duplicateGroups: groups, duplicateLastScanned: Math.floor(Date.now() / 1000) }); + set({ duplicateGroups: groups, duplicateLastScanned: Math.floor(Date.now() / 1000), duplicateScanFolderId: folderId }); void notifyTaskComplete( "Duplicate scan complete", groups.length === 1 ? "Found 1 duplicate group." : `Found ${groups.length.toLocaleString()} duplicate groups.`, -- 2.52.0 From 4e5923ba844dfa14077cfd17cfea2f07b1cd23c7 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 09:16:56 +0100 Subject: [PATCH 16/25] feat(folders): add rename, missing-folder recovery, and right-click context menu - Backend: new `rename_folder` command (updates display name only, not path) and `update_folder_path` command (rewrites image paths in DB before reindexing so thumbnails/embeddings are not regenerated unnecessarily) - DB: `rename_folder` and `update_folder_path` helpers; `scan_error` column on folders table to surface indexing failures without silently dropping images - Store: `renameFolder` and `updateFolderPath` actions - Sidebar: right-click context menu on each folder row (Reindex, Rename, Locate Folder, Remove from app); inline rename input on Enter/Escape/blur; missing- folder recovery banner with Locate/Remove; hover icon buttons (reindex + remove with Confirm/Cancel) preserved alongside the context menu --- src-tauri/src/commands.rs | 14 ++ src-tauri/src/db.rs | 8 + src-tauri/src/lib.rs | 1 + src/components/Sidebar.tsx | 354 ++++++++++++++++++++++++------------- src/store.ts | 6 + 5 files changed, 258 insertions(+), 125 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 434cfd4..53982b3 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -325,6 +325,20 @@ pub async fn reindex_folder( Ok(()) } +#[tauri::command] +pub async fn rename_folder( + db: State<'_, DbState>, + folder_id: i64, + new_name: String, +) -> Result<(), String> { + let new_name = new_name.trim().to_string(); + if new_name.is_empty() { + return Err("Folder name cannot be empty".to_string()); + } + let conn = db.get().map_err(|e| e.to_string())?; + db::rename_folder(&conn, folder_id, &new_name).map_err(|e| e.to_string()) +} + #[tauri::command] pub async fn update_folder_path( app: AppHandle, diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 549e59d..e102004 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -1356,6 +1356,14 @@ pub fn get_folders(conn: &Connection) -> Result> { Ok(rows.collect::>>()?) } +pub fn rename_folder(conn: &Connection, folder_id: i64, new_name: &str) -> Result<()> { + conn.execute( + "UPDATE folders SET name = ?2 WHERE id = ?1", + params![folder_id, new_name], + )?; + Ok(()) +} + pub fn update_folder_path(conn: &Connection, folder_id: i64, old_path: &str, new_path: &str, new_name: &str) -> Result<()> { conn.execute( "UPDATE folders SET path = ?2, name = ?3, scan_error = NULL WHERE id = ?1", diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 26522b5..165219f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -133,6 +133,7 @@ pub fn run() { commands::find_duplicates, commands::load_duplicate_scan_cache, commands::delete_images_from_disk, + commands::rename_folder, commands::update_folder_path, ]) .run(tauri::generate_context!()) diff --git a/src/components/Sidebar.tsx b/src/components/Sidebar.tsx index 84c039f..18f2fce 100644 --- a/src/components/Sidebar.tsx +++ b/src/components/Sidebar.tsx @@ -1,7 +1,73 @@ -import { useEffect, useState } from "react"; +import { useEffect, useRef, useState } from "react"; import { open } from "@tauri-apps/plugin-dialog"; import { useGalleryStore, Folder, IndexProgress } from "../store"; +interface ContextMenuState { + folderId: number; + x: number; + y: number; +} + +function FolderContextMenu({ + menu, + folder, + onClose, + onRename, + onReindex, + onLocate, + onRemove, +}: { + menu: ContextMenuState; + folder: Folder; + onClose: () => void; + onRename: () => void; + onReindex: () => void; + onLocate: () => void; + onRemove: () => void; +}) { + const ref = useRef(null); + + useEffect(() => { + const handleDown = (e: MouseEvent) => { + if (ref.current && !ref.current.contains(e.target as Node)) onClose(); + }; + const handleKey = (e: KeyboardEvent) => { if (e.key === "Escape") onClose(); }; + document.addEventListener("mousedown", handleDown); + document.addEventListener("keydown", handleKey); + return () => { + document.removeEventListener("mousedown", handleDown); + document.removeEventListener("keydown", handleKey); + }; + }, [onClose]); + + const item = (label: string, onClick: () => void, danger = false) => ( + + ); + + return ( +
+ {item("Reindex", onReindex)} + {item("Rename", onRename)} + {folder.scan_error && item("Locate Folder", onLocate)} +
+ {item("Remove from app", onRemove, true)} +
+ ); +} + function FolderItem({ folder, selected, @@ -11,148 +77,186 @@ function FolderItem({ selected: boolean; progress: IndexProgress | undefined; }) { - const { selectFolder, removeFolder, reindexFolder, updateFolderPath } = useGalleryStore(); + const { selectFolder, removeFolder, reindexFolder, renameFolder, updateFolderPath } = useGalleryStore(); const isIndexing = progress && !progress.done; - const [confirmingRemoval, setConfirmingRemoval] = useState(false); const isMissing = !!folder.scan_error && !isIndexing; - const handleLocateFolder = async (e: React.MouseEvent) => { + const [contextMenu, setContextMenu] = useState(null); + const [renaming, setRenaming] = useState(false); + const [renameValue, setRenameValue] = useState(folder.name); + const [confirmingRemoval, setConfirmingRemoval] = useState(false); + const renameInputRef = useRef(null); + + useEffect(() => { + if (renaming) { + setRenameValue(folder.name); + setTimeout(() => renameInputRef.current?.select(), 0); + } + }, [renaming, folder.name]); + + const handleContextMenu = (e: React.MouseEvent) => { + e.preventDefault(); e.stopPropagation(); - const selected = await open({ directory: true, multiple: false, title: `Locate "${folder.name}"` }); - if (selected && typeof selected === "string") { - await updateFolderPath(folder.id, selected); + // Keep menu inside viewport + const x = Math.min(e.clientX, window.innerWidth - 180); + const y = Math.min(e.clientY, window.innerHeight - 160); + setContextMenu({ folderId: folder.id, x, y }); + }; + + const handleLocateFolder = async () => { + const picked = await open({ directory: true, multiple: false, title: `Locate "${folder.name}"` }); + if (picked && typeof picked === "string") { + await updateFolderPath(folder.id, picked); } }; - useEffect(() => { - if (!confirmingRemoval) return; + const commitRename = async () => { + const trimmed = renameValue.trim(); + if (trimmed && trimmed !== folder.name) { + await renameFolder(folder.id, trimmed); + } + setRenaming(false); + }; - const timeout = window.setTimeout(() => { - setConfirmingRemoval(false); - }, 4000); - - return () => window.clearTimeout(timeout); - }, [confirmingRemoval]); - - const handleRemove = async (event: React.MouseEvent) => { - event.stopPropagation(); - await removeFolder(folder.id); + const handleRenameKey = (e: React.KeyboardEvent) => { + if (e.key === "Enter") { e.preventDefault(); void commitRename(); } + if (e.key === "Escape") { setRenaming(false); } }; return ( <> -
selectFolder(folder.id)} - > - {folder.scan_error ? ( - - - - - - ) : ( - - - - )} - -
-
- {folder.name} -
- {isIndexing ? ( - <> -
{progress.indexed}/{progress.total}
-
-
0 ? (progress.indexed / progress.total) * 100 : 0}%` }} - /> -
- +
!renaming && selectFolder(folder.id)} + onContextMenu={handleContextMenu} + > + {isMissing ? ( + + + + + ) : ( -
{folder.image_count.toLocaleString()}
+ + + + )} + +
+ {renaming ? ( + setRenameValue(e.target.value)} + onKeyDown={handleRenameKey} + onBlur={() => void commitRename()} + onClick={(e) => e.stopPropagation()} + /> + ) : ( +
+ {folder.name} +
+ )} + {isIndexing ? ( + <> +
{progress.indexed}/{progress.total}
+
+
0 ? (progress.indexed / progress.total) * 100 : 0}%` }} + /> +
+ + ) : ( +
{folder.image_count.toLocaleString()}
+ )} +
+ + {/* Hover action buttons */} + {!renaming && ( + confirmingRemoval ? ( +
e.stopPropagation()}> + + +
+ ) : ( +
+ + +
+ ) )}
- {confirmingRemoval ? ( -
event.stopPropagation()} - > - - -
- ) : ( -
- - + {isMissing && ( +
+

Folder not found

+

+ This folder may have been moved or renamed. Locate it to resume, or remove it from the app. +

+
+ + +
)} -
- {isMissing && ( -
-

Folder not found

-

- This folder may have been moved or renamed. Locate it to resume, or remove it from the app. -

-
- - -
-
- )} + + {contextMenu && contextMenu.folderId === folder.id && ( + setContextMenu(null)} + onRename={() => setRenaming(true)} + onReindex={() => void reindexFolder(folder.id)} + onLocate={() => void handleLocateFolder()} + onRemove={() => setConfirmingRemoval(true)} + /> + )} ); } diff --git a/src/store.ts b/src/store.ts index 1ebb358..d0c404f 100644 --- a/src/store.ts +++ b/src/store.ts @@ -287,6 +287,7 @@ interface GalleryState { addFolder: (path: string) => Promise; removeFolder: (folderId: number) => Promise; reindexFolder: (folderId: number) => Promise; + renameFolder: (folderId: number, newName: string) => Promise; updateFolderPath: (folderId: number, newPath: string) => Promise; selectFolder: (folderId: number | null) => void; loadImages: (reset?: boolean) => Promise; @@ -671,6 +672,11 @@ export const useGalleryStore = create((set, get) => ({ await loadBackgroundJobProgress(); }, + renameFolder: async (folderId, newName) => { + await invoke("rename_folder", { folderId, newName }); + await get().loadFolders(); + }, + updateFolderPath: async (folderId, newPath) => { const { loadFolders, loadBackgroundJobProgress } = get(); await invoke("update_folder_path", { folderId, newPath }); -- 2.52.0 From 183cdff6b1bc8be38abae5e5990d27fc172979e3 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 09:45:57 +0100 Subject: [PATCH 17/25] fix: address all 6 P2 review issues from PR #8 round 4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - BackgroundTasks: per-task Retry now dispatches to the correct backend — retryFailedEmbeddings for embedding failures, queueTaggingJobs for tagging failures, or both; previously always called embedding retry only - commands: tag-cloud cache key now hashes both image IDs and embedding bytes (xxh3) so re-embedding a file after reindex correctly invalidates the cache; removed now-dead fnv_hash_ids helper - db: update_folder_path uses SUBSTR instead of replace() to rewrite only the leading path prefix, preventing corruption when the old folder name also appears in a subdirectory name - db: add_user_tag promotes an existing AI-owned tag to user-owned on conflict instead of DO NOTHING, preventing the tagger from deleting a manually added tag that coincides with an AI tag - db/commands: add clear_duplicate_scan_cache / invalidate_duplicate_scan_cache so deletion removes the stale persisted cache entry; without this a restart would reload groups containing deleted images - store: setSimilarScope re-runs loadSimilarByRegion (with saved crop) when the active collection is Region Search Results, instead of silently switching to whole-image similarity and discarding the crop --- src-tauri/src/commands.rs | 43 ++++++++++++++++++++---------- src-tauri/src/db.rs | 20 +++++++++++--- src-tauri/src/lib.rs | 1 + src/components/BackgroundTasks.tsx | 6 ++++- src/store.ts | 12 ++++++--- 5 files changed, 61 insertions(+), 21 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 53982b3..0d92a10 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -753,15 +753,6 @@ pub struct TagCloudEntry { pub image_ids: Vec, } -fn fnv_hash_ids(ids: &[i64]) -> u64 { - let mut h: u64 = 0xcbf29ce484222325; - for &id in ids { - for b in id.to_le_bytes() { - h = h.wrapping_mul(0x100000001b3) ^ (b as u64); - } - } - h -} /// Clusters the library's image embeddings with k-means and returns one representative /// image per cluster — the member whose embedding is closest to its cluster centroid. @@ -782,11 +773,22 @@ pub async fn get_tag_cloud( return Ok(vec![]); } - // Compute a hash of the current embedded image IDs (sorted for stability) - let mut sorted_ids: Vec = embeddings_with_ids.iter().map(|(id, _)| *id).collect(); - sorted_ids.sort_unstable(); - let current_hash = fnv_hash_ids(&sorted_ids); - + // Sort by ID for stable ordering; hash both IDs and embedding bytes so that + // replacing a file and regenerating its embedding invalidates the cache even + // when the set of image IDs hasn't changed. + let mut sorted_pairs: Vec<_> = embeddings_with_ids.iter().collect(); + sorted_pairs.sort_unstable_by_key(|(id, _)| *id); + let current_hash = { + use xxhash_rust::xxh3::Xxh3; + let mut hasher = Xxh3::new(); + for (id, embedding) in &sorted_pairs { + hasher.update(&id.to_le_bytes()); + for val in embedding.iter() { + hasher.update(&val.to_le_bytes()); + } + } + hasher.digest() + }; let folder_scope = match folder_id { Some(id) => format!("folder_{}", id), None => "all".to_string(), @@ -1097,6 +1099,19 @@ pub async fn load_duplicate_scan_cache( } } +#[tauri::command] +pub async fn invalidate_duplicate_scan_cache( + db: State<'_, DbState>, + folder_id: Option, +) -> Result<(), String> { + let folder_scope = match folder_id { + Some(id) => format!("folder:{}", id), + None => "all".to_string(), + }; + let conn = db.get().map_err(|e| e.to_string())?; + db::clear_duplicate_scan_cache(&conn, &folder_scope).map_err(|e| e.to_string()) +} + #[tauri::command] pub async fn delete_images_from_disk( db: State<'_, DbState>, diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index e102004..e36c3bb 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -1371,9 +1371,10 @@ pub fn update_folder_path(conn: &Connection, folder_id: i64, old_path: &str, new )?; // Rewrite image paths so the indexer can match them by path and skip // re-generating thumbnails and embeddings for unchanged files. - // SQLite's replace() does a literal prefix substitution on each path. + // Use SUBSTR to replace only the leading prefix; SQLite's replace() + // would corrupt paths where the old folder name also appears deeper in the tree. conn.execute( - "UPDATE images SET path = replace(path, ?1, ?2) WHERE folder_id = ?3", + "UPDATE images SET path = ?2 || SUBSTR(path, LENGTH(?1) + 1) WHERE folder_id = ?3", params![old_path, new_path, folder_id], )?; Ok(()) @@ -1887,7 +1888,11 @@ pub fn add_user_tag(conn: &Connection, image_id: i64, tag: &str) -> Result Result<()> { + conn.execute( + "DELETE FROM duplicate_scan_cache WHERE folder_scope = ?1", + params![folder_scope], + )?; + Ok(()) +} + /// Upserts the duplicate scan cache for the given scope. pub fn set_duplicate_scan_cache( conn: &Connection, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 165219f..009851f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -132,6 +132,7 @@ pub fn run() { commands::search_tags_autocomplete, commands::find_duplicates, commands::load_duplicate_scan_cache, + commands::invalidate_duplicate_scan_cache, commands::delete_images_from_disk, commands::rename_folder, commands::update_folder_path, diff --git a/src/components/BackgroundTasks.tsx b/src/components/BackgroundTasks.tsx index 62ead7b..46d1861 100644 --- a/src/components/BackgroundTasks.tsx +++ b/src/components/BackgroundTasks.tsx @@ -57,6 +57,7 @@ export function BackgroundTasks() { const indexingProgress = useGalleryStore((state) => state.indexingProgress); const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings); + const queueTaggingJobs = useGalleryStore((state) => state.queueTaggingJobs); const duplicateScanning = useGalleryStore((state) => state.duplicateScanning); const duplicateScanProgress = useGalleryStore((state) => state.duplicateScanProgress); const [expanded, setExpanded] = useState(false); @@ -475,7 +476,10 @@ export function BackgroundTasks() { {taskHasFailed && ( diff --git a/src/store.ts b/src/store.ts index d0c404f..e5755e3 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1126,10 +1126,14 @@ export const useGalleryStore = create((set, get) => ({ setSimilarScope: (similarScope) => { set({ similarScope }); - const { similarSourceImageId, similarSourceFolderId, selectedFolderId } = get(); + const { similarSourceImageId, similarSourceFolderId, selectedFolderId, collectionTitle, similarCrop } = get(); if (similarSourceImageId === null) return; const folderId = similarScope === "current_folder" ? (similarSourceFolderId ?? selectedFolderId) : null; - void get().loadSimilarImages(similarSourceImageId, folderId, true, similarSourceFolderId); + if (collectionTitle === "Region Search Results" && similarCrop !== null) { + void get().loadSimilarByRegion(similarSourceImageId, similarCrop, folderId, similarSourceFolderId); + } else { + void get().loadSimilarImages(similarSourceImageId, folderId, true, similarSourceFolderId); + } }, suggestImageTags: async (imageId) => { @@ -1498,7 +1502,7 @@ export const useGalleryStore = create((set, get) => ({ clearDuplicateSelection: () => set({ duplicateSelectedIds: new Set() }), deleteSelectedDuplicates: async () => { - const { duplicateSelectedIds } = get(); + const { duplicateSelectedIds, duplicateScanFolderId } = get(); const ids = Array.from(duplicateSelectedIds); if (ids.length === 0) return 0; const deleted = await invoke("delete_images_from_disk", { params: { image_ids: ids } }); @@ -1509,6 +1513,8 @@ export const useGalleryStore = create((set, get) => ({ .map((g) => ({ ...g, images: g.images.filter((img) => !duplicateSelectedIds.has(img.id)) })) .filter((g) => g.images.length > 1), })); + // Invalidate the persisted cache so a restart doesn't reload stale entries + await invoke("invalidate_duplicate_scan_cache", { folderId: duplicateScanFolderId ?? null }); return deleted; }, -- 2.52.0 From 04ade3d5d67349c81e8ec3b4e7f13fd024ca51e2 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 10:20:28 +0100 Subject: [PATCH 18/25] fix: address 3 P2 review issues from PR #8 round 5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit captioner: set_caption_detail now marks CAPTION_SESSION_DIRTY so the cached FlorenceCaptioner session is torn down and rebuilt with the new prompt and token limit; previously only the file was updated on disk db: reset_generated_captions now cancels in-flight caption jobs instead of deleting them; the worker checks for a cancelled row before writing back, so deleting the row allowed results to be written immediately after a reset — now mirrors the clear_caption_jobs cancellation protocol BackgroundTasks: caption_pending/ready/failed counts are now included in pendingMediaWork, task stages, hasFailed checks, and the dismiss snapshot; previously a folder with only caption work caused the panel to disappear entirely and exposed no caption progress or failure status --- src-tauri/src/captioner.rs | 1 + src-tauri/src/db.rs | 24 +++++++++++++++++-- src/components/BackgroundTasks.tsx | 38 ++++++++++++++++++++++++++---- 3 files changed, 56 insertions(+), 7 deletions(-) diff --git a/src-tauri/src/captioner.rs b/src-tauri/src/captioner.rs index 27bb5e5..aea7833 100644 --- a/src-tauri/src/captioner.rs +++ b/src-tauri/src/captioner.rs @@ -238,6 +238,7 @@ pub fn set_caption_detail(app_data_dir: &Path, detail: CaptionDetail) -> Result< std::fs::create_dir_all(parent)?; } std::fs::write(path, detail.as_str())?; + CAPTION_SESSION_DIRTY.store(true, Ordering::Relaxed); Ok(detail) } diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index e36c3bb..980f871 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -693,9 +693,19 @@ pub fn reset_generated_captions(conn: &Connection, folder_id: Option) -> Re WHERE folder_id = ?1", [folder_id], )?; + // Mark in-flight jobs cancelled so the worker skips writing back; + // delete all others. + tx.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; tx.execute( "DELETE FROM caption_jobs - WHERE image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + WHERE status NOT IN ('processing', 'cancelled') + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", [folder_id], )?; } @@ -708,7 +718,17 @@ pub fn reset_generated_captions(conn: &Connection, folder_id: Option) -> Re caption_error = NULL", [], )?; - tx.execute("DELETE FROM caption_jobs", [])?; + // Same cancellation protocol as clear_caption_jobs. + tx.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing'", + [], + )?; + tx.execute( + "DELETE FROM caption_jobs WHERE status NOT IN ('processing', 'cancelled')", + [], + )?; } } diff --git a/src/components/BackgroundTasks.tsx b/src/components/BackgroundTasks.tsx index 46d1861..d34ceb2 100644 --- a/src/components/BackgroundTasks.tsx +++ b/src/components/BackgroundTasks.tsx @@ -24,6 +24,7 @@ interface Task { stages: TaskStage[]; hasFailedEmbeddings: boolean; hasFailedTagging: boolean; + hasFailedCaptions: boolean; pendingMediaWork: number; embeddingProcessed: number; embeddingTotal: number; @@ -148,16 +149,22 @@ export function BackgroundTasks() { const taggingPending = jobs?.tagging_pending ?? 0; const taggingReady = jobs?.tagging_ready ?? 0; const taggingFailed = jobs?.tagging_failed ?? 0; + const captionPending = jobs?.caption_pending ?? 0; + const captionReady = jobs?.caption_ready ?? 0; + const captionFailed = jobs?.caption_failed ?? 0; - const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending + taggingPending; + const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending + taggingPending + captionPending; const embeddingProcessed = embeddingReady + embeddingFailed; const embeddingTotal = embeddingProcessed + embeddingPending; const taggingProcessed = taggingReady + taggingFailed; const taggingTotal = taggingProcessed + taggingPending; + const captionProcessed = captionReady + captionFailed; + const captionTotal = captionProcessed + captionPending; const hasFailedEmbeddings = embeddingFailed > 0; const hasFailedTagging = taggingFailed > 0; + const hasFailedCaptions = captionFailed > 0; - if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings && !hasFailedTagging) return null; + if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings && !hasFailedTagging && !hasFailedCaptions) return null; const stages: TaskStage[] = []; @@ -209,6 +216,16 @@ export function BackgroundTasks() { }); } + if (captionPending > 0) { + const pct = captionTotal > 0 ? (captionProcessed / captionTotal) * 100 : 0; + stages.push({ + label: "Captions", + detail: `${captionProcessed.toLocaleString()} / ${captionTotal.toLocaleString()}`, + progress: pct, + failed: false, + }); + } + if (hasFailedEmbeddings && pendingMediaWork === 0) { stages.push({ label: "Failed", @@ -227,7 +244,16 @@ export function BackgroundTasks() { }); } - const snapshot = `${pendingMediaWork}:${embeddingFailed}:${taggingFailed}`; + if (hasFailedCaptions && pendingMediaWork === 0) { + stages.push({ + label: "Failed", + detail: `${captionFailed.toLocaleString()} captions`, + progress: null, + failed: true, + }); + } + + const snapshot = `${pendingMediaWork}:${embeddingFailed}:${taggingFailed}:${captionFailed}`; return { id: folder.id, @@ -235,6 +261,7 @@ export function BackgroundTasks() { stages, hasFailedEmbeddings, hasFailedTagging, + hasFailedCaptions, pendingMediaWork, embeddingProcessed, embeddingTotal, @@ -262,6 +289,7 @@ export function BackgroundTasks() { }], hasFailedEmbeddings: false, hasFailedTagging: false, + hasFailedCaptions: false, pendingMediaWork: 1, embeddingProcessed: 0, embeddingTotal: 0, @@ -275,7 +303,7 @@ export function BackgroundTasks() { const primary = allTasks[0]; const extraCount = allTasks.length - 1; - const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedTagging) && t.pendingMediaWork === 0); + const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedTagging || t.hasFailedCaptions) && t.pendingMediaWork === 0); // Best progress bar value: use embedding progress if available (most informative), // otherwise tagging progress, otherwise fall back to scanning progress, otherwise indeterminate. @@ -407,7 +435,7 @@ export function BackgroundTasks() { const taskTaggingStage = task.stages.find((s) => s.label === "Tags"); const taskScanningStage = task.stages.find((s) => s.label === "Scanning"); const taskBarProgress = taskEmbeddingStage?.progress ?? taskTaggingStage?.progress ?? taskScanningStage?.progress ?? null; - const taskHasFailed = (task.hasFailedEmbeddings || task.hasFailedTagging) && task.pendingMediaWork === 0; + const taskHasFailed = (task.hasFailedEmbeddings || task.hasFailedTagging || task.hasFailedCaptions) && task.pendingMediaWork === 0; return (
-- 2.52.0 From 7efb1879710fe7340a5d8bb8354ac865b5e7afde Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 22:15:48 +0100 Subject: [PATCH 19/25] fix: address 4 P2 review issues from PR #8 round 6 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit db/commands: paginate tag search — add offset + COUNT query to search_images_by_tag returning (images, total); TagSearchParams gains offset field and command returns TagSearchPage; store handles both the initial load and load-more path so results beyond PAGE_SIZE are reachable commands: semantic search over-fetches by 4× when any post-query filter (folder, media kind, favorites, rating) is active, then truncates to the requested limit; previously well-rated matches ranked just beyond the bare vector limit were silently omitted db: wrap update_folder_path in a transaction so a path-rewrite failure (e.g. uniqueness collision) rolls back the already-updated folder row instead of leaving the DB in an inconsistent half-migrated state store: invalidate exploreTagsFolderId when a tagging run completes so Explore reflects the updated tag distribution without requiring a manual folder switch; previously the cache was never busted by background tagging --- src-tauri/src/commands.rs | 34 ++++++++++++++++++++++++++++----- src-tauri/src/db.rs | 34 ++++++++++++++++++++++++++------- src/store.ts | 40 ++++++++++++++++++++++++++------------- 3 files changed, 83 insertions(+), 25 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 0d92a10..647d5d1 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -151,6 +151,15 @@ pub struct TagSearchParams { pub favorites_only: Option, pub rating_min: Option, pub limit: Option, + pub offset: Option, +} + +#[derive(Serialize)] +pub struct TagSearchPage { + pub images: Vec, + pub total: usize, + pub offset: usize, + pub limit: usize, } #[derive(Deserialize)] @@ -522,7 +531,17 @@ pub async fn semantic_search_images( let conn = db.get().map_err(|e| e.to_string())?; let limit = params.limit.unwrap_or(64); - let ids = vector::search_image_ids_by_embedding(&conn, &embedding, limit) + + // When post-query filters are active the vector search may discard many + // candidates ranked just beyond the bare limit. Over-fetch by 4× so that + // well-rated or folder-scoped matches are not silently omitted. + let has_filters = params.folder_id.is_some() + || params.media_kind.is_some() + || params.favorites_only.unwrap_or(false) + || params.rating_min.map_or(false, |r| r > 0); + let fetch_limit = if has_filters { limit * 4 } else { limit }; + + let ids = vector::search_image_ids_by_embedding(&conn, &embedding, fetch_limit) .map_err(|e| e.to_string())?; let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?; @@ -538,6 +557,7 @@ pub async fn semantic_search_images( if let Some(rating_min) = params.rating_min { images.retain(|image| image.rating >= rating_min); } + images.truncate(limit); Ok(images) } @@ -546,18 +566,22 @@ pub async fn semantic_search_images( pub async fn search_images_by_tag( db: State<'_, DbState>, params: TagSearchParams, -) -> Result, String> { +) -> Result { let conn = db.get().map_err(|e| e.to_string())?; - db::search_images_by_tag( + let limit = params.limit.unwrap_or(64); + let offset = params.offset.unwrap_or(0); + let (images, total) = db::search_images_by_tag( &conn, ¶ms.query, params.folder_id, params.media_kind.as_deref(), params.favorites_only.unwrap_or(false), params.rating_min.unwrap_or(0), - params.limit.unwrap_or(64), + limit, + offset, ) - .map_err(|e| e.to_string()) + .map_err(|e| e.to_string())?; + Ok(TagSearchPage { images, total, offset, limit }) } #[tauri::command] diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 980f871..83d9ac2 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -1385,7 +1385,10 @@ pub fn rename_folder(conn: &Connection, folder_id: i64, new_name: &str) -> Resul } pub fn update_folder_path(conn: &Connection, folder_id: i64, old_path: &str, new_path: &str, new_name: &str) -> Result<()> { - conn.execute( + // Both updates must be atomic: if the image path rewrite fails (e.g. a + // uniqueness collision) the folder row must not remain at the new location. + let tx = conn.unchecked_transaction()?; + tx.execute( "UPDATE folders SET path = ?2, name = ?3, scan_error = NULL WHERE id = ?1", params![folder_id, new_path, new_name], )?; @@ -1393,10 +1396,11 @@ pub fn update_folder_path(conn: &Connection, folder_id: i64, old_path: &str, new // re-generating thumbnails and embeddings for unchanged files. // Use SUBSTR to replace only the leading prefix; SQLite's replace() // would corrupt paths where the old folder name also appears deeper in the tree. - conn.execute( + tx.execute( "UPDATE images SET path = ?2 || SUBSTR(path, LENGTH(?1) + 1) WHERE folder_id = ?3", params![old_path, new_path, folder_id], )?; + tx.commit()?; Ok(()) } @@ -1522,13 +1526,28 @@ pub fn search_images_by_tag( favorites_only: bool, rating_min: i64, limit: usize, -) -> Result> { + offset: usize, +) -> Result<(Vec, usize)> { let normalized_query = query.trim().to_ascii_lowercase(); if normalized_query.is_empty() { - return Ok(Vec::new()); + return Ok((Vec::new(), 0)); } let favorites_flag = i64::from(favorites_only); + // Total count (for pagination) + let total: usize = conn.query_row( + "SELECT COUNT(DISTINCT i.id) + FROM images i + JOIN image_tags t ON t.image_id = i.id + WHERE (?1 IS NULL OR i.folder_id = ?1) + AND (?2 IS NULL OR i.media_kind = ?2) + AND (?3 = 0 OR i.favorite = 1) + AND i.rating >= ?4 + AND LOWER(TRIM(t.tag)) = ?5", + params![folder_id, media_kind, favorites_flag, rating_min, normalized_query], + |row| row.get::<_, i64>(0), + )? as usize; + let mut stmt = conn.prepare( "SELECT DISTINCT i.id FROM images i @@ -1539,7 +1558,7 @@ pub fn search_images_by_tag( AND i.rating >= ?4 AND LOWER(TRIM(t.tag)) = ?5 ORDER BY i.rating DESC, i.modified_at DESC NULLS LAST, i.filename ASC - LIMIT ?6", + LIMIT ?6 OFFSET ?7", )?; let image_ids = stmt @@ -1550,13 +1569,14 @@ pub fn search_images_by_tag( favorites_flag, rating_min, normalized_query, - limit as i64 + limit as i64, + offset as i64 ], |row| row.get::<_, i64>(0), )? .collect::>>()?; - get_images_by_ids(conn, &image_ids) + Ok((get_images_by_ids(conn, &image_ids)?, total)) } pub fn search_tags_autocomplete( diff --git a/src/store.ts b/src/store.ts index e5755e3..3d0bd28 100644 --- a/src/store.ts +++ b/src/store.ts @@ -725,7 +725,8 @@ export const useGalleryStore = create((set, get) => ({ } if (parsedSearch.mode === "tag" && parsedSearch.query) { - const images = await invoke("search_images_by_tag", { + const offset = reset ? 0 : loadedCount; + const result = await invoke<{ images: ImageRecord[]; total: number; offset: number; limit: number }>("search_images_by_tag", { params: { query: parsedSearch.query, folder_id: selectedFolderId, @@ -733,22 +734,32 @@ export const useGalleryStore = create((set, get) => ({ favorites_only: favoritesOnly, rating_min: minimumRating > 0 ? minimumRating : null, limit: PAGE_SIZE, + offset, }, }); if (requestToken !== galleryRequestToken) return; - set({ - images, - totalImages: images.length, - loadedCount: images.length, - loadingImages: false, - collectionTitle: `Tag search: ${parsedSearch.query}`, - selectedFolderId, - similarSourceImageId: null, - similarSourceFolderId: null, - similarHasMore: false, - similarFolderId: null, - }); + if (reset) { + set({ + images: result.images, + totalImages: result.total, + loadedCount: result.images.length, + loadingImages: false, + collectionTitle: `Tag search: ${parsedSearch.query}`, + selectedFolderId, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarFolderId: null, + }); + } else { + set((state) => ({ + images: [...state.images, ...result.images], + loadedCount: state.loadedCount + result.images.length, + totalImages: result.total, + loadingImages: false, + })); + } return; } @@ -1608,6 +1619,9 @@ export const useGalleryStore = create((set, get) => ({ "AI tagging complete", `${folderName} finished generating tags.${failureDetail}`, ); + // New tags are now in the DB — invalidate the Explore tag cache so + // reopening Explore reflects the updated tag distribution. + set({ exploreTagsFolderId: undefined }); } } -- 2.52.0 From 20c43d205614d675bf6d4512031abed8e0197833 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 22:30:02 +0100 Subject: [PATCH 20/25] fix: address P1 and P2 review issues from PR #8 round 7 indexer: clear duplicate scan cache (folder + global) at the end of every do_index so stale groups are never presented after a reindex; previously a file modified or replaced on disk would still appear as a duplicate candidate from the cached result until a manual rescan db/vector: replace timestamp-based HNSW revision with a monotonically incremented integer counter (app_kv.embedding_revision); mark_embedding_ready atomically increments the counter so two embeddings saved within the same second correctly advance the revision; get_embedding_revision now reads this counter instead of MAX(embedding_updated_at), preventing the HNSW cache from serving stale vectors after same-second batch updates --- src-tauri/src/db.rs | 14 ++++++++++++++ src-tauri/src/indexer.rs | 5 +++++ src-tauri/src/vector.rs | 18 +++++++++++------- 3 files changed, 30 insertions(+), 7 deletions(-) diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 83d9ac2..572212d 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -263,6 +263,13 @@ pub fn migrate(conn: &Connection) -> Result<()> { groups_json TEXT NOT NULL ); + CREATE TABLE IF NOT EXISTS app_kv ( + key TEXT PRIMARY KEY, + value INTEGER NOT NULL DEFAULT 0 + ); + + INSERT OR IGNORE INTO app_kv (key, value) VALUES ('embedding_revision', 0); + CREATE INDEX IF NOT EXISTS idx_images_folder_id ON images(folder_id); CREATE INDEX IF NOT EXISTS idx_images_modified_at ON images(modified_at); CREATE INDEX IF NOT EXISTS idx_embedding_jobs_status ON embedding_jobs(status); @@ -915,6 +922,13 @@ pub fn mark_embedding_ready(conn: &Connection, image_id: i64, model: &str) -> Re params![image_id, model], )?; conn.execute("DELETE FROM embedding_jobs WHERE image_id = ?1", [image_id])?; + // Advance the monotonic revision so the HNSW cache is always rebuilt after + // any embedding change, regardless of clock resolution. + conn.execute( + "INSERT INTO app_kv (key, value) VALUES ('embedding_revision', 1) + ON CONFLICT(key) DO UPDATE SET value = value + 1", + [], + )?; Ok(()) } diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index a5d062d..25b4c94 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -373,6 +373,11 @@ fn do_index(app: AppHandle, pool: &DbPool, folder_id: i64, folder_path: PathBuf) let _ = db::backfill_embedding_jobs(&conn)?; db::update_folder_count(&conn, folder_id)?; let _ = db::clear_folder_scan_error(&conn, folder_id); + // Invalidate duplicate scan cache — any reindex can change file contents + // or the set of files, making cached duplicate groups stale. + let folder_scope = format!("folder:{}", folder_id); + let _ = db::clear_duplicate_scan_cache(&conn, &folder_scope); + let _ = db::clear_duplicate_scan_cache(&conn, "all"); } emit_progress( diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 644ece8..25a3220 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -215,13 +215,17 @@ pub fn get_image_embedding(conn: &Connection, image_id: i64) -> Result Result { - let count: i64 = conn.query_row("SELECT COUNT(*) FROM image_vec", [], |row| row.get(0))?; - let max_updated_at: Option = conn.query_row( - "SELECT MAX(embedding_updated_at) FROM images WHERE embedding_status = 'ready'", - [], - |row| row.get(0), - )?; - Ok(format!("{}:{}", count, max_updated_at.unwrap_or_default())) + // Use the monotonically incremented app_kv counter so that two embeddings + // saved within the same clock second still advance the revision, preventing + // the HNSW cache from serving stale vectors. + let revision: i64 = conn + .query_row( + "SELECT COALESCE((SELECT value FROM app_kv WHERE key = 'embedding_revision'), 0)", + [], + |row| row.get(0), + ) + .unwrap_or(0); + Ok(revision.to_string()) } // fn image_ids_for_folder( -- 2.52.0 From ae3d6e10343e7484ed4da465e4f4da70ded9f3a3 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 22:46:05 +0100 Subject: [PATCH 21/25] fix(search): close three P2 correctness gaps from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - vector: increment embedding_revision whenever delete_embedding is called so the HNSW cache is invalidated after image/folder deletions, not just after inserts - hnsw_index: fix race in build_index — read the revision before fetching embeddings and again after the parallel_insert; retry the whole build if the revision advanced during construction, ensuring the cached index always reflects a consistent snapshot - commands: replace the fixed 4× over-fetch in semantic_search_images with a progressive doubling loop; start at exactly `limit` candidates and double the batch (capped at 8192) until the page is filled or the vector table is exhausted, preventing both under- and over-fetching --- src-tauri/src/commands.rs | 55 +++++++++++++++++---------- src-tauri/src/hnsw_index.rs | 76 +++++++++++++++++++++---------------- src-tauri/src/vector.rs | 6 +++ 3 files changed, 84 insertions(+), 53 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 647d5d1..cd239ad 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -532,34 +532,47 @@ pub async fn semantic_search_images( let conn = db.get().map_err(|e| e.to_string())?; let limit = params.limit.unwrap_or(64); - // When post-query filters are active the vector search may discard many - // candidates ranked just beyond the bare limit. Over-fetch by 4× so that - // well-rated or folder-scoped matches are not silently omitted. let has_filters = params.folder_id.is_some() || params.media_kind.is_some() || params.favorites_only.unwrap_or(false) || params.rating_min.map_or(false, |r| r > 0); - let fetch_limit = if has_filters { limit * 4 } else { limit }; - let ids = vector::search_image_ids_by_embedding(&conn, &embedding, fetch_limit) - .map_err(|e| e.to_string())?; - let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?; + if !has_filters { + // No post-query filtering — a single fetch of exactly `limit` results is optimal. + let ids = vector::search_image_ids_by_embedding(&conn, &embedding, limit) + .map_err(|e| e.to_string())?; + return db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string()); + } - if let Some(folder_id) = params.folder_id { - images.retain(|image| image.folder_id == folder_id); - } - if let Some(media_kind) = params.media_kind.as_deref() { - images.retain(|image| image.media_kind == media_kind); - } - if params.favorites_only.unwrap_or(false) { - images.retain(|image| image.favorite); - } - if let Some(rating_min) = params.rating_min { - images.retain(|image| image.rating >= rating_min); - } - images.truncate(limit); + // Post-query filters are active. Progressively double the fetch batch until we + // collect enough results or exhaust the vector table, so we never over-fetch + // more than necessary while still filling the requested page. + let mut batch = limit; + loop { + let ids = vector::search_image_ids_by_embedding(&conn, &embedding, batch) + .map_err(|e| e.to_string())?; + let exhausted = ids.len() < batch; + let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?; - Ok(images) + if let Some(folder_id) = params.folder_id { + images.retain(|image| image.folder_id == folder_id); + } + if let Some(media_kind) = params.media_kind.as_deref() { + images.retain(|image| image.media_kind == media_kind); + } + if params.favorites_only.unwrap_or(false) { + images.retain(|image| image.favorite); + } + if let Some(rating_min) = params.rating_min { + images.retain(|image| image.rating >= rating_min); + } + + if images.len() >= limit || exhausted { + images.truncate(limit); + return Ok(images); + } + batch = (batch * 2).min(8192); + } } #[tauri::command] diff --git a/src-tauri/src/hnsw_index.rs b/src-tauri/src/hnsw_index.rs index b7ef6c7..2cdc110 100644 --- a/src-tauri/src/hnsw_index.rs +++ b/src-tauri/src/hnsw_index.rs @@ -23,41 +23,53 @@ fn cache() -> &'static RwLock> { } fn build_index(conn: &Connection) -> Result { - let embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?; - let max_elements = embeddings.len().max(1); - let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).max(1); - let mut hnsw = Hnsw::::new( - HNSW_MAX_CONNECTIONS, - max_elements, - max_layer, - HNSW_EF_CONSTRUCTION, - DistCosine {}, - ); + // Read the revision *before* fetching embeddings so we can detect any write + // that races with the build. If the revision advances while we are building, + // the resulting index would be stale — retry until it is stable. + loop { + let revision_before = vector::get_embedding_revision(conn)?; - let image_ids_by_external = embeddings - .iter() - .map(|(image_id, _)| *image_id) - .collect::>(); - let external_by_image_id = image_ids_by_external - .iter() - .enumerate() - .map(|(external_id, image_id)| (*image_id, external_id)) - .collect::>(); - let data_with_id = embeddings - .iter() - .enumerate() - .map(|(external_id, (_, embedding))| (embedding, external_id)) - .collect::>(); + let embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?; + let max_elements = embeddings.len().max(1); + let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).max(1); + let mut hnsw = Hnsw::::new( + HNSW_MAX_CONNECTIONS, + max_elements, + max_layer, + HNSW_EF_CONSTRUCTION, + DistCosine {}, + ); - hnsw.parallel_insert(&data_with_id); - hnsw.set_searching_mode(true); + let image_ids_by_external = embeddings + .iter() + .map(|(image_id, _)| *image_id) + .collect::>(); + let external_by_image_id = image_ids_by_external + .iter() + .enumerate() + .map(|(external_id, image_id)| (*image_id, external_id)) + .collect::>(); + let data_with_id = embeddings + .iter() + .enumerate() + .map(|(external_id, (_, embedding))| (embedding, external_id)) + .collect::>(); - Ok(CachedHnswIndex { - revision: vector::get_embedding_revision(conn)?, - image_ids_by_external, - external_by_image_id, - hnsw, - }) + hnsw.parallel_insert(&data_with_id); + hnsw.set_searching_mode(true); + + // If the revision is unchanged the index reflects a consistent snapshot. + let revision_after = vector::get_embedding_revision(conn)?; + if revision_before == revision_after { + return Ok(CachedHnswIndex { + revision: revision_after, + image_ids_by_external, + external_by_image_id, + hnsw, + }); + } + // A concurrent write advanced the revision — discard this build and retry. + } } fn ensure_index(conn: &Connection) -> Result<()> { diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 25a3220..98b5349 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -33,6 +33,12 @@ pub fn migrate(conn: &Connection) -> Result<()> { #[allow(dead_code)] pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> { conn.execute("DELETE FROM image_vec WHERE image_id = ?1", [image_id])?; + // Advance the revision so any cached HNSW index is invalidated after deletions. + conn.execute( + "INSERT INTO app_kv (key, value) VALUES ('embedding_revision', 1) + ON CONFLICT(key) DO UPDATE SET value = value + 1", + [], + )?; Ok(()) } -- 2.52.0 From 1f6650021c8ad0c2bc66ed3442ffb9125d1543ac Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 23:02:51 +0100 Subject: [PATCH 22/25] fix(duplicates): close three more P2 correctness issues from review - commands: delete_images_from_disk now returns Vec of IDs that were actually removed from disk rather than a bare count, so callers can act precisely on what succeeded vs. what failed - store: deleteSelectedDuplicates uses the returned ID set to filter duplicate groups (files that failed to delete stay visible for retry) and invalidates both the global "all" cache and each affected folder cache, not just the currently-viewed scope - Lightbox: guard the getImageTags effect with a cancellation flag so a response that arrives after the user has navigated to another image cannot overwrite the new image's tag list --- src-tauri/src/commands.rs | 8 ++++---- src/components/Lightbox.tsx | 8 ++++++-- src/store.ts | 27 ++++++++++++++++++++------- 3 files changed, 30 insertions(+), 13 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index cd239ad..06f5031 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -1153,9 +1153,9 @@ pub async fn invalidate_duplicate_scan_cache( pub async fn delete_images_from_disk( db: State<'_, DbState>, params: DeleteImagesFromDiskParams, -) -> Result { +) -> Result, String> { if params.image_ids.is_empty() { - return Ok(0); + return Ok(Vec::new()); } let conn = db.get().map_err(|e| e.to_string())?; let records = db::get_all_image_paths(&conn, None).map_err(|e| e.to_string())?; @@ -1170,11 +1170,11 @@ pub async fn delete_images_from_disk( succeeded_ids.push(r.id); } } - let deleted = succeeded_ids.len(); if !succeeded_ids.is_empty() { db::delete_images_by_ids(&conn, &succeeded_ids).map_err(|e| e.to_string())?; } - Ok(deleted) + // Return the IDs that were actually removed so the caller can update state precisely. + Ok(succeeded_ids) } #[tauri::command] diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index 0ee42f5..054370a 100644 --- a/src/components/Lightbox.tsx +++ b/src/components/Lightbox.tsx @@ -175,9 +175,13 @@ export function Lightbox() { useEffect(() => { if (!selectedImage) return; + // Capture the ID so a stale response for image A cannot overwrite B's tags + // when the user navigates before the request resolves. + let cancelled = false; void getImageTags(selectedImage.id) - .then(setImageTags) - .catch(() => setImageTags([])); + .then((tags) => { if (!cancelled) setImageTags(tags); }) + .catch(() => { if (!cancelled) setImageTags([]); }); + return () => { cancelled = true; }; }, [selectedImage?.id, selectedImage?.ai_tagged_at, getImageTags]); // Reset the queued state once the worker finishes so the button is usable again diff --git a/src/store.ts b/src/store.ts index 3d0bd28..82be8d4 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1513,20 +1513,33 @@ export const useGalleryStore = create((set, get) => ({ clearDuplicateSelection: () => set({ duplicateSelectedIds: new Set() }), deleteSelectedDuplicates: async () => { - const { duplicateSelectedIds, duplicateScanFolderId } = get(); + const { duplicateSelectedIds, duplicateGroups } = get(); const ids = Array.from(duplicateSelectedIds); if (ids.length === 0) return 0; - const deleted = await invoke("delete_images_from_disk", { params: { image_ids: ids } }); - // Remove deleted images from groups and drop now-trivial groups + // Backend returns only the IDs that were actually removed from disk. + const succeededIds = await invoke("delete_images_from_disk", { params: { image_ids: ids } }); + const succeededSet = new Set(succeededIds); + // Only remove images confirmed deleted — failed files remain visible so the user can retry. set((state) => ({ duplicateSelectedIds: new Set(), duplicateGroups: state.duplicateGroups - .map((g) => ({ ...g, images: g.images.filter((img) => !duplicateSelectedIds.has(img.id)) })) + .map((g) => ({ ...g, images: g.images.filter((img) => !succeededSet.has(img.id)) })) .filter((g) => g.images.length > 1), })); - // Invalidate the persisted cache so a restart doesn't reload stale entries - await invoke("invalidate_duplicate_scan_cache", { folderId: duplicateScanFolderId ?? null }); - return deleted; + // Invalidate the persisted cache for every affected scope: + // - global "all" cache (always, since a folder-scoped deletion still makes the global result stale) + // - each folder that contained a deleted image (so a folder-scoped scan is also evicted) + const affectedFolderIds = new Set( + duplicateGroups + .flatMap((g) => g.images) + .filter((img) => succeededSet.has(img.id)) + .map((img) => img.folder_id), + ); + await invoke("invalidate_duplicate_scan_cache", { folderId: null }); // global + for (const folderId of affectedFolderIds) { + await invoke("invalidate_duplicate_scan_cache", { folderId }); + } + return succeededIds.length; }, retryFailedEmbeddings: async (folderId) => { -- 2.52.0 From 7ead26d7fbfcbb5191312f2d3b039054ab4175c0 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 23:10:34 +0100 Subject: [PATCH 23/25] fix: pre-empt review findings before next push - hnsw_index: fetch folder embedding IDs before acquiring the read lock in find_similar_image_matches; holding the lock across a SQLite query was delaying concurrent ensure_index write-lock requests - store/DuplicateFinder: add duplicateScanError state; scanDuplicates now catches errors and stores them rather than silently leaving the user with empty results and no feedback --- src-tauri/src/hnsw_index.rs | 19 ++++++++++++++++--- src/components/DuplicateFinder.tsx | 4 ++++ src/store.ts | 6 +++++- 3 files changed, 25 insertions(+), 4 deletions(-) diff --git a/src-tauri/src/hnsw_index.rs b/src-tauri/src/hnsw_index.rs index 2cdc110..0836a03 100644 --- a/src-tauri/src/hnsw_index.rs +++ b/src-tauri/src/hnsw_index.rs @@ -103,16 +103,29 @@ pub fn find_similar_image_matches( None => return Ok(Vec::new()), }; + // Fetch folder image IDs *before* acquiring the read lock so we don't hold + // the lock across a potentially slow SQLite query, which would delay any + // concurrent ensure_index call waiting for a write lock. + let folder_image_ids: Option> = if let Some(folder_id) = folder_id { + let ids = vector::get_all_image_embeddings_with_ids(conn, Some(folder_id))? + .into_iter() + .map(|(id, _)| id) + .collect(); + Some(ids) + } else { + None + }; + let guard = cache().read().expect("hnsw cache poisoned"); let Some(cached) = guard.as_ref() else { return Ok(Vec::new()); }; let knbn = (offset + limit).max(limit).saturating_add(32); - let neighbours: Vec = if let Some(folder_id) = folder_id { - let mut allowed_ids = vector::get_all_image_embeddings_with_ids(conn, Some(folder_id))? + let neighbours: Vec = if let Some(image_ids) = folder_image_ids { + let mut allowed_ids = image_ids .into_iter() - .filter_map(|(allowed_image_id, _)| { + .filter_map(|allowed_image_id| { cached.external_by_image_id.get(&allowed_image_id).copied() }) .collect::>(); diff --git a/src/components/DuplicateFinder.tsx b/src/components/DuplicateFinder.tsx index d01f05d..cae277f 100644 --- a/src/components/DuplicateFinder.tsx +++ b/src/components/DuplicateFinder.tsx @@ -116,6 +116,7 @@ export function DuplicateFinder() { const duplicateGroups = useGalleryStore((state) => state.duplicateGroups); const duplicateScanning = useGalleryStore((state) => state.duplicateScanning); const duplicateScanProgress = useGalleryStore((state) => state.duplicateScanProgress); + const duplicateScanError = useGalleryStore((state) => state.duplicateScanError); const duplicateSelectedIds = useGalleryStore((state) => state.duplicateSelectedIds); const duplicateLastScanned = useGalleryStore((state) => state.duplicateLastScanned); const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); @@ -228,6 +229,9 @@ export function DuplicateFinder() {
) : null} + {duplicateScanError ? ( +

{duplicateScanError}

+ ) : null} {deleteResult ? (

{deleteResult}

) : null} diff --git a/src/store.ts b/src/store.ts index 82be8d4..60f406a 100644 --- a/src/store.ts +++ b/src/store.ts @@ -278,6 +278,7 @@ interface GalleryState { duplicateGroups: DuplicateGroup[]; duplicateScanning: boolean; duplicateScanProgress: { scanned: number; total: number } | null; + duplicateScanError: string | null; duplicateSelectedIds: Set; duplicateLastScanned: number | null; // Unix timestamp (seconds) duplicateScanFolderId: number | null | undefined; // undefined = never scanned @@ -613,6 +614,7 @@ export const useGalleryStore = create((set, get) => ({ duplicateGroups: [], duplicateScanning: false, duplicateScanProgress: null, + duplicateScanError: null, duplicateSelectedIds: new Set(), duplicateLastScanned: null, duplicateScanFolderId: undefined, @@ -1466,7 +1468,7 @@ export const useGalleryStore = create((set, get) => ({ scanDuplicates: async (folderId = null) => { const { listen } = await import("@tauri-apps/api/event"); - set({ duplicateScanning: true, duplicateGroups: [], duplicateScanProgress: null, duplicateSelectedIds: new Set() }); + set({ duplicateScanning: true, duplicateGroups: [], duplicateScanProgress: null, duplicateScanError: null, duplicateSelectedIds: new Set() }); const unlisten = await listen<[number, number]>("duplicate_scan_progress", (event) => { const [scanned, total] = event.payload; set({ duplicateScanProgress: { scanned, total } }); @@ -1478,6 +1480,8 @@ export const useGalleryStore = create((set, get) => ({ "Duplicate scan complete", groups.length === 1 ? "Found 1 duplicate group." : `Found ${groups.length.toLocaleString()} duplicate groups.`, ); + } catch (e) { + set({ duplicateScanError: String(e) }); } finally { unlisten(); set({ duplicateScanning: false }); -- 2.52.0 From ba989b37b9d766fdb7c417cf0ea8fb3351bf5cd5 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 23:25:25 +0100 Subject: [PATCH 24/25] fix: four review issues (search cap, tagging pause, stale tag, walkdir) - commands: semantic search doubling loop now returns when batch reaches the 8192 cap rather than re-fetching the same IDs forever when filters yield fewer than limit matches (P1 infinite loop) - db/indexer: tagging worker checks is_tagging_job_processing instead of is_tagging_job_cancelled so paused jobs reset to 'pending' are also discarded; cancelled rows are deleted, paused rows kept for retry - Lightbox: capture image id at form-submit time and compare against currentImageIdRef in the addUserTag callback, matching the guard already applied to the getImageTags effect - indexer: WalkDir errors are no longer silently swallowed; paths under unreadable subtrees are excluded from the missing-file deletion pass using Path::starts_with (component-aware) so a transient permission error cannot cascade into data loss or affect sibling directories --- src-tauri/src/commands.rs | 5 ++++ src-tauri/src/db.rs | 12 ++++++++++ src-tauri/src/indexer.rs | 47 +++++++++++++++++++++++++++++++------ src/components/Lightbox.tsx | 10 +++++++- 4 files changed, 66 insertions(+), 8 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 06f5031..fbeef49 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -571,6 +571,11 @@ pub async fn semantic_search_images( images.truncate(limit); return Ok(images); } + // If we are already at the fetch cap, another iteration would fetch the + // exact same IDs — break out rather than looping forever. + if batch >= 8192 { + return Ok(images); + } batch = (batch * 2).min(8192); } } diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 572212d..a9caf14 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -2024,6 +2024,18 @@ pub fn is_tagging_job_cancelled(conn: &Connection, image_id: i64) -> Result 0) } +/// Returns `true` when the job row for `image_id` is still `processing`. +/// A `false` result means the row was reset to `pending` (pause) or `cancelled` +/// while inference was running — either way the result must be discarded. +pub fn is_tagging_job_processing(conn: &Connection, image_id: i64) -> Result { + let count: i64 = conn.query_row( + "SELECT COUNT(*) FROM tagging_jobs WHERE image_id = ?1 AND status = 'processing'", + [image_id], + |row| row.get(0), + )?; + Ok(count > 0) +} + pub fn requeue_tagging_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()> { for image_id in image_ids { conn.execute( diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index 25b4c94..7b1426d 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -284,12 +284,32 @@ fn do_index(app: AppHandle, pool: &DbPool, folder_id: i64, folder_path: PathBuf) .map(|entry| (entry.path.clone(), entry)) .collect::>(); + // Collect traversal errors separately. An unreadable subdirectory means we + // cannot distinguish absent files from inaccessible ones; tracking errors + // lets us skip the deletion step for paths under affected subtrees. + let mut unreadable_prefixes: Vec = Vec::new(); let media_paths: Vec = WalkDir::new(&folder_path) .follow_links(true) .into_iter() - .filter_map(|entry| entry.ok()) - .filter(|entry| entry.file_type().is_file() && is_supported_media(entry.path())) - .map(|entry| entry.path().to_path_buf()) + .filter_map(|entry| match entry { + Ok(e) if e.file_type().is_file() && is_supported_media(e.path()) => { + Some(e.path().to_path_buf()) + } + Ok(_) => None, + Err(err) => { + // Record the inaccessible path so we can protect its descendants + // from the missing-file deletion pass. + if let Some(path) = err.path() { + unreadable_prefixes.push(path.to_string_lossy().into_owned()); + } + eprintln!( + "WalkDir error while scanning folder {}: {}", + folder_path.display(), + err + ); + None + } + }) .collect(); let total = media_paths.len(); @@ -362,6 +382,17 @@ fn do_index(app: AppHandle, pool: &DbPool, folder_id: i64, folder_path: PathBuf) let missing_ids = existing_by_path .values() .filter(|entry| !seen_paths.contains(&entry.path)) + .filter(|entry| { + // If this path lives under a subtree that WalkDir couldn't read, + // we don't know whether the file is gone or just temporarily + // inaccessible — keep the record to avoid false deletions. + // Use Path::starts_with (component-aware) so a sibling directory + // whose name shares a prefix is not incorrectly protected. + let entry_path = std::path::Path::new(&entry.path); + unreadable_prefixes + .iter() + .all(|prefix| !entry_path.starts_with(prefix)) + }) .map(|entry| entry.id) .collect::>(); @@ -982,11 +1013,13 @@ fn process_tagging_batch( let mut updated_images = Vec::with_capacity(tag_results.len()); for (job, tag_result) in &tag_results { - // If the job was cancelled while inference was running, discard - // the result and delete the row — don't save tags or mark failed. - if db::is_tagging_job_cancelled(&tx, job.image_id)? { + // If the job is no longer in 'processing' state it was either + // cancelled or reset to 'pending' (pause) while inference was running. + // Discard the result in both cases — for cancelled rows also delete + // the row; for paused rows leave it so the job is retried later. + if !db::is_tagging_job_processing(&tx, job.image_id)? { tx.execute( - "DELETE FROM tagging_jobs WHERE image_id = ?1", + "DELETE FROM tagging_jobs WHERE image_id = ?1 AND status = 'cancelled'", [job.image_id], )?; continue; diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index 054370a..dadec21 100644 --- a/src/components/Lightbox.tsx +++ b/src/components/Lightbox.tsx @@ -129,6 +129,11 @@ export function Lightbox() { const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus); const queueTaggingForImage = useGalleryStore((state) => state.queueTaggingForImage); + // Tracks the image id that is currently displayed, used to discard async + // tag mutations that resolve after the user has navigated to another image. + const currentImageIdRef = useRef(null); + currentImageIdRef.current = selectedImage?.id ?? null; + const [zoom, setZoom] = useState(1); const [imageTags, setImageTags] = useState([]); const [tagInput, setTagInput] = useState(""); @@ -668,8 +673,11 @@ export function Lightbox() { const raw = tagInput.trim(); if (!raw || tagAdding) return; setTagAdding(true); - void addUserTag(selectedImage.id, raw) + const taggedImageId = selectedImage.id; + void addUserTag(taggedImageId, raw) .then((newTag) => { + // Discard if the user navigated away before the request resolved. + if (currentImageIdRef.current !== taggedImageId) return; setImageTags((prev) => [...prev, newTag]); setTagInput(""); }) -- 2.52.0 From fd1585b5e2a8b0bcc6694ef24c43f6096ea3ad78 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 23:35:14 +0100 Subject: [PATCH 25/25] fix: three pre-emptive issues caught before push MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - store: removeFolder now also resets exploreTagsFolderId so the Explore tag list is invalidated when a folder is removed, not just tagCloud - db: delete_folder moves vector::delete_embedding calls outside the transaction — sqlite-vec virtual-table DML is unreliable inside a transaction and caused remove_folder to fail for folders with embeddings - store: loadMoreImages now passes similarSourceFolderId as the fourth arg to loadSimilarImages; omitting it caused each paginated load to overwrite the state with similarFolderId, corrupting scope-switches that use similarSourceFolderId afterward --- src-tauri/src/db.rs | 13 +++++++++---- src/store.ts | 6 +++--- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index a9caf14..9b5caf8 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -2074,11 +2074,16 @@ pub fn delete_folder(conn: &Connection, folder_id: i64) -> Result<()> { rows }; - let tx = conn.unchecked_transaction()?; - for image_id in image_ids { - vector::delete_embedding(&tx, image_id)?; - vector::delete_caption_embedding(&tx, image_id)?; + // Delete sqlite-vec rows outside any transaction — DML on virtual tables is + // unreliable inside a transaction and will cause remove_folder to fail for + // folders that have generated embeddings. The folder cascade (images rows) + // is handled by the FK ON DELETE CASCADE when the folder row is deleted. + for image_id in &image_ids { + vector::delete_embedding(conn, *image_id)?; + vector::delete_caption_embedding(conn, *image_id)?; } + + let tx = conn.unchecked_transaction()?; tx.execute("DELETE FROM folders WHERE id = ?1", params![folder_id])?; tx.commit()?; Ok(()) diff --git a/src/store.ts b/src/store.ts index 60f406a..82d456d 100644 --- a/src/store.ts +++ b/src/store.ts @@ -657,8 +657,8 @@ export const useGalleryStore = create((set, get) => ({ const { selectedFolderId, loadFolders, loadImages, loadBackgroundJobProgress } = get(); await loadFolders(); await loadBackgroundJobProgress(); - // Invalidate tag cloud cache since library content changed - set({ tagCloudFolderId: undefined, tagCloudEntries: [] }); + // Invalidate tag cloud and explore-tags cache since library content changed + set({ tagCloudFolderId: undefined, tagCloudEntries: [], exploreTagsFolderId: undefined }); if (selectedFolderId === folderId) { set({ selectedFolderId: null }); await loadImages(true); @@ -810,7 +810,7 @@ export const useGalleryStore = create((set, get) => ({ if (collectionTitle === "Explore Cluster") return; if (collectionTitle === "Similar Images" && similarSourceImageId !== null) { if (!similarHasMore) return; - await get().loadSimilarImages(similarSourceImageId, similarFolderId, false); + await get().loadSimilarImages(similarSourceImageId, similarFolderId, false, get().similarSourceFolderId ?? null); return; } if (collectionTitle === "Region Search Results" && similarSourceImageId !== null && similarCrop !== null) { -- 2.52.0