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/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);
}