feat(discovery): AI tagging, semantic search, similarity, duplicate finder, folder management
Adds a full discovery and organisation layer to the gallery.
## AI Tagger
- WD tagger (ONNX via ort) with DirectML/CPU acceleration
- Per-image and per-folder tag queuing; configurable batch size and
confidence threshold; model downloaded on first use via ureq/zip
- Tags stored with source ('ai' | 'user'); user tags are never
overwritten by AI; AI tags protected from accidental promotion
- Tagger pause/resume per folder; in-flight batches discarded cleanly
on pause or cancellation without leaving jobs stuck in processing
## Semantic & Tag Search
- CLIP text-query embedding via candle (HuggingFace hub model)
- Progressive candidate doubling with filter post-processing so
folder/rating/media-kind filters do not silently under-return results
- Tag search with DB-backed pagination (offset + COUNT total)
- Filename search unchanged; prefix syntax: s:, t:, f: in search bar
## Similarity Search
- HNSW in-memory index (hnsw_rs) with monotonic embedding_revision
counter for cache invalidation; build_index retries if a concurrent
write advances the revision during construction
- Folder-scoped similar search uses brute-force cosine scan (no k limit)
- Region-based similarity: crop an arbitrary rectangle in the lightbox,
embed it with CLIP, find the nearest images globally or within folder
- Pagination for both similar and region results; scope toggle
(all media / current folder) re-runs the query without reopening image
## Duplicate Finder
- Three-phase detection: stat (size) → sample hash (4×16 KB windows)
→ full-file hash (xxh3, large files only) to eliminate false positives
- Filesystem-first deletion: only removes DB rows for files successfully
deleted from disk; failed files remain visible for retry
- Persisted scan cache (SQLite) with invalidation on reindex and deletion;
both global and per-folder scopes invalidated after any deletion
## Tag Cloud & Explore
- Visual cluster explore: k-means over CLIP embeddings, configurable k,
representative thumbnail per cluster; cache keyed on xxh3 of embedding set
- Tag explore: ranked tag list with image counts and representative
thumbnail per tag; invalidated when AI tagging completes or folder removed
- Tag autocomplete for search bar
## Folder Management
- Inline rename (display name only, not OS rename)
- Relocate: file-picker to update folder path; all image paths rewritten
using SUBSTR prefix replacement to avoid corrupting paths that contain
the folder name as a substring
- Missing-folder recovery banner: Locate or Remove, never silent deletion
- Right-click context menu on sidebar items: Reindex, Rename, Locate,
Remove; hover buttons preserved alongside context menu
## Background Tasks
- Per-folder progress panel with embedding, tagging, caption, and
thumbnail stage breakdown
- Retry button per task for failed embeddings and tagging jobs
- Completed-task notifications (system tray via tauri-plugin-notification)
- Scan errors shown inline; WalkDir permission errors protect existing
records from deletion rather than cascading data loss
## Backend correctness
- sqlite-vec DML performed outside transactions throughout (virtual-table
limitation); embedding revision incremented on delete as well as insert
- Tagging job discard checks status = 'processing' so paused jobs reset
to 'pending' are also skipped, not just cancelled ones
- Stale async responses in Lightbox guarded with cancellation flags and
currentImageIdRef for both tag-load and tag-add callbacks
- Duplicate cache cleared on reindex; folder-removal clears vector rows
outside transaction before deleting the folder row
This commit was merged in pull request #8.
This commit is contained in:
@@ -0,0 +1,12 @@
|
||||
[target.x86_64-pc-windows-msvc]
|
||||
rustflags = [
|
||||
# Disable MSVC incremental linking — avoids .ilk file corruption
|
||||
# and removes one source of link failure on restart.
|
||||
"-C", "link-arg=/INCREMENTAL:NO",
|
||||
# Skip PDB generation entirely in dev builds.
|
||||
# The .pdb file is the most common reason the linker fails after
|
||||
# an unclean shutdown: the previous phokus.exe process holds the
|
||||
# file open, so link.exe cannot write a new one → LNK error.
|
||||
# Rust backtraces still work without MSVC PDBs.
|
||||
"-C", "link-arg=/DEBUG:NONE",
|
||||
]
|
||||
Generated
+478
-1
@@ -61,6 +61,73 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anndists"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9a8396b473aa0bceed68fb32462505387ea39fa47c7029417e0a49f10592b036"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"cfg-if",
|
||||
"cpu-time",
|
||||
"env_logger",
|
||||
"lazy_static",
|
||||
"log",
|
||||
"num-traits",
|
||||
"num_cpus",
|
||||
"rayon",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anstream"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d"
|
||||
dependencies = [
|
||||
"anstyle",
|
||||
"anstyle-parse",
|
||||
"anstyle-query",
|
||||
"anstyle-wincon",
|
||||
"colorchoice",
|
||||
"is_terminal_polyfill",
|
||||
"utf8parse",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anstyle"
|
||||
version = "1.0.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
|
||||
|
||||
[[package]]
|
||||
name = "anstyle-parse"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e"
|
||||
dependencies = [
|
||||
"utf8parse",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anstyle-query"
|
||||
version = "1.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anstyle-wincon"
|
||||
version = "3.0.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d"
|
||||
dependencies = [
|
||||
"anstyle",
|
||||
"once_cell_polyfill",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anyhow"
|
||||
version = "1.0.102"
|
||||
@@ -266,6 +333,15 @@ version = "1.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
|
||||
|
||||
[[package]]
|
||||
name = "bincode"
|
||||
version = "1.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bit-set"
|
||||
version = "0.8.0"
|
||||
@@ -595,6 +671,12 @@ version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||
|
||||
[[package]]
|
||||
name = "cfg_aliases"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
|
||||
|
||||
[[package]]
|
||||
name = "chacha20"
|
||||
version = "0.10.0"
|
||||
@@ -626,6 +708,12 @@ version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b"
|
||||
|
||||
[[package]]
|
||||
name = "colorchoice"
|
||||
version = "1.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570"
|
||||
|
||||
[[package]]
|
||||
name = "combine"
|
||||
version = "4.6.7"
|
||||
@@ -757,6 +845,16 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cpu-time"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e9e393a7668fe1fad3075085b86c781883000b4ede868f43627b34a87c8b7ded"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cpufeatures"
|
||||
version = "0.2.17"
|
||||
@@ -874,6 +972,27 @@ dependencies = [
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "csv"
|
||||
version = "1.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938"
|
||||
dependencies = [
|
||||
"csv-core",
|
||||
"itoa",
|
||||
"ryu",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "csv-core"
|
||||
version = "0.1.13"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ctor"
|
||||
version = "0.2.9"
|
||||
@@ -1331,12 +1450,35 @@ dependencies = [
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "env_filter"
|
||||
version = "1.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef"
|
||||
dependencies = [
|
||||
"log",
|
||||
"regex",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "env_home"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe"
|
||||
|
||||
[[package]]
|
||||
name = "env_logger"
|
||||
version = "0.11.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a"
|
||||
dependencies = [
|
||||
"anstream",
|
||||
"anstyle",
|
||||
"env_filter",
|
||||
"jiff",
|
||||
"log",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
@@ -2319,6 +2461,8 @@ version = "0.15.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
|
||||
dependencies = [
|
||||
"allocator-api2",
|
||||
"equivalent",
|
||||
"foldhash 0.1.5",
|
||||
]
|
||||
|
||||
@@ -2389,6 +2533,37 @@ dependencies = [
|
||||
"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]]
|
||||
name = "hnsw_rs"
|
||||
version = "0.3.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "43a5258f079b97bf2e8311ff9579e903c899dcbac0d9a138d62e9a066778bd07"
|
||||
dependencies = [
|
||||
"anndists",
|
||||
"anyhow",
|
||||
"bincode",
|
||||
"cfg-if",
|
||||
"cpu-time",
|
||||
"env_logger",
|
||||
"hashbrown 0.15.5",
|
||||
"indexmap 2.13.1",
|
||||
"lazy_static",
|
||||
"log",
|
||||
"mmap-rs",
|
||||
"num-traits",
|
||||
"num_cpus",
|
||||
"parking_lot",
|
||||
"rand 0.9.2",
|
||||
"rayon",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "html5ever"
|
||||
version = "0.29.1"
|
||||
@@ -2792,6 +2967,12 @@ dependencies = [
|
||||
"once_cell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "is_terminal_polyfill"
|
||||
version = "1.70.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695"
|
||||
|
||||
[[package]]
|
||||
name = "itertools"
|
||||
version = "0.14.0"
|
||||
@@ -2830,6 +3011,30 @@ dependencies = [
|
||||
"system-deps",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jiff"
|
||||
version = "0.2.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359"
|
||||
dependencies = [
|
||||
"jiff-static",
|
||||
"log",
|
||||
"portable-atomic",
|
||||
"portable-atomic-util",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jiff-static"
|
||||
version = "0.2.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jni"
|
||||
version = "0.21.1"
|
||||
@@ -2931,6 +3136,12 @@ dependencies = [
|
||||
"selectors 0.24.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
|
||||
|
||||
[[package]]
|
||||
name = "leb128fmt"
|
||||
version = "0.1.0"
|
||||
@@ -3059,6 +3270,12 @@ version = "0.4.29"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
|
||||
|
||||
[[package]]
|
||||
name = "lzma-rust2"
|
||||
version = "0.15.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1670343e58806300d87950e3401e820b519b9384281bbabfb15e3636689ffd69"
|
||||
|
||||
[[package]]
|
||||
name = "lzma-sys"
|
||||
version = "0.1.20"
|
||||
@@ -3076,6 +3293,27 @@ version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
|
||||
|
||||
[[package]]
|
||||
name = "mac-notification-sys"
|
||||
version = "0.6.13"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50efa634682b3fc5a1ab6f3dd5b2bce7b848011fc485b53b063dc68f2f74feae"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"objc2",
|
||||
"objc2-foundation",
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mach2"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d640282b302c0bb0a2a8e0233ead9035e3bed871f0b7e81fe4a1ec829765db44"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "macro_rules_attribute"
|
||||
version = "0.2.2"
|
||||
@@ -3134,6 +3372,16 @@ version = "0.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
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]]
|
||||
name = "memchr"
|
||||
version = "2.8.0"
|
||||
@@ -3192,6 +3440,23 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mmap-rs"
|
||||
version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4ecce9d566cb9234ae3db9e249c8b55665feaaf32b0859ff1e27e310d2beb3d8"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"combine",
|
||||
"libc",
|
||||
"mach2",
|
||||
"nix",
|
||||
"sysctl",
|
||||
"thiserror 2.0.18",
|
||||
"widestring",
|
||||
"windows 0.48.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "monostate"
|
||||
version = "0.1.18"
|
||||
@@ -3262,6 +3527,21 @@ dependencies = [
|
||||
"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]]
|
||||
name = "ndk"
|
||||
version = "0.9.0"
|
||||
@@ -3298,6 +3578,18 @@ version = "1.0.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086"
|
||||
|
||||
[[package]]
|
||||
name = "nix"
|
||||
version = "0.30.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"cfg-if",
|
||||
"cfg_aliases",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nodrop"
|
||||
version = "0.1.14"
|
||||
@@ -3314,6 +3606,20 @@ dependencies = [
|
||||
"minimal-lexical",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "notify-rust"
|
||||
version = "4.17.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50ff2e74231b72c832d82982193b417f230945be6bdb5575b251d941d31adb00"
|
||||
dependencies = [
|
||||
"futures-lite",
|
||||
"log",
|
||||
"mac-notification-sys",
|
||||
"serde",
|
||||
"tauri-winrt-notification",
|
||||
"zbus",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ntapi"
|
||||
version = "0.4.3"
|
||||
@@ -3576,6 +3882,12 @@ version = "1.21.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
|
||||
|
||||
[[package]]
|
||||
name = "once_cell_polyfill"
|
||||
version = "1.70.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
|
||||
|
||||
[[package]]
|
||||
name = "onig"
|
||||
version = "6.5.1"
|
||||
@@ -3670,6 +3982,31 @@ dependencies = [
|
||||
"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]]
|
||||
name = "pango"
|
||||
version = "0.18.3"
|
||||
@@ -3947,11 +4284,15 @@ dependencies = [
|
||||
"candle-nn",
|
||||
"candle-transformers",
|
||||
"chrono",
|
||||
"csv",
|
||||
"fast_image_resize",
|
||||
"ffmpeg-sidecar",
|
||||
"hf-hub",
|
||||
"hnsw_rs",
|
||||
"image",
|
||||
"log",
|
||||
"memmap2",
|
||||
"ort",
|
||||
"r2d2",
|
||||
"r2d2_sqlite",
|
||||
"rayon",
|
||||
@@ -3964,12 +4305,15 @@ dependencies = [
|
||||
"tauri-build",
|
||||
"tauri-plugin-dialog",
|
||||
"tauri-plugin-fs",
|
||||
"tauri-plugin-notification",
|
||||
"tauri-plugin-opener",
|
||||
"tokenizers",
|
||||
"tokio",
|
||||
"ureq",
|
||||
"uuid",
|
||||
"walkdir",
|
||||
"xxhash-rust",
|
||||
"zip 4.6.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4009,7 +4353,7 @@ checksum = "740ebea15c5d1428f910cd1a5f52cebf8d25006245ed8ade92702f4943d91e07"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"indexmap 2.13.1",
|
||||
"quick-xml",
|
||||
"quick-xml 0.38.4",
|
||||
"serde",
|
||||
"time",
|
||||
]
|
||||
@@ -4060,6 +4404,15 @@ version = "1.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
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]]
|
||||
name = "potential_utf"
|
||||
version = "0.1.5"
|
||||
@@ -4217,6 +4570,15 @@ version = "2.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3"
|
||||
|
||||
[[package]]
|
||||
name = "quick-xml"
|
||||
version = "0.37.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "331e97a1af0bf59823e6eadffe373d7b27f485be8748f71471c662c1f269b7fb"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quick-xml"
|
||||
version = "0.38.4"
|
||||
@@ -4421,6 +4783,12 @@ version = "0.6.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"
|
||||
|
||||
[[package]]
|
||||
name = "rawpointer"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
||||
|
||||
[[package]]
|
||||
name = "rayon"
|
||||
version = "1.11.0"
|
||||
@@ -5684,6 +6052,25 @@ dependencies = [
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-plugin-notification"
|
||||
version = "2.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "01fc2c5ff41105bd1f7242d8201fdf3efd70749b82fa013a17f2126357d194cc"
|
||||
dependencies = [
|
||||
"log",
|
||||
"notify-rust",
|
||||
"rand 0.9.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_repr",
|
||||
"tauri",
|
||||
"tauri-plugin",
|
||||
"thiserror 2.0.18",
|
||||
"time",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-plugin-opener"
|
||||
version = "2.5.3"
|
||||
@@ -5806,6 +6193,18 @@ dependencies = [
|
||||
"toml 0.9.12+spec-1.1.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-winrt-notification"
|
||||
version = "0.7.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b1e66e07de489fe43a46678dd0b8df65e0c973909df1b60ba33874e297ba9b9"
|
||||
dependencies = [
|
||||
"quick-xml 0.37.5",
|
||||
"thiserror 2.0.18",
|
||||
"windows 0.61.3",
|
||||
"windows-version",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tempfile"
|
||||
version = "3.27.0"
|
||||
@@ -6473,6 +6872,12 @@ version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "utf8parse"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.23.0"
|
||||
@@ -6836,6 +7241,12 @@ dependencies = [
|
||||
"winsafe",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "widestring"
|
||||
version = "1.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471"
|
||||
|
||||
[[package]]
|
||||
name = "winapi"
|
||||
version = "0.3.9"
|
||||
@@ -6882,6 +7293,15 @@ dependencies = [
|
||||
"windows-version",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows"
|
||||
version = "0.48.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f"
|
||||
dependencies = [
|
||||
"windows-targets 0.48.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows"
|
||||
version = "0.61.3"
|
||||
@@ -7134,6 +7554,21 @@ dependencies = [
|
||||
"windows_x86_64_msvc 0.42.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-targets"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c"
|
||||
dependencies = [
|
||||
"windows_aarch64_gnullvm 0.48.5",
|
||||
"windows_aarch64_msvc 0.48.5",
|
||||
"windows_i686_gnu 0.48.5",
|
||||
"windows_i686_msvc 0.48.5",
|
||||
"windows_x86_64_gnu 0.48.5",
|
||||
"windows_x86_64_gnullvm 0.48.5",
|
||||
"windows_x86_64_msvc 0.48.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-targets"
|
||||
version = "0.52.6"
|
||||
@@ -7200,6 +7635,12 @@ version = "0.42.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_gnullvm"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_gnullvm"
|
||||
version = "0.52.6"
|
||||
@@ -7218,6 +7659,12 @@ version = "0.42.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_msvc"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_msvc"
|
||||
version = "0.52.6"
|
||||
@@ -7236,6 +7683,12 @@ version = "0.42.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnu"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnu"
|
||||
version = "0.52.6"
|
||||
@@ -7266,6 +7719,12 @@ version = "0.42.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_msvc"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_msvc"
|
||||
version = "0.52.6"
|
||||
@@ -7284,6 +7743,12 @@ version = "0.42.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnu"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnu"
|
||||
version = "0.52.6"
|
||||
@@ -7302,6 +7767,12 @@ version = "0.42.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnullvm"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnullvm"
|
||||
version = "0.52.6"
|
||||
@@ -7320,6 +7791,12 @@ version = "0.42.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_msvc"
|
||||
version = "0.48.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_msvc"
|
||||
version = "0.52.6"
|
||||
|
||||
@@ -27,6 +27,7 @@ rusqlite = { version = "0.32", features = ["bundled"] }
|
||||
r2d2 = "0.8"
|
||||
r2d2_sqlite = "0.25"
|
||||
sqlite-vec = "=0.1.9"
|
||||
hnsw_rs = "0.3.4"
|
||||
image = { version = "0.25", default-features = false, features = ["jpeg", "png", "webp", "gif", "bmp", "tiff"] }
|
||||
fast_image_resize = { version = "6.0.0", features = ["image"] }
|
||||
walkdir = "2"
|
||||
@@ -38,9 +39,63 @@ anyhow = "1"
|
||||
log = "0.4"
|
||||
ffmpeg-sidecar = "2.5.0"
|
||||
xxhash-rust = { version = "0.8", features = ["xxh3"] }
|
||||
memmap2 = "0.9"
|
||||
sysinfo = "0.38.4"
|
||||
candle-core = { version = "0.10.2", features = ["cuda"] }
|
||||
candle-nn = { version = "0.10.2", features = ["cuda"] }
|
||||
candle-transformers = { version = "0.10.2", features = ["cuda"] }
|
||||
hf-hub = { version = "0.5.0", default-features = false, features = ["ureq", "native-tls"] }
|
||||
tokenizers = "0.22.1"
|
||||
ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "download-binaries", "copy-dylibs", "load-dynamic", "directml", "api-24", "tls-native"] }
|
||||
ureq = { version = "3.3.0", default-features = false, features = ["native-tls"] }
|
||||
zip = { version = "4.6.1", default-features = false, features = ["deflate"] }
|
||||
csv = "1"
|
||||
tauri-plugin-notification = "2"
|
||||
|
||||
# ── Dev-mode performance ────────────────────────────────────────────────────
|
||||
# opt-level=1 on the main crate keeps incremental compile short.
|
||||
# Only the packages that are genuine hot-path bottlenecks in dev get opt-level=3
|
||||
# (using "*" caused cargo to recheck all dependency fingerprints too aggressively,
|
||||
# which caused spurious full rebuilds and linker conflicts).
|
||||
[profile.dev]
|
||||
opt-level = 1
|
||||
|
||||
# ML inference — without opt these run 20-50× slower than release
|
||||
[profile.dev.package.candle-core]
|
||||
opt-level = 3
|
||||
[profile.dev.package.candle-nn]
|
||||
opt-level = 3
|
||||
[profile.dev.package.candle-transformers]
|
||||
opt-level = 3
|
||||
|
||||
# ONNX runtime (WD tagger)
|
||||
[profile.dev.package.ort]
|
||||
opt-level = 3
|
||||
[profile.dev.package.ort-sys]
|
||||
opt-level = 3
|
||||
|
||||
# Image decode/resize workers
|
||||
[profile.dev.package.image]
|
||||
opt-level = 3
|
||||
[profile.dev.package.fast_image_resize]
|
||||
opt-level = 3
|
||||
|
||||
# Parallel work scheduler
|
||||
[profile.dev.package.rayon]
|
||||
opt-level = 3
|
||||
[profile.dev.package.rayon-core]
|
||||
opt-level = 3
|
||||
|
||||
# Tokenisation (embedding model pre-processing)
|
||||
[profile.dev.package.tokenizers]
|
||||
opt-level = 3
|
||||
|
||||
# Hashing (duplicate finder)
|
||||
[profile.dev.package.xxhash-rust]
|
||||
opt-level = 3
|
||||
|
||||
# SQLite (frequent db calls in workers)
|
||||
[profile.dev.package.rusqlite]
|
||||
opt-level = 3
|
||||
[profile.dev.package.libsqlite3-sys]
|
||||
opt-level = 3
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
"fs:allow-read-dir",
|
||||
"fs:read-files",
|
||||
"fs:read-dirs",
|
||||
"notification:default",
|
||||
"core:window:allow-minimize",
|
||||
"core:window:allow-close",
|
||||
"core:window:allow-toggle-maximize",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+1184
-48
File diff suppressed because it is too large
Load Diff
+1307
-36
File diff suppressed because it is too large
Load Diff
@@ -74,6 +74,51 @@ impl ClipImageEmbedder {
|
||||
Ok(self.embed_images(&[path.to_path_buf()])?.remove(0))
|
||||
}
|
||||
|
||||
/// Embed a cropped region of an image without writing a temp file to disk.
|
||||
/// `crop_x`, `crop_y`, `crop_w`, `crop_h` are normalized 0.0–1.0 coordinates.
|
||||
pub fn embed_image_crop(
|
||||
&self,
|
||||
path: &Path,
|
||||
crop_x: f32,
|
||||
crop_y: f32,
|
||||
crop_w: f32,
|
||||
crop_h: f32,
|
||||
) -> Result<Vec<f32>> {
|
||||
let img = image::ImageReader::open(path)?
|
||||
.with_guessed_format()?
|
||||
.decode()?;
|
||||
|
||||
let img_w = img.width() as f32;
|
||||
let img_h = img.height() as f32;
|
||||
|
||||
let x = ((crop_x * img_w) as u32).min(img.width().saturating_sub(1));
|
||||
let y = ((crop_y * img_h) as u32).min(img.height().saturating_sub(1));
|
||||
let w = ((crop_w * img_w) as u32).max(1).min(img.width() - x);
|
||||
let h = ((crop_h * img_h) as u32).max(1).min(img.height() - y);
|
||||
|
||||
let cropped = img.crop_imm(x, y, w, h);
|
||||
let resized = cropped.resize_to_fill(
|
||||
self.image_size as u32,
|
||||
self.image_size as u32,
|
||||
image::imageops::FilterType::Triangle,
|
||||
);
|
||||
|
||||
let raw = resized.to_rgb8().into_raw();
|
||||
let tensor = candle_core::Tensor::from_vec(
|
||||
raw,
|
||||
(self.image_size, self.image_size, 3),
|
||||
&candle_core::Device::Cpu,
|
||||
)?
|
||||
.permute((2, 0, 1))?
|
||||
.to_dtype(candle_core::DType::F32)?
|
||||
.affine(2.0 / 255.0, -1.0)?;
|
||||
|
||||
let batch = tensor.unsqueeze(0)?.to_device(&self.device)?;
|
||||
let features = self.model.get_image_features(&batch)?;
|
||||
let normalized = candle_transformers::models::clip::div_l2_norm(&features)?;
|
||||
Ok(normalized.get(0)?.flatten_all()?.to_vec1::<f32>()?)
|
||||
}
|
||||
|
||||
pub fn embed_images(&self, paths: &[PathBuf]) -> Result<Vec<Vec<f32>>> {
|
||||
let images = load_images(paths, self.image_size)?.to_device(&self.device)?;
|
||||
let features = self.model.get_image_features(&images)?;
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
use crate::vector;
|
||||
use anyhow::Result;
|
||||
use hnsw_rs::prelude::{DistCosine, Hnsw, Neighbour};
|
||||
use rusqlite::Connection;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{OnceLock, RwLock};
|
||||
|
||||
const HNSW_MAX_CONNECTIONS: usize = 24;
|
||||
const HNSW_EF_CONSTRUCTION: usize = 300;
|
||||
const HNSW_EF_SEARCH: usize = 96;
|
||||
|
||||
struct CachedHnswIndex {
|
||||
revision: String,
|
||||
image_ids_by_external: Vec<i64>,
|
||||
external_by_image_id: HashMap<i64, usize>,
|
||||
hnsw: Hnsw<'static, f32, DistCosine>,
|
||||
}
|
||||
|
||||
static IMAGE_HNSW_INDEX: OnceLock<RwLock<Option<CachedHnswIndex>>> = OnceLock::new();
|
||||
|
||||
fn cache() -> &'static RwLock<Option<CachedHnswIndex>> {
|
||||
IMAGE_HNSW_INDEX.get_or_init(|| RwLock::new(None))
|
||||
}
|
||||
|
||||
fn build_index(conn: &Connection) -> Result<CachedHnswIndex> {
|
||||
// Read the revision *before* fetching embeddings so we can detect any write
|
||||
// that races with the build. If the revision advances while we are building,
|
||||
// the resulting index would be stale — retry until it is stable.
|
||||
loop {
|
||||
let revision_before = vector::get_embedding_revision(conn)?;
|
||||
|
||||
let embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?;
|
||||
let max_elements = embeddings.len().max(1);
|
||||
let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).max(1);
|
||||
let mut hnsw = Hnsw::<f32, DistCosine>::new(
|
||||
HNSW_MAX_CONNECTIONS,
|
||||
max_elements,
|
||||
max_layer,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
DistCosine {},
|
||||
);
|
||||
|
||||
let image_ids_by_external = embeddings
|
||||
.iter()
|
||||
.map(|(image_id, _)| *image_id)
|
||||
.collect::<Vec<_>>();
|
||||
let external_by_image_id = image_ids_by_external
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(external_id, image_id)| (*image_id, external_id))
|
||||
.collect::<HashMap<_, _>>();
|
||||
let data_with_id = embeddings
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(external_id, (_, embedding))| (embedding, external_id))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
hnsw.parallel_insert(&data_with_id);
|
||||
hnsw.set_searching_mode(true);
|
||||
|
||||
// If the revision is unchanged the index reflects a consistent snapshot.
|
||||
let revision_after = vector::get_embedding_revision(conn)?;
|
||||
if revision_before == revision_after {
|
||||
return Ok(CachedHnswIndex {
|
||||
revision: revision_after,
|
||||
image_ids_by_external,
|
||||
external_by_image_id,
|
||||
hnsw,
|
||||
});
|
||||
}
|
||||
// A concurrent write advanced the revision — discard this build and retry.
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_index(conn: &Connection) -> Result<()> {
|
||||
let revision = vector::get_embedding_revision(conn)?;
|
||||
|
||||
{
|
||||
let guard = cache().read().expect("hnsw cache poisoned");
|
||||
if guard.as_ref().map(|cached| cached.revision.as_str()) == Some(revision.as_str()) {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
let next = build_index(conn)?;
|
||||
let mut guard = cache().write().expect("hnsw cache poisoned");
|
||||
*guard = Some(next);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn find_similar_image_matches(
|
||||
conn: &Connection,
|
||||
image_id: i64,
|
||||
folder_id: Option<i64>,
|
||||
threshold: f32,
|
||||
offset: usize,
|
||||
limit: usize,
|
||||
) -> Result<Vec<(i64, f32)>> {
|
||||
ensure_index(conn)?;
|
||||
|
||||
let query_embedding = match vector::get_image_embedding(conn, image_id)? {
|
||||
Some(embedding) => embedding,
|
||||
None => return Ok(Vec::new()),
|
||||
};
|
||||
|
||||
// Fetch folder image IDs *before* acquiring the read lock so we don't hold
|
||||
// the lock across a potentially slow SQLite query, which would delay any
|
||||
// concurrent ensure_index call waiting for a write lock.
|
||||
let folder_image_ids: Option<Vec<i64>> = if let Some(folder_id) = folder_id {
|
||||
let ids = vector::get_all_image_embeddings_with_ids(conn, Some(folder_id))?
|
||||
.into_iter()
|
||||
.map(|(id, _)| id)
|
||||
.collect();
|
||||
Some(ids)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let guard = cache().read().expect("hnsw cache poisoned");
|
||||
let Some(cached) = guard.as_ref() else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let knbn = (offset + limit).max(limit).saturating_add(32);
|
||||
let neighbours: Vec<Neighbour> = if let Some(image_ids) = folder_image_ids {
|
||||
let mut allowed_ids = image_ids
|
||||
.into_iter()
|
||||
.filter_map(|allowed_image_id| {
|
||||
cached.external_by_image_id.get(&allowed_image_id).copied()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
allowed_ids.sort_unstable();
|
||||
cached
|
||||
.hnsw
|
||||
.search_filter(&query_embedding, knbn, HNSW_EF_SEARCH, Some(&allowed_ids))
|
||||
} else {
|
||||
cached.hnsw.search(&query_embedding, knbn, HNSW_EF_SEARCH)
|
||||
};
|
||||
|
||||
let matches = neighbours
|
||||
.into_iter()
|
||||
.filter_map(|neighbour| {
|
||||
let image_id_match = cached.image_ids_by_external.get(neighbour.d_id).copied()?;
|
||||
if image_id_match == image_id || neighbour.distance > threshold {
|
||||
return None;
|
||||
}
|
||||
Some((image_id_match, neighbour.distance))
|
||||
})
|
||||
.skip(offset)
|
||||
.take(limit)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(matches)
|
||||
}
|
||||
+477
-51
@@ -1,7 +1,9 @@
|
||||
use crate::captioner::{self, FlorenceCaptioner};
|
||||
use crate::db::{self, DbPool, EmbeddingJob, FolderJobProgress, ImageRecord, IndexedMediaEntry};
|
||||
use crate::embedder::{embedding_source_path, ClipImageEmbedder};
|
||||
use crate::media::{probe_video_metadata, MediaTools};
|
||||
use crate::storage::{detect_storage_profile, RuntimeAdaptiveProfile, StorageProfile};
|
||||
use crate::tagger::{self, WdTagger};
|
||||
use crate::thumbnail;
|
||||
use crate::vector;
|
||||
use anyhow::Result;
|
||||
@@ -9,7 +11,6 @@ use rayon::prelude::*;
|
||||
use serde::Serialize;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
use tauri::{AppHandle, Emitter};
|
||||
@@ -22,29 +23,97 @@ const IMAGE_EXTENSIONS: &[&str] = &[
|
||||
const VIDEO_EXTENSIONS: &[&str] = &["mp4", "mov", "m4v", "webm"];
|
||||
|
||||
const JOB_PROGRESS_EMIT_INTERVAL: Duration = Duration::from_millis(750);
|
||||
const CAPTION_BATCH_SIZE: usize = 1;
|
||||
|
||||
static LAST_JOB_PROGRESS_EMIT: OnceLock<Mutex<HashMap<i64, Instant>>> = OnceLock::new();
|
||||
static ACTIVE_INDEXING_FOLDERS: OnceLock<Mutex<HashSet<i64>>> = OnceLock::new();
|
||||
|
||||
static THUMBNAIL_WORKER_PAUSED: AtomicBool = AtomicBool::new(false);
|
||||
static METADATA_WORKER_PAUSED: AtomicBool = AtomicBool::new(false);
|
||||
static EMBEDDING_WORKER_PAUSED: AtomicBool = AtomicBool::new(false);
|
||||
static PAUSED_WORKER_FOLDERS: OnceLock<Mutex<PausedWorkerFolders>> = OnceLock::new();
|
||||
|
||||
pub fn set_worker_paused(worker: &str, paused: bool) {
|
||||
match worker {
|
||||
"thumbnail" => THUMBNAIL_WORKER_PAUSED.store(paused, Ordering::Relaxed),
|
||||
"metadata" => METADATA_WORKER_PAUSED.store(paused, Ordering::Relaxed),
|
||||
"embedding" => EMBEDDING_WORKER_PAUSED.store(paused, Ordering::Relaxed),
|
||||
_ => {}
|
||||
#[derive(Default)]
|
||||
struct PausedWorkerFolders {
|
||||
thumbnail: HashSet<i64>,
|
||||
metadata: HashSet<i64>,
|
||||
embedding: HashSet<i64>,
|
||||
caption: HashSet<i64>,
|
||||
tagging: HashSet<i64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct FolderWorkerPausedState {
|
||||
pub thumbnail: bool,
|
||||
pub metadata: bool,
|
||||
pub embedding: bool,
|
||||
pub caption: bool,
|
||||
pub tagging: 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),
|
||||
"tagging" => Some(&mut paused_folders.tagging),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if let Some(folder_set) = folder_set {
|
||||
if paused {
|
||||
folder_set.insert(folder_id);
|
||||
} else {
|
||||
folder_set.remove(&folder_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_worker_paused_states() -> [bool; 3] {
|
||||
[
|
||||
THUMBNAIL_WORKER_PAUSED.load(Ordering::Relaxed),
|
||||
METADATA_WORKER_PAUSED.load(Ordering::Relaxed),
|
||||
EMBEDDING_WORKER_PAUSED.load(Ordering::Relaxed),
|
||||
]
|
||||
pub fn get_worker_paused_states(folder_ids: &[i64]) -> HashMap<i64, FolderWorkerPausedState> {
|
||||
let Ok(paused_folders) = PAUSED_WORKER_FOLDERS
|
||||
.get_or_init(|| Mutex::new(PausedWorkerFolders::default()))
|
||||
.lock()
|
||||
else {
|
||||
return HashMap::new();
|
||||
};
|
||||
|
||||
folder_ids
|
||||
.iter()
|
||||
.copied()
|
||||
.map(|folder_id| {
|
||||
(
|
||||
folder_id,
|
||||
FolderWorkerPausedState {
|
||||
thumbnail: paused_folders.thumbnail.contains(&folder_id),
|
||||
metadata: paused_folders.metadata.contains(&folder_id),
|
||||
embedding: paused_folders.embedding.contains(&folder_id),
|
||||
caption: paused_folders.caption.contains(&folder_id),
|
||||
tagging: paused_folders.tagging.contains(&folder_id),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn paused_folder_ids(worker: &str) -> HashSet<i64> {
|
||||
let Ok(paused_folders) = PAUSED_WORKER_FOLDERS
|
||||
.get_or_init(|| Mutex::new(PausedWorkerFolders::default()))
|
||||
.lock()
|
||||
else {
|
||||
return HashSet::new();
|
||||
};
|
||||
|
||||
match worker {
|
||||
"thumbnail" => paused_folders.thumbnail.clone(),
|
||||
"metadata" => paused_folders.metadata.clone(),
|
||||
"embedding" => paused_folders.embedding.clone(),
|
||||
"caption" => paused_folders.caption.clone(),
|
||||
"tagging" => paused_folders.tagging.clone(),
|
||||
_ => HashSet::new(),
|
||||
}
|
||||
}
|
||||
static FOLDER_STORAGE_PROFILES: OnceLock<Mutex<HashMap<i64, RuntimeAdaptiveProfile>>> =
|
||||
OnceLock::new();
|
||||
@@ -78,11 +147,50 @@ pub struct MediaJobProgressEvent {
|
||||
|
||||
pub fn index_folder(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) {
|
||||
std::thread::spawn(move || {
|
||||
set_folder_indexing_state(folder_id, true);
|
||||
|
||||
// If the folder path no longer exists on disk, record the error and
|
||||
// emit done. Images are intentionally kept in the DB so the user can
|
||||
// choose to relocate the folder or remove it explicitly — they should
|
||||
// not be silently destroyed.
|
||||
if !folder_path.is_dir() {
|
||||
let error_msg = format!("Folder not found: {}", folder_path.display());
|
||||
eprintln!("Indexing error for folder {}: {}", folder_id, error_msg);
|
||||
if let Ok(conn) = pool.get() {
|
||||
let _ = db::set_folder_scan_error(&conn, folder_id, &error_msg);
|
||||
}
|
||||
emit_progress(
|
||||
&app,
|
||||
&IndexProgress {
|
||||
folder_id,
|
||||
total: 0,
|
||||
indexed: 0,
|
||||
current_file: String::new(),
|
||||
done: true,
|
||||
},
|
||||
);
|
||||
set_folder_indexing_state(folder_id, false);
|
||||
return;
|
||||
}
|
||||
|
||||
let storage_profile = detect_storage_profile(&folder_path);
|
||||
set_folder_storage_profile(folder_id, RuntimeAdaptiveProfile::new(storage_profile));
|
||||
set_folder_indexing_state(folder_id, true);
|
||||
if let Err(error) = do_index(app, pool, folder_id, folder_path) {
|
||||
eprintln!("Indexing error: {}", error);
|
||||
if let Err(error) = do_index(app.clone(), &pool, folder_id, folder_path) {
|
||||
eprintln!("Indexing error for folder {}: {}", folder_id, error);
|
||||
if let Ok(conn) = pool.get() {
|
||||
let _ = db::set_folder_scan_error(&conn, folder_id, &error.to_string());
|
||||
}
|
||||
// Always emit done so the frontend reloads and recovers from partial state.
|
||||
emit_progress(
|
||||
&app,
|
||||
&IndexProgress {
|
||||
folder_id,
|
||||
total: 0,
|
||||
indexed: 0,
|
||||
current_file: String::new(),
|
||||
done: true,
|
||||
},
|
||||
);
|
||||
}
|
||||
set_folder_indexing_state(folder_id, false);
|
||||
});
|
||||
@@ -95,10 +203,6 @@ pub fn start_thumbnail_worker(
|
||||
cache_dir: PathBuf,
|
||||
) {
|
||||
std::thread::spawn(move || loop {
|
||||
if THUMBNAIL_WORKER_PAUSED.load(Ordering::Relaxed) {
|
||||
std::thread::sleep(std::time::Duration::from_millis(500));
|
||||
continue;
|
||||
}
|
||||
if let Err(error) = process_thumbnail_batch(&app, &pool, &media_tools, &cache_dir) {
|
||||
eprintln!("Thumbnail worker error: {}", error);
|
||||
}
|
||||
@@ -108,10 +212,6 @@ pub fn start_thumbnail_worker(
|
||||
|
||||
pub fn start_metadata_worker(app: AppHandle, pool: DbPool, media_tools: MediaTools) {
|
||||
std::thread::spawn(move || loop {
|
||||
if METADATA_WORKER_PAUSED.load(Ordering::Relaxed) {
|
||||
std::thread::sleep(std::time::Duration::from_millis(500));
|
||||
continue;
|
||||
}
|
||||
if let Err(error) = process_metadata_batch(&app, &pool, &media_tools) {
|
||||
eprintln!("Metadata worker error: {}", error);
|
||||
}
|
||||
@@ -124,10 +224,6 @@ pub fn start_embedding_worker(app: AppHandle, pool: DbPool) {
|
||||
let mut embedder: Option<ClipImageEmbedder> = None;
|
||||
println!("Embedding worker started.");
|
||||
loop {
|
||||
if EMBEDDING_WORKER_PAUSED.load(Ordering::Relaxed) {
|
||||
std::thread::sleep(std::time::Duration::from_millis(500));
|
||||
continue;
|
||||
}
|
||||
if let Err(error) = process_embedding_batch(&app, &pool, &mut embedder) {
|
||||
eprintln!("Embedding worker error: {}", error);
|
||||
}
|
||||
@@ -136,7 +232,49 @@ pub fn start_embedding_worker(app: AppHandle, pool: DbPool) {
|
||||
});
|
||||
}
|
||||
|
||||
fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> {
|
||||
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 the acceleration setting changed, drop the cached session so
|
||||
// the next batch picks it up with the new execution provider.
|
||||
if captioner::CAPTION_SESSION_DIRTY.swap(false, std::sync::atomic::Ordering::Relaxed) {
|
||||
println!("Caption worker: acceleration setting changed — resetting session.");
|
||||
captioner = None;
|
||||
}
|
||||
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));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub fn start_tagging_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) {
|
||||
std::thread::spawn(move || {
|
||||
let mut tagger_instance: Option<WdTagger> = None;
|
||||
println!("Tagging worker started.");
|
||||
loop {
|
||||
// If the acceleration setting changed, drop the cached session so
|
||||
// the next batch picks it up with the new execution provider.
|
||||
if tagger::TAGGER_SESSION_DIRTY.swap(false, std::sync::atomic::Ordering::Relaxed) {
|
||||
println!("Tagging worker: acceleration setting changed — resetting session.");
|
||||
tagger_instance = None;
|
||||
}
|
||||
if let Err(error) =
|
||||
process_tagging_batch(&app, &pool, &app_data_dir, &mut tagger_instance)
|
||||
{
|
||||
eprintln!("Tagging worker error: {}", error);
|
||||
tagger_instance = None;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(750));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn do_index(app: AppHandle, pool: &DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> {
|
||||
let existing_entries = {
|
||||
let conn = pool.get()?;
|
||||
db::get_folder_media_index(&conn, folder_id)?
|
||||
@@ -146,12 +284,32 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf)
|
||||
.map(|entry| (entry.path.clone(), entry))
|
||||
.collect::<HashMap<_, _>>();
|
||||
|
||||
// Collect traversal errors separately. An unreadable subdirectory means we
|
||||
// cannot distinguish absent files from inaccessible ones; tracking errors
|
||||
// lets us skip the deletion step for paths under affected subtrees.
|
||||
let mut unreadable_prefixes: Vec<String> = Vec::new();
|
||||
let media_paths: Vec<PathBuf> = WalkDir::new(&folder_path)
|
||||
.follow_links(true)
|
||||
.into_iter()
|
||||
.filter_map(|entry| entry.ok())
|
||||
.filter(|entry| entry.file_type().is_file() && is_supported_media(entry.path()))
|
||||
.map(|entry| entry.path().to_path_buf())
|
||||
.filter_map(|entry| match entry {
|
||||
Ok(e) if e.file_type().is_file() && is_supported_media(e.path()) => {
|
||||
Some(e.path().to_path_buf())
|
||||
}
|
||||
Ok(_) => None,
|
||||
Err(err) => {
|
||||
// Record the inaccessible path so we can protect its descendants
|
||||
// from the missing-file deletion pass.
|
||||
if let Some(path) = err.path() {
|
||||
unreadable_prefixes.push(path.to_string_lossy().into_owned());
|
||||
}
|
||||
eprintln!(
|
||||
"WalkDir error while scanning folder {}: {}",
|
||||
folder_path.display(),
|
||||
err
|
||||
);
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let total = media_paths.len();
|
||||
@@ -197,7 +355,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf)
|
||||
images: committed,
|
||||
},
|
||||
);
|
||||
emit_folder_job_progress(&app, &pool, &[folder_id]);
|
||||
emit_folder_job_progress(&app, &pool, &[folder_id], false);
|
||||
}
|
||||
|
||||
processed += path_chunk.len();
|
||||
@@ -224,6 +382,17 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf)
|
||||
let missing_ids = existing_by_path
|
||||
.values()
|
||||
.filter(|entry| !seen_paths.contains(&entry.path))
|
||||
.filter(|entry| {
|
||||
// If this path lives under a subtree that WalkDir couldn't read,
|
||||
// we don't know whether the file is gone or just temporarily
|
||||
// inaccessible — keep the record to avoid false deletions.
|
||||
// Use Path::starts_with (component-aware) so a sibling directory
|
||||
// whose name shares a prefix is not incorrectly protected.
|
||||
let entry_path = std::path::Path::new(&entry.path);
|
||||
unreadable_prefixes
|
||||
.iter()
|
||||
.all(|prefix| !entry_path.starts_with(prefix))
|
||||
})
|
||||
.map(|entry| entry.id)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
@@ -232,7 +401,14 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf)
|
||||
if !missing_ids.is_empty() {
|
||||
db::delete_images_by_ids(&conn, &missing_ids)?;
|
||||
}
|
||||
let _ = db::backfill_embedding_jobs(&conn)?;
|
||||
db::update_folder_count(&conn, folder_id)?;
|
||||
let _ = db::clear_folder_scan_error(&conn, folder_id);
|
||||
// Invalidate duplicate scan cache — any reindex can change file contents
|
||||
// or the set of files, making cached duplicate groups stale.
|
||||
let folder_scope = format!("folder:{}", folder_id);
|
||||
let _ = db::clear_duplicate_scan_cache(&conn, &folder_scope);
|
||||
let _ = db::clear_duplicate_scan_cache(&conn, "all");
|
||||
}
|
||||
|
||||
emit_progress(
|
||||
@@ -245,7 +421,7 @@ fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf)
|
||||
done: true,
|
||||
},
|
||||
);
|
||||
emit_folder_job_progress(&app, &pool, &[folder_id]);
|
||||
emit_folder_job_progress(&app, &pool, &[folder_id], true);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -302,6 +478,14 @@ fn build_record(
|
||||
embedding_model: Some(vector::CLIP_MODEL_NAME.to_string()),
|
||||
embedding_updated_at: None,
|
||||
embedding_error: None,
|
||||
generated_caption: None,
|
||||
caption_model: None,
|
||||
caption_updated_at: None,
|
||||
caption_error: None,
|
||||
ai_rating: None,
|
||||
ai_tagger_model: None,
|
||||
ai_tagged_at: None,
|
||||
ai_tagger_error: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -335,11 +519,13 @@ fn process_thumbnail_batch(
|
||||
with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
let active_folders = active_indexing_folders();
|
||||
let paused_folders = paused_folder_ids("thumbnail");
|
||||
let worker_batch_size = max_worker_batch_size(&active_folders);
|
||||
let worker_fetch_size = max_worker_fetch_size(&active_folders);
|
||||
db::claim_thumbnail_jobs(
|
||||
&mut conn,
|
||||
&active_folders,
|
||||
&paused_folders,
|
||||
worker_fetch_size,
|
||||
worker_batch_size,
|
||||
)
|
||||
@@ -428,7 +614,7 @@ fn process_thumbnail_batch(
|
||||
images: updated_images,
|
||||
},
|
||||
);
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>());
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -439,11 +625,13 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo
|
||||
with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
let active_folders = active_indexing_folders();
|
||||
let paused_folders = paused_folder_ids("metadata");
|
||||
let worker_batch_size = max_worker_batch_size(&active_folders);
|
||||
let worker_fetch_size = max_worker_fetch_size(&active_folders);
|
||||
db::claim_metadata_jobs(
|
||||
&mut conn,
|
||||
&active_folders,
|
||||
&paused_folders,
|
||||
worker_fetch_size,
|
||||
worker_batch_size,
|
||||
)
|
||||
@@ -506,7 +694,7 @@ fn process_metadata_batch(app: &AppHandle, pool: &DbPool, media_tools: &MediaToo
|
||||
images: updated_images,
|
||||
},
|
||||
);
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>());
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -518,14 +706,11 @@ fn process_embedding_batch(
|
||||
embedder: &mut Option<ClipImageEmbedder>,
|
||||
) -> Result<()> {
|
||||
let batch_started_at = Instant::now();
|
||||
if embedder.is_none() {
|
||||
*embedder = Some(ClipImageEmbedder::new()?);
|
||||
}
|
||||
|
||||
let claim_started_at = Instant::now();
|
||||
let paused_folders = paused_folder_ids("embedding");
|
||||
let jobs = with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
db::claim_embedding_jobs(&mut conn, EMBEDDING_BATCH_SIZE)
|
||||
db::claim_embedding_jobs(&mut conn, &paused_folders, EMBEDDING_BATCH_SIZE)
|
||||
})?;
|
||||
let claim_elapsed = claim_started_at.elapsed();
|
||||
|
||||
@@ -533,9 +718,18 @@ fn process_embedding_batch(
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if embedder.is_none() {
|
||||
*embedder = Some(ClipImageEmbedder::new()?);
|
||||
}
|
||||
|
||||
println!("Embedding batch claimed: {} items", jobs.len());
|
||||
let folder_ids = jobs.iter().map(|job| job.folder_id).collect::<HashSet<_>>();
|
||||
emit_folder_job_progress(app, pool, &folder_ids.iter().copied().collect::<Vec<_>>());
|
||||
emit_folder_job_progress(
|
||||
app,
|
||||
pool,
|
||||
&folder_ids.iter().copied().collect::<Vec<_>>(),
|
||||
false,
|
||||
);
|
||||
let embedder = embedder.as_ref().expect("embedder should be initialized");
|
||||
|
||||
let infer_started_at = Instant::now();
|
||||
@@ -639,7 +833,7 @@ fn process_embedding_batch(
|
||||
images: updated_images,
|
||||
},
|
||||
);
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>());
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
|
||||
}
|
||||
|
||||
let write_elapsed = write_started_at.elapsed();
|
||||
@@ -652,6 +846,237 @@ fn process_embedding_batch(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn process_caption_batch(
|
||||
app: &AppHandle,
|
||||
pool: &DbPool,
|
||||
app_data_dir: &Path,
|
||||
captioner: &mut Option<FlorenceCaptioner>,
|
||||
) -> Result<()> {
|
||||
if !captioner::caption_model_status(app_data_dir).ready {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let paused_folders = paused_folder_ids("caption");
|
||||
let jobs = with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
db::claim_caption_jobs(&mut conn, &paused_folders, CAPTION_BATCH_SIZE)
|
||||
})?;
|
||||
|
||||
if jobs.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if captioner.is_none() {
|
||||
match FlorenceCaptioner::new(app_data_dir) {
|
||||
Ok(model) => *captioner = Some(model),
|
||||
Err(error) => {
|
||||
with_db_write_lock(|| {
|
||||
let conn = pool.get()?;
|
||||
db::requeue_caption_jobs(
|
||||
&conn,
|
||||
&jobs.iter().map(|job| job.image_id).collect::<Vec<_>>(),
|
||||
)
|
||||
})?;
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let folder_ids = jobs.iter().map(|job| job.folder_id).collect::<HashSet<_>>();
|
||||
emit_folder_job_progress(
|
||||
app,
|
||||
pool,
|
||||
&folder_ids.iter().copied().collect::<Vec<_>>(),
|
||||
false,
|
||||
);
|
||||
|
||||
let captioner = captioner
|
||||
.as_mut()
|
||||
.expect("captioner should be initialized before caption batch processing");
|
||||
let caption_results = jobs
|
||||
.iter()
|
||||
.map(|job| (job.clone(), captioner.generate(Path::new(&job.path))))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let updated_images = with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
let tx = conn.transaction()?;
|
||||
let mut updated_images = Vec::with_capacity(caption_results.len());
|
||||
|
||||
for (job, caption_result) in &caption_results {
|
||||
if db::is_caption_job_cancelled(&tx, job.image_id)? {
|
||||
tx.execute(
|
||||
"DELETE FROM caption_jobs WHERE image_id = ?1",
|
||||
[job.image_id],
|
||||
)?;
|
||||
continue;
|
||||
}
|
||||
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 process_tagging_batch(
|
||||
app: &AppHandle,
|
||||
pool: &DbPool,
|
||||
app_data_dir: &Path,
|
||||
tagger_instance: &mut Option<WdTagger>,
|
||||
) -> Result<()> {
|
||||
if !tagger::tagger_model_status(app_data_dir).ready {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let paused_folders = paused_folder_ids("tagging");
|
||||
let batch_size = crate::tagger::tagger_batch_size(app_data_dir);
|
||||
let jobs = with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
db::claim_tagging_jobs(&mut conn, &paused_folders, batch_size)
|
||||
})?;
|
||||
|
||||
if jobs.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if tagger_instance.is_none() {
|
||||
match WdTagger::new(app_data_dir) {
|
||||
Ok(model) => *tagger_instance = Some(model),
|
||||
Err(error) => {
|
||||
with_db_write_lock(|| {
|
||||
let conn = pool.get()?;
|
||||
db::requeue_tagging_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 tagger_ref = tagger_instance
|
||||
.as_mut()
|
||||
.expect("tagger should be initialized before tagging batch processing");
|
||||
|
||||
let tag_results = jobs
|
||||
.iter()
|
||||
.map(|job| {
|
||||
(
|
||||
job.clone(),
|
||||
tagger_ref.run(Path::new(&job.path), tagger::DEFAULT_MAX_TAGS),
|
||||
)
|
||||
})
|
||||
.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(tag_results.len());
|
||||
|
||||
for (job, tag_result) in &tag_results {
|
||||
// If the job is no longer in 'processing' state it was either
|
||||
// cancelled or reset to 'pending' (pause) while inference was running.
|
||||
// Discard the result in both cases — for cancelled rows also delete
|
||||
// the row; for paused rows leave it so the job is retried later.
|
||||
if !db::is_tagging_job_processing(&tx, job.image_id)? {
|
||||
tx.execute(
|
||||
"DELETE FROM tagging_jobs WHERE image_id = ?1 AND status = 'cancelled'",
|
||||
[job.image_id],
|
||||
)?;
|
||||
continue;
|
||||
}
|
||||
match tag_result {
|
||||
Ok(output) => {
|
||||
let tag_pairs: Vec<(String, f64)> = output
|
||||
.tags
|
||||
.iter()
|
||||
.map(|t| (t.tag.clone(), t.confidence as f64))
|
||||
.collect();
|
||||
db::update_ai_tags(
|
||||
&tx,
|
||||
job.image_id,
|
||||
&tag_pairs,
|
||||
&output.rating,
|
||||
tagger::WD_TAGGER_MODEL_NAME,
|
||||
)?;
|
||||
}
|
||||
Err(error) => {
|
||||
db::mark_tagging_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)
|
||||
})
|
||||
.or_else(|db_err| {
|
||||
// The DB write failed. Try to requeue the claimed jobs so they aren't
|
||||
// left stuck in 'processing' until the next app restart.
|
||||
let image_ids: Vec<i64> = jobs.iter().map(|job| job.image_id).collect();
|
||||
let _ = with_db_write_lock(|| {
|
||||
let conn = pool.get()?;
|
||||
db::requeue_tagging_jobs(&conn, &image_ids)
|
||||
});
|
||||
Err(db_err)
|
||||
})?;
|
||||
|
||||
if !updated_images.is_empty() {
|
||||
let folder_ids = updated_images
|
||||
.iter()
|
||||
.map(|image| image.folder_id)
|
||||
.collect::<HashSet<_>>();
|
||||
emit_media_updates(
|
||||
app,
|
||||
&MediaUpdateBatch {
|
||||
images: updated_images,
|
||||
},
|
||||
);
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn active_indexing_folders() -> HashSet<i64> {
|
||||
ACTIVE_INDEXING_FOLDERS
|
||||
.get_or_init(|| Mutex::new(HashSet::new()))
|
||||
@@ -737,7 +1162,7 @@ fn emit_media_updates(app: &AppHandle, batch: &MediaUpdateBatch) {
|
||||
let _ = app.emit("media-updated", batch);
|
||||
}
|
||||
|
||||
fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64]) {
|
||||
pub fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64], force: bool) {
|
||||
let mut unique_folder_ids = folder_ids.iter().copied().collect::<Vec<_>>();
|
||||
unique_folder_ids.sort_unstable();
|
||||
unique_folder_ids.dedup();
|
||||
@@ -749,10 +1174,11 @@ fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64])
|
||||
Err(_) => return,
|
||||
};
|
||||
unique_folder_ids.retain(|folder_id| {
|
||||
let should_emit = tracker
|
||||
.get(folder_id)
|
||||
.map(|last_emit| now.duration_since(*last_emit) >= JOB_PROGRESS_EMIT_INTERVAL)
|
||||
.unwrap_or(true);
|
||||
let should_emit = force
|
||||
|| tracker
|
||||
.get(folder_id)
|
||||
.map(|last_emit| now.duration_since(*last_emit) >= JOB_PROGRESS_EMIT_INTERVAL)
|
||||
.unwrap_or(true);
|
||||
if should_emit {
|
||||
tracker.insert(*folder_id, now);
|
||||
}
|
||||
|
||||
+60
-3
@@ -1,14 +1,17 @@
|
||||
mod captioner;
|
||||
mod commands;
|
||||
mod db;
|
||||
mod embedder;
|
||||
mod hnsw_index;
|
||||
mod indexer;
|
||||
mod media;
|
||||
mod storage;
|
||||
mod tagger;
|
||||
mod thumbnail;
|
||||
mod vector;
|
||||
|
||||
use tauri::Manager;
|
||||
use crate::storage::StorageProfile;
|
||||
use tauri::Manager;
|
||||
|
||||
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
||||
pub fn run() {
|
||||
@@ -16,6 +19,7 @@ pub fn run() {
|
||||
.plugin(tauri_plugin_opener::init())
|
||||
.plugin(tauri_plugin_dialog::init())
|
||||
.plugin(tauri_plugin_fs::init())
|
||||
.plugin(tauri_plugin_notification::init())
|
||||
.setup(|app| {
|
||||
let app_dir = app
|
||||
.path()
|
||||
@@ -34,11 +38,19 @@ pub fn run() {
|
||||
let conn = pool.get().expect("Failed to get connection for migration");
|
||||
db::migrate(&conn).expect("Failed to run migrations");
|
||||
db::reset_inflight_jobs(&conn).expect("Failed to reset inflight jobs");
|
||||
let backfilled = db::backfill_embedding_jobs(&conn)
|
||||
.expect("Failed to backfill embedding jobs");
|
||||
let backfilled =
|
||||
db::backfill_embedding_jobs(&conn).expect("Failed to backfill embedding jobs");
|
||||
if backfilled > 0 {
|
||||
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");
|
||||
@@ -58,6 +70,9 @@ pub fn run() {
|
||||
}
|
||||
indexer::start_metadata_worker(app.handle().clone(), pool.clone(), media_tools.clone());
|
||||
indexer::start_embedding_worker(app.handle().clone(), pool.clone());
|
||||
// Caption worker disabled — UI removed; keeping backend code intact for future use.
|
||||
// indexer::start_caption_worker(app.handle().clone(), pool.clone(), app_dir.clone());
|
||||
indexer::start_tagging_worker(app.handle().clone(), pool.clone(), app_dir.clone());
|
||||
|
||||
app.manage(pool);
|
||||
app.manage(media_tools);
|
||||
@@ -73,12 +88,54 @@ pub fn run() {
|
||||
commands::reindex_folder,
|
||||
commands::update_image_details,
|
||||
commands::find_similar_images,
|
||||
commands::find_similar_by_region,
|
||||
commands::debug_similar_images,
|
||||
commands::retry_failed_embeddings,
|
||||
commands::semantic_search_images,
|
||||
commands::search_images_by_tag,
|
||||
commands::get_caption_model_status,
|
||||
commands::get_caption_acceleration,
|
||||
commands::set_caption_acceleration,
|
||||
commands::get_caption_detail,
|
||||
commands::set_caption_detail,
|
||||
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::clear_caption_jobs,
|
||||
commands::reset_generated_captions,
|
||||
commands::set_generated_caption,
|
||||
commands::suggest_image_tags,
|
||||
commands::set_worker_paused,
|
||||
commands::get_worker_states,
|
||||
commands::get_tag_cloud,
|
||||
commands::get_explore_tags,
|
||||
commands::get_images_by_ids,
|
||||
commands::get_failed_embedding_images,
|
||||
commands::get_tagger_model_status,
|
||||
commands::get_tagger_acceleration,
|
||||
commands::set_tagger_acceleration,
|
||||
commands::probe_tagger_runtime,
|
||||
commands::get_tagger_threshold,
|
||||
commands::set_tagger_threshold,
|
||||
commands::get_tagger_batch_size,
|
||||
commands::set_tagger_batch_size,
|
||||
commands::prepare_tagger_model,
|
||||
commands::delete_tagger_model,
|
||||
commands::queue_tagging_jobs,
|
||||
commands::clear_tagging_jobs,
|
||||
commands::get_image_tags,
|
||||
commands::add_user_tag,
|
||||
commands::remove_tag,
|
||||
commands::search_tags_autocomplete,
|
||||
commands::find_duplicates,
|
||||
commands::load_duplicate_scan_cache,
|
||||
commands::invalidate_duplicate_scan_cache,
|
||||
commands::delete_images_from_disk,
|
||||
commands::rename_folder,
|
||||
commands::update_folder_path,
|
||||
])
|
||||
.run(tauri::generate_context!())
|
||||
.expect("error while running tauri application");
|
||||
|
||||
@@ -0,0 +1,656 @@
|
||||
use anyhow::Result;
|
||||
use hf_hub::{api::sync::Api, Repo, RepoType};
|
||||
use image::{imageops::FilterType, DynamicImage, ImageReader};
|
||||
use ort::ep;
|
||||
use ort::session::{builder::GraphOptimizationLevel, Session};
|
||||
use ort::value::Tensor;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
pub const WD_TAGGER_MODEL_ID: &str = "SmilingWolf/wd-swinv2-tagger-v3";
|
||||
pub const WD_TAGGER_MODEL_NAME: &str = "wd-swinv2-tagger-v3";
|
||||
|
||||
const TAGGER_ACCELERATION_FILE: &str = "settings/tagger_acceleration.txt";
|
||||
const TAGGER_THRESHOLD_FILE: &str = "settings/tagger_threshold.txt";
|
||||
const TAGGER_BATCH_SIZE_FILE: &str = "settings/tagger_batch_size.txt";
|
||||
|
||||
// Files required on disk before the tagger can run. The ONNX runtime DLLs
|
||||
// are shared with the captioner and live in the same `onnxruntime/` directory.
|
||||
const TAGGER_REQUIRED_FILES: &[&str] = &[
|
||||
"onnxruntime/onnxruntime.dll",
|
||||
"onnxruntime/onnxruntime_providers_shared.dll",
|
||||
"onnxruntime/DirectML.dll",
|
||||
"model.onnx",
|
||||
"selected_tags.csv",
|
||||
];
|
||||
|
||||
// Tags in these Danbooru categories are kept in the output.
|
||||
// Category 0 = general, category 4 = character.
|
||||
// Category 9 = rating (explicit/questionable/sensitive/general) – used for
|
||||
// `ai_rating` but NOT emitted as individual tags.
|
||||
const GENERAL_CATEGORY: u8 = 0;
|
||||
const CHARACTER_CATEGORY: u8 = 4;
|
||||
const RATING_CATEGORY: u8 = 9;
|
||||
|
||||
pub const DEFAULT_THRESHOLD: f32 = 0.35;
|
||||
pub const DEFAULT_MAX_TAGS: usize = 30;
|
||||
|
||||
/// Set to `true` by `set_tagger_acceleration` so the tagging worker loop
|
||||
/// knows to drop its cached `WdTagger` and reload with the new EP.
|
||||
pub static TAGGER_SESSION_DIRTY: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Settings types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum TaggerAcceleration {
|
||||
Auto,
|
||||
Cpu,
|
||||
Directml,
|
||||
}
|
||||
|
||||
impl TaggerAcceleration {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Auto => "auto",
|
||||
Self::Cpu => "cpu",
|
||||
Self::Directml => "directml",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TaggerAcceleration {
|
||||
fn default() -> Self {
|
||||
Self::Auto
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status / probe types exposed to the frontend
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TaggerModelStatus {
|
||||
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 TaggerModelProgress {
|
||||
pub total_files: usize,
|
||||
pub completed_files: usize,
|
||||
pub current_file: Option<String>,
|
||||
pub done: bool,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Runtime probe types exposed to the frontend
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TaggerRuntimeProbe {
|
||||
pub ready: bool,
|
||||
pub acceleration: TaggerAcceleration,
|
||||
pub session: TaggerRuntimeSessionProbe,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TaggerRuntimeSessionProbe {
|
||||
pub file: &'static str,
|
||||
pub inputs: Vec<String>,
|
||||
pub outputs: Vec<String>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tag record returned to callers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TagResult {
|
||||
pub tag: String,
|
||||
pub confidence: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TaggerOutput {
|
||||
pub tags: Vec<TagResult>,
|
||||
/// Highest-scoring rating label: "general" | "sensitive" | "questionable" | "explicit"
|
||||
pub rating: String,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Internal label table built from selected_tags.csv
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct TagEntry {
|
||||
name: String,
|
||||
category: u8,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Path helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub fn model_dir(app_data_dir: &Path) -> PathBuf {
|
||||
app_data_dir.join("models").join("wd-swinv2-tagger-v3")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Settings persistence
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub fn tagger_acceleration(app_data_dir: &Path) -> TaggerAcceleration {
|
||||
let path = app_data_dir.join(TAGGER_ACCELERATION_FILE);
|
||||
let Ok(value) = std::fs::read_to_string(path) else {
|
||||
return TaggerAcceleration::default();
|
||||
};
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"cpu" => TaggerAcceleration::Cpu,
|
||||
"directml" => TaggerAcceleration::Directml,
|
||||
_ => TaggerAcceleration::Auto,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_tagger_acceleration(
|
||||
app_data_dir: &Path,
|
||||
acceleration: TaggerAcceleration,
|
||||
) -> Result<TaggerAcceleration> {
|
||||
let path = app_data_dir.join(TAGGER_ACCELERATION_FILE);
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::write(path, acceleration.as_str())?;
|
||||
TAGGER_SESSION_DIRTY.store(true, Ordering::Relaxed);
|
||||
Ok(acceleration)
|
||||
}
|
||||
|
||||
pub fn tagger_threshold(app_data_dir: &Path) -> f32 {
|
||||
let path = app_data_dir.join(TAGGER_THRESHOLD_FILE);
|
||||
let Ok(value) = std::fs::read_to_string(path) else {
|
||||
return DEFAULT_THRESHOLD;
|
||||
};
|
||||
value
|
||||
.trim()
|
||||
.parse::<f32>()
|
||||
.unwrap_or(DEFAULT_THRESHOLD)
|
||||
.clamp(0.01, 1.0)
|
||||
}
|
||||
|
||||
pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result<f32> {
|
||||
let clamped = threshold.clamp(0.01, 1.0);
|
||||
let path = app_data_dir.join(TAGGER_THRESHOLD_FILE);
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::write(path, clamped.to_string())?;
|
||||
TAGGER_SESSION_DIRTY.store(true, Ordering::Relaxed);
|
||||
Ok(clamped)
|
||||
}
|
||||
|
||||
pub fn tagger_batch_size(app_data_dir: &Path) -> usize {
|
||||
let path = app_data_dir.join(TAGGER_BATCH_SIZE_FILE);
|
||||
let Ok(value) = std::fs::read_to_string(path) else {
|
||||
return 8;
|
||||
};
|
||||
value
|
||||
.trim()
|
||||
.parse::<usize>()
|
||||
.unwrap_or(8)
|
||||
.clamp(1, 100)
|
||||
}
|
||||
|
||||
pub fn set_tagger_batch_size(app_data_dir: &Path, batch_size: usize) -> Result<usize> {
|
||||
let clamped = batch_size.clamp(1, 100);
|
||||
let path = app_data_dir.join(TAGGER_BATCH_SIZE_FILE);
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::write(path, clamped.to_string())?;
|
||||
Ok(clamped)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Model status / download
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub fn tagger_model_status(app_data_dir: &Path) -> TaggerModelStatus {
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
// The ONNX runtime DLLs live in the caption model dir; reuse them.
|
||||
let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft");
|
||||
let missing_files = TAGGER_REQUIRED_FILES
|
||||
.iter()
|
||||
.filter(|file| {
|
||||
let path = if file.starts_with("onnxruntime/") {
|
||||
caption_model_dir.join(file)
|
||||
} else {
|
||||
local_dir.join(file)
|
||||
};
|
||||
!path.exists()
|
||||
})
|
||||
.map(|file| (*file).to_string())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
TaggerModelStatus {
|
||||
model_id: WD_TAGGER_MODEL_ID,
|
||||
model_name: WD_TAGGER_MODEL_NAME,
|
||||
local_dir: local_dir.to_string_lossy().to_string(),
|
||||
ready: missing_files.is_empty(),
|
||||
missing_files,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn prepare_tagger_model_with_progress(
|
||||
app_data_dir: &Path,
|
||||
emit_progress: impl Fn(TaggerModelProgress),
|
||||
) -> Result<TaggerModelStatus> {
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
std::fs::create_dir_all(&local_dir)?;
|
||||
|
||||
// Ensure the shared ONNX runtime DLLs are present. The tagger shares
|
||||
// them with the captioner; install them here so the tagger is fully
|
||||
// functional even on a clean install where the caption model has never
|
||||
// been downloaded.
|
||||
let caption_model_dir = crate::captioner::model_dir(app_data_dir);
|
||||
std::fs::create_dir_all(&caption_model_dir)?;
|
||||
crate::captioner::ensure_onnx_runtime(&caption_model_dir)?;
|
||||
|
||||
// Download the two tagger-specific files.
|
||||
const DOWNLOAD_FILES: &[&str] = &["model.onnx", "selected_tags.csv"];
|
||||
|
||||
let api = Api::new()?;
|
||||
let repo = api.repo(Repo::new(WD_TAGGER_MODEL_ID.to_string(), RepoType::Model));
|
||||
|
||||
let mut completed_files = DOWNLOAD_FILES
|
||||
.iter()
|
||||
.filter(|file| local_dir.join(file).exists())
|
||||
.count();
|
||||
|
||||
emit_progress(TaggerModelProgress {
|
||||
total_files: DOWNLOAD_FILES.len(),
|
||||
completed_files,
|
||||
current_file: None,
|
||||
done: completed_files == DOWNLOAD_FILES.len(),
|
||||
});
|
||||
|
||||
for file in DOWNLOAD_FILES {
|
||||
let destination = local_dir.join(file);
|
||||
if destination.exists() {
|
||||
continue;
|
||||
}
|
||||
emit_progress(TaggerModelProgress {
|
||||
total_files: DOWNLOAD_FILES.len(),
|
||||
completed_files,
|
||||
current_file: Some((*file).to_string()),
|
||||
done: false,
|
||||
});
|
||||
let cached = repo.get(file)?;
|
||||
std::fs::copy(cached, destination)?;
|
||||
completed_files += 1;
|
||||
emit_progress(TaggerModelProgress {
|
||||
total_files: DOWNLOAD_FILES.len(),
|
||||
completed_files,
|
||||
current_file: Some((*file).to_string()),
|
||||
done: completed_files == DOWNLOAD_FILES.len(),
|
||||
});
|
||||
}
|
||||
|
||||
emit_progress(TaggerModelProgress {
|
||||
total_files: DOWNLOAD_FILES.len(),
|
||||
completed_files,
|
||||
current_file: None,
|
||||
done: true,
|
||||
});
|
||||
|
||||
Ok(tagger_model_status(app_data_dir))
|
||||
}
|
||||
|
||||
pub fn delete_tagger_model(app_data_dir: &Path) -> Result<TaggerModelStatus> {
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
if local_dir.exists() {
|
||||
std::fs::remove_dir_all(&local_dir)?;
|
||||
}
|
||||
Ok(tagger_model_status(app_data_dir))
|
||||
}
|
||||
|
||||
pub fn probe_tagger_runtime(app_data_dir: &Path) -> Result<TaggerRuntimeProbe> {
|
||||
let status = tagger_model_status(app_data_dir);
|
||||
if !status.ready {
|
||||
anyhow::bail!(
|
||||
"WD Tagger model is missing {} required file(s): {}",
|
||||
status.missing_files.len(),
|
||||
status.missing_files.join(", ")
|
||||
);
|
||||
}
|
||||
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft");
|
||||
crate::captioner::ensure_onnx_runtime(&caption_model_dir)?;
|
||||
|
||||
let acceleration = tagger_acceleration(app_data_dir);
|
||||
let model_path = local_dir.join("model.onnx");
|
||||
|
||||
// Verify that the model file exists and has non-zero size before trying
|
||||
// to create an ORT session (better error message on corruption).
|
||||
let metadata = std::fs::metadata(&model_path)?;
|
||||
if metadata.len() == 0 {
|
||||
anyhow::bail!("model.onnx is empty");
|
||||
}
|
||||
|
||||
// Actually create a session to verify the EP loads correctly.
|
||||
let loaded_acceleration = match acceleration {
|
||||
TaggerAcceleration::Cpu => {
|
||||
create_tagger_session(&model_path, TaggerAcceleration::Cpu)?;
|
||||
TaggerAcceleration::Cpu
|
||||
}
|
||||
TaggerAcceleration::Auto => {
|
||||
// Try DirectML explicitly; if it fails the real session would have
|
||||
// fallen back to CPU silently — report what would actually run.
|
||||
let directml_ok =
|
||||
create_tagger_session(&model_path, TaggerAcceleration::Directml).is_ok();
|
||||
if directml_ok {
|
||||
TaggerAcceleration::Directml
|
||||
} else {
|
||||
create_tagger_session(&model_path, TaggerAcceleration::Cpu)?;
|
||||
TaggerAcceleration::Cpu
|
||||
}
|
||||
}
|
||||
TaggerAcceleration::Directml => {
|
||||
create_tagger_session(&model_path, TaggerAcceleration::Directml)?;
|
||||
TaggerAcceleration::Directml
|
||||
}
|
||||
};
|
||||
|
||||
Ok(TaggerRuntimeProbe {
|
||||
ready: true,
|
||||
acceleration: loaded_acceleration,
|
||||
session: TaggerRuntimeSessionProbe {
|
||||
file: "model.onnx",
|
||||
inputs: vec!["pixel_values".to_string()],
|
||||
outputs: vec![format!("output [EP: {:?}]", loaded_acceleration)],
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Top-level inference entry point
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tagger implementation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub struct WdTagger {
|
||||
session: Session,
|
||||
labels: Vec<TagEntry>,
|
||||
threshold: f32,
|
||||
input_size: usize,
|
||||
}
|
||||
|
||||
impl WdTagger {
|
||||
pub fn new(app_data_dir: &Path) -> Result<Self> {
|
||||
let started_at = Instant::now();
|
||||
|
||||
let status = tagger_model_status(app_data_dir);
|
||||
if !status.ready {
|
||||
anyhow::bail!(
|
||||
"WD tagger model is missing {} required file(s): {}",
|
||||
status.missing_files.len(),
|
||||
status.missing_files.join(", ")
|
||||
);
|
||||
}
|
||||
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
|
||||
// The ONNX runtime DLLs are shared with the captioner; use the
|
||||
// captioner's shared ORT init lock to avoid double-initialisation.
|
||||
let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft");
|
||||
crate::captioner::ensure_onnx_runtime(&caption_model_dir).map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"ONNX Runtime not initialised — download the Florence-2 caption model first \
|
||||
to get the shared runtime DLLs. Original error: {e}"
|
||||
)
|
||||
})?;
|
||||
|
||||
let acceleration = tagger_acceleration(app_data_dir);
|
||||
let threshold = tagger_threshold(app_data_dir);
|
||||
let model_path = local_dir.join("model.onnx");
|
||||
let labels_path = local_dir.join("selected_tags.csv");
|
||||
|
||||
let session = create_tagger_session(&model_path, acceleration)?;
|
||||
|
||||
// Determine the input spatial size from the ONNX model graph.
|
||||
// WD v3 models use (1, H, W, 3) where H == W, typically 448.
|
||||
let input_size = {
|
||||
let inputs = session.inputs();
|
||||
let dim = inputs
|
||||
.first()
|
||||
.and_then(|inp| {
|
||||
if let ort::value::ValueType::Tensor { shape, .. } = inp.dtype() {
|
||||
shape
|
||||
.get(1)
|
||||
.and_then(|&d| if d > 0 { Some(d as usize) } else { None })
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.unwrap_or(448);
|
||||
dim
|
||||
};
|
||||
|
||||
let labels = load_labels(&labels_path)?;
|
||||
|
||||
println!(
|
||||
"WD tagger loaded in {:?} ({} labels, input {}x{}, {:?} acceleration)",
|
||||
started_at.elapsed(),
|
||||
labels.len(),
|
||||
input_size,
|
||||
input_size,
|
||||
acceleration,
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
session,
|
||||
labels,
|
||||
threshold,
|
||||
input_size,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn run(&mut self, image_path: &Path, max_tags: usize) -> Result<TaggerOutput> {
|
||||
let started_at = Instant::now();
|
||||
|
||||
let image_array = preprocess_image(image_path, self.input_size)?;
|
||||
let batch_size = 1usize;
|
||||
let input = Tensor::from_array((
|
||||
[batch_size, self.input_size, self.input_size, 3usize],
|
||||
image_array.into_boxed_slice(),
|
||||
))
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
let input_name: String = self.session.inputs()[0].name().to_string();
|
||||
let outputs = self
|
||||
.session
|
||||
.run(ort::inputs! { input_name.as_str() => input })
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
let (_, probabilities) = outputs[0]
|
||||
.try_extract_tensor::<f32>()
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
let probs: &[f32] = &probabilities;
|
||||
|
||||
if probs.len() != self.labels.len() {
|
||||
anyhow::bail!(
|
||||
"Model output length {} does not match label count {}",
|
||||
probs.len(),
|
||||
self.labels.len()
|
||||
);
|
||||
}
|
||||
|
||||
// Collect rating scores (category 9) - pick the argmax as the rating.
|
||||
let rating = self
|
||||
.labels
|
||||
.iter()
|
||||
.zip(probs.iter())
|
||||
.filter(|(entry, _)| entry.category == RATING_CATEGORY)
|
||||
.max_by(|(_, a), (_, b)| a.total_cmp(b))
|
||||
.map(|(entry, _)| entry.name.clone())
|
||||
.unwrap_or_else(|| "general".to_string());
|
||||
|
||||
// Collect general + character tags above threshold, sorted by confidence.
|
||||
let mut tags: Vec<TagResult> = self
|
||||
.labels
|
||||
.iter()
|
||||
.zip(probs.iter())
|
||||
.filter(|(entry, prob)| {
|
||||
(entry.category == GENERAL_CATEGORY || entry.category == CHARACTER_CATEGORY)
|
||||
&& **prob >= self.threshold
|
||||
})
|
||||
.map(|(entry, prob)| TagResult {
|
||||
tag: entry.name.clone(),
|
||||
confidence: *prob,
|
||||
})
|
||||
.collect();
|
||||
|
||||
tags.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
|
||||
tags.truncate(max_tags);
|
||||
|
||||
println!(
|
||||
"WD tagger: {} tags (threshold={}, rating={}) in {:?} for {}",
|
||||
tags.len(),
|
||||
self.threshold,
|
||||
rating,
|
||||
started_at.elapsed(),
|
||||
image_path.display(),
|
||||
);
|
||||
|
||||
Ok(TaggerOutput { tags, rating })
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Session creation – mirrors captioner's `create_session`
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn create_tagger_session(path: &Path, acceleration: TaggerAcceleration) -> 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 use_directml = matches!(
|
||||
acceleration,
|
||||
TaggerAcceleration::Auto | TaggerAcceleration::Directml
|
||||
);
|
||||
|
||||
let builder = builder
|
||||
.with_memory_pattern(!use_directml)
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
let builder = builder
|
||||
.with_parallel_execution(false)
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
let builder = builder
|
||||
.with_intra_threads(1)
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
let mut builder = match acceleration {
|
||||
TaggerAcceleration::Cpu => {
|
||||
println!("WD tagger: using CPU execution provider");
|
||||
builder
|
||||
}
|
||||
TaggerAcceleration::Auto => builder
|
||||
.with_execution_providers([ep::DirectML::default().build().fail_silently()])
|
||||
.unwrap_or_else(|error| {
|
||||
println!("WD tagger: DirectML unavailable, falling back to CPU");
|
||||
error.recover()
|
||||
}),
|
||||
TaggerAcceleration::Directml => builder
|
||||
.with_execution_providers([ep::DirectML::default().build().error_on_failure()])
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?,
|
||||
};
|
||||
|
||||
let session = builder
|
||||
.commit_from_file(path)
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
Ok(session)
|
||||
}
|
||||
// ---------------------------------------------------------------------------
|
||||
// Label loading
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn load_labels(path: &Path) -> Result<Vec<TagEntry>> {
|
||||
let mut reader = csv::Reader::from_path(path)?;
|
||||
let mut entries = Vec::new();
|
||||
for result in reader.records() {
|
||||
let record = result?;
|
||||
// CSV columns: tag_id, name, category, count
|
||||
let name = record
|
||||
.get(1)
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing name column in selected_tags.csv"))?
|
||||
.replace('_', " ");
|
||||
let category: u8 = record
|
||||
.get(2)
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing category column in selected_tags.csv"))?
|
||||
.parse()
|
||||
.unwrap_or(0);
|
||||
entries.push(TagEntry { name, category });
|
||||
}
|
||||
if entries.is_empty() {
|
||||
anyhow::bail!("selected_tags.csv is empty or could not be parsed");
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Image preprocessing
|
||||
// WD tagger expects: (1, H, W, 3) float32, raw [0,255] values, BGR channel
|
||||
// order, padded to square with white (255,255,255).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn preprocess_image(image_path: &Path, target_size: usize) -> Result<Vec<f32>> {
|
||||
let image = ImageReader::open(image_path)?.decode()?;
|
||||
|
||||
// Composite any alpha channel onto a white background.
|
||||
let image_rgba = image.to_rgba8();
|
||||
let (width, height) = image_rgba.dimensions();
|
||||
let mut canvas_rgba =
|
||||
image::RgbaImage::from_pixel(width, height, image::Rgba([255, 255, 255, 255]));
|
||||
image::imageops::overlay(&mut canvas_rgba, &image_rgba, 0, 0);
|
||||
let image_rgb = DynamicImage::ImageRgba8(canvas_rgba).to_rgb8();
|
||||
|
||||
// Pad to square.
|
||||
let max_dim = width.max(height);
|
||||
let pad_left = (max_dim - width) / 2;
|
||||
let pad_top = (max_dim - height) / 2;
|
||||
let mut square = image::RgbImage::from_pixel(max_dim, max_dim, image::Rgb([255, 255, 255]));
|
||||
image::imageops::overlay(&mut square, &image_rgb, pad_left as i64, pad_top as i64);
|
||||
|
||||
// Resize to model input size.
|
||||
let resized = image::imageops::resize(
|
||||
&square,
|
||||
target_size as u32,
|
||||
target_size as u32,
|
||||
FilterType::CatmullRom,
|
||||
);
|
||||
|
||||
// Flatten to (H, W, 3) float32 in BGR order, values in [0, 255].
|
||||
let mut pixel_values = vec![0.0f32; target_size * target_size * 3];
|
||||
for (x, y, pixel) in resized.enumerate_pixels() {
|
||||
let base = (y as usize * target_size + x as usize) * 3;
|
||||
// BGR order
|
||||
pixel_values[base] = f32::from(pixel[2]); // B
|
||||
pixel_values[base + 1] = f32::from(pixel[1]); // G
|
||||
pixel_values[base + 2] = f32::from(pixel[0]); // R
|
||||
}
|
||||
|
||||
Ok(pixel_values)
|
||||
}
|
||||
+317
-9
@@ -1,5 +1,5 @@
|
||||
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 std::sync::Once;
|
||||
|
||||
@@ -19,8 +19,13 @@ pub fn migrate(conn: &Connection) -> Result<()> {
|
||||
"CREATE VIRTUAL TABLE IF NOT EXISTS image_vec USING vec0(
|
||||
image_id INTEGER PRIMARY KEY,
|
||||
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(())
|
||||
}
|
||||
@@ -28,6 +33,18 @@ pub fn migrate(conn: &Connection) -> Result<()> {
|
||||
#[allow(dead_code)]
|
||||
pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> {
|
||||
conn.execute("DELETE FROM image_vec WHERE image_id = ?1", [image_id])?;
|
||||
// Advance the revision so any cached HNSW index is invalidated after deletions.
|
||||
conn.execute(
|
||||
"INSERT INTO app_kv (key, value) VALUES ('embedding_revision', 1)
|
||||
ON CONFLICT(key) DO UPDATE SET value = value + 1",
|
||||
[],
|
||||
)?;
|
||||
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(())
|
||||
}
|
||||
|
||||
@@ -50,26 +67,74 @@ pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) ->
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Result<Vec<i64>> {
|
||||
let embedding: Vec<u8> = conn.query_row(
|
||||
#[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,
|
||||
folder_id: Option<i64>,
|
||||
) -> Result<Vec<i64>> {
|
||||
let embedding: Vec<u8> = match conn.query_row(
|
||||
"SELECT embedding FROM image_vec WHERE image_id = ?1",
|
||||
[image_id],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
) {
|
||||
Ok(embedding) => embedding,
|
||||
Err(SqliteError::QueryReturnedNoRows) => return Ok(Vec::new()),
|
||||
Err(error) => return Err(error.into()),
|
||||
};
|
||||
|
||||
if let Some(folder_id) = folder_id {
|
||||
// Brute-force cosine scan scoped to the folder — avoids the KNN k=4096 limit
|
||||
// and returns exact nearest neighbours within the folder.
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT v.image_id
|
||||
FROM image_vec v
|
||||
JOIN images i ON i.id = v.image_id
|
||||
WHERE i.folder_id = ?2
|
||||
AND v.image_id != ?3
|
||||
ORDER BY vec_distance_cosine(v.embedding, vec_f32(?1)) ASC
|
||||
LIMIT ?4",
|
||||
)?;
|
||||
let rows = stmt.query_map((&embedding, folder_id, image_id, limit as i64), |row| {
|
||||
row.get::<_, i64>(0)
|
||||
})?;
|
||||
return Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?);
|
||||
}
|
||||
|
||||
// Global KNN search (no folder filter) — use the ANN index.
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT image_id
|
||||
FROM image_vec
|
||||
WHERE embedding MATCH vec_f32(?1)
|
||||
AND k = ?2",
|
||||
)?;
|
||||
let rows = stmt.query_map((&embedding, (limit as i64) + 1), |row| row.get::<_, i64>(0))?;
|
||||
let rows = stmt
|
||||
.query_map((&embedding, (limit + 1) as i64), |row| row.get::<_, i64>(0))?
|
||||
.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
|
||||
let mut ids = Vec::new();
|
||||
for row in rows {
|
||||
let candidate_id = row?;
|
||||
if candidate_id != image_id {
|
||||
ids.push(candidate_id);
|
||||
if row != image_id {
|
||||
ids.push(row);
|
||||
}
|
||||
if ids.len() >= limit {
|
||||
break;
|
||||
@@ -78,6 +143,106 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) ->
|
||||
Ok(ids)
|
||||
}
|
||||
|
||||
// pub fn find_similar_image_matches(
|
||||
// conn: &Connection,
|
||||
// image_id: i64,
|
||||
// folder_id: Option<i64>,
|
||||
// threshold: f32,
|
||||
// offset: usize,
|
||||
// limit: usize,
|
||||
// ) -> Result<Vec<(i64, f32)>> {
|
||||
// let embedding: Vec<u8> = match conn.query_row(
|
||||
// "SELECT embedding FROM image_vec WHERE image_id = ?1",
|
||||
// [image_id],
|
||||
// |row| row.get(0),
|
||||
// ) {
|
||||
// Ok(embedding) => embedding,
|
||||
// Err(SqliteError::QueryReturnedNoRows) => return Ok(Vec::new()),
|
||||
// Err(error) => return Err(error.into()),
|
||||
// };
|
||||
|
||||
// let query = match folder_id {
|
||||
// Some(_) => {
|
||||
// "SELECT v.image_id, vec_distance_cosine(v.embedding, vec_f32(?1)) AS distance
|
||||
// FROM image_vec v
|
||||
// JOIN images i ON i.id = v.image_id
|
||||
// WHERE i.folder_id = ?2
|
||||
// AND v.image_id != ?3
|
||||
// AND vec_distance_cosine(v.embedding, vec_f32(?1)) <= ?4
|
||||
// ORDER BY distance ASC
|
||||
// LIMIT ?5 OFFSET ?6"
|
||||
// }
|
||||
// None => {
|
||||
// "SELECT v.image_id, vec_distance_cosine(v.embedding, vec_f32(?1)) AS distance
|
||||
// FROM image_vec v
|
||||
// WHERE v.image_id != ?2
|
||||
// AND vec_distance_cosine(v.embedding, vec_f32(?1)) <= ?3
|
||||
// ORDER BY distance ASC
|
||||
// LIMIT ?4 OFFSET ?5"
|
||||
// }
|
||||
// };
|
||||
|
||||
// let mut stmt = conn.prepare(query)?;
|
||||
// match folder_id {
|
||||
// Some(folder_id) => Ok(stmt
|
||||
// .query_map(
|
||||
// (
|
||||
// &embedding,
|
||||
// folder_id,
|
||||
// image_id,
|
||||
// threshold,
|
||||
// limit as i64,
|
||||
// offset as i64,
|
||||
// ),
|
||||
// |row| Ok((row.get::<_, i64>(0)?, row.get::<_, f32>(1)?)),
|
||||
// )?
|
||||
// .collect::<rusqlite::Result<Vec<_>>>()?),
|
||||
// None => Ok(stmt
|
||||
// .query_map(
|
||||
// (&embedding, image_id, threshold, limit as i64, offset as i64),
|
||||
// |row| Ok((row.get::<_, i64>(0)?, row.get::<_, f32>(1)?)),
|
||||
// )?
|
||||
// .collect::<rusqlite::Result<Vec<_>>>()?),
|
||||
// }
|
||||
// }
|
||||
|
||||
pub fn get_image_embedding(conn: &Connection, image_id: i64) -> Result<Option<Vec<f32>>> {
|
||||
let embedding: Result<Vec<u8>, rusqlite::Error> = conn.query_row(
|
||||
"SELECT embedding FROM image_vec WHERE image_id = ?1",
|
||||
[image_id],
|
||||
|row| row.get(0),
|
||||
);
|
||||
|
||||
match embedding {
|
||||
Ok(bytes) => Ok(Some(unpack_f32(&bytes))),
|
||||
Err(SqliteError::QueryReturnedNoRows) => Ok(None),
|
||||
Err(error) => Err(error.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_embedding_revision(conn: &Connection) -> Result<String> {
|
||||
// Use the monotonically incremented app_kv counter so that two embeddings
|
||||
// saved within the same clock second still advance the revision, preventing
|
||||
// the HNSW cache from serving stale vectors.
|
||||
let revision: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COALESCE((SELECT value FROM app_kv WHERE key = 'embedding_revision'), 0)",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.unwrap_or(0);
|
||||
Ok(revision.to_string())
|
||||
}
|
||||
|
||||
// fn image_ids_for_folder(
|
||||
// conn: &Connection,
|
||||
// folder_id: i64,
|
||||
// ) -> Result<std::collections::HashSet<i64>> {
|
||||
// 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))?;
|
||||
// Ok(rows.collect::<rusqlite::Result<std::collections::HashSet<_>>>()?)
|
||||
// }
|
||||
|
||||
/// Returns all stored image embeddings with their image IDs, optionally filtered to one folder.
|
||||
/// Each entry is `(image_id, normalized_f32_embedding)`.
|
||||
pub fn get_all_image_embeddings_with_ids(
|
||||
@@ -155,6 +320,149 @@ pub fn search_image_ids_by_embedding(
|
||||
Ok(ids)
|
||||
}
|
||||
|
||||
/// Brute-force cosine search scoped to a single folder, ordered by ascending distance.
|
||||
/// Used for region-based similarity search where we want folder-scoped results.
|
||||
pub fn search_image_ids_by_embedding_in_folder(
|
||||
conn: &Connection,
|
||||
embedding: &[f32],
|
||||
folder_id: i64,
|
||||
exclude_image_id: Option<i64>,
|
||||
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 exclude_id = exclude_image_id.unwrap_or(-1);
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT v.image_id
|
||||
FROM image_vec v
|
||||
JOIN images i ON i.id = v.image_id
|
||||
WHERE i.folder_id = ?2
|
||||
AND v.image_id != ?3
|
||||
ORDER BY vec_distance_cosine(v.embedding, vec_f32(?1)) ASC
|
||||
LIMIT ?4",
|
||||
)?;
|
||||
let rows = stmt.query_map((&packed, folder_id, exclude_id, limit as i64), |row| {
|
||||
row.get::<_, i64>(0)
|
||||
})?;
|
||||
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
|
||||
}
|
||||
|
||||
#[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)]
|
||||
fn pack_f32(values: &[f32]) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(values.len() * std::mem::size_of::<f32>());
|
||||
|
||||
@@ -4,9 +4,9 @@
|
||||
"version": "0.1.0",
|
||||
"identifier": "wtf.jezz.phokus",
|
||||
"build": {
|
||||
"beforeDevCommand": "pnpm dev",
|
||||
"beforeDevCommand": "pnpm dev:vite",
|
||||
"devUrl": "http://localhost:1420",
|
||||
"beforeBuildCommand": "pnpm build",
|
||||
"beforeBuildCommand": "pnpm build:vite",
|
||||
"frontendDist": "../dist"
|
||||
},
|
||||
"app": {
|
||||
@@ -25,7 +25,9 @@
|
||||
"csp": null,
|
||||
"assetProtocol": {
|
||||
"enable": true,
|
||||
"scope": ["**"]
|
||||
"scope": [
|
||||
"**"
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -40,4 +42,4 @@
|
||||
"icons/icon.ico"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user