feat: add local AI captions and queue controls
This commit is contained in:
Generated
+78
@@ -2389,6 +2389,12 @@ dependencies = [
|
|||||||
"windows-sys 0.61.2",
|
"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]]
|
[[package]]
|
||||||
name = "html5ever"
|
name = "html5ever"
|
||||||
version = "0.29.1"
|
version = "0.29.1"
|
||||||
@@ -3059,6 +3065,12 @@ version = "0.4.29"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
|
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "lzma-rust2"
|
||||||
|
version = "0.15.7"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "1670343e58806300d87950e3401e820b519b9384281bbabfb15e3636689ffd69"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "lzma-sys"
|
name = "lzma-sys"
|
||||||
version = "0.1.20"
|
version = "0.1.20"
|
||||||
@@ -3134,6 +3146,16 @@ version = "0.1.10"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "2532096657941c2fea9c289d370a250971c689d4f143798ff67113ec042024a5"
|
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]]
|
[[package]]
|
||||||
name = "memchr"
|
name = "memchr"
|
||||||
version = "2.8.0"
|
version = "2.8.0"
|
||||||
@@ -3262,6 +3284,21 @@ dependencies = [
|
|||||||
"tempfile",
|
"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]]
|
[[package]]
|
||||||
name = "ndk"
|
name = "ndk"
|
||||||
version = "0.9.0"
|
version = "0.9.0"
|
||||||
@@ -3670,6 +3707,31 @@ dependencies = [
|
|||||||
"pin-project-lite",
|
"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]]
|
[[package]]
|
||||||
name = "pango"
|
name = "pango"
|
||||||
version = "0.18.3"
|
version = "0.18.3"
|
||||||
@@ -3952,6 +4014,7 @@ dependencies = [
|
|||||||
"hf-hub",
|
"hf-hub",
|
||||||
"image",
|
"image",
|
||||||
"log",
|
"log",
|
||||||
|
"ort",
|
||||||
"r2d2",
|
"r2d2",
|
||||||
"r2d2_sqlite",
|
"r2d2_sqlite",
|
||||||
"rayon",
|
"rayon",
|
||||||
@@ -4060,6 +4123,15 @@ version = "1.13.1"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
|
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]]
|
[[package]]
|
||||||
name = "potential_utf"
|
name = "potential_utf"
|
||||||
version = "0.1.5"
|
version = "0.1.5"
|
||||||
@@ -4421,6 +4493,12 @@ version = "0.6.2"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"
|
checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "rawpointer"
|
||||||
|
version = "0.2.1"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "rayon"
|
name = "rayon"
|
||||||
version = "1.11.0"
|
version = "1.11.0"
|
||||||
|
|||||||
@@ -44,3 +44,4 @@ candle-nn = { version = "0.10.2", features = ["cuda"] }
|
|||||||
candle-transformers = { 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"] }
|
hf-hub = { version = "0.5.0", default-features = false, features = ["ureq", "native-tls"] }
|
||||||
tokenizers = "0.22.1"
|
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"] }
|
||||||
|
|||||||
@@ -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],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
+265
-20
@@ -1,10 +1,11 @@
|
|||||||
|
use crate::captioner::{self, CaptionModelStatus, CaptionRuntimeProbe, CaptionVisionProbe};
|
||||||
use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord};
|
use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord};
|
||||||
use crate::embedder;
|
use crate::embedder;
|
||||||
use crate::indexer;
|
use crate::indexer;
|
||||||
use crate::vector;
|
use crate::vector;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
use tauri::{AppHandle, State};
|
use tauri::{AppHandle, Emitter, Manager, State};
|
||||||
|
|
||||||
pub type DbState = DbPool;
|
pub type DbState = DbPool;
|
||||||
|
|
||||||
@@ -41,11 +42,46 @@ pub struct FindSimilarImagesParams {
|
|||||||
pub limit: Option<usize>,
|
pub limit: Option<usize>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
pub struct DebugSimilarImagesParams {
|
||||||
|
pub image_id: i64,
|
||||||
|
pub limit: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
pub struct RetryFailedEmbeddingsParams {
|
pub struct RetryFailedEmbeddingsParams {
|
||||||
pub folder_id: i64,
|
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)]
|
#[derive(Deserialize)]
|
||||||
pub struct SemanticSearchParams {
|
pub struct SemanticSearchParams {
|
||||||
pub query: String,
|
pub query: String,
|
||||||
@@ -127,7 +163,14 @@ pub async fn get_images(
|
|||||||
let favorites_only = params.favorites_only.unwrap_or(false);
|
let favorites_only = params.favorites_only.unwrap_or(false);
|
||||||
let embedding_failed_only = params.embedding_failed_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)
|
let total = db::count_images(
|
||||||
|
&conn,
|
||||||
|
params.folder_id,
|
||||||
|
search,
|
||||||
|
media_kind,
|
||||||
|
favorites_only,
|
||||||
|
embedding_failed_only,
|
||||||
|
)
|
||||||
.map_err(|e| e.to_string())?;
|
.map_err(|e| e.to_string())?;
|
||||||
|
|
||||||
let images = db::get_images(
|
let images = db::get_images(
|
||||||
@@ -188,11 +231,45 @@ pub async fn find_similar_images(
|
|||||||
) -> Result<Vec<ImageRecord>, String> {
|
) -> Result<Vec<ImageRecord>, String> {
|
||||||
let conn = db.get().map_err(|e| e.to_string())?;
|
let conn = db.get().map_err(|e| e.to_string())?;
|
||||||
let limit = params.limit.unwrap_or(32);
|
let limit = params.limit.unwrap_or(32);
|
||||||
let image_ids = vector::find_similar_image_ids(&conn, params.image_id, limit)
|
if !vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())? {
|
||||||
.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())
|
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]
|
#[tauri::command]
|
||||||
pub async fn retry_failed_embeddings(
|
pub async fn retry_failed_embeddings(
|
||||||
db: State<'_, DbState>,
|
db: State<'_, DbState>,
|
||||||
@@ -211,7 +288,8 @@ pub async fn semantic_search_images(
|
|||||||
|
|
||||||
let conn = db.get().map_err(|e| e.to_string())?;
|
let conn = db.get().map_err(|e| e.to_string())?;
|
||||||
let limit = params.limit.unwrap_or(64);
|
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())?;
|
let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?;
|
||||||
|
|
||||||
if let Some(folder_id) = params.folder_id {
|
if let Some(folder_id) = params.folder_id {
|
||||||
@@ -227,6 +305,142 @@ pub async fn semantic_search_images(
|
|||||||
Ok(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, ¶ms.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)]
|
#[derive(Serialize, Deserialize)]
|
||||||
pub struct TagCloudEntry {
|
pub struct TagCloudEntry {
|
||||||
pub count: usize,
|
pub count: usize,
|
||||||
@@ -287,7 +501,10 @@ pub async fn get_tag_cloud(
|
|||||||
|
|
||||||
// Cache miss — run k-means
|
// Cache miss — run k-means
|
||||||
let ids: Vec<i64> = embeddings_with_ids.iter().map(|(id, _)| *id).collect();
|
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 k = (n / 20).clamp(5, 30);
|
||||||
let (centroids, cluster_counts, assignments) = kmeans_cosine(&points, k, 40);
|
let (centroids, cluster_counts, assignments) = kmeans_cosine(&points, k, 40);
|
||||||
@@ -361,7 +578,10 @@ fn kmeans_cosine(
|
|||||||
let next = points
|
let next = points
|
||||||
.iter()
|
.iter()
|
||||||
.map(|p| {
|
.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
|
1.0 - best_sim // distance = 1 - cosine_similarity
|
||||||
})
|
})
|
||||||
.enumerate()
|
.enumerate()
|
||||||
@@ -389,7 +609,9 @@ fn kmeans_cosine(
|
|||||||
changed = true;
|
changed = true;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if !changed { break; }
|
if !changed {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
// Update step: mean of assigned points, then normalize
|
// Update step: mean of assigned points, then normalize
|
||||||
let mut sums = vec![vec![0.0f32; dim]; k];
|
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);
|
sums[c].iter_mut().zip(p.iter()).for_each(|(s, v)| *s += v);
|
||||||
counts[c] += 1;
|
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 {
|
if count > 0 {
|
||||||
sum.iter_mut().for_each(|v| *v /= count as f32);
|
sum.iter_mut().for_each(|v| *v /= count as f32);
|
||||||
normalize(sum);
|
normalize(sum);
|
||||||
@@ -408,7 +632,9 @@ fn kmeans_cosine(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let mut counts = vec![0usize; k];
|
let mut counts = vec![0usize; k];
|
||||||
for &a in &assignments { counts[a] += 1; }
|
for &a in &assignments {
|
||||||
|
counts[a] += 1;
|
||||||
|
}
|
||||||
|
|
||||||
(centroids, counts, assignments)
|
(centroids, counts, assignments)
|
||||||
}
|
}
|
||||||
@@ -438,24 +664,43 @@ pub async fn get_failed_embedding_images(
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Serialize)]
|
#[derive(Serialize)]
|
||||||
pub struct WorkerStates {
|
pub struct FolderWorkerStates {
|
||||||
|
pub folder_id: i64,
|
||||||
pub thumbnail_paused: bool,
|
pub thumbnail_paused: bool,
|
||||||
pub metadata_paused: bool,
|
pub metadata_paused: bool,
|
||||||
pub embedding_paused: bool,
|
pub embedding_paused: bool,
|
||||||
|
pub caption_paused: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tauri::command]
|
#[tauri::command]
|
||||||
pub async fn set_worker_paused(worker: String, paused: bool) -> Result<(), String> {
|
pub async fn set_worker_paused(worker: String, folder_id: i64, paused: bool) -> Result<(), String> {
|
||||||
indexer::set_worker_paused(&worker, paused);
|
indexer::set_worker_paused(&worker, folder_id, paused);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tauri::command]
|
#[tauri::command]
|
||||||
pub async fn get_worker_states() -> Result<WorkerStates, String> {
|
pub async fn get_worker_states(folder_ids: Vec<i64>) -> Result<Vec<FolderWorkerStates>, String> {
|
||||||
let states = indexer::get_worker_paused_states();
|
let states = indexer::get_worker_paused_states(&folder_ids);
|
||||||
Ok(WorkerStates {
|
Ok(folder_ids
|
||||||
thumbnail_paused: states[0],
|
.into_iter()
|
||||||
metadata_paused: states[1],
|
.map(|folder_id| {
|
||||||
embedding_paused: states[2],
|
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())
|
||||||
}
|
}
|
||||||
|
|||||||
+438
-30
@@ -57,6 +57,10 @@ pub struct ImageRecord {
|
|||||||
pub embedding_model: Option<String>,
|
pub embedding_model: Option<String>,
|
||||||
pub embedding_updated_at: Option<String>,
|
pub embedding_updated_at: Option<String>,
|
||||||
pub embedding_error: 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)]
|
#[allow(dead_code)]
|
||||||
@@ -89,6 +93,13 @@ pub struct MetadataJob {
|
|||||||
pub path: String,
|
pub path: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct CaptionJob {
|
||||||
|
pub image_id: i64,
|
||||||
|
pub folder_id: i64,
|
||||||
|
pub path: String,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct IndexedMediaEntry {
|
pub struct IndexedMediaEntry {
|
||||||
pub id: i64,
|
pub id: i64,
|
||||||
@@ -106,6 +117,9 @@ pub struct FolderJobProgress {
|
|||||||
pub embedding_pending: i64,
|
pub embedding_pending: i64,
|
||||||
pub embedding_ready: i64,
|
pub embedding_ready: i64,
|
||||||
pub embedding_failed: 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> {
|
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'))
|
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 (
|
CREATE TABLE IF NOT EXISTS tag_cloud_cache (
|
||||||
folder_scope TEXT PRIMARY KEY,
|
folder_scope TEXT PRIMARY KEY,
|
||||||
image_ids_hash INTEGER NOT NULL,
|
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_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_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_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", "audio_codec", "TEXT")?;
|
||||||
ensure_column(conn, "images", "metadata_updated_at", "TEXT")?;
|
ensure_column(conn, "images", "metadata_updated_at", "TEXT")?;
|
||||||
ensure_column(conn, "images", "metadata_error", "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)?;
|
vector::migrate(conn)?;
|
||||||
Ok(())
|
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> {
|
pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result<i64> {
|
||||||
let id = conn.query_row(
|
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)
|
"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)
|
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
|
ON CONFLICT(path) DO UPDATE SET
|
||||||
folder_id = excluded.folder_id,
|
folder_id = excluded.folder_id,
|
||||||
filename = excluded.filename,
|
filename = excluded.filename,
|
||||||
@@ -250,7 +278,11 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result<i64> {
|
|||||||
embedding_status = excluded.embedding_status,
|
embedding_status = excluded.embedding_status,
|
||||||
embedding_model = excluded.embedding_model,
|
embedding_model = excluded.embedding_model,
|
||||||
embedding_updated_at = excluded.embedding_updated_at,
|
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",
|
RETURNING id",
|
||||||
params![
|
params![
|
||||||
img.folder_id,
|
img.folder_id,
|
||||||
@@ -275,6 +307,10 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result<i64> {
|
|||||||
img.embedding_model,
|
img.embedding_model,
|
||||||
img.embedding_updated_at,
|
img.embedding_updated_at,
|
||||||
img.embedding_error,
|
img.embedding_error,
|
||||||
|
img.generated_caption,
|
||||||
|
img.caption_model,
|
||||||
|
img.caption_updated_at,
|
||||||
|
img.caption_error,
|
||||||
],
|
],
|
||||||
|row| row.get(0),
|
|row| row.get(0),
|
||||||
)?;
|
)?;
|
||||||
@@ -307,6 +343,51 @@ pub fn backfill_embedding_jobs(conn: &Connection) -> Result<usize> {
|
|||||||
Ok(inserted)
|
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> {
|
pub fn retry_failed_embedding_jobs(conn: &Connection, folder_id: i64) -> Result<usize> {
|
||||||
// Only re-queue images that are actually embeddable right now.
|
// Only re-queue images that are actually embeddable right now.
|
||||||
// Videos without a thumbnail would just fail again immediately, so skip them —
|
// 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'",
|
"UPDATE embedding_jobs SET status = 'pending' WHERE status = 'processing'",
|
||||||
[],
|
[],
|
||||||
)?;
|
)?;
|
||||||
|
conn.execute(
|
||||||
|
"UPDATE caption_jobs SET status = 'pending' WHERE status = 'processing'",
|
||||||
|
[],
|
||||||
|
)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -377,16 +462,102 @@ pub fn enqueue_metadata_job(conn: &Connection, image_id: i64) -> Result<()> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get_pending_embedding_jobs(conn: &Connection, limit: usize) -> Result<Vec<EmbeddingJob>> {
|
pub fn enqueue_caption_job(conn: &Connection, image_id: i64) -> Result<()> {
|
||||||
let mut stmt = conn.prepare(
|
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,
|
"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
|
j.status, j.attempts, j.last_error, j.created_at, j.updated_at
|
||||||
FROM embedding_jobs j
|
FROM embedding_jobs j
|
||||||
JOIN images i ON i.id = j.image_id
|
JOIN images i ON i.id = j.image_id
|
||||||
WHERE status = 'pending'
|
WHERE status = 'pending'
|
||||||
|
{}
|
||||||
ORDER BY j.updated_at, j.image_id
|
ORDER BY j.updated_at, j.image_id
|
||||||
LIMIT ?1",
|
LIMIT ?1",
|
||||||
)?;
|
folder_exclusion_clause("i", excluded_folder_ids)
|
||||||
|
);
|
||||||
|
let mut stmt = conn.prepare(&sql)?;
|
||||||
let rows = stmt.query_map([limit as i64], |row| {
|
let rows = stmt.query_map([limit as i64], |row| {
|
||||||
Ok(EmbeddingJob {
|
Ok(EmbeddingJob {
|
||||||
image_id: row.get(0)?,
|
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<_>>>()?)
|
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 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);
|
let mut claimed = Vec::with_capacity(limit);
|
||||||
|
|
||||||
for job in candidates {
|
for job in candidates {
|
||||||
@@ -430,6 +605,64 @@ pub fn claim_embedding_jobs(conn: &mut Connection, limit: usize) -> Result<Vec<E
|
|||||||
Ok(claimed)
|
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)]
|
#[allow(dead_code)]
|
||||||
pub fn mark_embedding_ready(conn: &Connection, image_id: i64, model: &str) -> Result<()> {
|
pub fn mark_embedding_ready(conn: &Connection, image_id: i64, model: &str) -> Result<()> {
|
||||||
conn.execute(
|
conn.execute(
|
||||||
@@ -489,6 +722,7 @@ pub fn delete_images_by_ids(conn: &Connection, image_ids: &[i64]) -> Result<()>
|
|||||||
let tx = conn.unchecked_transaction()?;
|
let tx = conn.unchecked_transaction()?;
|
||||||
for image_id in image_ids {
|
for image_id in image_ids {
|
||||||
vector::delete_embedding(&tx, *image_id)?;
|
vector::delete_embedding(&tx, *image_id)?;
|
||||||
|
vector::delete_caption_embedding(&tx, *image_id)?;
|
||||||
tx.execute("DELETE FROM images WHERE id = ?1", [image_id])?;
|
tx.execute("DELETE FROM images WHERE id = ?1", [image_id])?;
|
||||||
}
|
}
|
||||||
tx.commit()?;
|
tx.commit()?;
|
||||||
@@ -539,6 +773,31 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result<Fold
|
|||||||
|row| row.get(0),
|
|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 {
|
Ok(FolderJobProgress {
|
||||||
folder_id,
|
folder_id,
|
||||||
thumbnail_pending,
|
thumbnail_pending,
|
||||||
@@ -546,6 +805,9 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result<Fold
|
|||||||
embedding_pending,
|
embedding_pending,
|
||||||
embedding_ready,
|
embedding_ready,
|
||||||
embedding_failed,
|
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)
|
Ok(progress)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get_pending_thumbnail_jobs(conn: &Connection, limit: usize) -> Result<Vec<ThumbnailJob>> {
|
fn get_pending_thumbnail_jobs_excluding(
|
||||||
let mut stmt = conn.prepare(
|
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
|
"SELECT j.image_id, i.folder_id, i.path, i.media_kind
|
||||||
FROM thumbnail_jobs j
|
FROM thumbnail_jobs j
|
||||||
JOIN images i ON i.id = j.image_id
|
JOIN images i ON i.id = j.image_id
|
||||||
WHERE j.status = 'pending'
|
WHERE j.status = 'pending'
|
||||||
|
{}
|
||||||
ORDER BY j.updated_at, j.image_id
|
ORDER BY j.updated_at, j.image_id
|
||||||
LIMIT ?1",
|
LIMIT ?1",
|
||||||
)?;
|
folder_exclusion_clause("i", excluded_folder_ids)
|
||||||
|
);
|
||||||
|
let mut stmt = conn.prepare(&sql)?;
|
||||||
let rows = stmt.query_map([limit as i64], |row| {
|
let rows = stmt.query_map([limit as i64], |row| {
|
||||||
Ok(ThumbnailJob {
|
Ok(ThumbnailJob {
|
||||||
image_id: row.get(0)?,
|
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(
|
pub fn claim_thumbnail_jobs(
|
||||||
conn: &mut Connection,
|
conn: &mut Connection,
|
||||||
active_folder_ids: &std::collections::HashSet<i64>,
|
active_folder_ids: &std::collections::HashSet<i64>,
|
||||||
|
paused_folder_ids: &std::collections::HashSet<i64>,
|
||||||
fetch_limit: usize,
|
fetch_limit: usize,
|
||||||
claim_limit: usize,
|
claim_limit: usize,
|
||||||
) -> Result<Vec<ThumbnailJob>> {
|
) -> Result<Vec<ThumbnailJob>> {
|
||||||
let tx = conn.transaction()?;
|
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);
|
let mut claimed = Vec::with_capacity(claim_limit);
|
||||||
|
|
||||||
for job in candidates {
|
for job in candidates {
|
||||||
if active_folder_ids.contains(&job.folder_id) {
|
debug_assert!(!excluded_folder_ids.contains(&job.folder_id));
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let updated = tx.execute(
|
let updated = tx.execute(
|
||||||
"UPDATE thumbnail_jobs
|
"UPDATE thumbnail_jobs
|
||||||
SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now')
|
SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now')
|
||||||
@@ -617,15 +888,22 @@ pub fn claim_thumbnail_jobs(
|
|||||||
Ok(claimed)
|
Ok(claimed)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn get_pending_metadata_jobs(conn: &Connection, limit: usize) -> Result<Vec<MetadataJob>> {
|
fn get_pending_metadata_jobs_excluding(
|
||||||
let mut stmt = conn.prepare(
|
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
|
"SELECT j.image_id, i.folder_id, i.path
|
||||||
FROM metadata_jobs j
|
FROM metadata_jobs j
|
||||||
JOIN images i ON i.id = j.image_id
|
JOIN images i ON i.id = j.image_id
|
||||||
WHERE j.status = 'pending' AND i.media_kind = 'video'
|
WHERE j.status = 'pending' AND i.media_kind = 'video'
|
||||||
|
{}
|
||||||
ORDER BY j.updated_at, j.image_id
|
ORDER BY j.updated_at, j.image_id
|
||||||
LIMIT ?1",
|
LIMIT ?1",
|
||||||
)?;
|
folder_exclusion_clause("i", excluded_folder_ids)
|
||||||
|
);
|
||||||
|
let mut stmt = conn.prepare(&sql)?;
|
||||||
let rows = stmt.query_map([limit as i64], |row| {
|
let rows = stmt.query_map([limit as i64], |row| {
|
||||||
Ok(MetadataJob {
|
Ok(MetadataJob {
|
||||||
image_id: row.get(0)?,
|
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(
|
pub fn claim_metadata_jobs(
|
||||||
conn: &mut Connection,
|
conn: &mut Connection,
|
||||||
active_folder_ids: &std::collections::HashSet<i64>,
|
active_folder_ids: &std::collections::HashSet<i64>,
|
||||||
|
paused_folder_ids: &std::collections::HashSet<i64>,
|
||||||
fetch_limit: usize,
|
fetch_limit: usize,
|
||||||
claim_limit: usize,
|
claim_limit: usize,
|
||||||
) -> Result<Vec<MetadataJob>> {
|
) -> Result<Vec<MetadataJob>> {
|
||||||
let tx = conn.transaction()?;
|
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);
|
let mut claimed = Vec::with_capacity(claim_limit);
|
||||||
|
|
||||||
for job in candidates {
|
for job in candidates {
|
||||||
if active_folder_ids.contains(&job.folder_id) {
|
debug_assert!(!excluded_folder_ids.contains(&job.folder_id));
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let updated = tx.execute(
|
let updated = tx.execute(
|
||||||
"UPDATE metadata_jobs
|
"UPDATE metadata_jobs
|
||||||
SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now')
|
SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now')
|
||||||
@@ -765,7 +1045,8 @@ pub fn update_image_details(
|
|||||||
conn.query_row(
|
conn.query_row(
|
||||||
"SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type,
|
"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,
|
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
|
FROM images
|
||||||
WHERE id = ?1",
|
WHERE id = ?1",
|
||||||
[image_id],
|
[image_id],
|
||||||
@@ -778,7 +1059,8 @@ pub fn get_image_by_id(conn: &Connection, image_id: i64) -> Result<ImageRecord>
|
|||||||
conn.query_row(
|
conn.query_row(
|
||||||
"SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type,
|
"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,
|
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
|
FROM images
|
||||||
WHERE id = ?1",
|
WHERE id = ?1",
|
||||||
[image_id],
|
[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>> {
|
pub fn get_images_by_ids(conn: &Connection, image_ids: &[i64]) -> Result<Vec<ImageRecord>> {
|
||||||
let mut images = Vec::with_capacity(image_ids.len());
|
let mut images = Vec::with_capacity(image_ids.len());
|
||||||
for image_id in image_ids {
|
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)
|
Ok(images)
|
||||||
}
|
}
|
||||||
@@ -839,7 +1125,8 @@ pub fn get_images(
|
|||||||
let sql = format!(
|
let sql = format!(
|
||||||
"SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type,
|
"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,
|
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
|
FROM images
|
||||||
WHERE (?1 IS NULL OR folder_id = ?1)
|
WHERE (?1 IS NULL OR folder_id = ?1)
|
||||||
AND (?2 IS NULL OR filename LIKE ?2)
|
AND (?2 IS NULL OR filename LIKE ?2)
|
||||||
@@ -918,11 +1205,104 @@ pub fn get_failed_embedding_images(
|
|||||||
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
|
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn delete_folder(conn: &Connection, folder_id: i64) -> Result<()> {
|
pub fn update_generated_caption(
|
||||||
conn.execute("DELETE FROM folders WHERE id = ?1", params![folder_id])?;
|
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(())
|
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> {
|
fn map_image_row(row: &Row<'_>) -> rusqlite::Result<ImageRecord> {
|
||||||
Ok(ImageRecord {
|
Ok(ImageRecord {
|
||||||
id: row.get(0)?,
|
id: row.get(0)?,
|
||||||
@@ -948,6 +1328,10 @@ fn map_image_row(row: &Row<'_>) -> rusqlite::Result<ImageRecord> {
|
|||||||
embedding_model: row.get(20)?,
|
embedding_model: row.get(20)?,
|
||||||
embedding_updated_at: row.get(21)?,
|
embedding_updated_at: row.get(21)?,
|
||||||
embedding_error: row.get(22)?,
|
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(())
|
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)
|
||||||
|
}
|
||||||
|
|||||||
+216
-42
@@ -1,3 +1,4 @@
|
|||||||
|
use crate::captioner::{self, FlorenceCaptioner};
|
||||||
use crate::db::{self, DbPool, EmbeddingJob, FolderJobProgress, ImageRecord, IndexedMediaEntry};
|
use crate::db::{self, DbPool, EmbeddingJob, FolderJobProgress, ImageRecord, IndexedMediaEntry};
|
||||||
use crate::embedder::{embedding_source_path, ClipImageEmbedder};
|
use crate::embedder::{embedding_source_path, ClipImageEmbedder};
|
||||||
use crate::media::{probe_video_metadata, MediaTools};
|
use crate::media::{probe_video_metadata, MediaTools};
|
||||||
@@ -9,7 +10,6 @@ use rayon::prelude::*;
|
|||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::atomic::{AtomicBool, Ordering};
|
|
||||||
use std::sync::{Mutex, OnceLock};
|
use std::sync::{Mutex, OnceLock};
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
use tauri::{AppHandle, Emitter};
|
use tauri::{AppHandle, Emitter};
|
||||||
@@ -22,30 +22,92 @@ const IMAGE_EXTENSIONS: &[&str] = &[
|
|||||||
const VIDEO_EXTENSIONS: &[&str] = &["mp4", "mov", "m4v", "webm"];
|
const VIDEO_EXTENSIONS: &[&str] = &["mp4", "mov", "m4v", "webm"];
|
||||||
|
|
||||||
const JOB_PROGRESS_EMIT_INTERVAL: Duration = Duration::from_millis(750);
|
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 LAST_JOB_PROGRESS_EMIT: OnceLock<Mutex<HashMap<i64, Instant>>> = OnceLock::new();
|
||||||
static ACTIVE_INDEXING_FOLDERS: OnceLock<Mutex<HashSet<i64>>> = OnceLock::new();
|
static ACTIVE_INDEXING_FOLDERS: OnceLock<Mutex<HashSet<i64>>> = OnceLock::new();
|
||||||
|
|
||||||
static THUMBNAIL_WORKER_PAUSED: AtomicBool = AtomicBool::new(false);
|
static PAUSED_WORKER_FOLDERS: OnceLock<Mutex<PausedWorkerFolders>> = OnceLock::new();
|
||||||
static METADATA_WORKER_PAUSED: AtomicBool = AtomicBool::new(false);
|
|
||||||
static EMBEDDING_WORKER_PAUSED: AtomicBool = AtomicBool::new(false);
|
#[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(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();
|
||||||
|
};
|
||||||
|
|
||||||
pub fn set_worker_paused(worker: &str, paused: bool) {
|
|
||||||
match worker {
|
match worker {
|
||||||
"thumbnail" => THUMBNAIL_WORKER_PAUSED.store(paused, Ordering::Relaxed),
|
"thumbnail" => paused_folders.thumbnail.clone(),
|
||||||
"metadata" => METADATA_WORKER_PAUSED.store(paused, Ordering::Relaxed),
|
"metadata" => paused_folders.metadata.clone(),
|
||||||
"embedding" => EMBEDDING_WORKER_PAUSED.store(paused, Ordering::Relaxed),
|
"embedding" => paused_folders.embedding.clone(),
|
||||||
_ => {}
|
_ => HashSet::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
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),
|
|
||||||
]
|
|
||||||
}
|
|
||||||
static FOLDER_STORAGE_PROFILES: OnceLock<Mutex<HashMap<i64, RuntimeAdaptiveProfile>>> =
|
static FOLDER_STORAGE_PROFILES: OnceLock<Mutex<HashMap<i64, RuntimeAdaptiveProfile>>> =
|
||||||
OnceLock::new();
|
OnceLock::new();
|
||||||
static DB_WRITE_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
static DB_WRITE_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||||
@@ -95,10 +157,6 @@ pub fn start_thumbnail_worker(
|
|||||||
cache_dir: PathBuf,
|
cache_dir: PathBuf,
|
||||||
) {
|
) {
|
||||||
std::thread::spawn(move || loop {
|
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) {
|
if let Err(error) = process_thumbnail_batch(&app, &pool, &media_tools, &cache_dir) {
|
||||||
eprintln!("Thumbnail worker error: {}", error);
|
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) {
|
pub fn start_metadata_worker(app: AppHandle, pool: DbPool, media_tools: MediaTools) {
|
||||||
std::thread::spawn(move || loop {
|
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) {
|
if let Err(error) = process_metadata_batch(&app, &pool, &media_tools) {
|
||||||
eprintln!("Metadata worker error: {}", error);
|
eprintln!("Metadata worker error: {}", error);
|
||||||
}
|
}
|
||||||
@@ -124,10 +178,6 @@ pub fn start_embedding_worker(app: AppHandle, pool: DbPool) {
|
|||||||
let mut embedder: Option<ClipImageEmbedder> = None;
|
let mut embedder: Option<ClipImageEmbedder> = None;
|
||||||
println!("Embedding worker started.");
|
println!("Embedding worker started.");
|
||||||
loop {
|
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) {
|
if let Err(error) = process_embedding_batch(&app, &pool, &mut embedder) {
|
||||||
eprintln!("Embedding worker error: {}", error);
|
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<()> {
|
fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> {
|
||||||
let existing_entries = {
|
let existing_entries = {
|
||||||
let conn = pool.get()?;
|
let conn = pool.get()?;
|
||||||
@@ -197,7 +261,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf)
|
|||||||
images: committed,
|
images: committed,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
emit_folder_job_progress(&app, &pool, &[folder_id]);
|
emit_folder_job_progress(&app, &pool, &[folder_id], false);
|
||||||
}
|
}
|
||||||
|
|
||||||
processed += path_chunk.len();
|
processed += path_chunk.len();
|
||||||
@@ -245,7 +309,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf)
|
|||||||
done: true,
|
done: true,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
emit_folder_job_progress(&app, &pool, &[folder_id]);
|
emit_folder_job_progress(&app, &pool, &[folder_id], true);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -302,6 +366,10 @@ fn build_record(
|
|||||||
embedding_model: Some(vector::CLIP_MODEL_NAME.to_string()),
|
embedding_model: Some(vector::CLIP_MODEL_NAME.to_string()),
|
||||||
embedding_updated_at: None,
|
embedding_updated_at: None,
|
||||||
embedding_error: 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(|| {
|
with_db_write_lock(|| {
|
||||||
let mut conn = pool.get()?;
|
let mut conn = pool.get()?;
|
||||||
let active_folders = active_indexing_folders();
|
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_batch_size = max_worker_batch_size(&active_folders);
|
||||||
let worker_fetch_size = max_worker_fetch_size(&active_folders);
|
let worker_fetch_size = max_worker_fetch_size(&active_folders);
|
||||||
db::claim_thumbnail_jobs(
|
db::claim_thumbnail_jobs(
|
||||||
&mut conn,
|
&mut conn,
|
||||||
&active_folders,
|
&active_folders,
|
||||||
|
&paused_folders,
|
||||||
worker_fetch_size,
|
worker_fetch_size,
|
||||||
worker_batch_size,
|
worker_batch_size,
|
||||||
)
|
)
|
||||||
@@ -428,7 +498,7 @@ fn process_thumbnail_batch(
|
|||||||
images: updated_images,
|
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(())
|
Ok(())
|
||||||
@@ -439,11 +509,13 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo
|
|||||||
with_db_write_lock(|| {
|
with_db_write_lock(|| {
|
||||||
let mut conn = pool.get()?;
|
let mut conn = pool.get()?;
|
||||||
let active_folders = active_indexing_folders();
|
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_batch_size = max_worker_batch_size(&active_folders);
|
||||||
let worker_fetch_size = max_worker_fetch_size(&active_folders);
|
let worker_fetch_size = max_worker_fetch_size(&active_folders);
|
||||||
db::claim_metadata_jobs(
|
db::claim_metadata_jobs(
|
||||||
&mut conn,
|
&mut conn,
|
||||||
&active_folders,
|
&active_folders,
|
||||||
|
&paused_folders,
|
||||||
worker_fetch_size,
|
worker_fetch_size,
|
||||||
worker_batch_size,
|
worker_batch_size,
|
||||||
)
|
)
|
||||||
@@ -506,7 +578,7 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo
|
|||||||
images: updated_images,
|
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(())
|
Ok(())
|
||||||
@@ -518,14 +590,11 @@ fn process_embedding_batch(
|
|||||||
embedder: &mut Option<ClipImageEmbedder>,
|
embedder: &mut Option<ClipImageEmbedder>,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let batch_started_at = Instant::now();
|
let batch_started_at = Instant::now();
|
||||||
if embedder.is_none() {
|
|
||||||
*embedder = Some(ClipImageEmbedder::new()?);
|
|
||||||
}
|
|
||||||
|
|
||||||
let claim_started_at = Instant::now();
|
let claim_started_at = Instant::now();
|
||||||
|
let paused_folders = paused_folder_ids("embedding");
|
||||||
let jobs = with_db_write_lock(|| {
|
let jobs = with_db_write_lock(|| {
|
||||||
let mut conn = pool.get()?;
|
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();
|
let claim_elapsed = claim_started_at.elapsed();
|
||||||
|
|
||||||
@@ -533,9 +602,18 @@ fn process_embedding_batch(
|
|||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if embedder.is_none() {
|
||||||
|
*embedder = Some(ClipImageEmbedder::new()?);
|
||||||
|
}
|
||||||
|
|
||||||
println!("Embedding batch claimed: {} items", jobs.len());
|
println!("Embedding batch claimed: {} items", jobs.len());
|
||||||
let folder_ids = jobs.iter().map(|job| job.folder_id).collect::<HashSet<_>>();
|
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 embedder = embedder.as_ref().expect("embedder should be initialized");
|
||||||
|
|
||||||
let infer_started_at = Instant::now();
|
let infer_started_at = Instant::now();
|
||||||
@@ -639,7 +717,7 @@ fn process_embedding_batch(
|
|||||||
images: updated_images,
|
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();
|
let write_elapsed = write_started_at.elapsed();
|
||||||
@@ -652,6 +730,101 @@ fn process_embedding_batch(
|
|||||||
Ok(())
|
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> {
|
fn active_indexing_folders() -> HashSet<i64> {
|
||||||
ACTIVE_INDEXING_FOLDERS
|
ACTIVE_INDEXING_FOLDERS
|
||||||
.get_or_init(|| Mutex::new(HashSet::new()))
|
.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);
|
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<_>>();
|
let mut unique_folder_ids = folder_ids.iter().copied().collect::<Vec<_>>();
|
||||||
unique_folder_ids.sort_unstable();
|
unique_folder_ids.sort_unstable();
|
||||||
unique_folder_ids.dedup();
|
unique_folder_ids.dedup();
|
||||||
@@ -749,7 +922,8 @@ fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64])
|
|||||||
Err(_) => return,
|
Err(_) => return,
|
||||||
};
|
};
|
||||||
unique_folder_ids.retain(|folder_id| {
|
unique_folder_ids.retain(|folder_id| {
|
||||||
let should_emit = tracker
|
let should_emit = force
|
||||||
|
|| tracker
|
||||||
.get(folder_id)
|
.get(folder_id)
|
||||||
.map(|last_emit| now.duration_since(*last_emit) >= JOB_PROGRESS_EMIT_INTERVAL)
|
.map(|last_emit| now.duration_since(*last_emit) >= JOB_PROGRESS_EMIT_INTERVAL)
|
||||||
.unwrap_or(true);
|
.unwrap_or(true);
|
||||||
|
|||||||
+23
-3
@@ -1,3 +1,4 @@
|
|||||||
|
mod captioner;
|
||||||
mod commands;
|
mod commands;
|
||||||
mod db;
|
mod db;
|
||||||
mod embedder;
|
mod embedder;
|
||||||
@@ -7,8 +8,8 @@ mod storage;
|
|||||||
mod thumbnail;
|
mod thumbnail;
|
||||||
mod vector;
|
mod vector;
|
||||||
|
|
||||||
use tauri::Manager;
|
|
||||||
use crate::storage::StorageProfile;
|
use crate::storage::StorageProfile;
|
||||||
|
use tauri::Manager;
|
||||||
|
|
||||||
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
||||||
pub fn run() {
|
pub fn run() {
|
||||||
@@ -34,11 +35,19 @@ pub fn run() {
|
|||||||
let conn = pool.get().expect("Failed to get connection for migration");
|
let conn = pool.get().expect("Failed to get connection for migration");
|
||||||
db::migrate(&conn).expect("Failed to run migrations");
|
db::migrate(&conn).expect("Failed to run migrations");
|
||||||
db::reset_inflight_jobs(&conn).expect("Failed to reset inflight jobs");
|
db::reset_inflight_jobs(&conn).expect("Failed to reset inflight jobs");
|
||||||
let backfilled = db::backfill_embedding_jobs(&conn)
|
let backfilled =
|
||||||
.expect("Failed to backfill embedding jobs");
|
db::backfill_embedding_jobs(&conn).expect("Failed to backfill embedding jobs");
|
||||||
if backfilled > 0 {
|
if backfilled > 0 {
|
||||||
println!("Backfilled {} embedding jobs.", backfilled);
|
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");
|
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_metadata_worker(app.handle().clone(), pool.clone(), media_tools.clone());
|
||||||
indexer::start_embedding_worker(app.handle().clone(), pool.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(pool);
|
||||||
app.manage(media_tools);
|
app.manage(media_tools);
|
||||||
@@ -73,8 +83,18 @@ pub fn run() {
|
|||||||
commands::reindex_folder,
|
commands::reindex_folder,
|
||||||
commands::update_image_details,
|
commands::update_image_details,
|
||||||
commands::find_similar_images,
|
commands::find_similar_images,
|
||||||
|
commands::debug_similar_images,
|
||||||
commands::retry_failed_embeddings,
|
commands::retry_failed_embeddings,
|
||||||
commands::semantic_search_images,
|
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::set_worker_paused,
|
||||||
commands::get_worker_states,
|
commands::get_worker_states,
|
||||||
commands::get_tag_cloud,
|
commands::get_tag_cloud,
|
||||||
|
|||||||
+147
-4
@@ -1,5 +1,5 @@
|
|||||||
use anyhow::{anyhow, Result};
|
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 sqlite_vec::sqlite3_vec_init;
|
||||||
use std::sync::Once;
|
use std::sync::Once;
|
||||||
|
|
||||||
@@ -19,8 +19,13 @@ pub fn migrate(conn: &Connection) -> Result<()> {
|
|||||||
"CREATE VIRTUAL TABLE IF NOT EXISTS image_vec USING vec0(
|
"CREATE VIRTUAL TABLE IF NOT EXISTS image_vec USING vec0(
|
||||||
image_id INTEGER PRIMARY KEY,
|
image_id INTEGER PRIMARY KEY,
|
||||||
embedding FLOAT[{}] distance_metric=cosine
|
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(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@@ -31,6 +36,12 @@ pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> {
|
|||||||
Ok(())
|
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)]
|
#[allow(dead_code)]
|
||||||
pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) -> Result<()> {
|
pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) -> Result<()> {
|
||||||
if embedding.len() != CLIP_VECTOR_DIM {
|
if embedding.len() != CLIP_VECTOR_DIM {
|
||||||
@@ -50,12 +61,35 @@ pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) ->
|
|||||||
Ok(())
|
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>> {
|
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",
|
"SELECT embedding FROM image_vec WHERE image_id = ?1",
|
||||||
[image_id],
|
[image_id],
|
||||||
|row| row.get(0),
|
|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(
|
let mut stmt = conn.prepare(
|
||||||
"SELECT image_id
|
"SELECT image_id
|
||||||
@@ -155,6 +189,115 @@ pub fn search_image_ids_by_embedding(
|
|||||||
Ok(ids)
|
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)]
|
#[allow(dead_code)]
|
||||||
fn pack_f32(values: &[f32]) -> Vec<u8> {
|
fn pack_f32(values: &[f32]) -> Vec<u8> {
|
||||||
let mut out = Vec::with_capacity(values.len() * std::mem::size_of::<f32>());
|
let mut out = Vec::with_capacity(values.len() * std::mem::size_of::<f32>());
|
||||||
|
|||||||
@@ -7,17 +7,20 @@ import { Gallery } from "./components/Gallery";
|
|||||||
import { Lightbox } from "./components/Lightbox";
|
import { Lightbox } from "./components/Lightbox";
|
||||||
import { TagCloud } from "./components/TagCloud";
|
import { TagCloud } from "./components/TagCloud";
|
||||||
import { TitleBar } from "./components/TitleBar";
|
import { TitleBar } from "./components/TitleBar";
|
||||||
|
import { SettingsModal } from "./components/SettingsModal";
|
||||||
|
|
||||||
export default function App() {
|
export default function App() {
|
||||||
const loadFolders = useGalleryStore((state) => state.loadFolders);
|
const loadFolders = useGalleryStore((state) => state.loadFolders);
|
||||||
const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress);
|
const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress);
|
||||||
const loadImages = useGalleryStore((state) => state.loadImages);
|
const loadImages = useGalleryStore((state) => state.loadImages);
|
||||||
|
const loadCaptionModelStatus = useGalleryStore((state) => state.loadCaptionModelStatus);
|
||||||
const subscribeToProgress = useGalleryStore((state) => state.subscribeToProgress);
|
const subscribeToProgress = useGalleryStore((state) => state.subscribeToProgress);
|
||||||
const activeView = useGalleryStore((state) => state.activeView);
|
const activeView = useGalleryStore((state) => state.activeView);
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
loadFolders().then(() => {
|
loadFolders().then(() => {
|
||||||
void loadBackgroundJobProgress();
|
void loadBackgroundJobProgress();
|
||||||
|
void loadCaptionModelStatus();
|
||||||
return loadImages(true);
|
return loadImages(true);
|
||||||
});
|
});
|
||||||
let unlisten: (() => void) | undefined;
|
let unlisten: (() => void) | undefined;
|
||||||
@@ -54,6 +57,7 @@ export default function App() {
|
|||||||
</div>
|
</div>
|
||||||
|
|
||||||
<Lightbox />
|
<Lightbox />
|
||||||
|
<SettingsModal />
|
||||||
</div>
|
</div>
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,12 +2,13 @@ import { useEffect, useMemo, useState } from "react";
|
|||||||
import { invoke } from "@tauri-apps/api/core";
|
import { invoke } from "@tauri-apps/api/core";
|
||||||
import { useGalleryStore } from "../store";
|
import { useGalleryStore } from "../store";
|
||||||
|
|
||||||
type WorkerKey = "thumbnail" | "metadata" | "embedding";
|
type WorkerKey = "thumbnail" | "metadata" | "embedding" | "caption";
|
||||||
|
|
||||||
const WORKER_FOR_STAGE: Record<string, WorkerKey> = {
|
const WORKER_FOR_STAGE: Record<string, WorkerKey> = {
|
||||||
Thumbnails: "thumbnail",
|
Thumbnails: "thumbnail",
|
||||||
Metadata: "metadata",
|
Metadata: "metadata",
|
||||||
Embeddings: "embedding",
|
Embeddings: "embedding",
|
||||||
|
Captions: "caption",
|
||||||
};
|
};
|
||||||
|
|
||||||
interface TaskStage {
|
interface TaskStage {
|
||||||
@@ -22,6 +23,7 @@ interface Task {
|
|||||||
name: string;
|
name: string;
|
||||||
stages: TaskStage[];
|
stages: TaskStage[];
|
||||||
hasFailedEmbeddings: boolean;
|
hasFailedEmbeddings: boolean;
|
||||||
|
hasFailedCaptions: boolean;
|
||||||
pendingMediaWork: number;
|
pendingMediaWork: number;
|
||||||
embeddingProcessed: number;
|
embeddingProcessed: number;
|
||||||
embeddingTotal: number;
|
embeddingTotal: number;
|
||||||
@@ -35,6 +37,21 @@ interface FailedEmbeddingItem {
|
|||||||
error: string | null;
|
error: string | null;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
interface FolderWorkerStates {
|
||||||
|
folder_id: number;
|
||||||
|
thumbnail_paused: boolean;
|
||||||
|
metadata_paused: boolean;
|
||||||
|
embedding_paused: boolean;
|
||||||
|
caption_paused: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
const DEFAULT_PAUSED_STATE: Record<WorkerKey, boolean> = {
|
||||||
|
thumbnail: false,
|
||||||
|
metadata: false,
|
||||||
|
embedding: false,
|
||||||
|
caption: false,
|
||||||
|
};
|
||||||
|
|
||||||
export function BackgroundTasks() {
|
export function BackgroundTasks() {
|
||||||
const folders = useGalleryStore((state) => state.folders);
|
const folders = useGalleryStore((state) => state.folders);
|
||||||
const indexingProgress = useGalleryStore((state) => state.indexingProgress);
|
const indexingProgress = useGalleryStore((state) => state.indexingProgress);
|
||||||
@@ -42,24 +59,32 @@ export function BackgroundTasks() {
|
|||||||
const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings);
|
const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings);
|
||||||
const [expanded, setExpanded] = useState(false);
|
const [expanded, setExpanded] = useState(false);
|
||||||
const [dismissed, setDismissed] = useState<Record<number, string>>({});
|
const [dismissed, setDismissed] = useState<Record<number, string>>({});
|
||||||
const [paused, setPaused] = useState<Record<WorkerKey, boolean>>({
|
const [paused, setPaused] = useState<Record<number, Record<WorkerKey, boolean>>>({});
|
||||||
thumbnail: false,
|
|
||||||
metadata: false,
|
|
||||||
embedding: false,
|
|
||||||
});
|
|
||||||
const [failedItems, setFailedItems] = useState<Record<number, FailedEmbeddingItem[]>>({});
|
const [failedItems, setFailedItems] = useState<Record<number, FailedEmbeddingItem[]>>({});
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
invoke<{ thumbnail_paused: boolean; metadata_paused: boolean; embedding_paused: boolean }>(
|
const folderIds = folders.map((folder) => folder.id);
|
||||||
"get_worker_states",
|
if (folderIds.length === 0) {
|
||||||
).then((states) => {
|
setPaused({});
|
||||||
setPaused({
|
return;
|
||||||
thumbnail: states.thumbnail_paused,
|
}
|
||||||
metadata: states.metadata_paused,
|
|
||||||
embedding: states.embedding_paused,
|
invoke<FolderWorkerStates[]>("get_worker_states", { folderIds }).then((states) => {
|
||||||
|
setPaused(
|
||||||
|
Object.fromEntries(
|
||||||
|
states.map((state) => [
|
||||||
|
state.folder_id,
|
||||||
|
{
|
||||||
|
thumbnail: state.thumbnail_paused,
|
||||||
|
metadata: state.metadata_paused,
|
||||||
|
embedding: state.embedding_paused,
|
||||||
|
caption: state.caption_paused,
|
||||||
|
},
|
||||||
|
]),
|
||||||
|
),
|
||||||
|
);
|
||||||
});
|
});
|
||||||
});
|
}, [folders]);
|
||||||
}, []);
|
|
||||||
|
|
||||||
// Fetch failed embedding filenames whenever the expanded panel opens or failure counts change.
|
// Fetch failed embedding filenames whenever the expanded panel opens or failure counts change.
|
||||||
const failedCounts = useMemo(
|
const failedCounts = useMemo(
|
||||||
@@ -83,10 +108,21 @@ export function BackgroundTasks() {
|
|||||||
}
|
}
|
||||||
}, [expanded, failedCounts]);
|
}, [expanded, failedCounts]);
|
||||||
|
|
||||||
const toggleWorker = (worker: WorkerKey) => {
|
const isWorkerPaused = (folderId: number, worker: WorkerKey) => {
|
||||||
const next = !paused[worker];
|
return paused[folderId]?.[worker] ?? DEFAULT_PAUSED_STATE[worker];
|
||||||
setPaused((prev) => ({ ...prev, [worker]: next }));
|
};
|
||||||
void invoke("set_worker_paused", { worker, paused: next });
|
|
||||||
|
const toggleWorker = (folderId: number, worker: WorkerKey) => {
|
||||||
|
const next = !isWorkerPaused(folderId, worker);
|
||||||
|
setPaused((prev) => ({
|
||||||
|
...prev,
|
||||||
|
[folderId]: {
|
||||||
|
...DEFAULT_PAUSED_STATE,
|
||||||
|
...prev[folderId],
|
||||||
|
[worker]: next,
|
||||||
|
},
|
||||||
|
}));
|
||||||
|
void invoke("set_worker_paused", { worker, folderId, paused: next });
|
||||||
};
|
};
|
||||||
|
|
||||||
const dismissTask = (id: number, snapshot: string) => {
|
const dismissTask = (id: number, snapshot: string) => {
|
||||||
@@ -105,13 +141,16 @@ export function BackgroundTasks() {
|
|||||||
const embeddingPending = jobs?.embedding_pending ?? 0;
|
const embeddingPending = jobs?.embedding_pending ?? 0;
|
||||||
const embeddingReady = jobs?.embedding_ready ?? 0;
|
const embeddingReady = jobs?.embedding_ready ?? 0;
|
||||||
const embeddingFailed = jobs?.embedding_failed ?? 0;
|
const embeddingFailed = jobs?.embedding_failed ?? 0;
|
||||||
|
const captionPending = jobs?.caption_pending ?? 0;
|
||||||
|
const captionFailed = jobs?.caption_failed ?? 0;
|
||||||
|
|
||||||
const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending;
|
const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending + captionPending;
|
||||||
const embeddingProcessed = embeddingReady + embeddingFailed;
|
const embeddingProcessed = embeddingReady + embeddingFailed;
|
||||||
const embeddingTotal = embeddingProcessed + embeddingPending;
|
const embeddingTotal = embeddingProcessed + embeddingPending;
|
||||||
const hasFailedEmbeddings = embeddingFailed > 0;
|
const hasFailedEmbeddings = embeddingFailed > 0;
|
||||||
|
const hasFailedCaptions = captionFailed > 0;
|
||||||
|
|
||||||
if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings) return null;
|
if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings && !hasFailedCaptions) return null;
|
||||||
|
|
||||||
const stages: TaskStage[] = [];
|
const stages: TaskStage[] = [];
|
||||||
|
|
||||||
@@ -153,6 +192,15 @@ export function BackgroundTasks() {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (captionPending > 0) {
|
||||||
|
stages.push({
|
||||||
|
label: "Captions",
|
||||||
|
detail: captionPending.toLocaleString(),
|
||||||
|
progress: null,
|
||||||
|
failed: false,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
if (hasFailedEmbeddings && pendingMediaWork === 0) {
|
if (hasFailedEmbeddings && pendingMediaWork === 0) {
|
||||||
stages.push({
|
stages.push({
|
||||||
label: "Failed",
|
label: "Failed",
|
||||||
@@ -162,13 +210,23 @@ export function BackgroundTasks() {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
const snapshot = `${pendingMediaWork}:${embeddingFailed}`;
|
if (hasFailedCaptions && pendingMediaWork === 0) {
|
||||||
|
stages.push({
|
||||||
|
label: "Failed",
|
||||||
|
detail: `${captionFailed.toLocaleString()} captions`,
|
||||||
|
progress: null,
|
||||||
|
failed: true,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
const snapshot = `${pendingMediaWork}:${embeddingFailed}:${captionFailed}`;
|
||||||
|
|
||||||
return {
|
return {
|
||||||
id: folder.id,
|
id: folder.id,
|
||||||
name: folder.name,
|
name: folder.name,
|
||||||
stages,
|
stages,
|
||||||
hasFailedEmbeddings,
|
hasFailedEmbeddings,
|
||||||
|
hasFailedCaptions,
|
||||||
pendingMediaWork,
|
pendingMediaWork,
|
||||||
embeddingProcessed,
|
embeddingProcessed,
|
||||||
embeddingTotal,
|
embeddingTotal,
|
||||||
@@ -184,7 +242,7 @@ export function BackgroundTasks() {
|
|||||||
|
|
||||||
const primary = tasks[0];
|
const primary = tasks[0];
|
||||||
const extraCount = tasks.length - 1;
|
const extraCount = tasks.length - 1;
|
||||||
const hasFailed = tasks.some((t) => t.hasFailedEmbeddings && t.pendingMediaWork === 0);
|
const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedCaptions) && t.pendingMediaWork === 0);
|
||||||
|
|
||||||
// Best progress bar value: use embedding progress if available (most informative),
|
// Best progress bar value: use embedding progress if available (most informative),
|
||||||
// otherwise fall back to scanning progress, otherwise indeterminate.
|
// otherwise fall back to scanning progress, otherwise indeterminate.
|
||||||
@@ -214,7 +272,7 @@ export function BackgroundTasks() {
|
|||||||
<div className="flex items-center gap-1.5 flex-1 min-w-0 overflow-hidden">
|
<div className="flex items-center gap-1.5 flex-1 min-w-0 overflow-hidden">
|
||||||
{primary.stages.map((stage) => {
|
{primary.stages.map((stage) => {
|
||||||
const workerKey = WORKER_FOR_STAGE[stage.label];
|
const workerKey = WORKER_FOR_STAGE[stage.label];
|
||||||
const isPaused = workerKey ? paused[workerKey] : false;
|
const isPaused = workerKey ? isWorkerPaused(primary.id, workerKey) : false;
|
||||||
return (
|
return (
|
||||||
<span
|
<span
|
||||||
key={stage.label}
|
key={stage.label}
|
||||||
@@ -234,7 +292,7 @@ export function BackgroundTasks() {
|
|||||||
<button
|
<button
|
||||||
className="ml-0.5 opacity-0 group-hover:opacity-100 hover:text-white transition-opacity"
|
className="ml-0.5 opacity-0 group-hover:opacity-100 hover:text-white transition-opacity"
|
||||||
title={isPaused ? `Resume ${stage.label}` : `Pause ${stage.label}`}
|
title={isPaused ? `Resume ${stage.label}` : `Pause ${stage.label}`}
|
||||||
onClick={(e) => { e.stopPropagation(); toggleWorker(workerKey); }}
|
onClick={(e) => { e.stopPropagation(); toggleWorker(primary.id, workerKey); }}
|
||||||
>
|
>
|
||||||
{isPaused ? (
|
{isPaused ? (
|
||||||
<svg className="h-2.5 w-2.5" fill="currentColor" viewBox="0 0 24 24">
|
<svg className="h-2.5 w-2.5" fill="currentColor" viewBox="0 0 24 24">
|
||||||
@@ -312,7 +370,7 @@ export function BackgroundTasks() {
|
|||||||
const taskEmbeddingStage = task.stages.find((s) => s.label === "Embeddings");
|
const taskEmbeddingStage = task.stages.find((s) => s.label === "Embeddings");
|
||||||
const taskScanningStage = task.stages.find((s) => s.label === "Scanning");
|
const taskScanningStage = task.stages.find((s) => s.label === "Scanning");
|
||||||
const taskBarProgress = taskEmbeddingStage?.progress ?? taskScanningStage?.progress ?? null;
|
const taskBarProgress = taskEmbeddingStage?.progress ?? taskScanningStage?.progress ?? null;
|
||||||
const taskHasFailed = task.hasFailedEmbeddings && task.pendingMediaWork === 0;
|
const taskHasFailed = (task.hasFailedEmbeddings || task.hasFailedCaptions) && task.pendingMediaWork === 0;
|
||||||
|
|
||||||
return (
|
return (
|
||||||
<div key={task.id}>
|
<div key={task.id}>
|
||||||
@@ -322,7 +380,7 @@ export function BackgroundTasks() {
|
|||||||
<div className="flex items-center gap-1.5 flex-1 min-w-0 overflow-hidden">
|
<div className="flex items-center gap-1.5 flex-1 min-w-0 overflow-hidden">
|
||||||
{task.stages.map((stage) => {
|
{task.stages.map((stage) => {
|
||||||
const workerKey = WORKER_FOR_STAGE[stage.label];
|
const workerKey = WORKER_FOR_STAGE[stage.label];
|
||||||
const isPaused = workerKey ? paused[workerKey] : false;
|
const isPaused = workerKey ? isWorkerPaused(task.id, workerKey) : false;
|
||||||
return (
|
return (
|
||||||
<span
|
<span
|
||||||
key={stage.label}
|
key={stage.label}
|
||||||
@@ -347,7 +405,7 @@ export function BackgroundTasks() {
|
|||||||
<button
|
<button
|
||||||
className="ml-0.5 text-gray-600 hover:text-white transition-colors"
|
className="ml-0.5 text-gray-600 hover:text-white transition-colors"
|
||||||
title={isPaused ? `Resume ${stage.label}` : `Pause ${stage.label}`}
|
title={isPaused ? `Resume ${stage.label}` : `Pause ${stage.label}`}
|
||||||
onClick={() => toggleWorker(workerKey)}
|
onClick={() => toggleWorker(task.id, workerKey)}
|
||||||
>
|
>
|
||||||
{isPaused ? (
|
{isPaused ? (
|
||||||
<svg className="h-2.5 w-2.5" fill="currentColor" viewBox="0 0 24 24">
|
<svg className="h-2.5 w-2.5" fill="currentColor" viewBox="0 0 24 24">
|
||||||
|
|||||||
@@ -251,6 +251,10 @@ export function Gallery() {
|
|||||||
const zoomPreset = useGalleryStore((state) => state.zoomPreset);
|
const zoomPreset = useGalleryStore((state) => state.zoomPreset);
|
||||||
const search = useGalleryStore((state) => state.search);
|
const search = useGalleryStore((state) => state.search);
|
||||||
const searchMode = useGalleryStore((state) => state.searchMode);
|
const searchMode = useGalleryStore((state) => state.searchMode);
|
||||||
|
const collectionTitle = useGalleryStore((state) => state.collectionTitle);
|
||||||
|
const imageLoadError = useGalleryStore((state) => state.imageLoadError);
|
||||||
|
const galleryScrollResetKey = useGalleryStore((state) => state.galleryScrollResetKey);
|
||||||
|
const isSimilarResults = collectionTitle === "Similar Images";
|
||||||
|
|
||||||
const parentRef = useRef<HTMLDivElement>(null);
|
const parentRef = useRef<HTMLDivElement>(null);
|
||||||
const [contextMenu, setContextMenu] = useState<{ x: number; y: number; image: ImageRecord } | null>(null);
|
const [contextMenu, setContextMenu] = useState<{ x: number; y: number; image: ImageRecord } | null>(null);
|
||||||
@@ -271,6 +275,10 @@ export function Gallery() {
|
|||||||
return () => element.removeEventListener("scroll", handleScroll);
|
return () => element.removeEventListener("scroll", handleScroll);
|
||||||
}, [handleScroll]);
|
}, [handleScroll]);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
parentRef.current?.scrollTo({ top: 0, left: 0 });
|
||||||
|
}, [galleryScrollResetKey]);
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
const close = (event: PointerEvent) => {
|
const close = (event: PointerEvent) => {
|
||||||
if ((event.target as HTMLElement | null)?.closest("[data-gallery-context-menu]")) return;
|
if ((event.target as HTMLElement | null)?.closest("[data-gallery-context-menu]")) return;
|
||||||
@@ -293,12 +301,16 @@ export function Gallery() {
|
|||||||
<div className="rounded-2xl border border-white/[0.06] bg-white/[0.02] p-8 min-w-72">
|
<div className="rounded-2xl border border-white/[0.06] bg-white/[0.02] p-8 min-w-72">
|
||||||
<div className="h-5 w-5 mx-auto rounded-full border-2 border-white/20 border-t-white/60 animate-spin" />
|
<div className="h-5 w-5 mx-auto rounded-full border-2 border-white/20 border-t-white/60 animate-spin" />
|
||||||
<p className="mt-4 text-sm text-white/40 font-medium">
|
<p className="mt-4 text-sm text-white/40 font-medium">
|
||||||
{searchMode === "semantic" && search.trim().length > 0
|
{isSimilarResults
|
||||||
|
? "Finding similar images"
|
||||||
|
: searchMode === "semantic" && search.trim().length > 0
|
||||||
? `Searching for matches to "${search}"`
|
? `Searching for matches to "${search}"`
|
||||||
: "Loading media"}
|
: "Loading media"}
|
||||||
</p>
|
</p>
|
||||||
<p className="text-xs text-white/20 mt-1">
|
<p className="text-xs text-white/20 mt-1">
|
||||||
{searchMode === "semantic" && search.trim().length > 0
|
{isSimilarResults
|
||||||
|
? "Comparing visual embeddings"
|
||||||
|
: searchMode === "semantic" && search.trim().length > 0
|
||||||
? "Semantic search can take a little longer than filename search"
|
? "Semantic search can take a little longer than filename search"
|
||||||
: "Fetching results"}
|
: "Fetching results"}
|
||||||
</p>
|
</p>
|
||||||
@@ -316,12 +328,20 @@ export function Gallery() {
|
|||||||
d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-6-6h.01M6 20h12a2 2 0 002-2V6a2 2 0 00-2-2H6a2 2 0 00-2 2v12a2 2 0 002 2z" />
|
d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-6-6h.01M6 20h12a2 2 0 002-2V6a2 2 0 00-2-2H6a2 2 0 00-2 2v12a2 2 0 002 2z" />
|
||||||
</svg>
|
</svg>
|
||||||
<p className="text-sm text-white/30 font-medium">
|
<p className="text-sm text-white/30 font-medium">
|
||||||
{searchMode === "semantic" && search.trim().length > 0
|
{imageLoadError
|
||||||
|
? "Could not load results"
|
||||||
|
: isSimilarResults
|
||||||
|
? "No similar images found"
|
||||||
|
: searchMode === "semantic" && search.trim().length > 0
|
||||||
? "No semantic matches found"
|
? "No semantic matches found"
|
||||||
: "No media found"}
|
: "No media found"}
|
||||||
</p>
|
</p>
|
||||||
<p className="text-xs text-white/15 mt-1">
|
<p className="text-xs text-white/15 mt-1">
|
||||||
{searchMode === "semantic" && search.trim().length > 0
|
{imageLoadError
|
||||||
|
? imageLoadError
|
||||||
|
: isSimilarResults
|
||||||
|
? "This item may be visually isolated, or more embeddings may need to finish processing"
|
||||||
|
: searchMode === "semantic" && search.trim().length > 0
|
||||||
? "Try a broader phrase, or wait for more embeddings to finish processing"
|
? "Try a broader phrase, or wait for more embeddings to finish processing"
|
||||||
: "Try adjusting your filters or add a new folder"}
|
: "Try adjusting your filters or add a new folder"}
|
||||||
</p>
|
</p>
|
||||||
|
|||||||
@@ -47,12 +47,27 @@ function embeddingLabel(status: string, model: string | null): string {
|
|||||||
|
|
||||||
export function Lightbox() {
|
export function Lightbox() {
|
||||||
const selectedImage = useGalleryStore((state) => state.selectedImage);
|
const selectedImage = useGalleryStore((state) => state.selectedImage);
|
||||||
|
const selectedFolderId = useGalleryStore((state) => state.selectedFolderId);
|
||||||
const closeImage = useGalleryStore((state) => state.closeImage);
|
const closeImage = useGalleryStore((state) => state.closeImage);
|
||||||
const images = useGalleryStore((state) => state.images);
|
const images = useGalleryStore((state) => state.images);
|
||||||
const openImage = useGalleryStore((state) => state.openImage);
|
const openImage = useGalleryStore((state) => state.openImage);
|
||||||
const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages);
|
const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages);
|
||||||
const updateImageDetails = useGalleryStore((state) => state.updateImageDetails);
|
const updateImageDetails = useGalleryStore((state) => state.updateImageDetails);
|
||||||
|
const suggestImageTags = useGalleryStore((state) => state.suggestImageTags);
|
||||||
|
const captionModelStatus = useGalleryStore((state) => state.captionModelStatus);
|
||||||
|
const captionModelPreparing = useGalleryStore((state) => state.captionModelPreparing);
|
||||||
|
const captionModelError = useGalleryStore((state) => state.captionModelError);
|
||||||
|
const captionModelProgress = useGalleryStore((state) => state.captionModelProgress);
|
||||||
|
const aiCaptionsEnabled = useGalleryStore((state) => state.aiCaptionsEnabled);
|
||||||
|
const setAiCaptionsEnabled = useGalleryStore((state) => state.setAiCaptionsEnabled);
|
||||||
|
const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen);
|
||||||
|
const prepareCaptionModel = useGalleryStore((state) => state.prepareCaptionModel);
|
||||||
|
const queueCaptionForImage = useGalleryStore((state) => state.queueCaptionForImage);
|
||||||
|
const queueCaptionJobs = useGalleryStore((state) => state.queueCaptionJobs);
|
||||||
const [zoom, setZoom] = useState(1);
|
const [zoom, setZoom] = useState(1);
|
||||||
|
const [suggestedTags, setSuggestedTags] = useState<string[]>([]);
|
||||||
|
const [captionQueueing, setCaptionQueueing] = useState(false);
|
||||||
|
const [captionQueueStatus, setCaptionQueueStatus] = useState<string | null>(null);
|
||||||
const imageViewportRef = useRef<HTMLDivElement>(null);
|
const imageViewportRef = useRef<HTMLDivElement>(null);
|
||||||
|
|
||||||
const currentIndex = selectedImage ? images.findIndex((image) => image.id === selectedImage.id) : -1;
|
const currentIndex = selectedImage ? images.findIndex((image) => image.id === selectedImage.id) : -1;
|
||||||
@@ -68,8 +83,17 @@ export function Lightbox() {
|
|||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
setZoom(1);
|
setZoom(1);
|
||||||
|
setSuggestedTags([]);
|
||||||
|
setCaptionQueueStatus(null);
|
||||||
}, [selectedImage?.id]);
|
}, [selectedImage?.id]);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
if (!selectedImage?.generated_caption) return;
|
||||||
|
void suggestImageTags(selectedImage.id)
|
||||||
|
.then(setSuggestedTags)
|
||||||
|
.catch(() => setSuggestedTags([]));
|
||||||
|
}, [selectedImage?.id, selectedImage?.generated_caption, suggestImageTags]);
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
const viewport = imageViewportRef.current;
|
const viewport = imageViewportRef.current;
|
||||||
if (!viewport || !selectedImage || selectedImage.media_kind !== "image") return;
|
if (!viewport || !selectedImage || selectedImage.media_kind !== "image") return;
|
||||||
@@ -316,6 +340,119 @@ export function Lightbox() {
|
|||||||
) : null}
|
) : null}
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
|
<div>
|
||||||
|
<p className="mb-1 text-xs uppercase tracking-wider text-gray-500">AI Caption</p>
|
||||||
|
{selectedImage.generated_caption ? (
|
||||||
|
<>
|
||||||
|
<p className="text-white">{selectedImage.generated_caption}</p>
|
||||||
|
{suggestedTags.length > 0 ? (
|
||||||
|
<div className="mt-2 flex flex-wrap gap-1.5">
|
||||||
|
{suggestedTags.map((tag) => (
|
||||||
|
<span
|
||||||
|
key={tag}
|
||||||
|
className="rounded-md border border-white/10 bg-white/5 px-2 py-0.5 text-xs text-gray-300"
|
||||||
|
>
|
||||||
|
{tag}
|
||||||
|
</span>
|
||||||
|
))}
|
||||||
|
</div>
|
||||||
|
) : null}
|
||||||
|
{selectedImage.caption_model ? (
|
||||||
|
<p className="mt-1 text-xs text-gray-600">{selectedImage.caption_model}</p>
|
||||||
|
) : null}
|
||||||
|
</>
|
||||||
|
) : (
|
||||||
|
<>
|
||||||
|
<p className="text-gray-600">
|
||||||
|
{selectedImage.caption_error ?? "Not generated"}
|
||||||
|
</p>
|
||||||
|
{captionModelStatus?.ready ? (
|
||||||
|
<>
|
||||||
|
<p className={`mt-2 text-xs ${aiCaptionsEnabled ? "text-emerald-400/70" : "text-gray-600"}`}>
|
||||||
|
{aiCaptionsEnabled ? "Florence-2 enabled locally" : "Florence-2 downloaded but disabled"}
|
||||||
|
</p>
|
||||||
|
<div className="mt-2 flex flex-wrap gap-2">
|
||||||
|
<button
|
||||||
|
className="rounded-md border border-white/10 bg-white/5 px-3 py-1.5 text-xs text-gray-300 transition-colors hover:bg-white/10 hover:text-white"
|
||||||
|
onClick={() => setAiCaptionsEnabled(!aiCaptionsEnabled)}
|
||||||
|
>
|
||||||
|
{aiCaptionsEnabled ? "Disable captions" : "Enable captions"}
|
||||||
|
</button>
|
||||||
|
{selectedImage.media_kind === "image" && aiCaptionsEnabled ? (
|
||||||
|
<button
|
||||||
|
className="rounded-md border border-white/10 bg-white/5 px-3 py-1.5 text-xs text-gray-300 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-50"
|
||||||
|
onClick={() => {
|
||||||
|
setCaptionQueueing(true);
|
||||||
|
setCaptionQueueStatus(null);
|
||||||
|
void queueCaptionForImage(selectedImage.id)
|
||||||
|
.then(() => setCaptionQueueStatus("Queued for background captions."))
|
||||||
|
.catch((error) => setCaptionQueueStatus(String(error)))
|
||||||
|
.finally(() => setCaptionQueueing(false));
|
||||||
|
}}
|
||||||
|
disabled={captionQueueing}
|
||||||
|
>
|
||||||
|
{captionQueueing ? "Queueing..." : "Queue caption"}
|
||||||
|
</button>
|
||||||
|
) : null}
|
||||||
|
{aiCaptionsEnabled ? (
|
||||||
|
<button
|
||||||
|
className="rounded-md border border-white/10 bg-white/5 px-3 py-1.5 text-xs text-gray-300 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-50"
|
||||||
|
onClick={() => {
|
||||||
|
setCaptionQueueing(true);
|
||||||
|
setCaptionQueueStatus(null);
|
||||||
|
void queueCaptionJobs(selectedFolderId)
|
||||||
|
.then((queued) =>
|
||||||
|
setCaptionQueueStatus(
|
||||||
|
queued === 0
|
||||||
|
? "No missing captions found."
|
||||||
|
: `Queued ${queued.toLocaleString()} image${queued === 1 ? "" : "s"} for background captions.`,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.catch((error) => setCaptionQueueStatus(String(error)))
|
||||||
|
.finally(() => setCaptionQueueing(false));
|
||||||
|
}}
|
||||||
|
disabled={captionQueueing}
|
||||||
|
>
|
||||||
|
{selectedFolderId === null ? "Queue all" : "Queue folder"}
|
||||||
|
</button>
|
||||||
|
) : null}
|
||||||
|
</div>
|
||||||
|
{captionQueueStatus ? (
|
||||||
|
<p className="mt-2 break-all text-xs text-gray-500">{captionQueueStatus}</p>
|
||||||
|
) : null}
|
||||||
|
</>
|
||||||
|
) : (
|
||||||
|
<button
|
||||||
|
className="mt-2 rounded-md border border-white/10 bg-white/5 px-3 py-1.5 text-xs text-gray-300 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-50"
|
||||||
|
onClick={() => void prepareCaptionModel()}
|
||||||
|
disabled={captionModelPreparing}
|
||||||
|
>
|
||||||
|
{captionModelProgress
|
||||||
|
? `Downloading ${captionModelProgress.completed_files}/${captionModelProgress.total_files}`
|
||||||
|
: captionModelPreparing
|
||||||
|
? "Preparing Florence-2..."
|
||||||
|
: "Enable local captions"}
|
||||||
|
</button>
|
||||||
|
)}
|
||||||
|
{captionModelProgress?.current_file ? (
|
||||||
|
<p className="mt-2 break-all text-xs text-gray-600">{captionModelProgress.current_file}</p>
|
||||||
|
) : null}
|
||||||
|
{captionModelError ? (
|
||||||
|
<p className="mt-2 text-xs text-amber-300">{captionModelError}</p>
|
||||||
|
) : null}
|
||||||
|
{!captionModelStatus?.ready && !captionModelError ? (
|
||||||
|
<p className="mt-2 text-xs text-gray-600">Downloads Florence-2 on demand for offline captions.</p>
|
||||||
|
) : null}
|
||||||
|
<button
|
||||||
|
className="mt-2 block text-xs text-gray-500 transition-colors hover:text-gray-300"
|
||||||
|
onClick={() => setSettingsOpen(true)}
|
||||||
|
>
|
||||||
|
Open settings
|
||||||
|
</button>
|
||||||
|
</>
|
||||||
|
)}
|
||||||
|
</div>
|
||||||
|
|
||||||
<div>
|
<div>
|
||||||
<p className="mb-1 text-xs uppercase tracking-wider text-gray-500">Path</p>
|
<p className="mb-1 text-xs uppercase tracking-wider text-gray-500">Path</p>
|
||||||
<p className="break-all text-xs text-gray-400">{selectedImage.path}</p>
|
<p className="break-all text-xs text-gray-400">{selectedImage.path}</p>
|
||||||
|
|||||||
@@ -0,0 +1,354 @@
|
|||||||
|
import { useEffect, useState } from "react";
|
||||||
|
import { useGalleryStore } from "../store";
|
||||||
|
|
||||||
|
type SettingsSection = "ai" | "library" | "display" | "storage";
|
||||||
|
|
||||||
|
const SECTIONS: { id: SettingsSection; label: string; detail: string }[] = [
|
||||||
|
{ id: "ai", label: "Local AI", detail: "Captions and suggestions" },
|
||||||
|
{ id: "library", label: "Library", detail: "Indexing and scanning" },
|
||||||
|
{ id: "display", label: "Display", detail: "Gallery preferences" },
|
||||||
|
{ id: "storage", label: "Storage", detail: "Cache and model files" },
|
||||||
|
];
|
||||||
|
|
||||||
|
function ToggleSwitch({
|
||||||
|
checked,
|
||||||
|
disabled,
|
||||||
|
onChange,
|
||||||
|
label,
|
||||||
|
}: {
|
||||||
|
checked: boolean;
|
||||||
|
disabled?: boolean;
|
||||||
|
onChange: (checked: boolean) => void;
|
||||||
|
label: string;
|
||||||
|
}) {
|
||||||
|
return (
|
||||||
|
<button
|
||||||
|
type="button"
|
||||||
|
role="switch"
|
||||||
|
aria-checked={checked}
|
||||||
|
aria-label={label}
|
||||||
|
disabled={disabled}
|
||||||
|
onClick={() => onChange(!checked)}
|
||||||
|
className={`relative h-7 w-12 shrink-0 rounded-full border transition-colors duration-150 focus:outline-none focus:ring-2 focus:ring-emerald-400/30 disabled:cursor-not-allowed disabled:opacity-45 ${
|
||||||
|
checked
|
||||||
|
? "border-emerald-400/50 bg-emerald-500/35"
|
||||||
|
: "border-white/12 bg-white/[0.06]"
|
||||||
|
}`}
|
||||||
|
>
|
||||||
|
<span
|
||||||
|
className={`absolute left-0.5 top-0.5 h-6 w-6 rounded-full bg-white shadow-lg shadow-black/40 transition-transform duration-150 ${
|
||||||
|
checked ? "translate-x-5" : "translate-x-0"
|
||||||
|
}`}
|
||||||
|
/>
|
||||||
|
</button>
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
function StatusPill({ children, tone }: { children: React.ReactNode; tone: "ready" | "muted" | "busy" }) {
|
||||||
|
const className =
|
||||||
|
tone === "ready"
|
||||||
|
? "border-emerald-400/25 bg-emerald-500/10 text-emerald-300"
|
||||||
|
: tone === "busy"
|
||||||
|
? "border-sky-400/25 bg-sky-500/10 text-sky-300"
|
||||||
|
: "border-white/10 bg-white/[0.04] text-gray-500";
|
||||||
|
|
||||||
|
return (
|
||||||
|
<span className={`inline-flex rounded-md border px-2 py-1 text-[11px] font-medium ${className}`}>
|
||||||
|
{children}
|
||||||
|
</span>
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
function SettingsRow({
|
||||||
|
title,
|
||||||
|
description,
|
||||||
|
children,
|
||||||
|
}: {
|
||||||
|
title: string;
|
||||||
|
description: string;
|
||||||
|
children: React.ReactNode;
|
||||||
|
}) {
|
||||||
|
return (
|
||||||
|
<div className="flex items-start justify-between gap-5 border-b border-white/[0.07] py-4 last:border-b-0">
|
||||||
|
<div className="min-w-0">
|
||||||
|
<p className="text-sm font-medium text-white">{title}</p>
|
||||||
|
<p className="mt-1 max-w-md text-xs leading-relaxed text-gray-500">{description}</p>
|
||||||
|
</div>
|
||||||
|
<div className="shrink-0">{children}</div>
|
||||||
|
</div>
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
function SectionShell({
|
||||||
|
eyebrow,
|
||||||
|
title,
|
||||||
|
children,
|
||||||
|
}: {
|
||||||
|
eyebrow: string;
|
||||||
|
title: string;
|
||||||
|
children: React.ReactNode;
|
||||||
|
}) {
|
||||||
|
return (
|
||||||
|
<div>
|
||||||
|
<p className="text-[10px] font-semibold uppercase tracking-[0.18em] text-gray-600">{eyebrow}</p>
|
||||||
|
<h3 className="mt-1 text-lg font-semibold text-white">{title}</h3>
|
||||||
|
<div className="mt-5 border-t border-white/[0.08]">{children}</div>
|
||||||
|
</div>
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
export function SettingsModal() {
|
||||||
|
const [activeSection, setActiveSection] = useState<SettingsSection>("ai");
|
||||||
|
const [captionQueueStatus, setCaptionQueueStatus] = useState<string | null>(null);
|
||||||
|
const [captionQueueing, setCaptionQueueing] = useState(false);
|
||||||
|
const settingsOpen = useGalleryStore((state) => state.settingsOpen);
|
||||||
|
const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen);
|
||||||
|
const folders = useGalleryStore((state) => state.folders);
|
||||||
|
const selectedFolderId = useGalleryStore((state) => state.selectedFolderId);
|
||||||
|
const captionModelStatus = useGalleryStore((state) => state.captionModelStatus);
|
||||||
|
const captionModelPreparing = useGalleryStore((state) => state.captionModelPreparing);
|
||||||
|
const captionModelProgress = useGalleryStore((state) => state.captionModelProgress);
|
||||||
|
const captionModelError = useGalleryStore((state) => state.captionModelError);
|
||||||
|
const captionRuntimeProbe = useGalleryStore((state) => state.captionRuntimeProbe);
|
||||||
|
const captionRuntimeChecking = useGalleryStore((state) => state.captionRuntimeChecking);
|
||||||
|
const aiCaptionsEnabled = useGalleryStore((state) => state.aiCaptionsEnabled);
|
||||||
|
const setAiCaptionsEnabled = useGalleryStore((state) => state.setAiCaptionsEnabled);
|
||||||
|
const prepareCaptionModel = useGalleryStore((state) => state.prepareCaptionModel);
|
||||||
|
const deleteCaptionModel = useGalleryStore((state) => state.deleteCaptionModel);
|
||||||
|
const probeCaptionRuntime = useGalleryStore((state) => state.probeCaptionRuntime);
|
||||||
|
const queueCaptionJobs = useGalleryStore((state) => state.queueCaptionJobs);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
if (!settingsOpen) return;
|
||||||
|
const handleKeyDown = (event: KeyboardEvent) => {
|
||||||
|
if (event.key === "Escape") setSettingsOpen(false);
|
||||||
|
};
|
||||||
|
window.addEventListener("keydown", handleKeyDown);
|
||||||
|
return () => window.removeEventListener("keydown", handleKeyDown);
|
||||||
|
}, [settingsOpen, setSettingsOpen]);
|
||||||
|
|
||||||
|
if (!settingsOpen) return null;
|
||||||
|
|
||||||
|
const modelReady = captionModelStatus?.ready ?? false;
|
||||||
|
const downloadLabel = captionModelProgress
|
||||||
|
? `Downloading ${captionModelProgress.completed_files}/${captionModelProgress.total_files}`
|
||||||
|
: captionModelPreparing
|
||||||
|
? "Preparing Florence-2..."
|
||||||
|
: modelReady
|
||||||
|
? "Downloaded"
|
||||||
|
: "Download Florence-2";
|
||||||
|
const downloadPercent = captionModelProgress
|
||||||
|
? Math.round((captionModelProgress.completed_files / Math.max(captionModelProgress.total_files, 1)) * 100)
|
||||||
|
: 0;
|
||||||
|
const selectedFolder = folders.find((folder) => folder.id === selectedFolderId);
|
||||||
|
const captionScopeLabel = selectedFolder ? selectedFolder.name : "all libraries";
|
||||||
|
|
||||||
|
return (
|
||||||
|
<div
|
||||||
|
className="fixed inset-0 z-[60] flex items-center justify-center bg-black/65 px-6 backdrop-blur-sm"
|
||||||
|
onClick={() => setSettingsOpen(false)}
|
||||||
|
>
|
||||||
|
<div
|
||||||
|
className="flex h-[min(680px,calc(100vh-56px))] w-full max-w-4xl overflow-hidden rounded-lg border border-white/10 bg-[#07080f] shadow-2xl shadow-black/60"
|
||||||
|
onClick={(event) => event.stopPropagation()}
|
||||||
|
>
|
||||||
|
<aside className="flex w-56 shrink-0 flex-col border-r border-white/[0.07] bg-white/[0.025]">
|
||||||
|
<div className="border-b border-white/[0.07] px-5 py-5">
|
||||||
|
<p className="text-base font-semibold text-white">Settings</p>
|
||||||
|
<p className="mt-1 text-xs text-gray-600">Phokus preferences</p>
|
||||||
|
</div>
|
||||||
|
<div className="flex-1 space-y-1 overflow-y-auto p-2">
|
||||||
|
{SECTIONS.map((section) => (
|
||||||
|
<button
|
||||||
|
key={section.id}
|
||||||
|
className={`w-full rounded-md px-3 py-2.5 text-left transition-colors ${
|
||||||
|
activeSection === section.id
|
||||||
|
? "bg-white/10 text-white"
|
||||||
|
: "text-gray-500 hover:bg-white/[0.055] hover:text-gray-200"
|
||||||
|
}`}
|
||||||
|
onClick={() => setActiveSection(section.id)}
|
||||||
|
>
|
||||||
|
<span className="block text-[13px] font-medium">{section.label}</span>
|
||||||
|
<span className="mt-0.5 block text-[11px] text-gray-600">{section.detail}</span>
|
||||||
|
</button>
|
||||||
|
))}
|
||||||
|
</div>
|
||||||
|
</aside>
|
||||||
|
|
||||||
|
<main className="flex min-w-0 flex-1 flex-col">
|
||||||
|
<div className="flex h-14 shrink-0 items-center justify-between border-b border-white/[0.07] px-6">
|
||||||
|
<div className="flex items-center gap-2">
|
||||||
|
{modelReady ? <StatusPill tone="ready">Florence-2 ready</StatusPill> : <StatusPill tone="muted">Optional</StatusPill>}
|
||||||
|
{captionModelPreparing ? <StatusPill tone="busy">Working</StatusPill> : null}
|
||||||
|
</div>
|
||||||
|
<button
|
||||||
|
className="rounded-md p-1.5 text-gray-500 transition-colors hover:bg-white/[0.06] hover:text-white"
|
||||||
|
onClick={() => setSettingsOpen(false)}
|
||||||
|
title="Close settings"
|
||||||
|
>
|
||||||
|
<svg className="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
|
||||||
|
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
|
||||||
|
</svg>
|
||||||
|
</button>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="flex-1 overflow-y-auto px-7 py-6">
|
||||||
|
{activeSection === "ai" ? (
|
||||||
|
<SectionShell eyebrow="Local AI" title="Captions and suggested tags">
|
||||||
|
<SettingsRow
|
||||||
|
title="AI captions"
|
||||||
|
description="Generate captions and suggested tags with the local Florence-2 model."
|
||||||
|
>
|
||||||
|
<ToggleSwitch
|
||||||
|
checked={aiCaptionsEnabled && modelReady}
|
||||||
|
disabled={!modelReady || captionModelPreparing}
|
||||||
|
onChange={setAiCaptionsEnabled}
|
||||||
|
label="AI captions"
|
||||||
|
/>
|
||||||
|
</SettingsRow>
|
||||||
|
|
||||||
|
<SettingsRow
|
||||||
|
title="Florence-2 model"
|
||||||
|
description={modelReady ? "Stored locally and available offline." : "Download the model before enabling local captions."}
|
||||||
|
>
|
||||||
|
<div className="flex flex-col items-end gap-2">
|
||||||
|
<button
|
||||||
|
className="relative overflow-hidden rounded-md border border-white/10 bg-white/[0.055] px-3 py-1.5 text-xs text-gray-300 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-45"
|
||||||
|
onClick={() => void prepareCaptionModel()}
|
||||||
|
disabled={captionModelPreparing || modelReady}
|
||||||
|
>
|
||||||
|
{captionModelProgress ? (
|
||||||
|
<span
|
||||||
|
className="absolute inset-y-0 left-0 bg-emerald-400/15 transition-[width] duration-200"
|
||||||
|
style={{ width: `${downloadPercent}%` }}
|
||||||
|
/>
|
||||||
|
) : null}
|
||||||
|
<span className="relative">{downloadLabel}</span>
|
||||||
|
</button>
|
||||||
|
{modelReady ? (
|
||||||
|
<>
|
||||||
|
<button
|
||||||
|
className="rounded-md border border-white/10 bg-white/[0.055] px-3 py-1.5 text-xs text-gray-300 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-45"
|
||||||
|
onClick={() => void probeCaptionRuntime()}
|
||||||
|
disabled={captionRuntimeChecking}
|
||||||
|
>
|
||||||
|
{captionRuntimeChecking ? "Checking runtime..." : "Check runtime"}
|
||||||
|
</button>
|
||||||
|
<button
|
||||||
|
className="rounded-md border border-red-400/20 bg-red-500/10 px-3 py-1.5 text-xs text-red-200 transition-colors hover:bg-red-500/15 disabled:cursor-not-allowed disabled:opacity-45"
|
||||||
|
onClick={() => void deleteCaptionModel()}
|
||||||
|
disabled={captionModelPreparing || captionRuntimeChecking}
|
||||||
|
>
|
||||||
|
Delete model files
|
||||||
|
</button>
|
||||||
|
</>
|
||||||
|
) : null}
|
||||||
|
</div>
|
||||||
|
</SettingsRow>
|
||||||
|
|
||||||
|
<SettingsRow
|
||||||
|
title="Caption queue"
|
||||||
|
description={`Generate missing captions in ${captionScopeLabel}. Captions update in the gallery as the local worker finishes each image.`}
|
||||||
|
>
|
||||||
|
<button
|
||||||
|
className="rounded-md border border-white/10 bg-white/[0.055] px-3 py-1.5 text-xs text-gray-300 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-45"
|
||||||
|
onClick={() => {
|
||||||
|
setCaptionQueueing(true);
|
||||||
|
setCaptionQueueStatus(null);
|
||||||
|
void queueCaptionJobs(selectedFolderId)
|
||||||
|
.then((queued) => {
|
||||||
|
setCaptionQueueStatus(
|
||||||
|
queued === 0
|
||||||
|
? "No missing captions found."
|
||||||
|
: `Queued ${queued.toLocaleString()} image${queued === 1 ? "" : "s"}.`,
|
||||||
|
);
|
||||||
|
})
|
||||||
|
.catch((error) => setCaptionQueueStatus(String(error)))
|
||||||
|
.finally(() => setCaptionQueueing(false));
|
||||||
|
}}
|
||||||
|
disabled={!modelReady || !aiCaptionsEnabled || captionQueueing}
|
||||||
|
>
|
||||||
|
{captionQueueing ? "Queueing..." : selectedFolder ? "Caption this library" : "Caption all libraries"}
|
||||||
|
</button>
|
||||||
|
</SettingsRow>
|
||||||
|
|
||||||
|
<div className="py-4">
|
||||||
|
<p className="text-xs font-medium text-gray-400">Model location</p>
|
||||||
|
<p className="mt-2 break-all rounded-md border border-white/[0.07] bg-black/20 px-3 py-2 text-xs text-gray-600">
|
||||||
|
{modelReady ? captionModelStatus?.local_dir : "Not downloaded"}
|
||||||
|
</p>
|
||||||
|
{captionModelProgress?.current_file ? (
|
||||||
|
<p className="mt-3 break-all text-xs text-gray-500">{captionModelProgress.current_file}</p>
|
||||||
|
) : null}
|
||||||
|
{captionModelError ? (
|
||||||
|
<p className="mt-3 text-xs text-amber-300">{captionModelError}</p>
|
||||||
|
) : null}
|
||||||
|
{captionQueueStatus ? (
|
||||||
|
<p className="mt-3 text-xs text-gray-500">{captionQueueStatus}</p>
|
||||||
|
) : null}
|
||||||
|
{captionRuntimeProbe ? (
|
||||||
|
<div className="mt-4 border-t border-white/[0.07] pt-4">
|
||||||
|
<div className="flex items-center justify-between gap-3">
|
||||||
|
<p className="text-xs font-medium text-gray-400">Runtime check</p>
|
||||||
|
<StatusPill tone="ready">Ready</StatusPill>
|
||||||
|
</div>
|
||||||
|
<p className="mt-2 text-xs text-gray-600">
|
||||||
|
Tokenizer vocabulary: {captionRuntimeProbe.tokenizer_vocab_size.toLocaleString()}
|
||||||
|
</p>
|
||||||
|
<div className="mt-3 space-y-2">
|
||||||
|
{captionRuntimeProbe.sessions.map((session) => (
|
||||||
|
<div key={session.file} className="rounded-md border border-white/[0.07] bg-black/20 px-3 py-2">
|
||||||
|
<p className="break-all text-xs text-gray-400">{session.file}</p>
|
||||||
|
<p className="mt-1 text-[11px] text-gray-600">
|
||||||
|
{session.inputs.length} input{session.inputs.length === 1 ? "" : "s"} · {session.outputs.length} output{session.outputs.length === 1 ? "" : "s"}
|
||||||
|
</p>
|
||||||
|
</div>
|
||||||
|
))}
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
) : null}
|
||||||
|
</div>
|
||||||
|
</SectionShell>
|
||||||
|
) : null}
|
||||||
|
|
||||||
|
{activeSection === "library" ? (
|
||||||
|
<SectionShell eyebrow="Library" title="Indexing and scanning">
|
||||||
|
<SettingsRow title="Background workers" description="Folder-level pause controls remain in the background tasks panel.">
|
||||||
|
<StatusPill tone="muted">Managed per folder</StatusPill>
|
||||||
|
</SettingsRow>
|
||||||
|
<SettingsRow title="Reindexing" description="Use the library sidebar to rescan a folder when files change.">
|
||||||
|
<StatusPill tone="muted">Available</StatusPill>
|
||||||
|
</SettingsRow>
|
||||||
|
</SectionShell>
|
||||||
|
) : null}
|
||||||
|
|
||||||
|
{activeSection === "display" ? (
|
||||||
|
<SectionShell eyebrow="Display" title="Gallery preferences">
|
||||||
|
<SettingsRow title="Grid density" description="Use the toolbar size control to change thumbnail density.">
|
||||||
|
<StatusPill tone="muted">Toolbar</StatusPill>
|
||||||
|
</SettingsRow>
|
||||||
|
<SettingsRow title="Result view" description="Similar results reset to the top on each new search.">
|
||||||
|
<StatusPill tone="ready">Enabled</StatusPill>
|
||||||
|
</SettingsRow>
|
||||||
|
</SectionShell>
|
||||||
|
) : null}
|
||||||
|
|
||||||
|
{activeSection === "storage" ? (
|
||||||
|
<SectionShell eyebrow="Storage" title="Local files">
|
||||||
|
<SettingsRow title="Florence-2" description="Remove the local model without changing the rest of the library.">
|
||||||
|
<button
|
||||||
|
className="rounded-md border border-red-400/20 bg-red-500/10 px-3 py-1.5 text-xs text-red-200 transition-colors hover:bg-red-500/15 disabled:cursor-not-allowed disabled:opacity-45"
|
||||||
|
onClick={() => void deleteCaptionModel()}
|
||||||
|
disabled={captionModelPreparing || !modelReady}
|
||||||
|
>
|
||||||
|
Delete model files
|
||||||
|
</button>
|
||||||
|
</SettingsRow>
|
||||||
|
</SectionShell>
|
||||||
|
) : null}
|
||||||
|
</div>
|
||||||
|
</main>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
);
|
||||||
|
}
|
||||||
@@ -1,5 +1,6 @@
|
|||||||
import { useState, useEffect } from "react";
|
import { useState, useEffect } from "react";
|
||||||
import { getCurrentWindow } from "@tauri-apps/api/window";
|
import { getCurrentWindow } from "@tauri-apps/api/window";
|
||||||
|
import { useGalleryStore } from "../store";
|
||||||
|
|
||||||
// SVG icons for window controls
|
// SVG icons for window controls
|
||||||
function MinimizeIcon() {
|
function MinimizeIcon() {
|
||||||
@@ -37,6 +38,7 @@ function CloseIcon() {
|
|||||||
|
|
||||||
export function TitleBar() {
|
export function TitleBar() {
|
||||||
const [isMaximized, setIsMaximized] = useState(false);
|
const [isMaximized, setIsMaximized] = useState(false);
|
||||||
|
const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen);
|
||||||
const appWindow = getCurrentWindow();
|
const appWindow = getCurrentWindow();
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
@@ -85,6 +87,17 @@ export function TitleBar() {
|
|||||||
className="flex items-stretch h-full"
|
className="flex items-stretch h-full"
|
||||||
style={{ WebkitAppRegion: "no-drag" } as React.CSSProperties}
|
style={{ WebkitAppRegion: "no-drag" } as React.CSSProperties}
|
||||||
>
|
>
|
||||||
|
<button
|
||||||
|
onClick={() => setSettingsOpen(true)}
|
||||||
|
title="Settings"
|
||||||
|
className="group flex h-full w-10 items-center justify-center text-gray-600 transition-colors duration-100 hover:bg-white/6 hover:text-gray-300"
|
||||||
|
>
|
||||||
|
<svg className="h-3.5 w-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor">
|
||||||
|
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.8} d="M10.325 4.317c.426-1.756 2.924-1.756 3.35 0a1.724 1.724 0 002.573 1.066c1.543-.94 3.31.826 2.37 2.37a1.724 1.724 0 001.065 2.572c1.756.426 1.756 2.924 0 3.35a1.724 1.724 0 00-1.066 2.573c.94 1.543-.826 3.31-2.37 2.37a1.724 1.724 0 00-2.572 1.065c-.426 1.756-2.924 1.756-3.35 0a1.724 1.724 0 00-2.573-1.066c-1.543.94-3.31-.826-2.37-2.37a1.724 1.724 0 00-1.065-2.572c-1.756-.426-1.756-2.924 0-3.35a1.724 1.724 0 001.066-2.573c-.94-1.543.826-3.31 2.37-2.37.996.607 2.296.07 2.572-1.065z" />
|
||||||
|
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.8} d="M15 12a3 3 0 11-6 0 3 3 0 016 0z" />
|
||||||
|
</svg>
|
||||||
|
</button>
|
||||||
|
|
||||||
{/* Minimize */}
|
{/* Minimize */}
|
||||||
<button
|
<button
|
||||||
onClick={handleMinimize}
|
onClick={handleMinimize}
|
||||||
|
|||||||
+237
-20
@@ -40,6 +40,10 @@ export interface ImageRecord {
|
|||||||
embedding_model: string | null;
|
embedding_model: string | null;
|
||||||
embedding_updated_at: string | null;
|
embedding_updated_at: string | null;
|
||||||
embedding_error: string | null;
|
embedding_error: string | null;
|
||||||
|
generated_caption: string | null;
|
||||||
|
caption_model: string | null;
|
||||||
|
caption_updated_at: string | null;
|
||||||
|
caption_error: string | null;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface IndexProgress {
|
export interface IndexProgress {
|
||||||
@@ -57,6 +61,9 @@ export interface FolderJobProgress {
|
|||||||
embedding_pending: number;
|
embedding_pending: number;
|
||||||
embedding_ready: number;
|
embedding_ready: number;
|
||||||
embedding_failed: number;
|
embedding_failed: number;
|
||||||
|
caption_pending: number;
|
||||||
|
caption_ready: number;
|
||||||
|
caption_failed: number;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface MediaJobProgressEvent {
|
export interface MediaJobProgressEvent {
|
||||||
@@ -80,6 +87,39 @@ export interface TagCloudEntry {
|
|||||||
thumbnail_path: string | null;
|
thumbnail_path: string | null;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export interface CaptionModelStatus {
|
||||||
|
model_id: string;
|
||||||
|
model_name: string;
|
||||||
|
local_dir: string;
|
||||||
|
ready: boolean;
|
||||||
|
missing_files: string[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface CaptionModelProgress {
|
||||||
|
total_files: number;
|
||||||
|
completed_files: number;
|
||||||
|
current_file: string | null;
|
||||||
|
done: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface CaptionRuntimeSessionProbe {
|
||||||
|
file: string;
|
||||||
|
inputs: string[];
|
||||||
|
outputs: string[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface CaptionRuntimeProbe {
|
||||||
|
ready: boolean;
|
||||||
|
tokenizer_vocab_size: number;
|
||||||
|
sessions: CaptionRuntimeSessionProbe[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface CaptionVisionProbe {
|
||||||
|
input_shape: number[];
|
||||||
|
output_shape: number[];
|
||||||
|
output_values: number;
|
||||||
|
}
|
||||||
|
|
||||||
export type SortOrder =
|
export type SortOrder =
|
||||||
| "date_desc"
|
| "date_desc"
|
||||||
| "date_asc"
|
| "date_asc"
|
||||||
@@ -97,6 +137,7 @@ interface GalleryState {
|
|||||||
totalImages: number;
|
totalImages: number;
|
||||||
loadedCount: number;
|
loadedCount: number;
|
||||||
loadingImages: boolean;
|
loadingImages: boolean;
|
||||||
|
imageLoadError: string | null;
|
||||||
search: string;
|
search: string;
|
||||||
searchMode: SearchMode;
|
searchMode: SearchMode;
|
||||||
sort: SortOrder;
|
sort: SortOrder;
|
||||||
@@ -106,6 +147,9 @@ interface GalleryState {
|
|||||||
zoomPreset: ZoomPreset;
|
zoomPreset: ZoomPreset;
|
||||||
selectedImage: ImageRecord | null;
|
selectedImage: ImageRecord | null;
|
||||||
collectionTitle: string | null;
|
collectionTitle: string | null;
|
||||||
|
similarSourceImageId: number | null;
|
||||||
|
similarHasMore: boolean;
|
||||||
|
galleryScrollResetKey: number;
|
||||||
activeView: ActiveView;
|
activeView: ActiveView;
|
||||||
tagCloudEntries: TagCloudEntry[];
|
tagCloudEntries: TagCloudEntry[];
|
||||||
tagCloudLoading: boolean;
|
tagCloudLoading: boolean;
|
||||||
@@ -113,6 +157,14 @@ interface GalleryState {
|
|||||||
indexingProgress: Record<number, IndexProgress>;
|
indexingProgress: Record<number, IndexProgress>;
|
||||||
mediaJobProgress: Record<number, FolderJobProgress>;
|
mediaJobProgress: Record<number, FolderJobProgress>;
|
||||||
cacheDir: string;
|
cacheDir: string;
|
||||||
|
captionModelStatus: CaptionModelStatus | null;
|
||||||
|
captionModelPreparing: boolean;
|
||||||
|
captionModelError: string | null;
|
||||||
|
captionModelProgress: CaptionModelProgress | null;
|
||||||
|
captionRuntimeProbe: CaptionRuntimeProbe | null;
|
||||||
|
captionRuntimeChecking: boolean;
|
||||||
|
aiCaptionsEnabled: boolean;
|
||||||
|
settingsOpen: boolean;
|
||||||
|
|
||||||
loadFolders: () => Promise<void>;
|
loadFolders: () => Promise<void>;
|
||||||
loadBackgroundJobProgress: () => Promise<void>;
|
loadBackgroundJobProgress: () => Promise<void>;
|
||||||
@@ -136,7 +188,18 @@ interface GalleryState {
|
|||||||
setView: (view: ActiveView) => void;
|
setView: (view: ActiveView) => void;
|
||||||
loadTagCloud: () => Promise<void>;
|
loadTagCloud: () => Promise<void>;
|
||||||
searchByTag: (imageId: number) => void;
|
searchByTag: (imageId: number) => void;
|
||||||
loadSimilarImages: (imageId: number) => Promise<void>;
|
loadSimilarImages: (imageId: number, folderId?: number | null, reset?: boolean) => Promise<void>;
|
||||||
|
suggestImageTags: (imageId: number) => Promise<string[]>;
|
||||||
|
loadCaptionModelStatus: () => Promise<void>;
|
||||||
|
prepareCaptionModel: () => Promise<void>;
|
||||||
|
deleteCaptionModel: () => Promise<void>;
|
||||||
|
probeCaptionRuntime: () => Promise<void>;
|
||||||
|
probeCaptionImage: (imageId: number) => Promise<CaptionVisionProbe>;
|
||||||
|
generateCaptionForImage: (imageId: number) => Promise<ImageRecord>;
|
||||||
|
queueCaptionJobs: (folderId?: number | null) => Promise<number>;
|
||||||
|
queueCaptionForImage: (imageId: number) => Promise<number>;
|
||||||
|
setAiCaptionsEnabled: (enabled: boolean) => void;
|
||||||
|
setSettingsOpen: (open: boolean) => void;
|
||||||
retryFailedEmbeddings: (folderId: number) => Promise<void>;
|
retryFailedEmbeddings: (folderId: number) => Promise<void>;
|
||||||
updateImageDetails: (imageId: number, updates: { favorite?: boolean; rating?: number }) => Promise<void>;
|
updateImageDetails: (imageId: number, updates: { favorite?: boolean; rating?: number }) => Promise<void>;
|
||||||
setCacheDir: (dir: string) => void;
|
setCacheDir: (dir: string) => void;
|
||||||
@@ -144,6 +207,12 @@ interface GalleryState {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const PAGE_SIZE = 200;
|
const PAGE_SIZE = 200;
|
||||||
|
const AI_CAPTIONS_ENABLED_KEY = "phokus.aiCaptionsEnabled";
|
||||||
|
|
||||||
|
function initialAiCaptionsEnabled(): boolean {
|
||||||
|
if (typeof window === "undefined") return false;
|
||||||
|
return window.localStorage.getItem(AI_CAPTIONS_ENABLED_KEY) === "true";
|
||||||
|
}
|
||||||
|
|
||||||
function mergeIntoVisibleWindow(
|
function mergeIntoVisibleWindow(
|
||||||
currentImages: ImageRecord[],
|
currentImages: ImageRecord[],
|
||||||
@@ -266,6 +335,7 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
totalImages: 0,
|
totalImages: 0,
|
||||||
loadedCount: 0,
|
loadedCount: 0,
|
||||||
loadingImages: false,
|
loadingImages: false,
|
||||||
|
imageLoadError: null,
|
||||||
search: "",
|
search: "",
|
||||||
searchMode: "filename",
|
searchMode: "filename",
|
||||||
sort: "date_desc",
|
sort: "date_desc",
|
||||||
@@ -275,6 +345,9 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
zoomPreset: "comfortable",
|
zoomPreset: "comfortable",
|
||||||
selectedImage: null,
|
selectedImage: null,
|
||||||
collectionTitle: null,
|
collectionTitle: null,
|
||||||
|
similarSourceImageId: null,
|
||||||
|
similarHasMore: false,
|
||||||
|
galleryScrollResetKey: 0,
|
||||||
activeView: "gallery",
|
activeView: "gallery",
|
||||||
tagCloudEntries: [],
|
tagCloudEntries: [],
|
||||||
tagCloudLoading: false,
|
tagCloudLoading: false,
|
||||||
@@ -282,6 +355,14 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
indexingProgress: {},
|
indexingProgress: {},
|
||||||
mediaJobProgress: {},
|
mediaJobProgress: {},
|
||||||
cacheDir: "",
|
cacheDir: "",
|
||||||
|
captionModelStatus: null,
|
||||||
|
captionModelPreparing: false,
|
||||||
|
captionModelError: null,
|
||||||
|
captionModelProgress: null,
|
||||||
|
captionRuntimeProbe: null,
|
||||||
|
captionRuntimeChecking: false,
|
||||||
|
aiCaptionsEnabled: initialAiCaptionsEnabled(),
|
||||||
|
settingsOpen: false,
|
||||||
|
|
||||||
setCacheDir: (cacheDir) => set({ cacheDir }),
|
setCacheDir: (cacheDir) => set({ cacheDir }),
|
||||||
|
|
||||||
@@ -327,13 +408,13 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
},
|
},
|
||||||
|
|
||||||
selectFolder: (folderId) => {
|
selectFolder: (folderId) => {
|
||||||
set({ selectedFolderId: folderId, images: [], loadedCount: 0, collectionTitle: null, activeView: "gallery", failedEmbeddingsOnly: false });
|
set({ selectedFolderId: folderId, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, activeView: "gallery", failedEmbeddingsOnly: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
loadImages: async (reset = false) => {
|
loadImages: async (reset = false) => {
|
||||||
const { selectedFolderId, search, searchMode, sort, loadedCount, mediaFilter, favoritesOnly, failedEmbeddingsOnly } = get();
|
const { selectedFolderId, search, searchMode, sort, loadedCount, mediaFilter, favoritesOnly, failedEmbeddingsOnly } = get();
|
||||||
set({ loadingImages: true });
|
set({ loadingImages: true, imageLoadError: null });
|
||||||
|
|
||||||
try {
|
try {
|
||||||
if (searchMode === "semantic" && search.trim()) {
|
if (searchMode === "semantic" && search.trim()) {
|
||||||
@@ -353,6 +434,8 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
loadedCount: images.length,
|
loadedCount: images.length,
|
||||||
loadingImages: false,
|
loadingImages: false,
|
||||||
collectionTitle: `Semantic search: ${search}`,
|
collectionTitle: `Semantic search: ${search}`,
|
||||||
|
similarSourceImageId: null,
|
||||||
|
similarHasMore: false,
|
||||||
});
|
});
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
@@ -382,56 +465,63 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
loadedCount: reset ? result.images.length : state.loadedCount + result.images.length,
|
loadedCount: reset ? result.images.length : state.loadedCount + result.images.length,
|
||||||
loadingImages: false,
|
loadingImages: false,
|
||||||
collectionTitle: reset ? null : state.collectionTitle,
|
collectionTitle: reset ? null : state.collectionTitle,
|
||||||
|
similarSourceImageId: null,
|
||||||
|
similarHasMore: false,
|
||||||
}));
|
}));
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error("Failed to load media:", error);
|
console.error("Failed to load media:", error);
|
||||||
set({ loadingImages: false });
|
set({ loadingImages: false, imageLoadError: String(error) });
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
|
||||||
loadMoreImages: async () => {
|
loadMoreImages: async () => {
|
||||||
const { loadedCount, totalImages, loadingImages } = get();
|
const { loadedCount, totalImages, loadingImages, collectionTitle, similarSourceImageId, similarHasMore, selectedFolderId } = get();
|
||||||
if (loadingImages || loadedCount >= totalImages) return;
|
if (loadingImages || loadedCount >= totalImages) return;
|
||||||
|
if (collectionTitle === "Similar Images" && similarSourceImageId !== null) {
|
||||||
|
if (!similarHasMore) return;
|
||||||
|
await get().loadSimilarImages(similarSourceImageId, selectedFolderId, false);
|
||||||
|
return;
|
||||||
|
}
|
||||||
await get().loadImages(false);
|
await get().loadImages(false);
|
||||||
},
|
},
|
||||||
|
|
||||||
setSearch: (search) => {
|
setSearch: (search) => {
|
||||||
set({ search, images: [], loadedCount: 0, collectionTitle: null });
|
set({ search, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
clearSearch: () => {
|
clearSearch: () => {
|
||||||
set({ search: "", images: [], loadedCount: 0, collectionTitle: null });
|
set({ search: "", images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
resetSearch: () => {
|
resetSearch: () => {
|
||||||
set({ search: "", searchMode: "filename", images: [], loadedCount: 0, collectionTitle: null });
|
set({ search: "", searchMode: "filename", images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
setSearchMode: (searchMode) => {
|
setSearchMode: (searchMode) => {
|
||||||
set({ searchMode, images: [], loadedCount: 0, collectionTitle: null });
|
set({ searchMode, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
setSort: (sort) => {
|
setSort: (sort) => {
|
||||||
set({ sort, images: [], loadedCount: 0, collectionTitle: null });
|
set({ sort, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
setMediaFilter: (mediaFilter) => {
|
setMediaFilter: (mediaFilter) => {
|
||||||
set({ mediaFilter, images: [], loadedCount: 0, collectionTitle: null });
|
set({ mediaFilter, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
setFavoritesOnly: (favoritesOnly) => {
|
setFavoritesOnly: (favoritesOnly) => {
|
||||||
set({ favoritesOnly, images: [], loadedCount: 0, collectionTitle: null });
|
set({ favoritesOnly, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
setFailedEmbeddingsOnly: (failedEmbeddingsOnly) => {
|
setFailedEmbeddingsOnly: (failedEmbeddingsOnly) => {
|
||||||
set({ failedEmbeddingsOnly, images: [], loadedCount: 0, collectionTitle: null });
|
set({ failedEmbeddingsOnly, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null });
|
||||||
void get().loadImages(true);
|
void get().loadImages(true);
|
||||||
},
|
},
|
||||||
|
|
||||||
@@ -461,26 +551,145 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
},
|
},
|
||||||
|
|
||||||
searchByTag: (imageId) => {
|
searchByTag: (imageId) => {
|
||||||
set({ activeView: "gallery", images: [], loadedCount: 0, loadingImages: true, collectionTitle: "Similar Images" });
|
const { selectedFolderId } = get();
|
||||||
void get().loadSimilarImages(imageId);
|
set((state) => ({ activeView: "gallery", images: [], loadedCount: 0, loadingImages: true, collectionTitle: "Similar Images", imageLoadError: null, galleryScrollResetKey: state.galleryScrollResetKey + 1 }));
|
||||||
|
void get().loadSimilarImages(imageId, selectedFolderId);
|
||||||
},
|
},
|
||||||
|
|
||||||
loadSimilarImages: async (imageId) => {
|
loadSimilarImages: async (imageId, folderId = get().selectedFolderId, reset = true) => {
|
||||||
set({ images: [], loadedCount: 0, loadingImages: true, collectionTitle: "Similar Images" });
|
const requestedLimit = reset ? PAGE_SIZE : get().loadedCount + PAGE_SIZE;
|
||||||
|
set((state) => ({
|
||||||
|
images: reset ? [] : get().images,
|
||||||
|
loadedCount: reset ? 0 : get().loadedCount,
|
||||||
|
loadingImages: true,
|
||||||
|
collectionTitle: "Similar Images",
|
||||||
|
imageLoadError: null,
|
||||||
|
similarSourceImageId: imageId,
|
||||||
|
galleryScrollResetKey: reset ? state.galleryScrollResetKey + 1 : state.galleryScrollResetKey,
|
||||||
|
}));
|
||||||
|
try {
|
||||||
const images = await invoke<ImageRecord[]>("find_similar_images", {
|
const images = await invoke<ImageRecord[]>("find_similar_images", {
|
||||||
params: { image_id: imageId, limit: PAGE_SIZE },
|
params: { image_id: imageId, limit: requestedLimit },
|
||||||
});
|
});
|
||||||
|
const hasMore = images.length >= requestedLimit;
|
||||||
set({
|
set({
|
||||||
images,
|
images,
|
||||||
totalImages: images.length,
|
totalImages: hasMore ? images.length + PAGE_SIZE : images.length,
|
||||||
loadedCount: images.length,
|
loadedCount: images.length,
|
||||||
loadingImages: false,
|
loadingImages: false,
|
||||||
|
imageLoadError: null,
|
||||||
collectionTitle: "Similar Images",
|
collectionTitle: "Similar Images",
|
||||||
selectedFolderId: null,
|
similarSourceImageId: imageId,
|
||||||
|
similarHasMore: hasMore,
|
||||||
|
selectedFolderId: folderId ?? null,
|
||||||
|
selectedImage: reset ? null : get().selectedImage,
|
||||||
|
});
|
||||||
|
} catch (error) {
|
||||||
|
console.error("Failed to load similar images:", error);
|
||||||
|
set({
|
||||||
|
images: [],
|
||||||
|
totalImages: 0,
|
||||||
|
loadedCount: 0,
|
||||||
|
loadingImages: false,
|
||||||
|
imageLoadError: String(error),
|
||||||
|
collectionTitle: "Similar Images",
|
||||||
|
similarSourceImageId: imageId,
|
||||||
|
similarHasMore: false,
|
||||||
|
selectedFolderId: folderId ?? null,
|
||||||
selectedImage: null,
|
selectedImage: null,
|
||||||
});
|
});
|
||||||
|
}
|
||||||
},
|
},
|
||||||
|
|
||||||
|
suggestImageTags: async (imageId) => {
|
||||||
|
return invoke<string[]>("suggest_image_tags", {
|
||||||
|
params: { image_id: imageId, limit: 2 },
|
||||||
|
});
|
||||||
|
},
|
||||||
|
|
||||||
|
loadCaptionModelStatus: async () => {
|
||||||
|
try {
|
||||||
|
const captionModelStatus = await invoke<CaptionModelStatus>("get_caption_model_status");
|
||||||
|
set({ captionModelStatus, captionModelError: null });
|
||||||
|
} catch (error) {
|
||||||
|
set({ captionModelError: String(error) });
|
||||||
|
}
|
||||||
|
},
|
||||||
|
|
||||||
|
prepareCaptionModel: async () => {
|
||||||
|
set({ captionModelPreparing: true, captionModelError: null, captionModelProgress: null });
|
||||||
|
try {
|
||||||
|
const captionModelStatus = await invoke<CaptionModelStatus>("prepare_caption_model");
|
||||||
|
window.localStorage.setItem(AI_CAPTIONS_ENABLED_KEY, String(captionModelStatus.ready));
|
||||||
|
set({ captionModelStatus, captionModelPreparing: false, captionModelError: null, captionModelProgress: null, aiCaptionsEnabled: captionModelStatus.ready });
|
||||||
|
} catch (error) {
|
||||||
|
set({ captionModelPreparing: false, captionModelError: String(error), captionModelProgress: null });
|
||||||
|
}
|
||||||
|
},
|
||||||
|
|
||||||
|
deleteCaptionModel: async () => {
|
||||||
|
set({ captionModelPreparing: true, captionModelError: null, captionModelProgress: null });
|
||||||
|
try {
|
||||||
|
const captionModelStatus = await invoke<CaptionModelStatus>("delete_caption_model");
|
||||||
|
window.localStorage.setItem(AI_CAPTIONS_ENABLED_KEY, "false");
|
||||||
|
set({ captionModelStatus, captionModelPreparing: false, captionModelError: null, captionModelProgress: null, captionRuntimeProbe: null, aiCaptionsEnabled: false });
|
||||||
|
} catch (error) {
|
||||||
|
set({ captionModelPreparing: false, captionModelError: String(error), captionModelProgress: null });
|
||||||
|
}
|
||||||
|
},
|
||||||
|
|
||||||
|
probeCaptionRuntime: async () => {
|
||||||
|
set({ captionRuntimeChecking: true, captionModelError: null });
|
||||||
|
try {
|
||||||
|
const captionRuntimeProbe = await invoke<CaptionRuntimeProbe>("probe_caption_runtime");
|
||||||
|
set({ captionRuntimeProbe, captionRuntimeChecking: false, captionModelError: null });
|
||||||
|
} catch (error) {
|
||||||
|
set({ captionRuntimeChecking: false, captionModelError: String(error), captionRuntimeProbe: null });
|
||||||
|
}
|
||||||
|
},
|
||||||
|
|
||||||
|
probeCaptionImage: async (imageId) => {
|
||||||
|
return invoke<CaptionVisionProbe>("probe_caption_image", {
|
||||||
|
params: { image_id: imageId },
|
||||||
|
});
|
||||||
|
},
|
||||||
|
|
||||||
|
generateCaptionForImage: async (imageId) => {
|
||||||
|
const updatedImage = await invoke<ImageRecord>("generate_caption_for_image", {
|
||||||
|
params: { image_id: imageId },
|
||||||
|
});
|
||||||
|
|
||||||
|
set((state) => ({
|
||||||
|
images: replaceImage(state.images, updatedImage, state.sort),
|
||||||
|
selectedImage: state.selectedImage?.id === updatedImage.id ? updatedImage : state.selectedImage,
|
||||||
|
}));
|
||||||
|
|
||||||
|
return updatedImage;
|
||||||
|
},
|
||||||
|
|
||||||
|
queueCaptionJobs: async (folderId = get().selectedFolderId) => {
|
||||||
|
const queued = await invoke<number>("queue_caption_jobs", {
|
||||||
|
params: { folder_id: folderId ?? null, image_id: null },
|
||||||
|
});
|
||||||
|
await get().loadBackgroundJobProgress();
|
||||||
|
return queued;
|
||||||
|
},
|
||||||
|
|
||||||
|
queueCaptionForImage: async (imageId) => {
|
||||||
|
const queued = await invoke<number>("queue_caption_jobs", {
|
||||||
|
params: { folder_id: null, image_id: imageId },
|
||||||
|
});
|
||||||
|
await get().loadBackgroundJobProgress();
|
||||||
|
return queued;
|
||||||
|
},
|
||||||
|
|
||||||
|
setAiCaptionsEnabled: (aiCaptionsEnabled) => {
|
||||||
|
window.localStorage.setItem(AI_CAPTIONS_ENABLED_KEY, String(aiCaptionsEnabled));
|
||||||
|
set({ aiCaptionsEnabled });
|
||||||
|
},
|
||||||
|
|
||||||
|
setSettingsOpen: (settingsOpen) => set({ settingsOpen }),
|
||||||
|
|
||||||
retryFailedEmbeddings: async (folderId) => {
|
retryFailedEmbeddings: async (folderId) => {
|
||||||
await invoke("retry_failed_embeddings", { params: { folder_id: folderId } });
|
await invoke("retry_failed_embeddings", { params: { folder_id: folderId } });
|
||||||
await get().loadBackgroundJobProgress();
|
await get().loadBackgroundJobProgress();
|
||||||
@@ -536,6 +745,13 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
const unlistenCaptionModelProgress = await listen<CaptionModelProgress>("caption-model-progress", (event) => {
|
||||||
|
set({
|
||||||
|
captionModelProgress: event.payload.done ? null : event.payload,
|
||||||
|
captionModelPreparing: !event.payload.done,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
const unlistenImages = await listen<IndexedImagesBatch>("indexed-images", (event) => {
|
const unlistenImages = await listen<IndexedImagesBatch>("indexed-images", (event) => {
|
||||||
const batch = event.payload;
|
const batch = event.payload;
|
||||||
|
|
||||||
@@ -600,6 +816,7 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
return () => {
|
return () => {
|
||||||
unlistenProgress();
|
unlistenProgress();
|
||||||
unlistenMediaJobs();
|
unlistenMediaJobs();
|
||||||
|
unlistenCaptionModelProgress();
|
||||||
unlistenImages();
|
unlistenImages();
|
||||||
unlistenThumbnails();
|
unlistenThumbnails();
|
||||||
};
|
};
|
||||||
|
|||||||
Reference in New Issue
Block a user