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:
Generated
+1565
-18
File diff suppressed because it is too large
Load Diff
@@ -35,3 +35,7 @@ log = "0.4"
|
||||
ffmpeg-sidecar = "2.5.0"
|
||||
xxhash-rust = { version = "0.8", features = ["xxh3"] }
|
||||
sysinfo = "0.38.4"
|
||||
candle-core = { version = "0.10.2", features = ["cuda"] }
|
||||
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"] }
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
+148
-25
@@ -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 {
|
||||
let rows = stmt.query_map([limit as i64], |row| {
|
||||
Ok(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)?,
|
||||
}))
|
||||
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(
|
||||
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",
|
||||
[image_id],
|
||||
WHERE image_id = ?1 AND status = 'pending'",
|
||||
[job.image_id],
|
||||
)?;
|
||||
Ok(())
|
||||
|
||||
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")?;
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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()))
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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>());
|
||||
|
||||
+5
-1
@@ -9,11 +9,15 @@ import { Lightbox } from "./components/Lightbox";
|
||||
|
||||
export default function App() {
|
||||
const loadFolders = useGalleryStore((state) => state.loadFolders);
|
||||
const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress);
|
||||
const loadImages = useGalleryStore((state) => state.loadImages);
|
||||
const subscribeToProgress = useGalleryStore((state) => state.subscribeToProgress);
|
||||
|
||||
useEffect(() => {
|
||||
loadFolders().then(() => loadImages(true));
|
||||
loadFolders().then(() => {
|
||||
void loadBackgroundJobProgress();
|
||||
return loadImages(true);
|
||||
});
|
||||
let unlisten: (() => void) | undefined;
|
||||
subscribeToProgress().then((fn) => {
|
||||
unlisten = fn;
|
||||
|
||||
@@ -16,19 +16,27 @@ export function BackgroundTasks() {
|
||||
const folders = useGalleryStore((state) => state.folders);
|
||||
const indexingProgress = useGalleryStore((state) => state.indexingProgress);
|
||||
const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress);
|
||||
const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings);
|
||||
|
||||
const tasks = useMemo(() => {
|
||||
return folders
|
||||
.map((folder) => {
|
||||
const index = indexingProgress[folder.id];
|
||||
const jobs = mediaJobProgress[folder.id];
|
||||
const pendingMediaWork = (jobs?.thumbnail_pending ?? 0) + (jobs?.metadata_pending ?? 0);
|
||||
const pendingMediaWork =
|
||||
(jobs?.thumbnail_pending ?? 0) +
|
||||
(jobs?.metadata_pending ?? 0) +
|
||||
(jobs?.embedding_pending ?? 0);
|
||||
const embeddingProcessed = (jobs?.embedding_ready ?? 0) + (jobs?.embedding_failed ?? 0);
|
||||
const embeddingTotal = embeddingProcessed + (jobs?.embedding_pending ?? 0);
|
||||
const hasFailedEmbeddings = (jobs?.embedding_failed ?? 0) > 0;
|
||||
|
||||
if (!index && pendingMediaWork === 0) {
|
||||
if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const indexPercent = index && index.total > 0 ? (index.indexed / index.total) * 100 : 0;
|
||||
const embeddingPercent = embeddingTotal > 0 ? (embeddingProcessed / embeddingTotal) * 100 : 0;
|
||||
return {
|
||||
id: folder.id,
|
||||
name: folder.name,
|
||||
@@ -36,6 +44,10 @@ export function BackgroundTasks() {
|
||||
jobs,
|
||||
pendingMediaWork,
|
||||
indexPercent,
|
||||
embeddingProcessed,
|
||||
embeddingTotal,
|
||||
embeddingPercent,
|
||||
hasFailedEmbeddings,
|
||||
};
|
||||
})
|
||||
.filter((task) => task !== null);
|
||||
@@ -60,12 +72,20 @@ export function BackgroundTasks() {
|
||||
<p className="text-[11px] text-gray-500">
|
||||
{task.index && !task.index.done
|
||||
? `${task.index.indexed.toLocaleString()} of ${task.index.total.toLocaleString()} scanned`
|
||||
: task.hasFailedEmbeddings && task.pendingMediaWork === 0
|
||||
? `Embedding failures require attention`
|
||||
: `${task.pendingMediaWork.toLocaleString()} media jobs remaining`}
|
||||
</p>
|
||||
</div>
|
||||
<div className="text-right text-[11px] text-gray-400">
|
||||
{task.jobs?.thumbnail_pending ? <div>{task.jobs.thumbnail_pending.toLocaleString()} thumbnails</div> : null}
|
||||
{task.jobs?.metadata_pending ? <div>{task.jobs.metadata_pending.toLocaleString()} metadata</div> : null}
|
||||
{task.embeddingTotal > 0 ? (
|
||||
<div>
|
||||
{task.embeddingProcessed.toLocaleString()} / {task.embeddingTotal.toLocaleString()} embeddings
|
||||
</div>
|
||||
) : null}
|
||||
{task.jobs?.embedding_failed ? <div>{task.jobs.embedding_failed.toLocaleString()} failed</div> : null}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -74,10 +94,32 @@ export function BackgroundTasks() {
|
||||
<ProgressBar value={task.indexPercent} />
|
||||
<p className="truncate text-[11px] text-gray-500">{task.index.current_file || "Scanning..."}</p>
|
||||
</div>
|
||||
) : task.embeddingTotal > 0 && (task.jobs?.embedding_pending ?? 0) > 0 ? (
|
||||
<div className="mt-2 space-y-1">
|
||||
<ProgressBar value={task.embeddingPercent} />
|
||||
<p className="text-[11px] text-gray-500">
|
||||
{task.embeddingProcessed.toLocaleString()} completed, {task.jobs?.embedding_pending?.toLocaleString() ?? 0} remaining
|
||||
</p>
|
||||
</div>
|
||||
) : task.hasFailedEmbeddings ? (
|
||||
<div className="mt-2 space-y-1">
|
||||
<ProgressBar value={100} />
|
||||
<div className="flex items-center justify-between gap-3">
|
||||
<p className="text-[11px] text-amber-300">
|
||||
{task.jobs?.embedding_failed?.toLocaleString() ?? 0} embedding failures need attention
|
||||
</p>
|
||||
<button
|
||||
className="rounded-full border border-amber-400/20 bg-amber-500/10 px-2.5 py-1 text-[11px] text-amber-200 hover:bg-amber-500/20"
|
||||
onClick={() => void retryFailedEmbeddings(task.id)}
|
||||
>
|
||||
Retry
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
) : task.pendingMediaWork > 0 ? (
|
||||
<div className="mt-2 space-y-1">
|
||||
<ProgressBar value={0} />
|
||||
<p className="text-[11px] text-gray-500">Processing thumbnails and metadata</p>
|
||||
<p className="text-[11px] text-gray-500">Processing thumbnails, metadata, and embeddings</p>
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
|
||||
+58
-10
@@ -38,6 +38,19 @@ function formatDuration(durationMs: number | null): string | null {
|
||||
return `${minutes}:${seconds.toString().padStart(2, "0")}`;
|
||||
}
|
||||
|
||||
function embeddingLabel(image: ImageRecord): string {
|
||||
if (image.embedding_status === "ready") {
|
||||
return image.embedding_model ? `Embeddings ready` : "Embeddings ready";
|
||||
}
|
||||
if (image.embedding_status === "failed") {
|
||||
return "Embeddings failed";
|
||||
}
|
||||
if (image.embedding_status === "processing") {
|
||||
return "Embedding...";
|
||||
}
|
||||
return "Embedding queued";
|
||||
}
|
||||
|
||||
function ContextMenu({
|
||||
x,
|
||||
y,
|
||||
@@ -49,10 +62,13 @@ function ContextMenu({
|
||||
image: ImageRecord;
|
||||
onClose: () => void;
|
||||
}) {
|
||||
const { openImage, updateImageDetails } = useGalleryStore();
|
||||
const openImage = useGalleryStore((state) => state.openImage);
|
||||
const updateImageDetails = useGalleryStore((state) => state.updateImageDetails);
|
||||
const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages);
|
||||
|
||||
return (
|
||||
<div
|
||||
data-gallery-context-menu
|
||||
className="fixed z-40 min-w-56 rounded-2xl border border-white/10 bg-gray-950/95 p-2 shadow-2xl backdrop-blur"
|
||||
style={{ left: x, top: y }}
|
||||
onClick={(event) => event.stopPropagation()}
|
||||
@@ -75,21 +91,37 @@ function ContextMenu({
|
||||
>
|
||||
{image.favorite ? "Remove Favorite" : "Add to Favorites"}
|
||||
</button>
|
||||
<button
|
||||
className="w-full rounded-xl px-3 py-2 text-left text-sm text-gray-200 hover:bg-white/5"
|
||||
onClick={async () => {
|
||||
await loadSimilarImages(image.id);
|
||||
onClose();
|
||||
}}
|
||||
>
|
||||
Find Similar
|
||||
</button>
|
||||
<div className="my-2 h-px bg-white/5" />
|
||||
<div className="px-3 pb-1 pt-1 text-[11px] uppercase tracking-[0.2em] text-gray-500">Rating</div>
|
||||
<div className="flex gap-1 px-2 pb-1">
|
||||
<div className="flex items-center gap-1 px-2 pb-1">
|
||||
{Array.from({ length: 5 }, (_, index) => {
|
||||
const rating = index + 1;
|
||||
return (
|
||||
<button
|
||||
key={rating}
|
||||
className={`rounded-lg px-2 py-1 text-sm ${rating <= image.rating ? "bg-amber-400/15 text-amber-300" : "bg-white/5 text-gray-400 hover:text-white"}`}
|
||||
className="rounded-md p-1"
|
||||
onClick={async () => {
|
||||
await updateImageDetails(image.id, { rating });
|
||||
onClose();
|
||||
}}
|
||||
title={`Set ${rating} star rating`}
|
||||
>
|
||||
{rating}
|
||||
<svg
|
||||
className={`h-5 w-5 ${rating <= image.rating ? "text-amber-300" : "text-white/20 hover:text-white/50"}`}
|
||||
fill="currentColor"
|
||||
viewBox="0 0 20 20"
|
||||
>
|
||||
<path d="M9.049 2.927c.3-.921 1.603-.921 1.902 0l1.07 3.292a1 1 0 00.95.69h3.462c.969 0 1.371 1.24.588 1.81l-2.8 2.034a1 1 0 00-.364 1.118l1.07 3.292c.3.921-.755 1.688-1.54 1.118l-2.8-2.034a1 1 0 00-1.176 0l-2.8 2.034c-.784.57-1.838-.197-1.539-1.118l1.07-3.292a1 1 0 00-.364-1.118L2.98 8.72c-.783-.57-.38-1.81.588-1.81H7.03a1 1 0 00.951-.69l1.07-3.292z" />
|
||||
</svg>
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
@@ -123,6 +155,7 @@ function ImageTile({
|
||||
}) {
|
||||
const [loaded, setLoaded] = useState(false);
|
||||
const [errored, setErrored] = useState(false);
|
||||
const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages);
|
||||
|
||||
const src = image.thumbnail_path
|
||||
? convertFileSrc(image.thumbnail_path)
|
||||
@@ -207,7 +240,19 @@ function ImageTile({
|
||||
<p className="truncate text-sm font-medium text-white">{image.filename}</p>
|
||||
<div className="mt-1 flex items-center justify-between gap-2 text-xs text-white/70">
|
||||
<RatingStars rating={image.rating} />
|
||||
<span>{image.embedding_status === "ready" ? "indexed" : image.embedding_status}</span>
|
||||
<span>{embeddingLabel(image)}</span>
|
||||
</div>
|
||||
<div className="mt-2 flex items-center gap-2 opacity-0 transition-opacity group-hover:opacity-100">
|
||||
<button
|
||||
className="rounded-full border border-white/10 bg-black/40 px-2.5 py-1 text-[11px] text-white/85 hover:bg-black/60"
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
void loadSimilarImages(image.id);
|
||||
}}
|
||||
>
|
||||
Find Similar
|
||||
</button>
|
||||
<span className="text-[11px] text-white/50">Right-click for more</span>
|
||||
</div>
|
||||
</div>
|
||||
</button>
|
||||
@@ -242,18 +287,21 @@ export function Gallery() {
|
||||
}, [handleScroll]);
|
||||
|
||||
useEffect(() => {
|
||||
const close = () => setContextMenu(null);
|
||||
const close = (event: PointerEvent) => {
|
||||
if ((event.target as HTMLElement | null)?.closest("[data-gallery-context-menu]")) {
|
||||
return;
|
||||
}
|
||||
setContextMenu(null);
|
||||
};
|
||||
const handleKeyDown = (event: KeyboardEvent) => {
|
||||
if (event.key === "Escape") {
|
||||
setContextMenu(null);
|
||||
}
|
||||
};
|
||||
window.addEventListener("click", close);
|
||||
window.addEventListener("contextmenu", close);
|
||||
window.addEventListener("pointerdown", close);
|
||||
window.addEventListener("keydown", handleKeyDown);
|
||||
return () => {
|
||||
window.removeEventListener("click", close);
|
||||
window.removeEventListener("contextmenu", close);
|
||||
window.removeEventListener("pointerdown", close);
|
||||
window.removeEventListener("keydown", handleKeyDown);
|
||||
};
|
||||
}, []);
|
||||
|
||||
@@ -32,11 +32,25 @@ function formatDuration(durationMs: number | null): string {
|
||||
return `${minutes}:${seconds.toString().padStart(2, "0")}`;
|
||||
}
|
||||
|
||||
function embeddingLabel(status: string, model: string | null): string {
|
||||
if (status === "ready") {
|
||||
return model ? `Ready (${model})` : "Ready";
|
||||
}
|
||||
if (status === "failed") {
|
||||
return "Failed";
|
||||
}
|
||||
if (status === "processing") {
|
||||
return "Processing";
|
||||
}
|
||||
return "Queued";
|
||||
}
|
||||
|
||||
export function Lightbox() {
|
||||
const selectedImage = useGalleryStore((state) => state.selectedImage);
|
||||
const closeImage = useGalleryStore((state) => state.closeImage);
|
||||
const images = useGalleryStore((state) => state.images);
|
||||
const openImage = useGalleryStore((state) => state.openImage);
|
||||
const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages);
|
||||
const updateImageDetails = useGalleryStore((state) => state.updateImageDetails);
|
||||
const [zoom, setZoom] = useState(1);
|
||||
const imageViewportRef = useRef<HTMLDivElement>(null);
|
||||
@@ -184,6 +198,12 @@ export function Lightbox() {
|
||||
<path d="M3.172 5.172a4 4 0 015.656 0L10 6.343l1.172-1.171a4 4 0 115.656 5.656L10 17.657l-6.828-6.829a4 4 0 010-5.656z" />
|
||||
</svg>
|
||||
</button>
|
||||
<button
|
||||
className="rounded-full border border-white/10 bg-white/5 px-3 py-1.5 text-xs text-gray-300 hover:text-white"
|
||||
onClick={() => void loadSimilarImages(selectedImage.id)}
|
||||
>
|
||||
Similar
|
||||
</button>
|
||||
</div>
|
||||
<button className="rounded p-1 text-gray-400 hover:text-white" onClick={closeImage}>
|
||||
<svg className="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
|
||||
@@ -281,7 +301,10 @@ export function Lightbox() {
|
||||
|
||||
<div>
|
||||
<p className="mb-1 text-xs uppercase tracking-wider text-gray-500">Embedding</p>
|
||||
<p className="text-white">{selectedImage.embedding_status}</p>
|
||||
<p className="text-white">{embeddingLabel(selectedImage.embedding_status, selectedImage.embedding_model)}</p>
|
||||
{selectedImage.embedding_error ? (
|
||||
<p className="mt-1 text-xs text-amber-300">{selectedImage.embedding_error}</p>
|
||||
) : null}
|
||||
</div>
|
||||
|
||||
<div>
|
||||
|
||||
@@ -41,6 +41,7 @@ export function Toolbar() {
|
||||
const totalImages = useGalleryStore((state) => state.totalImages);
|
||||
const loadedCount = useGalleryStore((state) => state.loadedCount);
|
||||
const selectedFolderId = useGalleryStore((state) => state.selectedFolderId);
|
||||
const collectionTitle = useGalleryStore((state) => state.collectionTitle);
|
||||
const folders = useGalleryStore((state) => state.folders);
|
||||
const mediaFilter = useGalleryStore((state) => state.mediaFilter);
|
||||
const setMediaFilter = useGalleryStore((state) => state.setMediaFilter);
|
||||
@@ -53,7 +54,7 @@ export function Toolbar() {
|
||||
const debounceRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
const selectedFolder = folders.find((folder) => folder.id === selectedFolderId);
|
||||
const title = selectedFolder ? selectedFolder.name : "All Media";
|
||||
const title = collectionTitle ?? (selectedFolder ? selectedFolder.name : "All Media");
|
||||
const tileSize = tileSizeForZoom(zoomPreset);
|
||||
|
||||
useEffect(() => {
|
||||
|
||||
+47
-8
@@ -53,6 +53,9 @@ export interface FolderJobProgress {
|
||||
folder_id: number;
|
||||
thumbnail_pending: number;
|
||||
metadata_pending: number;
|
||||
embedding_pending: number;
|
||||
embedding_ready: number;
|
||||
embedding_failed: number;
|
||||
}
|
||||
|
||||
export interface MediaJobProgressEvent {
|
||||
@@ -89,11 +92,13 @@ interface GalleryState {
|
||||
favoritesOnly: boolean;
|
||||
zoomPreset: ZoomPreset;
|
||||
selectedImage: ImageRecord | null;
|
||||
collectionTitle: string | null;
|
||||
indexingProgress: Record<number, IndexProgress>;
|
||||
mediaJobProgress: Record<number, FolderJobProgress>;
|
||||
cacheDir: string;
|
||||
|
||||
loadFolders: () => Promise<void>;
|
||||
loadBackgroundJobProgress: () => Promise<void>;
|
||||
addFolder: (path: string) => Promise<void>;
|
||||
removeFolder: (folderId: number) => Promise<void>;
|
||||
reindexFolder: (folderId: number) => Promise<void>;
|
||||
@@ -107,6 +112,8 @@ interface GalleryState {
|
||||
setZoomPreset: (zoomPreset: ZoomPreset) => void;
|
||||
openImage: (image: ImageRecord) => void;
|
||||
closeImage: () => void;
|
||||
loadSimilarImages: (imageId: number) => Promise<void>;
|
||||
retryFailedEmbeddings: (folderId: number) => Promise<void>;
|
||||
updateImageDetails: (imageId: number, updates: { favorite?: boolean; rating?: number }) => Promise<void>;
|
||||
setCacheDir: (dir: string) => void;
|
||||
subscribeToProgress: () => Promise<UnlistenFn>;
|
||||
@@ -235,6 +242,7 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
favoritesOnly: false,
|
||||
zoomPreset: "comfortable",
|
||||
selectedImage: null,
|
||||
collectionTitle: null,
|
||||
indexingProgress: {},
|
||||
mediaJobProgress: {},
|
||||
cacheDir: "",
|
||||
@@ -246,16 +254,25 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
set({ folders });
|
||||
},
|
||||
|
||||
loadBackgroundJobProgress: async () => {
|
||||
const progress = await invoke<FolderJobProgress[]>("get_background_job_progress");
|
||||
set(() => ({
|
||||
mediaJobProgress: Object.fromEntries(progress.map((entry) => [entry.folder_id, entry])),
|
||||
}));
|
||||
},
|
||||
|
||||
addFolder: async (path) => {
|
||||
const { loadFolders } = get();
|
||||
const { loadFolders, loadBackgroundJobProgress } = get();
|
||||
await invoke("add_folder", { path });
|
||||
await loadFolders();
|
||||
await loadBackgroundJobProgress();
|
||||
},
|
||||
|
||||
removeFolder: async (folderId) => {
|
||||
await invoke("remove_folder", { folderId });
|
||||
const { selectedFolderId, loadFolders, loadImages } = get();
|
||||
const { selectedFolderId, loadFolders, loadImages, loadBackgroundJobProgress } = get();
|
||||
await loadFolders();
|
||||
await loadBackgroundJobProgress();
|
||||
if (selectedFolderId === folderId) {
|
||||
set({ selectedFolderId: null });
|
||||
await loadImages(true);
|
||||
@@ -263,13 +280,14 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
},
|
||||
|
||||
reindexFolder: async (folderId) => {
|
||||
const { loadFolders } = get();
|
||||
const { loadFolders, loadBackgroundJobProgress } = get();
|
||||
await invoke("reindex_folder", { folderId });
|
||||
await loadFolders();
|
||||
await loadBackgroundJobProgress();
|
||||
},
|
||||
|
||||
selectFolder: (folderId) => {
|
||||
set({ selectedFolderId: folderId, images: [], loadedCount: 0 });
|
||||
set({ selectedFolderId: folderId, images: [], loadedCount: 0, collectionTitle: null });
|
||||
void get().loadImages(true);
|
||||
},
|
||||
|
||||
@@ -301,6 +319,7 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
totalImages: result.total,
|
||||
loadedCount: reset ? result.images.length : state.loadedCount + result.images.length,
|
||||
loadingImages: false,
|
||||
collectionTitle: reset ? null : state.collectionTitle,
|
||||
}));
|
||||
} catch (error) {
|
||||
console.error("Failed to load media:", error);
|
||||
@@ -315,22 +334,22 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
},
|
||||
|
||||
setSearch: (search) => {
|
||||
set({ search, images: [], loadedCount: 0 });
|
||||
set({ search, images: [], loadedCount: 0, collectionTitle: null });
|
||||
void get().loadImages(true);
|
||||
},
|
||||
|
||||
setSort: (sort) => {
|
||||
set({ sort, images: [], loadedCount: 0 });
|
||||
set({ sort, images: [], loadedCount: 0, collectionTitle: null });
|
||||
void get().loadImages(true);
|
||||
},
|
||||
|
||||
setMediaFilter: (mediaFilter) => {
|
||||
set({ mediaFilter, images: [], loadedCount: 0 });
|
||||
set({ mediaFilter, images: [], loadedCount: 0, collectionTitle: null });
|
||||
void get().loadImages(true);
|
||||
},
|
||||
|
||||
setFavoritesOnly: (favoritesOnly) => {
|
||||
set({ favoritesOnly, images: [], loadedCount: 0 });
|
||||
set({ favoritesOnly, images: [], loadedCount: 0, collectionTitle: null });
|
||||
void get().loadImages(true);
|
||||
},
|
||||
|
||||
@@ -339,6 +358,25 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
openImage: (image) => set({ selectedImage: image }),
|
||||
closeImage: () => set({ selectedImage: null }),
|
||||
|
||||
loadSimilarImages: async (imageId) => {
|
||||
const images = await invoke<ImageRecord[]>("find_similar_images", {
|
||||
params: { image_id: imageId, limit: PAGE_SIZE },
|
||||
});
|
||||
set({
|
||||
images,
|
||||
totalImages: images.length,
|
||||
loadedCount: images.length,
|
||||
loadingImages: false,
|
||||
collectionTitle: "Similar Images",
|
||||
selectedFolderId: null,
|
||||
});
|
||||
},
|
||||
|
||||
retryFailedEmbeddings: async (folderId) => {
|
||||
await invoke("retry_failed_embeddings", { params: { folder_id: folderId } });
|
||||
await get().loadBackgroundJobProgress();
|
||||
},
|
||||
|
||||
updateImageDetails: async (imageId, updates) => {
|
||||
const updatedImage = await invoke<ImageRecord>("update_image_details", {
|
||||
params: {
|
||||
@@ -366,6 +404,7 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
|
||||
if (progress.done) {
|
||||
void get().loadFolders();
|
||||
void get().loadBackgroundJobProgress();
|
||||
void get().loadImages(true);
|
||||
|
||||
setTimeout(() => {
|
||||
|
||||
Reference in New Issue
Block a user