JezzWTF/phokus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add local AI captions and queue controls

Browse Source
This commit is contained in:
LyAhn
2026-04-07 06:11:01 +01:00
parent 8905baf4a5
commit bee6adc61a
15 changed files with 2653 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/Cargo.lock
Generated
+78
View File
@@ -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"
src-tauri/Cargo.toml
+1
View File
@@ -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"] }
src-tauri/src/captioner.rs
+620
View File
@@ -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<String>,
}
#[derive(Clone, Serialize)]
pub struct CaptionModelProgress {
pub total_files: usize,
pub completed_files: usize,
pub current_file: Option<String>,
pub done: bool,
}
#[derive(Serialize)]
pub struct CaptionRuntimeProbe {
pub ready: bool,
pub tokenizer_vocab_size: usize,
pub sessions: Vec<CaptionRuntimeSessionProbe>,
}
#[derive(Serialize)]
pub struct CaptionRuntimeSessionProbe {
pub file: &'static str,
pub inputs: Vec<String>,
pub outputs: Vec<String>,
}
#[derive(Serialize)]
pub struct CaptionVisionProbe {
pub input_shape: Vec<i64>,
pub output_shape: Vec<i64>,
pub output_values: usize,
}
#[derive(Clone)]
struct TensorData {
shape: Vec<i64>,
values: Vec<f32>,
}
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::<Vec<_>>();
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<CaptionModelStatus> {
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<CaptionModelStatus> {
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<CaptionRuntimeProbe> {
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::<Result<Vec<_>>>()?;
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<CaptionVisionProbe> {
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::<f32>()
.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<String> {
let mut captioner = FlorenceCaptioner::new(app_data_dir)?;
captioner.generate(image_path)
}
impl FlorenceCaptioner {
pub fn new(app_data_dir: &Path) -> Result<Self> {
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<String> {
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::<Vec<_>>();
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::<Vec<_>>();
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<CaptionRuntimeSessionProbe> {
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<TensorData> {
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<TensorData> {
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<TensorData> {
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<Vec<i64>> {
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<TensorData> = 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<Session> {
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<TensorData> {
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<TensorData> {
let (shape, values) = value
.try_extract_tensor::<f32>()
.map_err(|error| anyhow::anyhow!("{error}"))?;
Ok(TensorData {
shape: shape.to_vec(),
values: values.to_vec(),
})
}
fn tensor_from_data(data: &TensorData) -> Result<Tensor<f32>> {
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<i64> {
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<Vec<f32>> {
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<i64>) -> Self {
let values_len = shape.iter().map(|value| (*value).max(0) as usize).product();
Self {
shape,
values: vec![0.0; values_len],
}
}
}
src-tauri/src/commands.rs
+266 -21
View File
@@ -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<usize>,
}
#[derive(Deserialize)]
pub struct DebugSimilarImagesParams {
pub image_id: i64,
pub limit: Option<usize>,
}
#[derive(Deserialize)]
pub struct RetryFailedEmbeddingsParams {
pub folder_id: i64,
}
#[derive(Deserialize)]
pub struct SetGeneratedCaptionParams {
pub image_id: i64,
pub caption: String,
pub model: Option<String>,
}
#[derive(Deserialize)]
pub struct SuggestImageTagsParams {
pub image_id: i64,
pub limit: Option<usize>,
}
#[derive(Deserialize)]
pub struct QueueCaptionJobsParams {
pub folder_id: Option<i64>,
pub image_id: Option<i64>,
}
#[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<Vec<ImageRecord>, 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<i64>,
}
#[tauri::command]
pub async fn debug_similar_images(
db: State<'_, DbState>,
params: DebugSimilarImagesParams,
) -> Result<SimilarImagesDebug, String> {
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<ImageRecord, String> {
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, &params.caption, model)
.map_err(|e| e.to_string())
}
#[tauri::command]
pub async fn suggest_image_tags(
db: State<'_, DbState>,
params: SuggestImageTagsParams,
) -> Result<Vec<String>, 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<CaptionModelStatus, String> {
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<CaptionModelStatus, String> {
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<CaptionModelStatus, String> {
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<CaptionRuntimeProbe, String> {
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<CaptionVisionProbe, String> {
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<ImageRecord, String> {
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<usize, String> {
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<i64> = embeddings_with_ids.iter().map(|(id, _)| *id).collect();
let points: Vec<Vec<f32>> = embeddings_with_ids.into_iter().map(|(_, emb)| emb).collect();
let points: Vec<Vec<f32>> = 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<WorkerStates, String> {
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<i64>) -> Result<Vec<FolderWorkerStates>, 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())
}
src-tauri/src/db.rs
+438 -30
View File
@@ -57,6 +57,10 @@ pub struct ImageRecord {
pub embedding_model: Option<String>,
pub embedding_updated_at: Option<String>,
pub embedding_error: Option<String>,
pub generated_caption: Option<String>,
pub caption_model: Option<String>,
pub caption_updated_at: Option<String>,
pub caption_error: Option<String>,
}
#[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<DbPool> {
@@ -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<i64> {
pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result<i64> {
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<i64> {
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<i64> {
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<usize> {
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::<rusqlite::Result<Vec<_>>>()?;
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<usize> {
// 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<Vec<EmbeddingJob>> {
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<usize> {
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<usize> {
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<i64>,
limit: usize,
) -> Result<Vec<EmbeddingJob>> {
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<Vec
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
pub fn claim_embedding_jobs(conn: &mut Connection, limit: usize) -> Result<Vec<EmbeddingJob>> {
pub fn claim_embedding_jobs(
conn: &mut Connection,
paused_folder_ids: &std::collections::HashSet<i64>,
limit: usize,
) -> Result<Vec<EmbeddingJob>> {
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<Vec<E
Ok(claimed)
}
fn get_pending_caption_jobs_excluding(
conn: &Connection,
excluded_folder_ids: &std::collections::HashSet<i64>,
limit: usize,
) -> Result<Vec<CaptionJob>> {
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::<rusqlite::Result<Vec<_>>>()?)
}
pub fn claim_caption_jobs(
conn: &mut Connection,
paused_folder_ids: &std::collections::HashSet<i64>,
limit: usize,
) -> Result<Vec<CaptionJob>> {
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<Fold
|row| row.get(0),
)?;
let caption_pending = conn.query_row(
"SELECT COUNT(*)
FROM caption_jobs j
JOIN images i ON i.id = j.image_id
WHERE i.folder_id = ?1 AND j.status IN ('pending', 'processing')",
[folder_id],
|row| row.get(0),
)?;
let caption_ready = conn.query_row(
"SELECT COUNT(*)
FROM images
WHERE folder_id = ?1 AND generated_caption IS NOT NULL",
[folder_id],
|row| row.get(0),
)?;
let caption_failed = conn.query_row(
"SELECT COUNT(*)
FROM images
WHERE folder_id = ?1 AND caption_error IS NOT NULL",
[folder_id],
|row| row.get(0),
)?;
Ok(FolderJobProgress {
folder_id,
thumbnail_pending,
@@ -546,6 +805,9 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result<Fold
embedding_pending,
embedding_ready,
embedding_failed,
caption_pending,
caption_ready,
caption_failed,
})
}
@@ -562,15 +824,22 @@ pub fn get_all_folder_job_progress(conn: &Connection) -> Result<Vec<FolderJobPro
Ok(progress)
}
pub fn get_pending_thumbnail_jobs(conn: &Connection, limit: usize) -> Result<Vec<ThumbnailJob>> {
let mut stmt = conn.prepare(
fn get_pending_thumbnail_jobs_excluding(
conn: &Connection,
excluded_folder_ids: &std::collections::HashSet<i64>,
limit: usize,
) -> Result<Vec<ThumbnailJob>> {
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<Vec
pub fn claim_thumbnail_jobs(
conn: &mut Connection,
active_folder_ids: &std::collections::HashSet<i64>,
paused_folder_ids: &std::collections::HashSet<i64>,
fetch_limit: usize,
claim_limit: usize,
) -> Result<Vec<ThumbnailJob>> {
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::<std::collections::HashSet<_>>();
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<Vec<MetadataJob>> {
let mut stmt = conn.prepare(
fn get_pending_metadata_jobs_excluding(
conn: &Connection,
excluded_folder_ids: &std::collections::HashSet<i64>,
limit: usize,
) -> Result<Vec<MetadataJob>> {
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<Vec<
pub fn claim_metadata_jobs(
conn: &mut Connection,
active_folder_ids: &std::collections::HashSet<i64>,
paused_folder_ids: &std::collections::HashSet<i64>,
fetch_limit: usize,
claim_limit: usize,
) -> Result<Vec<MetadataJob>> {
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::<std::collections::HashSet<_>>();
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<ImageRecord>
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<ImageRecord>
pub fn get_images_by_ids(conn: &Connection, image_ids: &[i64]) -> Result<Vec<ImageRecord>> {
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::<rusqlite::Result<Vec<_>>>()?)
}
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<ImageRecord> {
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<Vec<String>> {
let caption = conn.query_row(
"SELECT generated_caption FROM images WHERE id = ?1",
[image_id],
|row| row.get::<_, Option<String>>(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::<rusqlite::Result<Vec<_>>>()?;
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<String> {
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::<std::collections::HashSet<_>>();
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<ImageRecord> {
Ok(ImageRecord {
id: row.get(0)?,
@@ -948,6 +1328,10 @@ fn map_image_row(row: &Row<'_>) -> rusqlite::Result<ImageRecord> {
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::<rusqlite::Error>()
.is_some_and(|error| matches!(error, rusqlite::Error::QueryReturnedNoRows))
}
fn folder_exclusion_clause(
image_alias: &str,
excluded_folder_ids: &std::collections::HashSet<i64>,
) -> String {
if excluded_folder_ids.is_empty() {
return String::new();
}
let mut ids = excluded_folder_ids.iter().copied().collect::<Vec<_>>();
ids.sort_unstable();
let id_list = ids
.into_iter()
.map(|id| id.to_string())
.collect::<Vec<_>>()
.join(",");
format!("AND {}.folder_id NOT IN ({})", image_alias, id_list)
}
src-tauri/src/indexer.rs
+218 -44
View File
@@ -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<Mutex<HashMap<i64, Instant>>> = OnceLock::new();
static ACTIVE_INDEXING_FOLDERS: OnceLock<Mutex<HashSet<i64>>> = 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<Mutex<PausedWorkerFolders>> = 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<i64>,
metadata: HashSet<i64>,
embedding: HashSet<i64>,
caption: HashSet<i64>,
}
#[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<i64, FolderWorkerPausedState> {
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<i64> {
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<Mutex<HashMap<i64, RuntimeAdaptiveProfile>>> =
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<ClipImageEmbedder> = 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<FlorenceCaptioner> = 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::<Vec<_>>());
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), 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::<Vec<_>>());
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
}
Ok(())
@@ -518,14 +590,11 @@ fn process_embedding_batch(
embedder: &mut Option<ClipImageEmbedder>,
) -> 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::<HashSet<_>>();
emit_folder_job_progress(app, pool, &folder_ids.iter().copied().collect::<Vec<_>>());
emit_folder_job_progress(
app,
pool,
&folder_ids.iter().copied().collect::<Vec<_>>(),
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::<Vec<_>>());
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), 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<FlorenceCaptioner>,
) -> 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::<Vec<_>>(),
)
})?;
return Err(error);
}
}
}
let folder_ids = jobs.iter().map(|job| job.folder_id).collect::<HashSet<_>>();
emit_folder_job_progress(
app,
pool,
&folder_ids.iter().copied().collect::<Vec<_>>(),
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::<Vec<_>>();
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::<HashSet<_>>();
emit_media_updates(
app,
&MediaUpdateBatch {
images: updated_images,
},
);
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
}
Ok(())
}
fn active_indexing_folders() -> HashSet<i64> {
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::<Vec<_>>();
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);
}
src-tauri/src/lib.rs
+23 -3
View File
@@ -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,
src-tauri/src/vector.rs
+147 -4
View File
@@ -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<Vec<i64>> {
let embedding: Vec<u8> = conn.query_row(
let embedding: Vec<u8> = 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<Vec<i64>> {
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<i64> {
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<i64> {
conn.query_row("SELECT COUNT(*) FROM caption_vec", [], |row| row.get(0))
.map_err(Into::into)
}
pub fn delete_orphaned_embeddings(conn: &Connection) -> Result<usize> {
let image_ids = {
let mut stmt = conn.prepare("SELECT id FROM images")?;
let rows = stmt
.query_map([], |row| row.get::<_, i64>(0))?
.collect::<rusqlite::Result<std::collections::HashSet<_>>>()?;
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::<rusqlite::Result<Vec<_>>>()?;
rows
};
let orphaned_ids = vector_ids
.into_iter()
.filter(|image_id| !image_ids.contains(image_id))
.collect::<Vec<_>>();
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<usize> {
let image_ids = {
let mut stmt = conn.prepare("SELECT id FROM images")?;
let rows = stmt
.query_map([], |row| row.get::<_, i64>(0))?
.collect::<rusqlite::Result<std::collections::HashSet<_>>>()?;
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::<rusqlite::Result<Vec<_>>>()?;
rows
};
let orphaned_ids = vector_ids
.into_iter()
.filter(|image_id| !image_ids.contains(image_id))
.collect::<Vec<_>>();
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<bool> {
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<u8> {
let mut out = Vec::with_capacity(values.len() * std::mem::size_of::<f32>());