Add CLIP embeddings and similar-image search

- add Candle + HF Hub CLIP image embedding pipeline with background embedding worker
- write image embeddings into sqlite-vec and expose similar-image lookup through a new backend command
- surface embedding progress and recovery in the UI, including retries for failed embeddings
- improve gallery/lightbox embedding UX and make similar-image actions directly accessible

Refs: #3, #4
This commit is contained in:
2026-04-06 00:54:14 +01:00
parent 76ec424167
commit 51e4c2c1f7
14 changed files with 2219 additions and 73 deletions
+42 -1
View File
@@ -1,5 +1,6 @@
use crate::db::{self, DbPool, Folder, ImageRecord};
use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord};
use crate::indexer;
use crate::vector;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use tauri::{AppHandle, State};
@@ -32,6 +33,17 @@ pub struct UpdateImageDetailsParams {
pub rating: Option<i64>,
}
#[derive(Deserialize)]
pub struct FindSimilarImagesParams {
pub image_id: i64,
pub limit: Option<usize>,
}
#[derive(Deserialize)]
pub struct RetryFailedEmbeddingsParams {
pub folder_id: i64,
}
#[tauri::command]
pub async fn add_folder(
app: AppHandle,
@@ -75,6 +87,14 @@ pub async fn get_folders(db: State<'_, DbState>) -> Result<Vec<Folder>, String>
db::get_folders(&conn).map_err(|e| e.to_string())
}
#[tauri::command]
pub async fn get_background_job_progress(
db: State<'_, DbState>,
) -> Result<Vec<FolderJobProgress>, String> {
let conn = db.get().map_err(|e| e.to_string())?;
db::get_all_folder_job_progress(&conn).map_err(|e| e.to_string())
}
#[tauri::command]
pub async fn remove_folder(db: State<'_, DbState>, folder_id: i64) -> Result<(), String> {
let conn = db.get().map_err(|e| e.to_string())?;
@@ -147,3 +167,24 @@ pub async fn reindex_folder(
indexer::index_folder(app, db.inner().clone(), folder_id, folder_path);
Ok(())
}
#[tauri::command]
pub async fn find_similar_images(
db: State<'_, DbState>,
params: FindSimilarImagesParams,
) -> 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())?;
db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string())
}
#[tauri::command]
pub async fn retry_failed_embeddings(
db: State<'_, DbState>,
params: RetryFailedEmbeddingsParams,
) -> Result<usize, String> {
let conn = db.get().map_err(|e| e.to_string())?;
db::retry_failed_embedding_jobs(&conn, params.folder_id).map_err(|e| e.to_string())
}
+152 -29
View File
@@ -63,6 +63,10 @@ pub struct ImageRecord {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingJob {
pub image_id: i64,
pub folder_id: i64,
pub path: String,
pub thumbnail_path: Option<String>,
pub media_kind: String,
pub status: String,
pub attempts: i64,
pub last_error: Option<String>,
@@ -99,6 +103,9 @@ pub struct FolderJobProgress {
pub folder_id: i64,
pub thumbnail_pending: i64,
pub metadata_pending: i64,
pub embedding_pending: i64,
pub embedding_ready: i64,
pub embedding_failed: i64,
}
pub fn create_pool(db_path: &Path) -> Result<DbPool> {
@@ -280,6 +287,56 @@ pub fn enqueue_embedding_job(conn: &Connection, image_id: i64) -> Result<()> {
Ok(())
}
pub fn backfill_embedding_jobs(conn: &Connection) -> Result<usize> {
let inserted = conn.execute(
"INSERT INTO embedding_jobs (image_id, status, attempts, last_error, created_at, updated_at)
SELECT i.id, 'pending', 0, NULL, datetime('now'), datetime('now')
FROM images i
LEFT JOIN embedding_jobs j ON j.image_id = i.id
WHERE i.embedding_status != 'ready'
AND j.image_id IS NULL",
[],
)?;
Ok(inserted)
}
pub fn retry_failed_embedding_jobs(conn: &Connection, folder_id: i64) -> Result<usize> {
let updated = conn.execute(
"INSERT INTO embedding_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 embedding_status = 'failed'
ON CONFLICT(image_id) DO UPDATE SET
status = 'pending',
last_error = NULL,
updated_at = datetime('now')",
[folder_id],
)?;
conn.execute(
"UPDATE images
SET embedding_status = 'pending', embedding_error = NULL
WHERE folder_id = ?1 AND embedding_status = 'failed'",
[folder_id],
)?;
Ok(updated)
}
pub fn reset_inflight_jobs(conn: &Connection) -> Result<()> {
conn.execute(
"UPDATE thumbnail_jobs SET status = 'pending' WHERE status = 'processing'",
[],
)?;
conn.execute(
"UPDATE metadata_jobs SET status = 'pending' WHERE status = 'processing'",
[],
)?;
conn.execute(
"UPDATE embedding_jobs SET status = 'pending' WHERE status = 'processing'",
[],
)?;
Ok(())
}
pub fn enqueue_thumbnail_job(conn: &Connection, image_id: i64) -> Result<()> {
conn.execute(
"INSERT INTO thumbnail_jobs (image_id, status, attempts, last_error, created_at, updated_at)
@@ -306,40 +363,57 @@ pub fn enqueue_metadata_job(conn: &Connection, image_id: i64) -> Result<()> {
Ok(())
}
#[allow(dead_code)]
pub fn get_next_embedding_job(conn: &Connection) -> Result<Option<EmbeddingJob>> {
pub fn get_pending_embedding_jobs(conn: &Connection, limit: usize) -> Result<Vec<EmbeddingJob>> {
let mut stmt = conn.prepare(
"SELECT image_id, status, attempts, last_error, created_at, updated_at
FROM embedding_jobs
"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 updated_at, image_id
LIMIT 1",
ORDER BY j.updated_at, j.image_id
LIMIT ?1",
)?;
let mut rows = stmt.query([])?;
let Some(row) = rows.next()? else {
return Ok(None);
};
Ok(Some(EmbeddingJob {
image_id: row.get(0)?,
status: row.get(1)?,
attempts: row.get(2)?,
last_error: row.get(3)?,
created_at: row.get(4)?,
updated_at: row.get(5)?,
}))
let rows = stmt.query_map([limit as i64], |row| {
Ok(EmbeddingJob {
image_id: row.get(0)?,
folder_id: row.get(1)?,
path: row.get(2)?,
thumbnail_path: row.get(3)?,
media_kind: row.get(4)?,
status: row.get(5)?,
attempts: row.get(6)?,
last_error: row.get(7)?,
created_at: row.get(8)?,
updated_at: row.get(9)?,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
#[allow(dead_code)]
pub fn mark_embedding_job_processing(conn: &Connection, image_id: i64) -> Result<()> {
conn.execute(
"UPDATE embedding_jobs
SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now')
WHERE image_id = ?1",
[image_id],
)?;
Ok(())
pub fn claim_embedding_jobs(conn: &mut Connection, limit: usize) -> Result<Vec<EmbeddingJob>> {
let tx = conn.transaction()?;
let candidates = get_pending_embedding_jobs(&tx, limit * 2)?;
let mut claimed = Vec::with_capacity(limit);
for job in candidates {
let updated = tx.execute(
"UPDATE embedding_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)]
@@ -426,13 +500,54 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result<Fold
|row| row.get(0),
)?;
let embedding_pending = conn.query_row(
"SELECT COUNT(*)
FROM embedding_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 embedding_ready = conn.query_row(
"SELECT COUNT(*)
FROM images
WHERE folder_id = ?1 AND embedding_status = 'ready'",
[folder_id],
|row| row.get(0),
)?;
let embedding_failed = conn.query_row(
"SELECT COUNT(*)
FROM images
WHERE folder_id = ?1 AND embedding_status = 'failed'",
[folder_id],
|row| row.get(0),
)?;
Ok(FolderJobProgress {
folder_id,
thumbnail_pending,
metadata_pending,
embedding_pending,
embedding_ready,
embedding_failed,
})
}
pub fn get_all_folder_job_progress(conn: &Connection) -> Result<Vec<FolderJobProgress>> {
let folder_ids = get_folders(conn)?
.into_iter()
.map(|folder| folder.id)
.collect::<Vec<_>>();
let mut progress = Vec::with_capacity(folder_ids.len());
for folder_id in folder_ids {
progress.push(get_folder_job_progress(conn, folder_id)?);
}
Ok(progress)
}
pub fn get_pending_thumbnail_jobs(conn: &Connection, limit: usize) -> Result<Vec<ThumbnailJob>> {
let mut stmt = conn.prepare(
"SELECT j.image_id, i.folder_id, i.path, i.media_kind
@@ -658,6 +773,14 @@ pub fn get_image_by_id(conn: &Connection, image_id: i64) -> Result<ImageRecord>
.map_err(Into::into)
}
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)?);
}
Ok(images)
}
pub fn get_folders(conn: &Connection) -> Result<Vec<Folder>> {
let mut stmt =
conn.prepare("SELECT id, path, name, image_count, indexed_at FROM folders ORDER BY name")?;
+118
View File
@@ -0,0 +1,118 @@
use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::clip::{self, ClipModel};
use hf_hub::{api::sync::Api, Repo, RepoType};
use std::path::{Path, PathBuf};
pub struct ClipImageEmbedder {
model: ClipModel,
device: Device,
image_size: usize,
}
impl ClipImageEmbedder {
pub fn new() -> Result<Self> {
println!("Initializing CLIP image embedder...");
let api = Api::new()?;
let repo = api.repo(Repo::new(
"laion/CLIP-ViT-B-32-laion2B-s34B-b79K".to_string(),
RepoType::Model,
));
println!("Resolving CLIP model weights from Hugging Face cache...");
let model_path = repo.get("model.safetensors")?;
let config = clip::ClipConfig::vit_base_patch32();
let device = resolve_device()?;
let vb = unsafe {
VarBuilder::from_mmaped_safetensors(
std::slice::from_ref(&model_path),
DType::F32,
&device,
)?
};
let model = ClipModel::new(vb, &config)?;
println!("CLIP image embedder ready.");
Ok(Self {
model,
device,
image_size: config.image_size,
})
}
pub fn embed_image(&self, path: &Path) -> Result<Vec<f32>> {
Ok(self.embed_images(&[path.to_path_buf()])?.remove(0))
}
pub fn embed_images(&self, paths: &[PathBuf]) -> Result<Vec<Vec<f32>>> {
let images = load_images(paths, self.image_size)?.to_device(&self.device)?;
let features = self.model.get_image_features(&images)?;
let normalized = clip::div_l2_norm(&features)?;
let mut embeddings = Vec::with_capacity(paths.len());
for index in 0..paths.len() {
embeddings.push(normalized.get(index)?.flatten_all()?.to_vec1::<f32>()?);
}
Ok(embeddings)
}
}
fn resolve_device() -> Result<Device> {
let cuda_device = Device::cuda_if_available(0)?;
if cuda_device.is_cuda() {
println!("CLIP embedder using CUDA device.");
return Ok(cuda_device);
}
#[cfg(target_os = "macos")]
{
let metal_device = Device::metal_if_available(0)?;
if metal_device.is_metal() {
println!("CLIP embedder using Metal device.");
return Ok(metal_device);
}
}
println!("CLIP embedder using CPU device.");
Ok(Device::Cpu)
}
fn load_image(path: &Path, image_size: usize) -> Result<Tensor> {
let image = image::ImageReader::open(path)?
.with_guessed_format()?
.decode()?;
let image = image.resize_to_fill(
image_size as u32,
image_size as u32,
image::imageops::FilterType::Triangle,
);
let image = image.to_rgb8().into_raw();
let tensor = Tensor::from_vec(image, (image_size, image_size, 3), &Device::Cpu)?
.permute((2, 0, 1))?
.to_dtype(DType::F32)?
.affine(2.0 / 255.0, -1.0)?;
Ok(tensor)
}
fn load_images(paths: &[PathBuf], image_size: usize) -> Result<Tensor> {
let mut images = Vec::with_capacity(paths.len());
for path in paths {
images.push(load_image(path, image_size)?);
}
Ok(Tensor::stack(&images, 0)?)
}
pub fn embedding_source_path(
path: &str,
thumbnail_path: Option<&str>,
media_kind: &str,
) -> PathBuf {
if media_kind == "video" {
thumbnail_path
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from(path))
} else {
PathBuf::from(path)
}
}
+117
View File
@@ -1,4 +1,5 @@
use crate::db::{self, DbPool, FolderJobProgress, ImageRecord, IndexedMediaEntry};
use crate::embedder::{embedding_source_path, ClipImageEmbedder};
use crate::media::{probe_video_metadata, MediaTools};
use crate::storage::{detect_storage_profile, RuntimeAdaptiveProfile, StorageProfile};
use crate::thumbnail;
@@ -26,6 +27,7 @@ static ACTIVE_INDEXING_FOLDERS: OnceLock<Mutex<HashSet<i64>>> = OnceLock::new();
static FOLDER_STORAGE_PROFILES: OnceLock<Mutex<HashMap<i64, RuntimeAdaptiveProfile>>> =
OnceLock::new();
static DB_WRITE_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
const EMBEDDING_BATCH_SIZE: usize = 8;
#[derive(Clone, Serialize)]
pub struct IndexProgress {
@@ -89,6 +91,20 @@ pub fn start_metadata_worker(app: AppHandle, pool: DbPool, media_tools: MediaToo
});
}
pub fn start_embedding_worker(app: AppHandle, pool: DbPool) {
std::thread::spawn(move || {
let mut embedder: Option<ClipImageEmbedder> = None;
println!("Embedding worker started.");
loop {
if let Err(error) = process_embedding_batch(&app, &pool, &mut embedder) {
eprintln!("Embedding worker error: {}", error);
}
std::thread::sleep(std::time::Duration::from_millis(500));
}
});
}
fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> {
let existing_entries = {
let conn = pool.get()?;
@@ -303,6 +319,8 @@ fn process_thumbnail_batch(
return Ok(());
}
println!("Embedding batch claimed: {} items", jobs.len());
let (image_jobs, video_jobs): (Vec<_>, Vec<_>) =
jobs.into_iter().partition(|job| job.media_kind == "image");
@@ -463,6 +481,105 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo
Ok(())
}
fn process_embedding_batch(
app: &AppHandle,
pool: &DbPool,
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 jobs = with_db_write_lock(|| {
let mut conn = pool.get()?;
db::claim_embedding_jobs(&mut conn, EMBEDDING_BATCH_SIZE)
})?;
let claim_elapsed = claim_started_at.elapsed();
if jobs.is_empty() {
return Ok(());
}
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<_>>());
let embedder = embedder.as_ref().expect("embedder should be initialized");
let infer_started_at = Instant::now();
let source_paths = jobs
.iter()
.map(|job| embedding_source_path(&job.path, job.thumbnail_path.as_deref(), &job.media_kind))
.collect::<Vec<_>>();
let results = match embedder.embed_images(&source_paths) {
Ok(embeddings) => jobs
.into_iter()
.zip(embeddings.into_iter().map(Ok))
.collect::<Vec<_>>(),
Err(batch_error) => {
eprintln!(
"Embedding batch fallback to per-image mode: {}",
batch_error
);
jobs.into_iter()
.zip(source_paths.into_iter())
.map(|(job, source_path)| (job, embedder.embed_image(&source_path)))
.collect::<Vec<_>>()
}
};
let infer_elapsed = infer_started_at.elapsed();
let write_started_at = Instant::now();
let updated_images = with_db_write_lock(|| {
let mut conn = pool.get()?;
let tx = conn.transaction()?;
let mut updated_images = Vec::new();
for (job, embedding_result) in results {
match embedding_result {
Ok(embedding) => {
vector::upsert_embedding(&tx, job.image_id, &embedding)?;
db::mark_embedding_ready(&tx, job.image_id, vector::CLIP_MODEL_NAME)?;
}
Err(error) => {
db::mark_embedding_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() {
println!("Embedding batch completed: {} items", updated_images.len());
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<_>>());
}
let write_elapsed = write_started_at.elapsed();
let batch_elapsed = batch_started_at.elapsed();
println!(
"Embedding batch timing: claimed {} in {:?}, infer {:?}, write {:?}, total {:?}",
EMBEDDING_BATCH_SIZE, claim_elapsed, infer_elapsed, write_elapsed, batch_elapsed
);
Ok(())
}
fn active_indexing_folders() -> HashSet<i64> {
ACTIVE_INDEXING_FOLDERS
.get_or_init(|| Mutex::new(HashSet::new()))
+11
View File
@@ -1,5 +1,6 @@
mod commands;
mod db;
mod embedder;
mod indexer;
mod media;
mod storage;
@@ -32,6 +33,12 @@ 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");
if backfilled > 0 {
println!("Backfilled {} embedding jobs.", backfilled);
}
}
let thumb_dir = app_dir.join("thumbnails");
@@ -50,6 +57,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());
app.manage(pool);
app.manage(media_tools);
@@ -59,10 +67,13 @@ pub fn run() {
.invoke_handler(tauri::generate_handler![
commands::add_folder,
commands::get_folders,
commands::get_background_job_progress,
commands::remove_folder,
commands::get_images,
commands::reindex_folder,
commands::update_image_details,
commands::find_similar_images,
commands::retry_failed_embeddings,
])
.run(tauri::generate_context!())
.expect("error while running tauri application");
+28
View File
@@ -50,6 +50,34 @@ pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) ->
Ok(())
}
pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Result<Vec<i64>> {
let embedding: Vec<u8> = conn.query_row(
"SELECT embedding FROM image_vec WHERE image_id = ?1",
[image_id],
|row| row.get(0),
)?;
let mut stmt = conn.prepare(
"SELECT image_id
FROM image_vec
WHERE embedding MATCH vec_f32(?1)
AND k = ?2",
)?;
let rows = stmt.query_map((&embedding, (limit as i64) + 1), |row| row.get::<_, i64>(0))?;
let mut ids = Vec::new();
for row in rows {
let candidate_id = row?;
if candidate_id != image_id {
ids.push(candidate_id);
}
if ids.len() >= limit {
break;
}
}
Ok(ids)
}
#[allow(dead_code)]
fn pack_f32(values: &[f32]) -> Vec<u8> {
let mut out = Vec::with_capacity(values.len() * std::mem::size_of::<f32>());