Add tag cloud feature with k-means clustering and CUDA support

Introduces an Explore view with a tag cloud that clusters image embeddings
using cosine k-means and labels clusters via vocabulary-nearest-neighbour
CLIP matching. Vocabulary embeddings are disk-cached (FNV hash-keyed) to
avoid redundant inference. Enables CUDA for candle dependencies and adds a
build.rs check that surfaces a clear error when the toolkit is missing.
This commit is contained in:
2026-04-06 12:43:44 +01:00
parent 35c1dafd65
commit 6c3fd449ce
11 changed files with 1248 additions and 37 deletions
+157 -4
View File
@@ -1,10 +1,10 @@
use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord};
use crate::embedder::ClipImageEmbedder;
use crate::embedder;
use crate::indexer;
use crate::vector;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use tauri::{AppHandle, State};
use tauri::{AppHandle, Manager, State};
pub type DbState = DbPool;
@@ -204,8 +204,7 @@ pub async fn semantic_search_images(
db: State<'_, DbState>,
params: SemanticSearchParams,
) -> Result<Vec<ImageRecord>, String> {
let embedder = ClipImageEmbedder::new().map_err(|e| e.to_string())?;
let embedding = embedder.embed_text(&params.query).map_err(|e| e.to_string())?;
let embedding = embedder::embed_text_query(&params.query).map_err(|e| e.to_string())?;
let conn = db.get().map_err(|e| e.to_string())?;
let limit = params.limit.unwrap_or(64);
@@ -225,6 +224,160 @@ pub async fn semantic_search_images(
Ok(images)
}
#[derive(Serialize)]
pub struct TagCloudEntry {
pub label: String,
pub count: usize,
}
/// Clusters the library's image embeddings with k-means, then labels each cluster by
/// finding the closest word in the vocabulary. The vocabulary is loaded from
/// `{app_data_dir}/vocabulary.txt` if present, otherwise from the bundled default.
/// Vocabulary embeddings are cached to disk — only recomputed when the vocabulary changes.
#[tauri::command]
pub async fn get_tag_cloud(
app: AppHandle,
db: State<'_, DbState>,
folder_id: Option<i64>,
) -> Result<Vec<TagCloudEntry>, String> {
let app_data_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
let image_embeddings = {
let conn = db.get().map_err(|e| e.to_string())?;
vector::get_all_image_embeddings(&conn, folder_id).map_err(|e| e.to_string())?
};
let n = image_embeddings.len();
if n < 5 {
return Ok(vec![]);
}
// Load vocabulary (custom file > bundled default)
let vocab = embedder::load_vocabulary(&app_data_dir);
// Embed vocabulary — disk-cached, only recomputes when vocabulary changes
let vocab_refs: Vec<&str> = vocab.iter().map(|s| s.as_str()).collect();
let vocab_embeddings = embedder::embed_vocab_cached(&vocab, &app_data_dir)
.map_err(|e| e.to_string())?;
// Choose k proportional to library size, capped at 30
let k = (n / 20).clamp(5, 30);
// Cluster image embeddings
let (centroids, cluster_counts) = kmeans_cosine(&image_embeddings, k, 40);
// Label each cluster with the nearest vocabulary word
let mut entries: Vec<TagCloudEntry> = Vec::new();
let mut used_labels = std::collections::HashSet::new();
let mut order: Vec<usize> = (0..k).collect();
order.sort_unstable_by(|&a, &b| cluster_counts[b].cmp(&cluster_counts[a]));
for ci in order {
let count = cluster_counts[ci];
if count == 0 { continue; }
let centroid = &centroids[ci];
let label = vocab_refs
.iter()
.zip(vocab_embeddings.iter())
.map(|(&word, emb)| {
let sim: f32 = centroid.iter().zip(emb.iter()).map(|(a, b)| a * b).sum();
(word, sim)
})
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
.map(|(word, _)| word.to_string())
.unwrap_or_else(|| "other".to_string());
if used_labels.insert(label.clone()) {
entries.push(TagCloudEntry { label, count });
}
}
entries.sort_unstable_by(|a, b| b.count.cmp(&a.count));
Ok(entries)
}
// ── k-means with cosine similarity (all vectors assumed to be unit-normalized) ──
fn dot(a: &[f32], b: &[f32]) -> f32 {
a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}
fn normalize(v: &mut Vec<f32>) {
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 1e-10 {
v.iter_mut().for_each(|x| *x /= norm);
}
}
fn kmeans_cosine(
points: &[Vec<f32>],
k: usize,
max_iter: usize,
) -> (Vec<Vec<f32>>, Vec<usize>) {
let n = points.len();
let dim = points[0].len();
// Deterministic k-means++ init: spread centroids as far apart as possible
let mut centroids: Vec<Vec<f32>> = Vec::with_capacity(k);
centroids.push(points[n / 2].clone());
for _ in 1..k {
let next = points
.iter()
.map(|p| {
let best_sim = centroids.iter().map(|c| dot(p, c)).fold(f32::NEG_INFINITY, f32::max);
1.0 - best_sim // distance = 1 - cosine_similarity
})
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
.map(|(i, _)| i)
.unwrap_or(0);
centroids.push(points[next].clone());
}
let mut assignments = vec![0usize; n];
for _ in 0..max_iter {
// Assignment step
let mut changed = false;
for (i, p) in points.iter().enumerate() {
let best = centroids
.iter()
.enumerate()
.map(|(j, c)| (j, dot(p, c)))
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
.map(|(j, _)| j)
.unwrap_or(0);
if assignments[i] != best {
assignments[i] = best;
changed = true;
}
}
if !changed { break; }
// Update step: mean of assigned points, then normalize
let mut sums = vec![vec![0.0f32; dim]; k];
let mut counts = vec![0usize; k];
for (p, &c) in points.iter().zip(assignments.iter()) {
sums[c].iter_mut().zip(p.iter()).for_each(|(s, v)| *s += v);
counts[c] += 1;
}
for (centroid, (sum, &count)) in centroids.iter_mut().zip(sums.iter_mut().zip(counts.iter())) {
if count > 0 {
sum.iter_mut().for_each(|v| *v /= count as f32);
normalize(sum);
*centroid = sum.clone();
}
}
}
let mut counts = vec![0usize; k];
for &a in &assignments { counts[a] += 1; }
(centroids, counts)
}
#[derive(Serialize)]
pub struct WorkerStates {
pub thumbnail_paused: bool,
+216 -22
View File
@@ -1,11 +1,180 @@
use anyhow::Result;
use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::clip::{self, ClipModel};
use hf_hub::{api::sync::Api, Repo, RepoType};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::sync::{Mutex, OnceLock};
use tokenizers::Tokenizer;
static TEXT_SEARCH_EMBEDDER: OnceLock<Mutex<Option<ClipImageEmbedder>>> = OnceLock::new();
/// In-process cache so the disk file is only read/written once per session.
static VOCAB_EMBED_CACHE: OnceLock<Mutex<Option<Vec<Vec<f32>>>>> = OnceLock::new();
/// Default vocabulary bundled with the binary.
pub const DEFAULT_VOCABULARY: &str = include_str!("../data/vocabulary.txt");
/// Embed a text query using a lazily-initialized, cached CLIP embedder.
pub fn embed_text_query(query: &str) -> Result<Vec<f32>> {
with_text_embedder(|e| e.embed_text(query))
}
/// Parse vocabulary text: strip comment lines (starting with `#`) and blank lines.
pub fn parse_vocabulary(text: &str) -> Vec<String> {
text.lines()
.map(|l| l.trim())
.filter(|l| !l.is_empty() && !l.starts_with('#'))
.map(|l| l.to_string())
.collect()
}
/// Load the vocabulary from `{cache_dir}/vocabulary.txt` if it exists,
/// otherwise fall back to the binary-bundled default.
pub fn load_vocabulary(cache_dir: &Path) -> Vec<String> {
let custom_path = cache_dir.join("vocabulary.txt");
if custom_path.exists() {
if let Ok(text) = std::fs::read_to_string(&custom_path) {
let words = parse_vocabulary(&text);
if !words.is_empty() {
println!("Using custom vocabulary from {:?} ({} words)", custom_path, words.len());
return words;
}
}
}
parse_vocabulary(DEFAULT_VOCABULARY)
}
/// Embed the vocabulary, using a disk cache so CLIP is only called when the vocabulary
/// actually changes. Cache file: `{cache_dir}/vocab_embeddings.bin`.
///
/// Format: `[u64 hash][u32 n_words][u32 dim][(n_words × dim) × f32 LE]`
pub fn embed_vocab_cached(vocab: &[String], cache_dir: &Path) -> Result<Vec<Vec<f32>>> {
// In-process cache hit
{
let guard = VOCAB_EMBED_CACHE
.get_or_init(|| Mutex::new(None))
.lock()
.map_err(|_| anyhow::anyhow!("Vocab cache lock poisoned"))?;
if let Some(cached) = guard.as_ref() {
return Ok(cached.clone());
}
}
let vocab_hash = fnv_hash(vocab);
let disk_path = cache_dir.join("vocab_embeddings.bin");
// Try loading from disk
if let Ok(embeddings) = load_disk_cache(&disk_path, vocab_hash) {
let mut guard = VOCAB_EMBED_CACHE
.get_or_init(|| Mutex::new(None))
.lock()
.map_err(|_| anyhow::anyhow!("Vocab cache lock poisoned"))?;
*guard = Some(embeddings.clone());
println!("Vocabulary embeddings loaded from disk cache ({} words).", embeddings.len());
return Ok(embeddings);
}
// Compute embeddings
println!("Computing vocabulary embeddings ({} words) — this is cached after the first run.", vocab.len());
let prompts: Vec<String> = vocab.iter().map(|w| format!("a photo of {}", w)).collect();
let prompt_refs: Vec<&str> = prompts.iter().map(|s| s.as_str()).collect();
let embeddings = with_text_embedder(|e| {
let mut all = Vec::with_capacity(vocab.len());
for chunk in prompt_refs.chunks(64) {
all.extend(e.embed_texts_batch(chunk)?);
}
Ok(all)
})?;
// Save to disk
if let Err(e) = save_disk_cache(&disk_path, vocab_hash, &embeddings) {
eprintln!("Warning: could not write vocab cache: {e}");
}
let mut guard = VOCAB_EMBED_CACHE
.get_or_init(|| Mutex::new(None))
.lock()
.map_err(|_| anyhow::anyhow!("Vocab cache lock poisoned"))?;
*guard = Some(embeddings.clone());
Ok(embeddings)
}
fn fnv_hash(words: &[String]) -> u64 {
let mut h: u64 = 0xcbf29ce484222325;
for w in words {
for b in w.bytes() {
h = h.wrapping_mul(0x100000001b3) ^ (b as u64);
}
h = h.wrapping_mul(0x100000001b3) ^ 0x0A; // newline separator
}
h
}
fn load_disk_cache(path: &Path, expected_hash: u64) -> Result<Vec<Vec<f32>>> {
let mut f = std::fs::File::open(path).context("no cache file")?;
let mut buf = Vec::new();
f.read_to_end(&mut buf)?;
if buf.len() < 16 {
anyhow::bail!("cache too short");
}
let hash = u64::from_le_bytes(buf[0..8].try_into()?);
if hash != expected_hash {
anyhow::bail!("vocab hash mismatch — recomputing");
}
let n = u32::from_le_bytes(buf[8..12].try_into()?) as usize;
let dim = u32::from_le_bytes(buf[12..16].try_into()?) as usize;
let expected_len = 16 + n * dim * 4;
if buf.len() != expected_len {
anyhow::bail!("cache size mismatch");
}
let mut embeddings = Vec::with_capacity(n);
let data = &buf[16..];
for i in 0..n {
let start = i * dim * 4;
let emb: Vec<f32> = data[start..start + dim * 4]
.chunks_exact(4)
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect();
embeddings.push(emb);
}
Ok(embeddings)
}
fn save_disk_cache(path: &Path, hash: u64, embeddings: &[Vec<f32>]) -> Result<()> {
let n = embeddings.len();
let dim = embeddings.first().map(|e| e.len()).unwrap_or(0);
let mut buf = Vec::with_capacity(16 + n * dim * 4);
buf.extend_from_slice(&hash.to_le_bytes());
buf.extend_from_slice(&(n as u32).to_le_bytes());
buf.extend_from_slice(&(dim as u32).to_le_bytes());
for emb in embeddings {
for &v in emb {
buf.extend_from_slice(&v.to_le_bytes());
}
}
std::fs::write(path, buf)?;
println!("Vocabulary embeddings cached to disk ({n} words, {dim} dims).");
Ok(())
}
fn with_text_embedder<T>(f: impl FnOnce(&ClipImageEmbedder) -> Result<T>) -> Result<T> {
let lock = TEXT_SEARCH_EMBEDDER.get_or_init(|| Mutex::new(None));
let mut guard = lock.lock().map_err(|_| anyhow::anyhow!("Text embedder lock poisoned"))?;
if guard.is_none() {
println!("Initializing CLIP text embedder...");
*guard = Some(ClipImageEmbedder::new()?);
}
f(guard.as_ref().unwrap())
}
pub struct ClipImageEmbedder {
model: ClipModel,
tokenizer: Tokenizer,
@@ -67,36 +236,61 @@ impl ClipImageEmbedder {
}
pub fn embed_text(&self, query: &str) -> Result<Vec<f32>> {
let encoding = self
.tokenizer
.encode(query, true)
.map_err(anyhow::Error::msg)?;
let token_ids = encoding
.get_ids()
Ok(self.embed_texts_batch(&[query])?.remove(0))
}
/// Embed multiple text queries in a single CLIP forward pass.
/// All sequences are padded to the same length (max 77 tokens).
pub fn embed_texts_batch(&self, queries: &[&str]) -> Result<Vec<Vec<f32>>> {
if queries.is_empty() {
return Ok(vec![]);
}
let encodings: Vec<_> = queries
.iter()
.map(|token| *token as u32)
.collect::<Vec<_>>();
let input_ids = Tensor::new(vec![token_ids], &self.device)?;
.map(|q| self.tokenizer.encode(*q, true).map_err(anyhow::Error::msg))
.collect::<Result<_>>()?;
let max_len = encodings
.iter()
.map(|e| e.get_ids().len())
.max()
.unwrap_or(10)
.min(77);
let n = queries.len();
let mut flat = vec![0u32; n * max_len];
for (i, enc) in encodings.iter().enumerate() {
let ids = enc.get_ids();
let len = ids.len().min(max_len);
for j in 0..len {
flat[i * max_len + j] = ids[j] as u32;
}
}
let input_ids = Tensor::new(flat, &self.device)?.reshape((n, max_len))?;
let features = self.model.get_text_features(&input_ids)?;
let normalized = clip::div_l2_norm(&features)?;
Ok(normalized.flatten_all()?.to_vec1::<f32>()?)
(0..n)
.map(|i| Ok(normalized.get(i)?.flatten_all()?.to_vec1::<f32>()?))
.collect()
}
}
fn resolve_device() -> Result<Device> {
#[cfg(feature = "candle-cuda")]
{
let cuda_device = Device::cuda_if_available(0)?;
if cuda_device.is_cuda() {
println!("CLIP embedder using CUDA device.");
return Ok(cuda_device);
match Device::cuda_if_available(0) {
Ok(device) if device.is_cuda() => {
println!("CLIP embedder: using CUDA GPU (device 0).");
return Ok(device);
}
Ok(_) => {
println!("CLIP embedder: no compatible CUDA GPU found — using CPU.");
}
Err(e) => {
println!("CLIP embedder: CUDA init failed ({e}) — using CPU.");
}
}
#[cfg(not(feature = "candle-cuda"))]
println!("CLIP embedder built without CUDA support.");
println!("CLIP embedder using CPU device.");
Ok(Device::Cpu)
}
+1
View File
@@ -77,6 +77,7 @@ pub fn run() {
commands::semantic_search_images,
commands::set_worker_paused,
commands::get_worker_states,
commands::get_tag_cloud,
])
.run(tauri::generate_context!())
.expect("error while running tauri application");
+35
View File
@@ -78,6 +78,41 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) ->
Ok(ids)
}
/// Returns all stored image embeddings, optionally filtered to one folder.
/// Each embedding is returned as a normalized f32 vector.
pub fn get_all_image_embeddings(conn: &Connection, folder_id: Option<i64>) -> Result<Vec<Vec<f32>>> {
let packed_rows: Vec<Vec<u8>> = match folder_id {
Some(fid) => {
let mut stmt = conn.prepare(
"SELECT embedding FROM image_vec
WHERE image_id IN (SELECT id FROM images WHERE folder_id = ?1)",
)?;
let rows: Vec<Vec<u8>> = stmt
.query_map([fid], |row| row.get::<_, Vec<u8>>(0))?
.filter_map(|r| r.ok())
.collect();
rows
}
None => {
let mut stmt = conn.prepare("SELECT embedding FROM image_vec")?;
let rows: Vec<Vec<u8>> = stmt
.query_map([], |row| row.get::<_, Vec<u8>>(0))?
.filter_map(|r| r.ok())
.collect();
rows
}
};
Ok(packed_rows.iter().map(|b| unpack_f32(b)).collect())
}
fn unpack_f32(bytes: &[u8]) -> Vec<f32> {
bytes
.chunks_exact(4)
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect()
}
pub fn search_image_ids_by_embedding(
conn: &Connection,
embedding: &[f32],