From bee6adc61aa94065247e48d7c3892832cfc15074 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Tue, 7 Apr 2026 06:11:01 +0100 Subject: [PATCH] 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 */}