From 0ca4d142d8187b2fc246579d7ecfea190e204d35 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 22:43:16 +0000 Subject: [PATCH] feat(discovery): AI tagging, semantic search, similarity, duplicate finder, folder management MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .gitignore | 6 + CLAUDE.md | 94 ++ README.md | 126 +-- index.html | 4 +- package.json | 9 +- pnpm-lock.yaml | 49 + src-tauri/.cargo/config.toml | 12 + src-tauri/Cargo.lock | 479 +++++++++- src-tauri/Cargo.toml | 55 ++ src-tauri/capabilities/default.json | 1 + src-tauri/src/captioner.rs | 1011 ++++++++++++++++++++ src-tauri/src/commands.rs | 1232 +++++++++++++++++++++++- src-tauri/src/db.rs | 1343 ++++++++++++++++++++++++++- src-tauri/src/embedder.rs | 45 + src-tauri/src/hnsw_index.rs | 154 +++ src-tauri/src/indexer.rs | 528 ++++++++++- src-tauri/src/lib.rs | 63 +- src-tauri/src/tagger.rs | 656 +++++++++++++ src-tauri/src/vector.rs | 326 ++++++- src-tauri/tauri.conf.json | 10 +- src/App.tsx | 14 + src/components/BackgroundTasks.tsx | 242 +++-- src/components/DuplicateFinder.tsx | 278 ++++++ src/components/Gallery.tsx | 170 ++-- src/components/Lightbox.tsx | 446 ++++++++- src/components/SettingsModal.tsx | 568 +++++++++++ src/components/Sidebar.tsx | 309 +++++- src/components/TagCloud.tsx | 512 ++++++---- src/components/TitleBar.tsx | 13 + src/components/Toolbar.tsx | 288 ++++-- src/notifications.ts | 31 + src/store.ts | 1225 +++++++++++++++++++++++- 32 files changed, 9539 insertions(+), 760 deletions(-) create mode 100644 CLAUDE.md create mode 100644 src-tauri/.cargo/config.toml create mode 100644 src-tauri/src/captioner.rs create mode 100644 src-tauri/src/hnsw_index.rs create mode 100644 src-tauri/src/tagger.rs create mode 100644 src/components/DuplicateFinder.tsx create mode 100644 src/components/SettingsModal.tsx create mode 100644 src/notifications.ts diff --git a/.gitignore b/.gitignore index 8b9cd9f..7e7224f 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,9 @@ dist-ssr # Local staging area /staging + +# Misc +*.py +*.json + +*.pyc diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..e9febb9 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,94 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project overview + +**Phokus** — a Tauri v2 desktop image gallery app with a React/TypeScript frontend and a Rust backend. The app indexes local media folders, generates thumbnails via FFmpeg, computes visual embeddings for semantic search and similarity, and AI-tags images using the WD tagger (ONNX via `ort`). AI captioning code exists in the backend but the UI surface has been removed; the worker is commented out in `lib.rs`. + +## Commands + +```bash +# Development (Vite hot-reload + Rust auto-rebuild) +pnpm dev:app + +# Frontend only (no Tauri window) +pnpm dev:vite + +# Production build +pnpm build:app + +# Type-check frontend +pnpm build:vite +``` + +Use **pnpm** — never npm. + +There are no test suites configured. + +## Architecture + +### Frontend (`src/`) + +- **`store.ts`** — single Zustand store (`useGalleryStore`) that owns all app state and all `invoke()` calls to the Tauri backend. Every feature (folders, images, search, similar images, tags, captions, tagger, duplicates) is implemented as store actions here. React components are thin consumers. +- **`App.tsx`** — sets up Tauri event listeners (`subscribeToProgress`) and renders the top-level layout (sidebar + active view). +- **`src/components/`** — UI components: `Gallery`, `Lightbox`, `Sidebar`, `Toolbar`, `TagCloud`, `DuplicateFinder`, `BackgroundTasks`, `SettingsModal`, `MenuBar`, `TitleBar`. +- State management: Zustand v5, no selectors library — components call `useGalleryStore(s => s.field)` directly. +- Styling: Tailwind CSS v4 (Vite plugin, no config file). +- Virtualized gallery grid: `@tanstack/react-virtual`. +- Animation: `framer-motion`. + +### Search modes + +The search bar supports prefix syntax parsed by `parseSearchValue` in `store.ts`: +- No prefix / `f:` — filename search (paginated, DB-backed) +- `/s ` or `s: ` — semantic (embedding) search +- `/t ` or `t: ` — tag search + +### Backend (`src-tauri/src/`) + +Workers are started in `lib.rs` and run as background threads throughout the app lifetime: +- **thumbnail worker** (multiple threads, count from `StorageProfile::Balanced`) +- **metadata worker** — FFmpeg probe for video files +- **embedding worker** — generates CLIP-style visual embeddings (candle, HuggingFace hub) +- **tagging worker** — WD tagger via ONNX Runtime (`ort`), DirectML/CPU acceleration + +Key modules: +| File | Purpose | +|------|---------| +| `db.rs` | SQLite pool (r2d2 + rusqlite), schema migrations, all query functions | +| `commands.rs` | All `#[tauri::command]` handlers — one-to-one with frontend `invoke()` calls | +| `indexer.rs` | Worker thread launchers and job dispatch | +| `embedder.rs` | Visual embedding generation (candle + HF hub models) | +| `vector.rs` | sqlite-vec integration + HNSW index for ANN search | +| `hnsw_index.rs` | In-memory HNSW index wrapper (hnsw_rs) | +| `tagger.rs` | WD tagger: ONNX model download, inference, CSV tag loading | +| `captioner.rs` | AI captioning (ONNX, disabled in workers but code intact) | +| `thumbnail.rs` | Thumbnail generation (image crate + fast_image_resize, FFmpeg for video) | +| `media.rs` | FFmpeg sidecar provisioning and probing | +| `storage.rs` | `StorageProfile` for tuning worker counts | + +Database: SQLite with WAL mode, stored in the Tauri app data directory as `gallery.db`. Thumbnails stored alongside as `thumbnails/`. + +### Tauri events (backend → frontend) + +| Event | Payload | +|-------|---------| +| `index-progress` | `IndexProgress` | +| `media-job-progress` | `MediaJobProgressEvent` | +| `indexed-images` | `IndexedImagesBatch` | +| `media-updated` | `ThumbnailBatch` | +| `caption-model-progress` | `CaptionModelProgress` | +| `tagger-model-progress` | `TaggerModelProgress` | +| `duplicate_scan_progress` | `[scanned, total]` | + +### Key types + +`ImageRecord` (mirrored in `store.ts` and `db.rs`) is the central data type. It carries embedding status, tagging status, caption data, and media metadata. The frontend type must stay in sync with the Rust struct serialization. + +## Development notes + +- Hot Reload is active during `dev:app` — do not restart the server for frontend changes. Restart only when adding Rust crates or changing Vite config. +- ML inference crates (`candle-*`, `ort`, `image`, `rayon`, `tokenizers`, `xxhash-rust`, `rusqlite`) use `opt-level = 3` in dev profile to keep inference performance acceptable. +- The caption worker is intentionally disabled (`lib.rs:73`) — the backend code is intact for future re-enabling. +- **Never use `any` type** in TypeScript — look up correct types. diff --git a/README.md b/README.md index a4b5578..784cc6b 100644 --- a/README.md +++ b/README.md @@ -1,116 +1,58 @@ # Phokus -## Overview +A local-first desktop media library for browsing, filtering, and curating image and video folders. -Phokus is a Tauri desktop app for building a fast, local media library from folders on disk. It indexes images and videos, stores metadata in SQLite, and gives you a dense browsing workflow with filtering, favorites, ratings, and a lightbox preview. +## Features -The current app is optimized for: - -- local folders instead of cloud import flows -- large visual libraries -- quick review and curation -- mixed image and video browsing - -## Current features - -- Add and remove media folders -- Background indexing with progress updates -- Browse all media or filter by folder -- Search by filename -- Filter by images, videos, or favorites -- Sort by modified date, name, or file size -- Grid density controls -- Lightbox preview with keyboard navigation -- Favorite and star-rating metadata saved in SQLite -- Virtualized/local-first architecture built on Tauri + React +- Add and remove media folders; background indexing with live progress +- Browse all media or filter by folder, type (image/video), favorites, or star rating +- **Filename search**, **semantic search** (`/s query`), and **tag search** (`/t tag`) +- **Similar image search** — find visually similar media by image or selected region +- **AI tagging** via WD tagger (ONNX, CPU/DirectML) with confidence threshold control +- **Explore view** — visual cluster map and tag cloud for browsing by theme +- **Duplicate finder** — scan for exact duplicates by file hash with bulk delete +- Lightbox preview with keyboard navigation, inline tag editing, and rating controls +- Sort by date, name, size, rating, or duration +- Grid density controls (compact / comfortable / detail) ## Supported formats -Images: - -- `jpg` -- `jpeg` -- `png` -- `gif` -- `bmp` -- `tiff` -- `tif` -- `webp` -- `avif` -- `heic` -- `heif` - -Videos: - -- `mp4` -- `mov` -- `m4v` -- `webm` +| Images | Videos | +|--------|--------| +| jpg, jpeg, png, gif, bmp | mp4, mov, m4v | +| tiff, tif, webp, avif, heic, heif | webm | ## Stack -- Tauri 2 -- React 19 -- TypeScript -- Zustand -- Rust -- SQLite + `sqlite-vec` -- Vite - -## Project structure - -- `src/`: React UI, state, and components -- `src-tauri/src/commands.rs`: Tauri command surface -- `src-tauri/src/db.rs`: SQLite schema and queries -- `src-tauri/src/indexer.rs`: folder crawling and batch indexing -- `src-tauri/src/vector.rs`: vector table setup for future semantic workflows +- Tauri 2 + Rust backend +- React 19 + TypeScript + Zustand +- SQLite + `sqlite-vec` (vector search) +- ONNX Runtime (`ort`) for AI tagging +- Candle (Rust ML) for visual embeddings +- FFmpeg sidecar for video thumbnails and metadata +- Vite + Tailwind CSS v4 ## Development -### Prerequisites - -- Node.js 20+ -- `pnpm` -- Rust toolchain -- Tauri system prerequisites for Windows - -### Install +**Prerequisites:** Node.js 20+, pnpm, Rust toolchain, Tauri system prerequisites for Windows. ```bash pnpm install -``` -### Run in development +# Run with hot-reload (frontend + Rust) +pnpm dev:app -```bash -pnpm tauri dev -``` +# Frontend only +pnpm dev:vite -### Build - -```bash -pnpm tauri build +# Production build +pnpm build:app ``` ## How it works -1. Add a folder from the sidebar or Library menu. -2. The Rust indexer walks the directory recursively. -3. Supported files are written into SQLite with metadata such as path, size, dimensions, media type, rating, and favorite state. +1. Add a folder from the sidebar — the Rust indexer walks it recursively. +2. Supported files are written to SQLite with metadata (path, dimensions, media type, etc.). +3. Background workers generate thumbnails, compute visual embeddings, and run AI tagging. 4. Progress events stream back to the UI while the gallery updates incrementally. -5. The gallery view loads media in pages and opens items in a lightbox for review. - -## Notes - -- This is currently a local desktop library, not a sync product. -- Search is filename-based right now. -- The vector table and embedding fields exist, but semantic search is not wired into the UI yet. -- Some visible UI copy may still use the old working name until the frontend text is updated. - -## Positioning - -The clearest product description today is: - -> A local-first desktop media library for browsing, filtering, and curating image and video folders. - -That description is more accurate than "gallery" alone and gives you a better base for future branding, onboarding copy, and a landing page. +5. Embeddings power semantic search and the similar images feature via an HNSW index. diff --git a/index.html b/index.html index ff93803..3354975 100644 --- a/index.html +++ b/index.html @@ -2,11 +2,9 @@ - - Tauri + React + Typescript + Phokus -
diff --git a/package.json b/package.json index c90f202..1836830 100644 --- a/package.json +++ b/package.json @@ -4,8 +4,10 @@ "version": "0.1.0", "type": "module", "scripts": { - "dev": "vite", - "build": "tsc && vite build", + "build:app": "tauri build", + "build:vite": "tsc && vite build", + "dev:app": "tauri dev", + "dev:vite": "vite", "preview": "vite preview", "tauri": "tauri" }, @@ -14,7 +16,9 @@ "@tauri-apps/api": "^2", "@tauri-apps/plugin-dialog": "^2.7.0", "@tauri-apps/plugin-fs": "^2.5.0", + "@tauri-apps/plugin-notification": "^2.3.3", "@tauri-apps/plugin-opener": "^2", + "d3-force": "^3.0.0", "framer-motion": "^12.38.0", "react": "^19.1.0", "react-dom": "^19.1.0", @@ -23,6 +27,7 @@ "devDependencies": { "@tailwindcss/vite": "^4.2.2", "@tauri-apps/cli": "^2", + "@types/d3-force": "^3.0.10", "@types/react": "^19.1.8", "@types/react-dom": "^19.1.6", "@vitejs/plugin-react": "^4.6.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4fd524e..ec0add6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -20,9 +20,15 @@ importers: '@tauri-apps/plugin-fs': specifier: ^2.5.0 version: 2.5.0 + '@tauri-apps/plugin-notification': + specifier: ^2.3.3 + version: 2.3.3 '@tauri-apps/plugin-opener': specifier: ^2 version: 2.5.3 + d3-force: + specifier: ^3.0.0 + version: 3.0.0 framer-motion: specifier: ^12.38.0 version: 12.38.0(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -42,6 +48,9 @@ importers: '@tauri-apps/cli': specifier: ^2 version: 2.10.1 + '@types/d3-force': + specifier: ^3.0.10 + version: 3.0.10 '@types/react': specifier: ^19.1.8 version: 19.2.14 @@ -647,6 +656,9 @@ packages: '@tauri-apps/plugin-fs@2.5.0': resolution: {integrity: sha512-c83kbz61AK+rKjhS+je9+stIO27nXj7p9cqeg36TwkIUtxpCFTttlHHtqon6h6FN54cXjyAjlMPOJcW3mwE5XQ==} + '@tauri-apps/plugin-notification@2.3.3': + resolution: {integrity: sha512-Zw+ZH18RJb41G4NrfHgIuofJiymusqN+q8fGUIIV7vyCH+5sSn5coqRv/MWB9qETsUs97vmU045q7OyseCV3Qg==} + '@tauri-apps/plugin-opener@2.5.3': resolution: {integrity: sha512-CCcUltXMOfUEArbf3db3kCE7Ggy1ExBEBl51Ko2ODJ6GDYHRp1nSNlQm5uNCFY5k7/ufaK5Ib3Du/Zir19IYQQ==} @@ -662,6 +674,9 @@ packages: '@types/babel__traverse@7.28.0': resolution: {integrity: sha512-8PvcXf70gTDZBgt9ptxJ8elBeBjcLOAcOtoO/mPJjtji1+CdGbHgm77om1GrsPxsiE+uXIpNSK64UYaIwQXd4Q==} + '@types/d3-force@3.0.10': + resolution: {integrity: sha512-ZYeSaCF3p73RdOKcjj+swRlZfnYpK1EbaDiYICEEp5Q6sUiqFaFQ9qgoshp5CzIyyb/yD09kD9o2zEltCexlgw==} + '@types/estree@1.0.8': resolution: {integrity: sha512-dWHzHa2WqEXI/O1E9OjrocMTKJl2mSrEolh1Iomrv6U+JuNwaHXsXx9bLu5gG7BUWFIN0skIQJQ/L1rIex4X6w==} @@ -698,6 +713,22 @@ packages: csstype@3.2.3: resolution: {integrity: sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==} + d3-dispatch@3.0.1: + resolution: {integrity: sha512-rzUyPU/S7rwUflMyLc1ETDeBj0NRuHKKAcvukozwhshr6g6c5d8zh4c2gQjY2bZ0dXeGLWc1PF174P2tVvKhfg==} + engines: {node: '>=12'} + + d3-force@3.0.0: + resolution: {integrity: sha512-zxV/SsA+U4yte8051P4ECydjD/S+qeYtnaIyAs9tgHCqfguma/aAQDjo85A9Z6EKhBirHRJHXIgJUlffT4wdLg==} + engines: {node: '>=12'} + + d3-quadtree@3.0.1: + resolution: {integrity: sha512-04xDrxQTDTCFwP5H6hRhsRcb9xxv2RzkcsygFzmkSIOJy3PeRJP7sNk3VRIbKXcog561P9oU0/rVH6vDROAgUw==} + engines: {node: '>=12'} + + d3-timer@3.0.1: + resolution: {integrity: sha512-ndfJ/JxxMd3nw31uyKoY2naivF+r29V+Lc0svZxe1JvvIRmi8hUsrMvdOwgS1o6uBHmiz91geQ0ylPP0aj1VUA==} + engines: {node: '>=12'} + debug@4.4.3: resolution: {integrity: sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==} engines: {node: '>=6.0'} @@ -1423,6 +1454,10 @@ snapshots: dependencies: '@tauri-apps/api': 2.10.1 + '@tauri-apps/plugin-notification@2.3.3': + dependencies: + '@tauri-apps/api': 2.10.1 + '@tauri-apps/plugin-opener@2.5.3': dependencies: '@tauri-apps/api': 2.10.1 @@ -1448,6 +1483,8 @@ snapshots: dependencies: '@babel/types': 7.29.0 + '@types/d3-force@3.0.10': {} + '@types/estree@1.0.8': {} '@types/react-dom@19.2.3(@types/react@19.2.14)': @@ -1486,6 +1523,18 @@ snapshots: csstype@3.2.3: {} + d3-dispatch@3.0.1: {} + + d3-force@3.0.0: + dependencies: + d3-dispatch: 3.0.1 + d3-quadtree: 3.0.1 + d3-timer: 3.0.1 + + d3-quadtree@3.0.1: {} + + d3-timer@3.0.1: {} + debug@4.4.3: dependencies: ms: 2.1.3 diff --git a/src-tauri/.cargo/config.toml b/src-tauri/.cargo/config.toml new file mode 100644 index 0000000..22422f2 --- /dev/null +++ b/src-tauri/.cargo/config.toml @@ -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", +] diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 2871f9e..edd9078 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -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" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index b9daa8a..2bd2dde 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -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 diff --git a/src-tauri/capabilities/default.json b/src-tauri/capabilities/default.json index c3a561a..a526b2c 100644 --- a/src-tauri/capabilities/default.json +++ b/src-tauri/capabilities/default.json @@ -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", diff --git a/src-tauri/src/captioner.rs b/src-tauri/src/captioner.rs new file mode 100644 index 0000000..aea7833 --- /dev/null +++ b/src-tauri/src/captioner.rs @@ -0,0 +1,1011 @@ +use anyhow::Result; +use hf_hub::{api::sync::Api, Repo, RepoType}; +use image::{imageops::FilterType, ImageReader}; +use ort::ep; +use ort::session::SessionInputValue; +use ort::session::SessionOutputs; +use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::value::{Shape, Tensor}; +use serde::{Deserialize, Serialize}; +use std::borrow::Cow; +use std::io::{Cursor, Read}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::OnceLock; +use std::time::Instant; +use tokenizers::Tokenizer; + +pub const FLORENCE_MODEL_ID: &str = "onnx-community/Florence-2-base-ft"; +pub const FLORENCE_CAPTION_MODEL_NAME: &str = "florence-2-base-ft-onnx-q4"; +const ONNX_RUNTIME_NUGET_URL: &str = + "https://www.nuget.org/api/v2/package/Microsoft.ML.OnnxRuntime.DirectML/1.24.2"; +const DIRECTML_NUGET_URL: &str = + "https://www.nuget.org/api/v2/package/Microsoft.AI.DirectML/1.15.4"; +const ONNX_RUNTIME_DLL_FILE: &str = "onnxruntime/onnxruntime.dll"; +const ONNX_RUNTIME_PROVIDERS_DLL_FILE: &str = "onnxruntime/onnxruntime_providers_shared.dll"; +const DIRECTML_DLL_FILE: &str = "onnxruntime/DirectML.dll"; +const CAPTION_ACCELERATION_FILE: &str = "settings/caption_acceleration.txt"; +const CAPTION_DETAIL_FILE: &str = "settings/caption_detail.txt"; + +const ONNX_RUNTIME_FILES: &[(&str, &str, &str)] = &[ + ( + ONNX_RUNTIME_DLL_FILE, + ONNX_RUNTIME_NUGET_URL, + "runtimes/win-x64/native/onnxruntime.dll", + ), + ( + ONNX_RUNTIME_PROVIDERS_DLL_FILE, + ONNX_RUNTIME_NUGET_URL, + "runtimes/win-x64/native/onnxruntime_providers_shared.dll", + ), + ( + DIRECTML_DLL_FILE, + DIRECTML_NUGET_URL, + "bin/x64-win/DirectML.dll", + ), +]; + +const REQUIRED_FILES: &[&str] = &[ + ONNX_RUNTIME_DLL_FILE, + ONNX_RUNTIME_PROVIDERS_DLL_FILE, + DIRECTML_DLL_FILE, + "config.json", + "generation_config.json", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "onnx/vision_encoder_fp16.onnx", + "onnx/encoder_model_q4.onnx", + "onnx/decoder_model_q4.onnx", + "onnx/decoder_model_merged_q4.onnx", + "onnx/embed_tokens_fp16.onnx", +]; + +static ORT_RUNTIME_INIT: OnceLock> = OnceLock::new(); + +/// Set to `true` by `set_caption_acceleration` so the caption worker loop +/// knows to drop its cached `FlorenceCaptioner` and reload with the new EP. +pub static CAPTION_SESSION_DIRTY: AtomicBool = AtomicBool::new(false); + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum CaptionAcceleration { + Auto, + Cpu, + Directml, +} + +impl CaptionAcceleration { + fn as_str(self) -> &'static str { + match self { + Self::Auto => "auto", + Self::Cpu => "cpu", + Self::Directml => "directml", + } + } +} + +impl Default for CaptionAcceleration { + fn default() -> Self { + Self::Auto + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum CaptionDetail { + Short, + Detailed, + Paragraph, +} + +impl CaptionDetail { + fn as_str(self) -> &'static str { + match self { + Self::Short => "short", + Self::Detailed => "detailed", + Self::Paragraph => "paragraph", + } + } + + fn prompt(self) -> &'static str { + match self { + Self::Short => "What does the image describe?", + Self::Detailed => "Describe in detail what is shown in the image.", + Self::Paragraph => "Describe with a paragraph what is shown in the image.", + } + } + + fn max_new_tokens(self) -> usize { + match self { + Self::Short => 32, + Self::Detailed => 56, + Self::Paragraph => 96, + } + } +} + +impl Default for CaptionDetail { + fn default() -> Self { + Self::Paragraph + } +} + +#[derive(Serialize)] +pub struct CaptionModelStatus { + pub model_id: &'static str, + pub model_name: &'static str, + pub local_dir: String, + pub ready: bool, + pub missing_files: Vec, +} + +#[derive(Clone, Serialize)] +pub struct CaptionModelProgress { + pub total_files: usize, + pub completed_files: usize, + pub current_file: Option, + pub done: bool, +} + +#[derive(Serialize)] +pub struct CaptionRuntimeProbe { + pub ready: bool, + pub acceleration: CaptionAcceleration, + pub detail: CaptionDetail, + pub tokenizer_vocab_size: usize, + pub sessions: Vec, +} + +#[derive(Serialize)] +pub struct CaptionRuntimeSessionProbe { + pub file: &'static str, + pub inputs: Vec, + pub outputs: Vec, +} + +#[derive(Serialize)] +pub struct CaptionVisionProbe { + pub input_shape: Vec, + pub output_shape: Vec, + pub output_values: usize, + pub acceleration: CaptionAcceleration, +} + +#[derive(Clone)] +struct TensorData { + shape: Vec, + values: Vec, +} + +pub struct FlorenceCaptioner { + tokenizer: Tokenizer, + caption_detail: CaptionDetail, + vision_session: Session, + embed_session: Session, + encoder_session: Session, + decoder_prefill_session: Session, + decoder_session: Session, +} + +pub fn model_dir(app_data_dir: &Path) -> PathBuf { + app_data_dir.join("models").join("florence-2-base-ft") +} + +pub fn caption_acceleration(app_data_dir: &Path) -> CaptionAcceleration { + let path = app_data_dir.join(CAPTION_ACCELERATION_FILE); + let Ok(value) = std::fs::read_to_string(path) else { + return CaptionAcceleration::default(); + }; + + match value.trim().to_ascii_lowercase().as_str() { + "cpu" => CaptionAcceleration::Cpu, + "directml" => CaptionAcceleration::Directml, + _ => CaptionAcceleration::Auto, + } +} + +pub fn set_caption_acceleration( + app_data_dir: &Path, + acceleration: CaptionAcceleration, +) -> Result { + let path = app_data_dir.join(CAPTION_ACCELERATION_FILE); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, acceleration.as_str())?; + CAPTION_SESSION_DIRTY.store(true, Ordering::Relaxed); + Ok(acceleration) +} + +pub fn caption_detail(app_data_dir: &Path) -> CaptionDetail { + let path = app_data_dir.join(CAPTION_DETAIL_FILE); + let Ok(value) = std::fs::read_to_string(path) else { + return CaptionDetail::default(); + }; + + match value.trim().to_ascii_lowercase().as_str() { + "short" => CaptionDetail::Short, + "paragraph" => CaptionDetail::Paragraph, + _ => CaptionDetail::Detailed, + } +} + +pub fn set_caption_detail(app_data_dir: &Path, detail: CaptionDetail) -> Result { + let path = app_data_dir.join(CAPTION_DETAIL_FILE); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(path, detail.as_str())?; + CAPTION_SESSION_DIRTY.store(true, Ordering::Relaxed); + Ok(detail) +} + +pub fn caption_model_status(app_data_dir: &Path) -> CaptionModelStatus { + let local_dir = model_dir(app_data_dir); + let missing_files = REQUIRED_FILES + .iter() + .filter(|file| !local_dir.join(file).exists()) + .map(|file| (*file).to_string()) + .collect::>(); + + CaptionModelStatus { + model_id: FLORENCE_MODEL_ID, + model_name: FLORENCE_CAPTION_MODEL_NAME, + local_dir: local_dir.to_string_lossy().to_string(), + ready: missing_files.is_empty(), + missing_files, + } +} + +pub fn prepare_caption_model_with_progress( + app_data_dir: &Path, + emit_progress: impl Fn(CaptionModelProgress), +) -> Result { + let local_dir = model_dir(app_data_dir); + std::fs::create_dir_all(&local_dir)?; + + let api = Api::new()?; + let repo = api.repo(Repo::new(FLORENCE_MODEL_ID.to_string(), RepoType::Model)); + let mut completed_files = REQUIRED_FILES + .iter() + .filter(|file| local_dir.join(file).exists()) + .count(); + + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: None, + done: completed_files == REQUIRED_FILES.len(), + }); + + for file in REQUIRED_FILES { + let destination = local_dir.join(file); + if destination.exists() { + continue; + } + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: Some((*file).to_string()), + done: false, + }); + if let Some(parent) = destination.parent() { + std::fs::create_dir_all(parent)?; + } + if ONNX_RUNTIME_FILES + .iter() + .any(|(runtime_file, _, _)| runtime_file == file) + { + download_onnx_runtime_files(&local_dir)?; + completed_files = REQUIRED_FILES + .iter() + .filter(|file| local_dir.join(file).exists()) + .count(); + } else { + let cached = repo.get(file)?; + std::fs::copy(cached, destination)?; + completed_files += 1; + } + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: Some((*file).to_string()), + done: completed_files == REQUIRED_FILES.len(), + }); + } + + emit_progress(CaptionModelProgress { + total_files: REQUIRED_FILES.len(), + completed_files, + current_file: None, + done: true, + }); + + Ok(caption_model_status(app_data_dir)) +} + +pub fn delete_caption_model(app_data_dir: &Path) -> Result { + let local_dir = model_dir(app_data_dir); + if local_dir.exists() { + std::fs::remove_dir_all(&local_dir)?; + } + Ok(caption_model_status(app_data_dir)) +} + +pub fn probe_caption_runtime(app_data_dir: &Path) -> Result { + let status = caption_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "Florence-2 model is missing {} required file(s)", + status.missing_files.len() + ); + } + + let local_dir = model_dir(app_data_dir); + ensure_onnx_runtime(&local_dir)?; + let tokenizer = + Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + + let acceleration = caption_acceleration(app_data_dir); + + // For the vision encoder (the only session that runs on the GPU EP) we + // create an actual ORT session so we can verify which EP was actually + // loaded, rather than just echoing the settings file. + let vision_session = probe_vision_session( + &local_dir.join("onnx/vision_encoder_fp16.onnx"), + acceleration, + )?; + + // The remaining entries are DLLs and CPU-only text/decoder models — probe + // them with the lightweight file-size check as before. + let other_files = [ + "onnxruntime/onnxruntime.dll", + "onnxruntime/onnxruntime_providers_shared.dll", + "onnxruntime/DirectML.dll", + "onnx/embed_tokens_fp16.onnx", + "onnx/encoder_model_q4.onnx", + "onnx/decoder_model_q4.onnx", + "onnx/decoder_model_merged_q4.onnx", + ]; + let mut sessions = vec![vision_session]; + for file in other_files { + sessions.push(probe_session(file, &local_dir.join(file))?); + } + + Ok(CaptionRuntimeProbe { + ready: true, + acceleration, + detail: caption_detail(app_data_dir), + tokenizer_vocab_size: tokenizer.get_vocab_size(false), + sessions, + }) +} + +pub fn probe_caption_vision(app_data_dir: &Path, image_path: &Path) -> Result { + let status = caption_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "Florence-2 model is missing {} required file(s)", + status.missing_files.len() + ); + } + + let local_dir = model_dir(app_data_dir); + ensure_onnx_runtime(&local_dir)?; + let pixels = preprocess_image(image_path)?; + let input_shape = vec![1, 3, 768, 768]; + let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice())) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let acceleration = caption_acceleration(app_data_dir); + let mut session = create_session( + &local_dir.join("onnx/vision_encoder_fp16.onnx"), + acceleration, + true, + )?; + let outputs = session + .run(ort::inputs! { + "pixel_values" => input + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let (output_shape, output_values) = outputs[0] + .try_extract_tensor::() + .map_err(|error| anyhow::anyhow!("{error}"))?; + + Ok(CaptionVisionProbe { + input_shape, + output_shape: output_shape.to_vec(), + output_values: output_values.len(), + acceleration, + }) +} + +pub fn generate_caption(app_data_dir: &Path, image_path: &Path) -> Result { + let mut captioner = FlorenceCaptioner::new(app_data_dir)?; + captioner.generate(image_path) +} + +impl FlorenceCaptioner { + pub fn new(app_data_dir: &Path) -> Result { + let started_at = Instant::now(); + let status = caption_model_status(app_data_dir); + if !status.ready { + anyhow::bail!( + "Florence-2 model is missing {} required file(s)", + status.missing_files.len() + ); + } + + let local_dir = model_dir(app_data_dir); + ensure_onnx_runtime(&local_dir)?; + let tokenizer = + Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + let caption_detail = caption_detail(app_data_dir); + + let sessions_started_at = Instant::now(); + let acceleration = caption_acceleration(app_data_dir); + let vision_session = create_session( + &local_dir.join("onnx/vision_encoder_fp16.onnx"), + acceleration, + true, + )?; + let embed_session = create_session( + &local_dir.join("onnx/embed_tokens_fp16.onnx"), + acceleration, + false, + )?; + let encoder_session = create_session( + &local_dir.join("onnx/encoder_model_q4.onnx"), + acceleration, + false, + )?; + let decoder_prefill_session = create_session( + &local_dir.join("onnx/decoder_model_q4.onnx"), + acceleration, + false, + )?; + let decoder_session = create_session( + &local_dir.join("onnx/decoder_model_merged_q4.onnx"), + acceleration, + false, + )?; + println!( + "Florence sessions loaded in {:?} with {:?} acceleration (total init {:?})", + sessions_started_at.elapsed(), + acceleration, + started_at.elapsed() + ); + + Ok(Self { + tokenizer, + caption_detail, + vision_session, + embed_session, + encoder_session, + decoder_prefill_session, + decoder_session, + }) + } + + pub fn generate(&mut self, image_path: &Path) -> Result { + let started_at = Instant::now(); + println!("Florence caption started: {}", image_path.display()); + let image_features = run_vision_encoder(&mut self.vision_session, image_path)?; + println!("Florence vision encoder done in {:?}", started_at.elapsed()); + let prompt_ids = self + .tokenizer + .encode(self.caption_detail.prompt(), false) + .map_err(anyhow::Error::msg)? + .get_ids() + .iter() + .map(|id| i64::from(*id)) + .collect::>(); + let prompt_embeds = run_token_embedder(&mut self.embed_session, &prompt_ids)?; + println!( + "Florence token embeddings done in {:?}", + started_at.elapsed() + ); + let encoder_embeds = concatenate_sequence_embeddings(&image_features, &prompt_embeds)?; + let encoder_attention_mask = vec![1_i64; encoder_embeds.shape[1] as usize]; + let encoder_hidden_states = run_encoder( + &mut self.encoder_session, + &encoder_embeds, + &encoder_attention_mask, + )?; + println!("Florence encoder done in {:?}", started_at.elapsed()); + + let generated_ids = run_decoder( + &mut self.decoder_prefill_session, + &mut self.decoder_session, + &mut self.embed_session, + &encoder_hidden_states, + &encoder_attention_mask, + self.caption_detail.max_new_tokens(), + )?; + println!("Florence decoder done in {:?}", started_at.elapsed()); + + let generated_u32 = generated_ids + .into_iter() + .map(|id| id as u32) + .collect::>(); + let caption = self + .tokenizer + .decode(&generated_u32, true) + .map_err(anyhow::Error::msg)? + .trim() + .to_string(); + + if caption.is_empty() { + anyhow::bail!("Florence-2 generated an empty caption"); + } + + Ok(clean_caption(&caption)) + } +} + +fn probe_session(file: &'static str, path: &Path) -> Result { + let metadata = std::fs::metadata(path)?; + if metadata.len() == 0 { + anyhow::bail!("{} is empty", path.display()); + } + + let (inputs, outputs) = match file { + "onnx/vision_encoder_fp16.onnx" => ( + vec!["pixel_values".to_string()], + vec!["image_features".to_string()], + ), + "onnx/embed_tokens_fp16.onnx" => ( + vec!["input_ids".to_string()], + vec!["inputs_embeds".to_string()], + ), + "onnx/encoder_model_q4.onnx" => ( + vec!["attention_mask".to_string(), "inputs_embeds".to_string()], + vec!["last_hidden_state".to_string()], + ), + "onnx/decoder_model_merged_q4.onnx" => ( + vec![ + "encoder_attention_mask".to_string(), + "encoder_hidden_states".to_string(), + "inputs_embeds".to_string(), + "past_key_values".to_string(), + "use_cache_branch".to_string(), + ], + vec!["logits".to_string(), "present_key_values".to_string()], + ), + "onnxruntime/onnxruntime.dll" => ( + vec!["OrtGetApiBase".to_string()], + vec!["ORT DirectML 1.24.2".to_string()], + ), + "onnxruntime/onnxruntime_providers_shared.dll" => ( + vec!["providers".to_string()], + vec!["shared provider bridge".to_string()], + ), + "onnxruntime/DirectML.dll" => ( + vec!["DirectX 12".to_string()], + vec!["DirectML 1.15.4".to_string()], + ), + "onnx/decoder_model_q4.onnx" => ( + vec![ + "encoder_attention_mask".to_string(), + "encoder_hidden_states".to_string(), + "inputs_embeds".to_string(), + ], + vec!["logits".to_string(), "present_key_values".to_string()], + ), + _ => (Vec::new(), Vec::new()), + }; + + Ok(CaptionRuntimeSessionProbe { + file, + inputs, + outputs, + }) +} + +/// Create an actual ORT session for the vision encoder and return a probe +/// entry that reflects which execution provider was successfully loaded. +/// This is the only session that can legitimately run on DirectML; all +/// text/decoder sessions intentionally use CPU only. +fn probe_vision_session( + path: &Path, + acceleration: CaptionAcceleration, +) -> Result { + let metadata = std::fs::metadata(path)?; + if metadata.len() == 0 { + anyhow::bail!("{} is empty", path.display()); + } + + // Try to create the session with the requested acceleration. If DirectML + // was requested but fails we report what actually loaded. + let loaded_acceleration = match acceleration { + CaptionAcceleration::Cpu => { + create_session(path, CaptionAcceleration::Cpu, false)?; + CaptionAcceleration::Cpu + } + CaptionAcceleration::Auto => { + // `fail_silently` — session will fall back to CPU if DirectML + // is unavailable; we detect this by trying Directml explicitly. + let directml_ok = create_session(path, CaptionAcceleration::Directml, true).is_ok(); + if directml_ok { + CaptionAcceleration::Directml + } else { + create_session(path, CaptionAcceleration::Cpu, false)?; + CaptionAcceleration::Cpu + } + } + CaptionAcceleration::Directml => { + create_session(path, CaptionAcceleration::Directml, true)?; + CaptionAcceleration::Directml + } + }; + + Ok(CaptionRuntimeSessionProbe { + file: "onnx/vision_encoder_fp16.onnx", + inputs: vec!["pixel_values".to_string()], + outputs: vec![format!("image_features [EP: {:?}]", loaded_acceleration)], + }) +} + +pub fn ensure_onnx_runtime(local_dir: &Path) -> Result<()> { + let dll_path = local_dir.join(ONNX_RUNTIME_DLL_FILE); + ORT_RUNTIME_INIT + .get_or_init(|| { + if !dll_path.exists() { + return Err(format!( + "ONNX Runtime DLL is missing: {}", + dll_path.display() + )); + } + ort::environment::init_from(&dll_path) + .map_err(|error| error.to_string())? + .with_name("phokus-florence") + .commit(); + Ok(()) + }) + .clone() + .map_err(anyhow::Error::msg) +} + +fn download_onnx_runtime_files(local_dir: &Path) -> Result<()> { + if !cfg!(target_os = "windows") { + anyhow::bail!( + "Florence-2 ONNX Runtime download is currently configured for Windows builds" + ); + } + + for (destination_file, source_url, archive_path) in ONNX_RUNTIME_FILES { + let destination = local_dir.join(destination_file); + download_nuget_file(source_url, archive_path, &destination)?; + } + + Ok(()) +} + +fn download_nuget_file(source_url: &str, archive_path: &str, destination: &Path) -> Result<()> { + let mut response = ureq::get(source_url) + .call() + .map_err(|error| anyhow::anyhow!("{error}"))?; + let mut bytes = Vec::new(); + response + .body_mut() + .as_reader() + .read_to_end(&mut bytes) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?; + let mut dll = archive.by_name(archive_path)?; + if let Some(parent) = destination.parent() { + std::fs::create_dir_all(parent)?; + } + + let temp_destination = destination.with_extension("tmp"); + { + let mut file = std::fs::File::create(&temp_destination)?; + std::io::copy(&mut dll, &mut file)?; + } + std::fs::rename(temp_destination, destination)?; + Ok(()) +} + +fn run_vision_encoder(session: &mut Session, image_path: &Path) -> Result { + let pixels = preprocess_image(image_path)?; + let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice())) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let outputs = session + .run(ort::inputs! { + "pixel_values" => input + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + tensor_data(&outputs[0]) +} + +fn run_token_embedder(session: &mut Session, token_ids: &[i64]) -> Result { + let input_ids = Tensor::from_array(( + [1usize, token_ids.len()], + token_ids.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let outputs = session + .run(ort::inputs! { + "input_ids" => input_ids + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + tensor_data(&outputs[0]) +} + +fn run_encoder( + session: &mut Session, + inputs_embeds: &TensorData, + attention_mask: &[i64], +) -> Result { + let attention_mask_tensor = Tensor::from_array(( + [1usize, attention_mask.len()], + attention_mask.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let inputs_embeds_tensor = tensor_from_data(inputs_embeds)?; + let outputs = session + .run(ort::inputs! { + "attention_mask" => attention_mask_tensor, + "inputs_embeds" => inputs_embeds_tensor + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + tensor_data(&outputs[0]) +} + +fn run_decoder( + decoder_prefill_session: &mut Session, + decoder_session: &mut Session, + embed_session: &mut Session, + encoder_hidden_states: &TensorData, + encoder_attention_mask: &[i64], + max_new_tokens: usize, +) -> Result> { + const DECODER_LAYERS: usize = 6; + const DECODER_START_TOKEN_ID: i64 = 2; + const FORCED_BOS_TOKEN_ID: i64 = 0; + const EOS_TOKEN_ID: i64 = 2; + + let mut generated = Vec::new(); + let encoder_attention_mask_tensor = Tensor::from_array(( + [1usize, encoder_attention_mask.len()], + encoder_attention_mask.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let encoder_hidden_states_tensor = tensor_from_data(encoder_hidden_states)?; + let prefill_inputs_embeds = run_token_embedder(embed_session, &[DECODER_START_TOKEN_ID])?; + let prefill_inputs_embeds_tensor = tensor_from_data(&prefill_inputs_embeds)?; + let prefill_started_at = Instant::now(); + let prefill_outputs = decoder_prefill_session + .run(ort::inputs! { + "encoder_attention_mask" => encoder_attention_mask_tensor, + "encoder_hidden_states" => encoder_hidden_states_tensor, + "inputs_embeds" => prefill_inputs_embeds_tensor + }) + .map_err(|error| anyhow::anyhow!("{error}"))?; + println!( + "Florence decoder prefill done in {:?}", + prefill_started_at.elapsed() + ); + let _prefill_logits = tensor_data(&prefill_outputs[0])?; + let mut next_input_id = FORCED_BOS_TOKEN_ID; + let encoder_kv = collect_present_key_values(&prefill_outputs, DECODER_LAYERS)?; + let mut decoder_kv = encoder_kv.clone(); + + for _ in 0..max_new_tokens { + if next_input_id == EOS_TOKEN_ID { + break; + } + generated.push(next_input_id); + + let inputs_embeds = run_token_embedder(embed_session, &[next_input_id])?; + let encoder_attention_mask_tensor = Tensor::from_array(( + [1usize, encoder_attention_mask.len()], + encoder_attention_mask.to_vec().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let encoder_hidden_states_tensor = tensor_from_data(encoder_hidden_states)?; + let inputs_embeds_tensor = tensor_from_data(&inputs_embeds)?; + let use_cache_branch_tensor = Tensor::from_array(([1usize], vec![true].into_boxed_slice())) + .map_err(|error| anyhow::anyhow!("{error}"))?; + + let mut inputs = ort::inputs! { + "encoder_attention_mask" => encoder_attention_mask_tensor, + "encoder_hidden_states" => encoder_hidden_states_tensor, + "inputs_embeds" => inputs_embeds_tensor, + "use_cache_branch" => use_cache_branch_tensor + }; + + for layer in 0..DECODER_LAYERS { + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.decoder.key"), + decoder_kv[layer * 4].clone(), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.decoder.value"), + decoder_kv[layer * 4 + 1].clone(), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.encoder.key"), + encoder_kv[layer * 4 + 2].clone(), + )?; + push_tensor_input( + &mut inputs, + format!("past_key_values.{layer}.encoder.value"), + encoder_kv[layer * 4 + 3].clone(), + )?; + } + + let outputs = decoder_session + .run(inputs) + .map_err(|error| anyhow::anyhow!("{error}"))?; + let logits = tensor_data(&outputs["logits"])?; + next_input_id = argmax_last_token(&logits)?; + decoder_kv = collect_present_key_values(&outputs, DECODER_LAYERS)?; + } + + println!("Florence decoder produced {} token(s)", generated.len()); + Ok(generated) +} + +fn create_session( + path: &Path, + acceleration: CaptionAcceleration, + allow_directml: bool, +) -> Result { + 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 = allow_directml + && matches!( + acceleration, + CaptionAcceleration::Auto | CaptionAcceleration::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 { + CaptionAcceleration::Cpu => builder, + CaptionAcceleration::Auto if use_directml => builder + .with_execution_providers([ep::DirectML::default().build().fail_silently()]) + .unwrap_or_else(|error| error.recover()), + CaptionAcceleration::Directml if use_directml => builder + .with_execution_providers([ep::DirectML::default().build().error_on_failure()]) + .map_err(|error| anyhow::anyhow!("{error}"))?, + CaptionAcceleration::Auto | CaptionAcceleration::Directml => { + // `allow_directml` is false for the 4 text/decoder sessions — + // they intentionally run on CPU regardless of the setting. + println!( + "Florence: using CPU for {} (DirectML disabled for this session type)", + path.display() + ); + builder + } + }; + let session = builder + .commit_from_file(path) + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok(session) +} + +fn concatenate_sequence_embeddings(text: &TensorData, image: &TensorData) -> Result { + if text.shape.len() != 3 || image.shape.len() != 3 { + anyhow::bail!("Expected 3D text and image embeddings"); + } + if text.shape[0] != image.shape[0] || text.shape[2] != image.shape[2] { + anyhow::bail!("Text and image embedding dimensions do not match"); + } + + let text_tokens = text.shape[1] as usize; + let image_tokens = image.shape[1] as usize; + let dim = text.shape[2] as usize; + let mut values = Vec::with_capacity((text_tokens + image_tokens) * dim); + values.extend_from_slice(&text.values); + values.extend_from_slice(&image.values); + + Ok(TensorData { + shape: vec![text.shape[0], text.shape[1] + image.shape[1], text.shape[2]], + values, + }) +} + +fn collect_present_key_values( + outputs: &SessionOutputs<'_>, + layers: usize, +) -> Result> { + let expected = layers * 4; + if outputs.len() < expected + 1 { + anyhow::bail!( + "Decoder returned {} output(s), expected at least {}", + outputs.len(), + expected + 1 + ); + } + + (1..=expected) + .map(|index| tensor_data(&outputs[index])) + .collect::>>() +} + +fn tensor_data(value: &ort::value::DynValue) -> Result { + let (shape, values) = value + .try_extract_tensor::() + .map_err(|error| anyhow::anyhow!("{error}"))?; + Ok(TensorData { + shape: shape.to_vec(), + values: values.to_vec(), + }) +} + +fn tensor_from_data(data: &TensorData) -> Result> { + Tensor::from_array(( + Shape::new(data.shape.clone()), + data.values.clone().into_boxed_slice(), + )) + .map_err(|error| anyhow::anyhow!("{error}")) +} + +fn push_tensor_input( + inputs: &mut Vec<(Cow<'_, str>, SessionInputValue<'_>)>, + name: String, + data: TensorData, +) -> Result<()> { + inputs.push((Cow::Owned(name), tensor_from_data(&data)?.into())); + Ok(()) +} + +fn argmax_last_token(logits: &TensorData) -> Result { + if logits.shape.len() != 3 { + anyhow::bail!("Expected decoder logits to be 3D"); + } + let sequence_length = logits.shape[1] as usize; + let vocab_size = logits.shape[2] as usize; + if sequence_length == 0 || vocab_size == 0 { + anyhow::bail!("Decoder logits are empty"); + } + let start = (sequence_length - 1) * vocab_size; + let (index, _) = logits.values[start..start + vocab_size] + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.total_cmp(b)) + .ok_or_else(|| anyhow::anyhow!("Decoder logits are empty"))?; + Ok(index as i64) +} + +fn clean_caption(caption: &str) -> String { + caption + .trim() + .trim_matches(|ch| matches!(ch, '<' | '>' | '|' | ' ')) + .to_string() +} + +fn preprocess_image(image_path: &Path) -> Result> { + let image = ImageReader::open(image_path)?.decode()?.to_rgb8(); + let resized = image::imageops::resize(&image, 768, 768, FilterType::CatmullRom); + let mut pixel_values = vec![0.0f32; 3 * 768 * 768]; + let mean = [0.485f32, 0.456, 0.406]; + let std = [0.229f32, 0.224, 0.225]; + + for (x, y, pixel) in resized.enumerate_pixels() { + let x = x as usize; + let y = y as usize; + let base = y * 768 + x; + for channel in 0..3 { + let value = f32::from(pixel[channel]) / 255.0; + pixel_values[channel * 768 * 768 + base] = (value - mean[channel]) / std[channel]; + } + } + + Ok(pixel_values) +} diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index fb16359..fbeef49 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -1,10 +1,16 @@ -use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord}; +use crate::captioner::{ + self, CaptionAcceleration, CaptionDetail, CaptionModelStatus, CaptionRuntimeProbe, + CaptionVisionProbe, +}; +use crate::db::{self, DbPool, ExploreTagEntry, Folder, FolderJobProgress, ImageRecord, ImageTag}; use crate::embedder; +use crate::hnsw_index; use crate::indexer; +use crate::tagger::{self, TaggerAcceleration, TaggerModelStatus, TaggerRuntimeProbe}; use crate::vector; use serde::{Deserialize, Serialize}; use std::path::PathBuf; -use tauri::{AppHandle, State}; +use tauri::{AppHandle, Emitter, Manager, State}; pub type DbState = DbPool; @@ -16,12 +22,21 @@ pub struct ImagesPage { pub limit: i64, } +#[derive(Serialize)] +pub struct SimilarImagesPage { + pub images: Vec, + pub offset: usize, + pub limit: usize, + pub has_more: bool, +} + #[derive(Deserialize)] pub struct GetImagesParams { pub folder_id: Option, pub search: Option, pub media_kind: Option, pub favorites_only: Option, + pub rating_min: Option, pub embedding_failed_only: Option, pub sort: Option, pub offset: Option, @@ -38,6 +53,29 @@ pub struct UpdateImageDetailsParams { #[derive(Deserialize)] pub struct FindSimilarImagesParams { pub image_id: i64, + pub folder_id: Option, + pub offset: Option, + pub limit: Option, + pub threshold: Option, +} + +#[derive(Deserialize)] +pub struct FindSimilarByRegionParams { + pub image_id: i64, + /// Normalized crop rect (0.0–1.0). + pub crop_x: f32, + pub crop_y: f32, + pub crop_w: f32, + pub crop_h: f32, + pub folder_id: Option, + pub offset: Option, + pub limit: Option, +} + +#[derive(Deserialize)] +pub struct DebugSimilarImagesParams { + pub image_id: i64, + pub folder_id: Option, pub limit: Option, } @@ -46,15 +84,120 @@ pub struct RetryFailedEmbeddingsParams { pub folder_id: i64, } +#[derive(Deserialize)] +pub struct SetGeneratedCaptionParams { + pub image_id: i64, + pub caption: String, + pub model: Option, +} + +#[derive(Deserialize)] +pub struct SuggestImageTagsParams { + pub image_id: i64, + pub limit: Option, +} + +#[derive(Deserialize)] +pub struct QueueCaptionJobsParams { + pub folder_id: Option, + pub image_id: Option, +} + +#[derive(Deserialize)] +pub struct ClearCaptionJobsParams { + pub folder_id: Option, +} + +#[derive(Deserialize)] +pub struct ResetCaptionsParams { + pub folder_id: Option, +} + +#[derive(Deserialize)] +pub struct SetCaptionAccelerationParams { + pub acceleration: CaptionAcceleration, +} + +#[derive(Deserialize)] +pub struct SetCaptionDetailParams { + pub detail: CaptionDetail, +} + +#[derive(Deserialize)] +pub struct ProbeCaptionImageParams { + pub image_id: i64, +} + +#[derive(Deserialize)] +pub struct GenerateCaptionParams { + pub image_id: i64, +} + #[derive(Deserialize)] pub struct SemanticSearchParams { pub query: String, pub folder_id: Option, pub media_kind: Option, pub favorites_only: Option, + pub rating_min: Option, pub limit: Option, } +#[derive(Deserialize)] +pub struct TagSearchParams { + pub query: String, + pub folder_id: Option, + pub media_kind: Option, + pub favorites_only: Option, + pub rating_min: Option, + pub limit: Option, + pub offset: Option, +} + +#[derive(Serialize)] +pub struct TagSearchPage { + pub images: Vec, + pub total: usize, + pub offset: usize, + pub limit: usize, +} + +#[derive(Deserialize)] +pub struct GetExploreTagsParams { + pub folder_id: Option, + pub limit: Option, +} + +#[derive(Deserialize)] +pub struct SearchTagsAutocompleteParams { + pub query: String, + pub folder_id: Option, + pub limit: Option, +} + +#[derive(Serialize, Deserialize)] +pub struct DuplicateGroup { + pub file_hash: String, + pub file_size: u64, + pub images: Vec, +} + +#[derive(Serialize)] +pub struct DuplicateScanCache { + pub groups: Vec, + pub scanned_at: i64, +} + +#[derive(Deserialize)] +pub struct DeleteImagesFromDiskParams { + pub image_ids: Vec, +} + +#[derive(Deserialize)] +pub struct GetImagesByIdsParams { + pub image_ids: Vec, +} + #[tauri::command] pub async fn add_folder( app: AppHandle, @@ -125,10 +268,19 @@ pub async fn get_images( let search = params.search.as_deref(); let media_kind = params.media_kind.as_deref(); let favorites_only = params.favorites_only.unwrap_or(false); + let rating_min = params.rating_min.unwrap_or(0); let embedding_failed_only = params.embedding_failed_only.unwrap_or(false); - let total = db::count_images(&conn, params.folder_id, search, media_kind, favorites_only, embedding_failed_only) - .map_err(|e| e.to_string())?; + let total = db::count_images( + &conn, + params.folder_id, + search, + media_kind, + favorites_only, + rating_min, + embedding_failed_only, + ) + .map_err(|e| e.to_string())?; let images = db::get_images( &conn, @@ -136,6 +288,7 @@ pub async fn get_images( search, media_kind, favorites_only, + rating_min, embedding_failed_only, sort, offset, @@ -181,16 +334,183 @@ pub async fn reindex_folder( Ok(()) } +#[tauri::command] +pub async fn rename_folder( + db: State<'_, DbState>, + folder_id: i64, + new_name: String, +) -> Result<(), String> { + let new_name = new_name.trim().to_string(); + if new_name.is_empty() { + return Err("Folder name cannot be empty".to_string()); + } + let conn = db.get().map_err(|e| e.to_string())?; + db::rename_folder(&conn, folder_id, &new_name).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn update_folder_path( + app: AppHandle, + db: State<'_, DbState>, + folder_id: i64, + new_path: String, +) -> Result<(), String> { + let new_path_buf = PathBuf::from(&new_path); + if !new_path_buf.is_dir() { + return Err(format!("Path is not a valid directory: {}", new_path)); + } + let new_name = new_path_buf + .file_name() + .map(|n| n.to_string_lossy().to_string()) + .unwrap_or_else(|| new_path.clone()); + { + let conn = db.get().map_err(|e| e.to_string())?; + // Fetch the old path before updating so image paths can be rewritten. + let old_path = db::get_folders(&conn) + .map_err(|e| e.to_string())? + .into_iter() + .find(|f| f.id == folder_id) + .map(|f| f.path) + .ok_or("Folder not found")?; + db::update_folder_path(&conn, folder_id, &old_path, &new_path, &new_name) + .map_err(|e| e.to_string())?; + } + indexer::index_folder(app, db.inner().clone(), folder_id, new_path_buf); + Ok(()) +} + #[tauri::command] pub async fn find_similar_images( db: State<'_, DbState>, params: FindSimilarImagesParams, -) -> Result, String> { +) -> Result { let conn = db.get().map_err(|e| e.to_string())?; let limit = params.limit.unwrap_or(32); - let image_ids = vector::find_similar_image_ids(&conn, params.image_id, limit) + let offset = params.offset.unwrap_or(0); + let threshold = params.threshold.unwrap_or(0.24); + if !vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())? { + db::repair_embedding_consistency(&conn).map_err(|e| e.to_string())?; + return Ok(SimilarImagesPage { + images: Vec::new(), + offset, + limit, + has_more: false, + }); + } + let matches = hnsw_index::find_similar_image_matches( + &conn, + params.image_id, + params.folder_id, + threshold, + offset, + limit + 1, + ) .map_err(|e| e.to_string())?; - db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string()) + let has_more = matches.len() > limit; + let image_ids = matches + .into_iter() + .take(limit) + .map(|(image_id, _)| image_id) + .collect::>(); + let images = db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string())?; + Ok(SimilarImagesPage { + images, + offset, + limit, + has_more, + }) +} + +#[tauri::command] +pub async fn find_similar_by_region( + db: State<'_, DbState>, + params: FindSimilarByRegionParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let limit = params.limit.unwrap_or(32); + let offset = params.offset.unwrap_or(0); + + // Look up the source image path + let image = db::get_image_by_id(&conn, params.image_id).map_err(|e| e.to_string())?; + let image_path = std::path::Path::new(&image.path); + + // Embed the cropped region in-memory (no temp file needed) + let embedder = embedder::ClipImageEmbedder::new().map_err(|e| e.to_string())?; + let embedding = embedder + .embed_image_crop( + image_path, + params.crop_x, + params.crop_y, + params.crop_w, + params.crop_h, + ) + .map_err(|e| e.to_string())?; + + // Search for similar images using the crop embedding + let image_ids = match params.folder_id { + Some(folder_id) => vector::search_image_ids_by_embedding_in_folder( + &conn, + &embedding, + folder_id, + Some(params.image_id), + offset + limit + 1, + ) + .map_err(|e| e.to_string())?, + None => { + // Fetch one extra candidate to compensate for the source image that + // will be removed, so has_more is accurate and results span multiple pages. + let mut ids = vector::search_image_ids_by_embedding(&conn, &embedding, offset + limit + 2) + .map_err(|e| e.to_string())?; + ids.retain(|&id| id != params.image_id); + ids + } + }; + + let has_more = image_ids.len() > offset + limit; + let page_ids = image_ids + .into_iter() + .skip(offset) + .take(limit) + .collect::>(); + + let images = db::get_images_by_ids(&conn, &page_ids).map_err(|e| e.to_string())?; + Ok(SimilarImagesPage { + images, + offset, + limit, + has_more, + }) +} + +#[derive(Serialize)] +pub struct SimilarImagesDebug { + pub image_id: i64, + pub vector_count: i64, + pub has_vector: bool, + pub similar_ids: Vec, +} + +#[tauri::command] +pub async fn debug_similar_images( + db: State<'_, DbState>, + params: DebugSimilarImagesParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let limit = params.limit.unwrap_or(32); + let vector_count = vector::count_image_vectors(&conn).map_err(|e| e.to_string())?; + let has_vector = vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())?; + let similar_ids = if has_vector { + vector::find_similar_image_ids(&conn, params.image_id, limit, params.folder_id) + .map_err(|e| e.to_string())? + } else { + Vec::new() + }; + Ok(SimilarImagesDebug { + image_id: params.image_id, + vector_count, + has_vector, + similar_ids, + }) } #[tauri::command] @@ -211,20 +531,259 @@ pub async fn semantic_search_images( let conn = db.get().map_err(|e| e.to_string())?; let limit = params.limit.unwrap_or(64); - let ids = vector::search_image_ids_by_embedding(&conn, &embedding, limit).map_err(|e| e.to_string())?; - let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?; - if let Some(folder_id) = params.folder_id { - images.retain(|image| image.folder_id == folder_id); - } - if let Some(media_kind) = params.media_kind.as_deref() { - images.retain(|image| image.media_kind == media_kind); - } - if params.favorites_only.unwrap_or(false) { - images.retain(|image| image.favorite); + let has_filters = params.folder_id.is_some() + || params.media_kind.is_some() + || params.favorites_only.unwrap_or(false) + || params.rating_min.map_or(false, |r| r > 0); + + if !has_filters { + // No post-query filtering — a single fetch of exactly `limit` results is optimal. + let ids = vector::search_image_ids_by_embedding(&conn, &embedding, limit) + .map_err(|e| e.to_string())?; + return db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string()); } - Ok(images) + // Post-query filters are active. Progressively double the fetch batch until we + // collect enough results or exhaust the vector table, so we never over-fetch + // more than necessary while still filling the requested page. + let mut batch = limit; + loop { + let ids = vector::search_image_ids_by_embedding(&conn, &embedding, batch) + .map_err(|e| e.to_string())?; + let exhausted = ids.len() < batch; + let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?; + + if let Some(folder_id) = params.folder_id { + images.retain(|image| image.folder_id == folder_id); + } + if let Some(media_kind) = params.media_kind.as_deref() { + images.retain(|image| image.media_kind == media_kind); + } + if params.favorites_only.unwrap_or(false) { + images.retain(|image| image.favorite); + } + if let Some(rating_min) = params.rating_min { + images.retain(|image| image.rating >= rating_min); + } + + if images.len() >= limit || exhausted { + images.truncate(limit); + return Ok(images); + } + // If we are already at the fetch cap, another iteration would fetch the + // exact same IDs — break out rather than looping forever. + if batch >= 8192 { + return Ok(images); + } + batch = (batch * 2).min(8192); + } +} + +#[tauri::command] +pub async fn search_images_by_tag( + db: State<'_, DbState>, + params: TagSearchParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let limit = params.limit.unwrap_or(64); + let offset = params.offset.unwrap_or(0); + let (images, total) = db::search_images_by_tag( + &conn, + ¶ms.query, + params.folder_id, + params.media_kind.as_deref(), + params.favorites_only.unwrap_or(false), + params.rating_min.unwrap_or(0), + limit, + offset, + ) + .map_err(|e| e.to_string())?; + Ok(TagSearchPage { images, total, offset, limit }) +} + +#[tauri::command] +pub async fn set_generated_caption( + db: State<'_, DbState>, + params: SetGeneratedCaptionParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let model = params.model.as_deref().unwrap_or("manual"); + db::update_generated_caption(&conn, params.image_id, ¶ms.caption, model) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn suggest_image_tags( + db: State<'_, DbState>, + params: SuggestImageTagsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::suggest_tags_from_caption(&conn, params.image_id, params.limit.unwrap_or(2)) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn get_caption_model_status(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(captioner::caption_model_status(&app_dir)) +} + +#[tauri::command] +pub async fn get_caption_acceleration(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(captioner::caption_acceleration(&app_dir)) +} + +#[tauri::command] +pub async fn set_caption_acceleration( + app: AppHandle, + params: SetCaptionAccelerationParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + captioner::set_caption_acceleration(&app_dir, params.acceleration).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn get_caption_detail(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(captioner::caption_detail(&app_dir)) +} + +#[tauri::command] +pub async fn set_caption_detail( + app: AppHandle, + params: SetCaptionDetailParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + captioner::set_caption_detail(&app_dir, params.detail).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn prepare_caption_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || { + let app = app.clone(); + captioner::prepare_caption_model_with_progress(&app_dir, move |progress| { + let _ = app.emit("caption-model-progress", progress); + }) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn delete_caption_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || captioner::delete_caption_model(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn probe_caption_runtime(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || captioner::probe_caption_runtime(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn probe_caption_image( + app: AppHandle, + db: State<'_, DbState>, + params: ProbeCaptionImageParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + let image_path = { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_image_by_id(&conn, params.image_id) + .map(|image| image.path) + .map_err(|e| e.to_string())? + }; + tauri::async_runtime::spawn_blocking(move || { + captioner::probe_caption_vision(&app_dir, std::path::Path::new(&image_path)) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn generate_caption_for_image( + app: AppHandle, + db: State<'_, DbState>, + params: GenerateCaptionParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + let image_path = { + let conn = db.get().map_err(|e| e.to_string())?; + let image = db::get_image_by_id(&conn, params.image_id).map_err(|e| e.to_string())?; + if image.media_kind != "image" { + return Err("AI captions can only be generated for images".to_string()); + } + image.path + }; + + let caption = tauri::async_runtime::spawn_blocking(move || { + captioner::generate_caption(&app_dir, std::path::Path::new(&image_path)) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|error| { + if let Ok(conn) = db.get() { + let _ = db::mark_caption_failed(&conn, params.image_id, &error.to_string()); + } + error.to_string() + })?; + + let conn = db.get().map_err(|e| e.to_string())?; + db::update_generated_caption( + &conn, + params.image_id, + &caption, + captioner::FLORENCE_CAPTION_MODEL_NAME, + ) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn queue_caption_jobs( + db: State<'_, DbState>, + params: QueueCaptionJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + match (params.folder_id, params.image_id) { + (_, Some(image_id)) => { + db::enqueue_caption_job(&conn, image_id).map_err(|e| e.to_string())?; + Ok(1) + } + (Some(folder_id), None) => { + db::enqueue_missing_caption_jobs_for_folder(&conn, folder_id).map_err(|e| e.to_string()) + } + (None, None) => db::enqueue_missing_caption_jobs(&conn).map_err(|e| e.to_string()), + } +} + +#[tauri::command] +pub async fn clear_caption_jobs( + db: State<'_, DbState>, + params: ClearCaptionJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + db::clear_caption_jobs(&conn, params.folder_id).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn reset_generated_captions( + db: State<'_, DbState>, + params: ResetCaptionsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + db::reset_generated_captions(&conn, params.folder_id).map_err(|e| e.to_string()) } #[derive(Serialize, Deserialize)] @@ -232,17 +791,10 @@ pub struct TagCloudEntry { pub count: usize, pub representative_image_id: i64, pub thumbnail_path: Option, + #[serde(default)] + pub image_ids: Vec, } -fn fnv_hash_ids(ids: &[i64]) -> u64 { - let mut h: u64 = 0xcbf29ce484222325; - for &id in ids { - for b in id.to_le_bytes() { - h = h.wrapping_mul(0x100000001b3) ^ (b as u64); - } - } - h -} /// Clusters the library's image embeddings with k-means and returns one representative /// image per cluster — the member whose embedding is closest to its cluster centroid. @@ -263,11 +815,22 @@ pub async fn get_tag_cloud( return Ok(vec![]); } - // Compute a hash of the current embedded image IDs (sorted for stability) - let mut sorted_ids: Vec = embeddings_with_ids.iter().map(|(id, _)| *id).collect(); - sorted_ids.sort_unstable(); - let current_hash = fnv_hash_ids(&sorted_ids); - + // Sort by ID for stable ordering; hash both IDs and embedding bytes so that + // replacing a file and regenerating its embedding invalidates the cache even + // when the set of image IDs hasn't changed. + let mut sorted_pairs: Vec<_> = embeddings_with_ids.iter().collect(); + sorted_pairs.sort_unstable_by_key(|(id, _)| *id); + let current_hash = { + use xxhash_rust::xxh3::Xxh3; + let mut hasher = Xxh3::new(); + for (id, embedding) in &sorted_pairs { + hasher.update(&id.to_le_bytes()); + for val in embedding.iter() { + hasher.update(&val.to_le_bytes()); + } + } + hasher.digest() + }; let folder_scope = match folder_id { Some(id) => format!("folder_{}", id), None => "all".to_string(), @@ -280,14 +843,21 @@ pub async fn get_tag_cloud( .map_err(|e| e.to_string())? { if let Ok(entries) = serde_json::from_str::>(&json) { - return Ok(entries); + // Reject cache entries written before image_ids were tracked — they all + // have empty image_ids which causes "No media found" when a cluster is opened. + if entries.iter().all(|e| !e.image_ids.is_empty()) { + return Ok(entries); + } } } } // Cache miss — run k-means let ids: Vec = embeddings_with_ids.iter().map(|(id, _)| *id).collect(); - let points: Vec> = embeddings_with_ids.into_iter().map(|(_, emb)| emb).collect(); + let points: Vec> = embeddings_with_ids + .into_iter() + .map(|(_, emb)| emb) + .collect(); let k = (n / 20).clamp(5, 30); let (centroids, cluster_counts, assignments) = kmeans_cosine(&points, k, 40); @@ -305,6 +875,13 @@ pub async fn get_tag_cloud( } let centroid = ¢roids[ci]; + let cluster_ids = points + .iter() + .enumerate() + .filter(|(i, _)| assignments[*i] == ci) + .map(|(i, _)| ids[i]) + .collect::>(); + let best_id = points .iter() .enumerate() @@ -322,6 +899,7 @@ pub async fn get_tag_cloud( count, representative_image_id: best_id, thumbnail_path, + image_ids: cluster_ids, }); } @@ -333,6 +911,286 @@ pub async fn get_tag_cloud( Ok(entries) } +#[tauri::command] +pub async fn get_explore_tags( + db: State<'_, DbState>, + params: GetExploreTagsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_explore_tags(&conn, params.folder_id, params.limit.unwrap_or(48)).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn search_tags_autocomplete( + db: State<'_, DbState>, + params: SearchTagsAutocompleteParams, +) -> Result, String> { + if params.query.trim().is_empty() { + return Ok(vec![]); + } + let conn = db.get().map_err(|e| e.to_string())?; + db::search_tags_autocomplete(&conn, ¶ms.query, params.folder_id, params.limit.unwrap_or(10)) + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn find_duplicates( + app: AppHandle, + db: State<'_, DbState>, + folder_id: Option, +) -> Result, String> { + let records = { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_all_image_paths(&conn, folder_id).map_err(|e| e.to_string())? + }; + + let total = records.len(); + let _ = app.emit("duplicate_scan_progress", (0usize, total)); + + // Three-phase detection. + // + // Phase 1 — stat: + // Read only file metadata (size). Files with a unique size cannot be + // duplicates; discard them immediately. Zero file content read. + // + // Phase 2 — sample hash: + // For each size-matched candidate, read four evenly-spaced 16 KB windows + // and hash them. For small files (≤ 64 KB) the windows cover the whole + // file so the hash is exact. For large files this is a fast shortlist. + // + // Phase 3 — full hash (large files only): + // For sample-hash groups that contain files > 64 KB, compute a full-file + // hash to confirm the match before presenting them as deletable duplicates. + let app_hash = app.clone(); + let pairs: Vec<(u64, u64, i64, String)> = tokio::task::spawn_blocking(move || { + use memmap2::Mmap; + use rayon::prelude::*; + use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; + use xxhash_rust::xxh3::{xxh3_64, Xxh3}; + + const WINDOW: usize = 16 * 1024; // 16 KB per sample window + const N: usize = 4; // windows at 0%, 33%, 66%, 100% + const COVERED: usize = WINDOW * N; // 64 KB total; also full-coverage threshold + + // Hash four evenly-spaced 16 KB windows. For files ≤ 64 KB this is + // equivalent to hashing the entire file. + fn sample(mmap: &[u8]) -> u64 { + if mmap.len() <= COVERED { + return xxh3_64(mmap); + } + let mut h = Xxh3::new(); + for i in 0..N { + let pos = (mmap.len() - WINDOW) * i / (N - 1); + h.update(&mmap[pos..pos + WINDOW]); + } + h.digest() + } + + // ── Phase 1: stat ───────────────────────────────────────────────── + let sized: Vec<(i64, String, u64)> = records + .par_iter() + .filter_map(|r| { + let size = std::fs::metadata(&r.path).ok()?.len(); + if size == 0 { return None; } + Some((r.id, r.path.clone(), size)) + }) + .collect(); + + let mut by_size: HashMap> = HashMap::new(); + for (i, (_, _, size)) in sized.iter().enumerate() { + by_size.entry(*size).or_default().push(i); + } + let candidates: Vec = by_size + .into_values() + .filter(|g| g.len() > 1) + .flatten() + .collect(); + + if candidates.is_empty() { + return vec![]; + } + + // ── Phase 2: sample hash ────────────────────────────────────────── + let c_total = candidates.len(); + let _ = app_hash.emit("duplicate_scan_progress", (0usize, c_total)); + let counter = AtomicUsize::new(0); + + candidates + .par_iter() + .filter_map(|&idx| { + let (id, path, size) = &sized[idx]; + let file = std::fs::File::open(path).ok()?; + // SAFETY: read-only; no external truncation expected during scan. + let mmap = unsafe { Mmap::map(&file).ok()? }; + let hash = sample(&mmap); + let done = counter.fetch_add(1, Ordering::Relaxed) + 1; + if done % 100 == 0 || done == c_total { + let _ = app_hash.emit("duplicate_scan_progress", (done, c_total)); + } + Some((hash, *size, *id, path.clone())) + }) + .collect() + }) + .await + .map_err(|e| e.to_string())?; + + // Group by (sample_hash, file_size). + let mut size_hash_map: std::collections::HashMap<(u64, u64), Vec<(i64, String)>> = + std::collections::HashMap::new(); + for (hash, file_size, id, path) in pairs { + size_hash_map.entry((hash, file_size)).or_default().push((id, path)); + } + + // Phase 3: for large-file groups (> 64 KB) the sample hash is not exact; + // re-hash the full file contents to confirm matches. Each distinct full + // hash becomes its own DuplicateGroup so disjoint sets (A,A vs B,B that + // happened to share size and sample hash) are never merged. + const COVERED: u64 = (16 * 1024 * 4) as u64; + let confirmed: Vec<(u64, u64, Vec)> = + tokio::task::spawn_blocking(move || { + use memmap2::Mmap; + use rayon::prelude::*; + use xxhash_rust::xxh3::xxh3_64; + + size_hash_map + .into_iter() + .filter(|(_, entries)| entries.len() > 1) + .flat_map(|((sample_hash, file_size), entries)| { + if file_size <= COVERED { + // Sample was already a full hash — one group, no re-read needed. + return vec![(sample_hash, file_size, entries.into_iter().map(|(id, _)| id).collect())]; + } + // Full-file hash pass. Group by full hash so that two sets of + // files that collide only in the sample produce separate groups. + let full_hashes: Vec<(i64, Option)> = entries + .par_iter() + .map(|(id, path)| { + let hash = std::fs::File::open(path) + .ok() + .and_then(|f| unsafe { Mmap::map(&f).ok() }) + .map(|mmap| xxh3_64(&mmap)); + (*id, hash) + }) + .collect(); + let mut by_full: std::collections::HashMap> = + std::collections::HashMap::new(); + for (id, maybe_hash) in full_hashes { + if let Some(h) = maybe_hash { + by_full.entry(h).or_default().push(id); + } + } + // Emit one entry per distinct full hash that has ≥ 2 members. + by_full + .into_iter() + .filter(|(_, ids)| ids.len() > 1) + .map(|(full_hash, ids)| (full_hash, file_size, ids)) + .collect() + }) + .collect() + }) + .await + .map_err(|e| e.to_string())?; + + // Resolve image records for each duplicate group + let conn = db.get().map_err(|e| e.to_string())?; + let mut groups: Vec = confirmed + .into_iter() + .filter_map(|(hash, file_size, ids)| { + let images = db::get_images_by_ids(&conn, &ids).ok()?; + Some(DuplicateGroup { + file_hash: format!("{:016x}", hash), + file_size, + images, + }) + }) + .collect(); + + // Largest duplicates first — wastes the most space + groups.sort_by(|a, b| b.file_size.cmp(&a.file_size)); + + // Persist results so they survive restart — best-effort, ignore errors. + let folder_scope = match folder_id { + Some(id) => format!("folder:{}", id), + None => "all".to_string(), + }; + if let Ok(json) = serde_json::to_string(&groups) { + let _ = db::set_duplicate_scan_cache(&conn, &folder_scope, &json); + } + + Ok(groups) +} + +#[tauri::command] +pub async fn load_duplicate_scan_cache( + db: State<'_, DbState>, + folder_id: Option, +) -> Result, String> { + let folder_scope = match folder_id { + Some(id) => format!("folder:{}", id), + None => "all".to_string(), + }; + let conn = db.get().map_err(|e| e.to_string())?; + match db::get_duplicate_scan_cache(&conn, &folder_scope).map_err(|e| e.to_string())? { + Some((json, scanned_at)) => { + let groups: Vec = + serde_json::from_str(&json).map_err(|e| e.to_string())?; + Ok(Some(DuplicateScanCache { groups, scanned_at })) + } + None => Ok(None), + } +} + +#[tauri::command] +pub async fn invalidate_duplicate_scan_cache( + db: State<'_, DbState>, + folder_id: Option, +) -> Result<(), String> { + let folder_scope = match folder_id { + Some(id) => format!("folder:{}", id), + None => "all".to_string(), + }; + let conn = db.get().map_err(|e| e.to_string())?; + db::clear_duplicate_scan_cache(&conn, &folder_scope).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn delete_images_from_disk( + db: State<'_, DbState>, + params: DeleteImagesFromDiskParams, +) -> Result, String> { + if params.image_ids.is_empty() { + return Ok(Vec::new()); + } + let conn = db.get().map_err(|e| e.to_string())?; + let records = db::get_all_image_paths(&conn, None).map_err(|e| e.to_string())?; + let id_set: std::collections::HashSet = params.image_ids.iter().copied().collect(); + + // Attempt filesystem removal first; only delete DB rows for files that + // were successfully removed. This prevents entries from disappearing from + // the library when a file is locked or permission-denied. + let mut succeeded_ids: Vec = Vec::new(); + for r in records.into_iter().filter(|r| id_set.contains(&r.id)) { + if std::fs::remove_file(&r.path).is_ok() { + succeeded_ids.push(r.id); + } + } + if !succeeded_ids.is_empty() { + db::delete_images_by_ids(&conn, &succeeded_ids).map_err(|e| e.to_string())?; + } + // Return the IDs that were actually removed so the caller can update state precisely. + Ok(succeeded_ids) +} + +#[tauri::command] +pub async fn get_images_by_ids( + db: State<'_, DbState>, + params: GetImagesByIdsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_images_by_ids(&conn, ¶ms.image_ids).map_err(|e| e.to_string()) +} + // ── k-means with cosine similarity (all vectors assumed to be unit-normalized) ── fn dot(a: &[f32], b: &[f32]) -> f32 { @@ -361,7 +1219,10 @@ fn kmeans_cosine( let next = points .iter() .map(|p| { - let best_sim = centroids.iter().map(|c| dot(p, c)).fold(f32::NEG_INFINITY, f32::max); + let best_sim = centroids + .iter() + .map(|c| dot(p, c)) + .fold(f32::NEG_INFINITY, f32::max); 1.0 - best_sim // distance = 1 - cosine_similarity }) .enumerate() @@ -389,7 +1250,9 @@ fn kmeans_cosine( changed = true; } } - if !changed { break; } + if !changed { + break; + } // Update step: mean of assigned points, then normalize let mut sums = vec![vec![0.0f32; dim]; k]; @@ -398,7 +1261,9 @@ fn kmeans_cosine( sums[c].iter_mut().zip(p.iter()).for_each(|(s, v)| *s += v); counts[c] += 1; } - for (centroid, (sum, &count)) in centroids.iter_mut().zip(sums.iter_mut().zip(counts.iter())) { + for (centroid, (sum, &count)) in + centroids.iter_mut().zip(sums.iter_mut().zip(counts.iter())) + { if count > 0 { sum.iter_mut().for_each(|v| *v /= count as f32); normalize(sum); @@ -408,7 +1273,9 @@ fn kmeans_cosine( } let mut counts = vec![0usize; k]; - for &a in &assignments { counts[a] += 1; } + for &a in &assignments { + counts[a] += 1; + } (centroids, counts, assignments) } @@ -438,24 +1305,293 @@ pub async fn get_failed_embedding_images( } #[derive(Serialize)] -pub struct WorkerStates { +pub struct FolderWorkerStates { + pub folder_id: i64, pub thumbnail_paused: bool, pub metadata_paused: bool, pub embedding_paused: bool, + pub caption_paused: bool, + pub tagging_paused: bool, } #[tauri::command] -pub async fn set_worker_paused(worker: String, paused: bool) -> Result<(), String> { - indexer::set_worker_paused(&worker, paused); +pub async fn set_worker_paused( + db: State<'_, DbState>, + worker: String, + folder_id: i64, + paused: bool, +) -> Result<(), String> { + indexer::set_worker_paused(&worker, folder_id, paused); + if worker == "caption" && paused { + let conn = db.get().map_err(|e| e.to_string())?; + db::requeue_processing_caption_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } + if worker == "tagging" && paused { + let conn = db.get().map_err(|e| e.to_string())?; + db::requeue_processing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } Ok(()) } #[tauri::command] -pub async fn get_worker_states() -> Result { - let states = indexer::get_worker_paused_states(); - Ok(WorkerStates { - thumbnail_paused: states[0], - metadata_paused: states[1], - embedding_paused: states[2], - }) +pub async fn get_worker_states(folder_ids: Vec) -> Result, String> { + let states = indexer::get_worker_paused_states(&folder_ids); + Ok(folder_ids + .into_iter() + .map(|folder_id| { + let state = states + .get(&folder_id) + .copied() + .unwrap_or(indexer::FolderWorkerPausedState { + thumbnail: false, + metadata: false, + embedding: false, + caption: false, + tagging: false, + }); + FolderWorkerStates { + folder_id, + thumbnail_paused: state.thumbnail, + metadata_paused: state.metadata, + embedding_paused: state.embedding, + caption_paused: state.caption, + tagging_paused: state.tagging, + } + }) + .collect()) +} + +// --------------------------------------------------------------------------- +// Tagger commands +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +pub struct SetTaggerAccelerationParams { + pub acceleration: TaggerAcceleration, +} + +#[derive(Deserialize)] +pub struct SetTaggerThresholdParams { + pub threshold: f32, +} + +#[derive(Deserialize)] +pub struct SetTaggerBatchSizeParams { + pub batch_size: usize, +} + +#[derive(Deserialize)] +pub struct QueueTaggingJobsParams { + pub folder_id: Option, + pub folder_ids: Option>, + pub image_id: Option, +} + +#[derive(Deserialize)] +pub struct ClearTaggingJobsParams { + pub folder_id: Option, + pub folder_ids: Option>, +} + +#[derive(Deserialize)] +pub struct GetImageTagsParams { + pub image_id: i64, +} + +#[derive(Deserialize)] +pub struct AddUserTagParams { + pub image_id: i64, + pub tag: String, +} + +#[derive(Deserialize)] +pub struct RemoveTagParams { + pub tag_id: i64, +} + +#[tauri::command] +pub async fn get_tagger_model_status(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_model_status(&app_dir)) +} + +#[tauri::command] +pub async fn get_tagger_acceleration(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_acceleration(&app_dir)) +} + +#[tauri::command] +pub async fn set_tagger_acceleration( + app: AppHandle, + params: SetTaggerAccelerationParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tagger::set_tagger_acceleration(&app_dir, params.acceleration).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn probe_tagger_runtime(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || tagger::probe_tagger_runtime(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn get_tagger_threshold(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_threshold(&app_dir)) +} + +#[tauri::command] +pub async fn set_tagger_threshold( + app: AppHandle, + params: SetTaggerThresholdParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tagger::set_tagger_threshold(&app_dir, params.threshold).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn get_tagger_batch_size(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + Ok(tagger::tagger_batch_size(&app_dir)) +} + +#[tauri::command] +pub async fn set_tagger_batch_size( + app: AppHandle, + params: SetTaggerBatchSizeParams, +) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tagger::set_tagger_batch_size(&app_dir, params.batch_size).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn prepare_tagger_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || { + let app = app.clone(); + tagger::prepare_tagger_model_with_progress(&app_dir, move |progress| { + let _ = app.emit("tagger-model-progress", progress); + }) + }) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn delete_tagger_model(app: AppHandle) -> Result { + let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; + tauri::async_runtime::spawn_blocking(move || tagger::delete_tagger_model(&app_dir)) + .await + .map_err(|e| e.to_string())? + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn queue_tagging_jobs( + app: AppHandle, + db: State<'_, DbState>, + params: QueueTaggingJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let requested_folder_ids = params.folder_ids.unwrap_or_default(); + let (total, folder_ids) = match (params.folder_id, params.image_id, requested_folder_ids.is_empty()) { + (_, Some(image_id), _) => { + db::enqueue_tagging_job(&conn, image_id).map_err(|e| e.to_string())?; + // Look up just this image's folder_id rather than fetching all folders + let image = db::get_image_by_id(&conn, image_id).map_err(|e| e.to_string())?; + (1usize, vec![image.folder_id]) + } + (Some(folder_id), None, _) => { + let n = db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + (n, vec![folder_id]) + } + (None, None, false) => { + let mut total = 0usize; + for &folder_id in &requested_folder_ids { + total += db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } + (total, requested_folder_ids) + } + (None, None, true) => { + let folders = db::get_folders(&conn).map_err(|e| e.to_string())?; + let folder_ids: Vec = folders.iter().map(|f| f.id).collect(); + let mut total = 0usize; + for &folder_id in &folder_ids { + total += db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id) + .map_err(|e| e.to_string())?; + } + (total, folder_ids) + } + }; + drop(conn); + indexer::emit_folder_job_progress(&app, db.inner(), &folder_ids, true); + Ok(total) +} + +#[tauri::command] +pub async fn clear_tagging_jobs( + app: AppHandle, + db: State<'_, DbState>, + params: ClearTaggingJobsParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + let requested_folder_ids = params.folder_ids.unwrap_or_default(); + let (n, folder_ids): (usize, Vec) = match (params.folder_id, requested_folder_ids.is_empty()) { + (Some(id), _) => ( + db::clear_tagging_jobs(&conn, Some(id)).map_err(|e| e.to_string())?, + vec![id], + ), + (None, false) => { + let mut total = 0usize; + for &folder_id in &requested_folder_ids { + total += db::clear_tagging_jobs(&conn, Some(folder_id)).map_err(|e| e.to_string())?; + } + (total, requested_folder_ids) + } + (None, true) => ( + db::clear_tagging_jobs(&conn, None).map_err(|e| e.to_string())?, + db::get_folders(&conn) + .map_err(|e| e.to_string())? + .into_iter() + .map(|f| f.id) + .collect(), + ), + }; + drop(conn); + indexer::emit_folder_job_progress(&app, db.inner(), &folder_ids, true); + Ok(n) +} + +#[tauri::command] +pub async fn get_image_tags( + db: State<'_, DbState>, + params: GetImageTagsParams, +) -> Result, String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::get_image_tags(&conn, params.image_id).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn add_user_tag( + db: State<'_, DbState>, + params: AddUserTagParams, +) -> Result { + let conn = db.get().map_err(|e| e.to_string())?; + db::add_user_tag(&conn, params.image_id, ¶ms.tag).map_err(|e| e.to_string()) +} + +#[tauri::command] +pub async fn remove_tag(db: State<'_, DbState>, params: RemoveTagParams) -> Result<(), String> { + let conn = db.get().map_err(|e| e.to_string())?; + db::remove_tag(&conn, params.tag_id).map_err(|e| e.to_string()) } diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 7fe7ead..9b5caf8 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -30,6 +30,7 @@ pub struct Folder { pub name: String, pub image_count: i64, pub indexed_at: Option, + pub scan_error: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -57,6 +58,14 @@ pub struct ImageRecord { pub embedding_model: Option, pub embedding_updated_at: Option, pub embedding_error: Option, + pub generated_caption: Option, + pub caption_model: Option, + pub caption_updated_at: Option, + pub caption_error: Option, + pub ai_rating: Option, + pub ai_tagger_model: Option, + pub ai_tagged_at: Option, + pub ai_tagger_error: Option, } #[allow(dead_code)] @@ -89,6 +98,39 @@ pub struct MetadataJob { pub path: String, } +#[derive(Debug, Clone)] +pub struct CaptionJob { + pub image_id: i64, + pub folder_id: i64, + pub path: String, +} + +#[derive(Debug, Clone)] +pub struct TaggingJob { + pub image_id: i64, + pub folder_id: i64, + pub path: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageTag { + pub id: i64, + pub image_id: i64, + pub tag: String, + pub source: String, + pub ai_model: Option, + pub confidence: Option, + pub created_at: String, +} + +#[derive(Debug, Clone, Serialize)] +pub struct ExploreTagEntry { + pub tag: String, + pub count: i64, + pub representative_image_id: i64, + pub thumbnail_path: Option, +} + #[derive(Debug, Clone)] pub struct IndexedMediaEntry { pub id: i64, @@ -106,6 +148,12 @@ pub struct FolderJobProgress { pub embedding_pending: i64, pub embedding_ready: i64, pub embedding_failed: i64, + pub caption_pending: i64, + pub caption_ready: i64, + pub caption_failed: i64, + pub tagging_pending: i64, + pub tagging_ready: i64, + pub tagging_failed: i64, } pub fn create_pool(db_path: &Path) -> Result { @@ -125,7 +173,8 @@ pub fn migrate(conn: &Connection) -> Result<()> { path TEXT NOT NULL UNIQUE, name TEXT NOT NULL, image_count INTEGER NOT NULL DEFAULT 0, - indexed_at TEXT + indexed_at TEXT, + scan_error TEXT ); CREATE TABLE IF NOT EXISTS images ( @@ -172,6 +221,35 @@ pub fn migrate(conn: &Connection) -> Result<()> { updated_at TEXT NOT NULL DEFAULT (datetime('now')) ); + CREATE TABLE IF NOT EXISTS caption_jobs ( + image_id INTEGER PRIMARY KEY REFERENCES images(id) ON DELETE CASCADE, + status TEXT NOT NULL DEFAULT 'pending', + attempts INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) + ); + + CREATE TABLE IF NOT EXISTS tagging_jobs ( + image_id INTEGER PRIMARY KEY REFERENCES images(id) ON DELETE CASCADE, + status TEXT NOT NULL DEFAULT 'pending', + attempts INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) + ); + + CREATE TABLE IF NOT EXISTS image_tags ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + image_id INTEGER NOT NULL REFERENCES images(id) ON DELETE CASCADE, + tag TEXT NOT NULL, + source TEXT NOT NULL DEFAULT 'user', + ai_model TEXT, + confidence REAL, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE(image_id, tag) + ); + CREATE TABLE IF NOT EXISTS tag_cloud_cache ( folder_scope TEXT PRIMARY KEY, image_ids_hash INTEGER NOT NULL, @@ -179,11 +257,29 @@ pub fn migrate(conn: &Connection) -> Result<()> { created_at TEXT NOT NULL DEFAULT (datetime('now')) ); + CREATE TABLE IF NOT EXISTS duplicate_scan_cache ( + folder_scope TEXT PRIMARY KEY, + scanned_at INTEGER NOT NULL, + groups_json TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS app_kv ( + key TEXT PRIMARY KEY, + value INTEGER NOT NULL DEFAULT 0 + ); + + INSERT OR IGNORE INTO app_kv (key, value) VALUES ('embedding_revision', 0); + CREATE INDEX IF NOT EXISTS idx_images_folder_id ON images(folder_id); CREATE INDEX IF NOT EXISTS idx_images_modified_at ON images(modified_at); CREATE INDEX IF NOT EXISTS idx_embedding_jobs_status ON embedding_jobs(status); CREATE INDEX IF NOT EXISTS idx_thumbnail_jobs_status ON thumbnail_jobs(status); CREATE INDEX IF NOT EXISTS idx_metadata_jobs_status ON metadata_jobs(status); + CREATE INDEX IF NOT EXISTS idx_caption_jobs_status ON caption_jobs(status); + CREATE INDEX IF NOT EXISTS idx_tagging_jobs_status ON tagging_jobs(status); + CREATE INDEX IF NOT EXISTS idx_image_tags_image_id ON image_tags(image_id); + CREATE INDEX IF NOT EXISTS idx_image_tags_source ON image_tags(source); + CREATE INDEX IF NOT EXISTS idx_image_tags_tag ON image_tags(tag); ", )?; @@ -209,6 +305,15 @@ pub fn migrate(conn: &Connection) -> Result<()> { ensure_column(conn, "images", "audio_codec", "TEXT")?; ensure_column(conn, "images", "metadata_updated_at", "TEXT")?; ensure_column(conn, "images", "metadata_error", "TEXT")?; + ensure_column(conn, "images", "generated_caption", "TEXT")?; + ensure_column(conn, "images", "caption_model", "TEXT")?; + ensure_column(conn, "images", "caption_updated_at", "TEXT")?; + ensure_column(conn, "images", "caption_error", "TEXT")?; + ensure_column(conn, "images", "ai_rating", "TEXT")?; + ensure_column(conn, "images", "ai_tagger_model", "TEXT")?; + ensure_column(conn, "images", "ai_tagged_at", "TEXT")?; + ensure_column(conn, "images", "ai_tagger_error", "TEXT")?; + ensure_column(conn, "folders", "scan_error", "TEXT")?; vector::migrate(conn)?; Ok(()) @@ -229,8 +334,8 @@ pub fn insert_folder(conn: &Connection, path: &str, name: &str) -> Result { pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { let id = conn.query_row( - "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22) + "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, generated_caption, caption_model, caption_updated_at, caption_error, ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26, ?27, ?28, ?29, ?30) ON CONFLICT(path) DO UPDATE SET folder_id = excluded.folder_id, filename = excluded.filename, @@ -250,7 +355,15 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { embedding_status = excluded.embedding_status, embedding_model = excluded.embedding_model, embedding_updated_at = excluded.embedding_updated_at, - embedding_error = excluded.embedding_error + embedding_error = excluded.embedding_error, + generated_caption = excluded.generated_caption, + caption_model = excluded.caption_model, + caption_updated_at = excluded.caption_updated_at, + caption_error = excluded.caption_error, + ai_rating = NULL, + ai_tagger_model = NULL, + ai_tagged_at = NULL, + ai_tagger_error = NULL RETURNING id", params![ img.folder_id, @@ -275,9 +388,23 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result { img.embedding_model, img.embedding_updated_at, img.embedding_error, + img.generated_caption, + img.caption_model, + img.caption_updated_at, + img.caption_error, + img.ai_rating, + img.ai_tagger_model, + img.ai_tagged_at, + img.ai_tagger_error, ], |row| row.get(0), )?; + // Remove stale AI tags when content changes so they aren't shown for the + // new file. User-sourced tags are preserved. + conn.execute( + "DELETE FROM image_tags WHERE image_id = ?1 AND source = 'ai'", + [id], + )?; Ok(id) } @@ -307,6 +434,51 @@ pub fn backfill_embedding_jobs(conn: &Connection) -> Result { Ok(inserted) } +pub fn repair_embedding_consistency(conn: &Connection) -> Result<(usize, usize)> { + let orphaned_vectors = vector::delete_orphaned_embeddings(conn)?; + let ready_ids = { + let mut stmt = conn.prepare( + "SELECT id + FROM images + WHERE embedding_status = 'ready'", + )?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + + let mut missing_vector_ids = Vec::new(); + for image_id in ready_ids { + if !vector::has_image_vector(conn, image_id)? { + missing_vector_ids.push(image_id); + } + } + + let tx = conn.unchecked_transaction()?; + for image_id in &missing_vector_ids { + tx.execute( + "INSERT INTO embedding_jobs (image_id, status, attempts, last_error, created_at, updated_at) + VALUES (?1, 'pending', 0, NULL, datetime('now'), datetime('now')) + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [image_id], + )?; + tx.execute( + "UPDATE images + SET embedding_status = 'pending', + embedding_error = NULL + WHERE id = ?1", + [image_id], + )?; + } + tx.commit()?; + + Ok((orphaned_vectors, missing_vector_ids.len())) +} + pub fn retry_failed_embedding_jobs(conn: &Connection, folder_id: i64) -> Result { // Only re-queue images that are actually embeddable right now. // Videos without a thumbnail would just fail again immediately, so skip them — @@ -348,6 +520,18 @@ pub fn reset_inflight_jobs(conn: &Connection) -> Result<()> { "UPDATE embedding_jobs SET status = 'pending' WHERE status = 'processing'", [], )?; + conn.execute( + "UPDATE caption_jobs SET status = 'pending' WHERE status = 'processing'", + [], + )?; + conn.execute( + "UPDATE tagging_jobs SET status = 'pending' WHERE status = 'processing'", + [], + )?; + // Delete any rows that were cancelled before the previous shutdown so + // they don't silently linger in the DB across restarts. + conn.execute("DELETE FROM tagging_jobs WHERE status = 'cancelled'", [])?; + conn.execute("DELETE FROM caption_jobs WHERE status = 'cancelled'", [])?; Ok(()) } @@ -377,16 +561,253 @@ pub fn enqueue_metadata_job(conn: &Connection, image_id: i64) -> Result<()> { Ok(()) } -pub fn get_pending_embedding_jobs(conn: &Connection, limit: usize) -> Result> { - let mut stmt = conn.prepare( +pub fn enqueue_caption_job(conn: &Connection, image_id: i64) -> Result<()> { + conn.execute( + "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at) + VALUES (?1, 'pending', 0, NULL, datetime('now'), datetime('now')) + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [image_id], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE id = ?1", + [image_id], + )?; + Ok(()) +} + +pub fn requeue_caption_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()> { + for image_id in image_ids { + conn.execute( + "UPDATE caption_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'processing'", + [image_id], + )?; + } + Ok(()) +} + +pub fn requeue_processing_caption_jobs_for_folder( + conn: &Connection, + folder_id: i64, +) -> Result { + conn.execute( + "UPDATE caption_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN ( + SELECT id FROM images WHERE folder_id = ?1 + )", + [folder_id], + ) + .map_err(Into::into) +} + +pub fn clear_caption_jobs(conn: &Connection, folder_id: Option) -> Result { + // Mirror the tagging worker pattern: mark in-flight rows as cancelled so + // the worker can detect cancellation before writing results, then delete + // every non-processing row immediately (pending, failed, etc.). + match folder_id { + Some(folder_id) => { + conn.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; + let deleted = conn.execute( + "DELETE FROM caption_jobs + WHERE status NOT IN ('processing', 'cancelled') + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE folder_id = ?1 + AND generated_caption IS NULL", + [folder_id], + )?; + Ok(deleted) + } + None => { + conn.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing'", + [], + )?; + let deleted = conn.execute( + "DELETE FROM caption_jobs WHERE status NOT IN ('processing', 'cancelled')", + [], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE generated_caption IS NULL", + [], + )?; + Ok(deleted) + } + } +} + +pub fn is_caption_job_cancelled(conn: &Connection, image_id: i64) -> Result { + let count: i64 = conn.query_row( + "SELECT COUNT(*) FROM caption_jobs WHERE image_id = ?1 AND status = 'cancelled'", + [image_id], + |row| row.get(0), + )?; + Ok(count > 0) +} + +pub fn reset_generated_captions(conn: &Connection, folder_id: Option) -> Result { + let image_ids = match folder_id { + Some(folder_id) => { + let mut stmt = conn.prepare( + "SELECT id FROM images WHERE folder_id = ?1 AND generated_caption IS NOT NULL", + )?; + let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?; + rows.collect::>>()? + } + None => { + let mut stmt = + conn.prepare("SELECT id FROM images WHERE generated_caption IS NOT NULL")?; + let rows = stmt.query_map([], |row| row.get::<_, i64>(0))?; + rows.collect::>>()? + } + }; + + let tx = conn.unchecked_transaction()?; + for image_id in &image_ids { + vector::delete_caption_embedding(&tx, *image_id)?; + } + + match folder_id { + Some(folder_id) => { + tx.execute( + "UPDATE images + SET generated_caption = NULL, + caption_model = NULL, + caption_updated_at = NULL, + caption_error = NULL + WHERE folder_id = ?1", + [folder_id], + )?; + // Mark in-flight jobs cancelled so the worker skips writing back; + // delete all others. + tx.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; + tx.execute( + "DELETE FROM caption_jobs + WHERE status NOT IN ('processing', 'cancelled') + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; + } + None => { + tx.execute( + "UPDATE images + SET generated_caption = NULL, + caption_model = NULL, + caption_updated_at = NULL, + caption_error = NULL", + [], + )?; + // Same cancellation protocol as clear_caption_jobs. + tx.execute( + "UPDATE caption_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing'", + [], + )?; + tx.execute( + "DELETE FROM caption_jobs WHERE status NOT IN ('processing', 'cancelled')", + [], + )?; + } + } + + tx.commit()?; + Ok(image_ids.len()) +} + +pub fn enqueue_missing_caption_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result { + let inserted = conn.execute( + "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at) + SELECT id, 'pending', 0, NULL, datetime('now'), datetime('now') + FROM images + WHERE folder_id = ?1 + AND media_kind = 'image' + AND generated_caption IS NULL + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [folder_id], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE folder_id = ?1 + AND media_kind = 'image' + AND generated_caption IS NULL", + [folder_id], + )?; + Ok(inserted) +} + +pub fn enqueue_missing_caption_jobs(conn: &Connection) -> Result { + let inserted = conn.execute( + "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at) + SELECT id, 'pending', 0, NULL, datetime('now'), datetime('now') + FROM images + WHERE media_kind = 'image' + AND generated_caption IS NULL + ON CONFLICT(image_id) DO UPDATE SET + status = 'pending', + last_error = NULL, + updated_at = datetime('now')", + [], + )?; + conn.execute( + "UPDATE images + SET caption_error = NULL + WHERE media_kind = 'image' + AND generated_caption IS NULL", + [], + )?; + Ok(inserted) +} + +pub fn get_pending_embedding_jobs( + conn: &Connection, + excluded_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let sql = format!( "SELECT j.image_id, i.folder_id, i.path, i.thumbnail_path, i.media_kind, j.status, j.attempts, j.last_error, j.created_at, j.updated_at FROM embedding_jobs j JOIN images i ON i.id = j.image_id WHERE status = 'pending' + {} ORDER BY j.updated_at, j.image_id LIMIT ?1", - )?; + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; let rows = stmt.query_map([limit as i64], |row| { Ok(EmbeddingJob { image_id: row.get(0)?, @@ -404,9 +825,13 @@ pub fn get_pending_embedding_jobs(conn: &Connection, limit: usize) -> Result>>()?) } -pub fn claim_embedding_jobs(conn: &mut Connection, limit: usize) -> Result> { +pub fn claim_embedding_jobs( + conn: &mut Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { let tx = conn.transaction()?; - let candidates = get_pending_embedding_jobs(&tx, limit * 2)?; + let candidates = get_pending_embedding_jobs(&tx, paused_folder_ids, limit * 2)?; let mut claimed = Vec::with_capacity(limit); for job in candidates { @@ -430,6 +855,64 @@ pub fn claim_embedding_jobs(conn: &mut Connection, limit: usize) -> Result, + limit: usize, +) -> Result> { + let sql = format!( + "SELECT j.image_id, i.folder_id, i.path + FROM caption_jobs j + JOIN images i ON i.id = j.image_id + WHERE j.status = 'pending' + AND i.media_kind = 'image' + AND i.generated_caption IS NULL + {} + ORDER BY j.updated_at, j.image_id + LIMIT ?1", + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map([limit as i64], |row| { + Ok(CaptionJob { + image_id: row.get(0)?, + folder_id: row.get(1)?, + path: row.get(2)?, + }) + })?; + Ok(rows.collect::>>()?) +} + +pub fn claim_caption_jobs( + conn: &mut Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let tx = conn.transaction()?; + let candidates = get_pending_caption_jobs_excluding(&tx, paused_folder_ids, limit * 2)?; + let mut claimed = Vec::with_capacity(limit); + + for job in candidates { + let updated = tx.execute( + "UPDATE caption_jobs + SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'pending'", + [job.image_id], + )?; + + if updated == 1 { + claimed.push(job); + } + + if claimed.len() >= limit { + break; + } + } + + tx.commit()?; + Ok(claimed) +} + #[allow(dead_code)] pub fn mark_embedding_ready(conn: &Connection, image_id: i64, model: &str) -> Result<()> { conn.execute( @@ -439,6 +922,13 @@ pub fn mark_embedding_ready(conn: &Connection, image_id: i64, model: &str) -> Re params![image_id, model], )?; conn.execute("DELETE FROM embedding_jobs WHERE image_id = ?1", [image_id])?; + // Advance the monotonic revision so the HNSW cache is always rebuilt after + // any embedding change, regardless of clock resolution. + conn.execute( + "INSERT INTO app_kv (key, value) VALUES ('embedding_revision', 1) + ON CONFLICT(key) DO UPDATE SET value = value + 1", + [], + )?; Ok(()) } @@ -486,9 +976,14 @@ pub fn get_folder_media_index(conn: &Connection, folder_id: i64) -> Result Result<()> { + // Delete from sqlite-vec virtual tables outside the transaction — sqlite-vec + // has known issues with transactional DML in early versions. + for image_id in image_ids { + vector::delete_embedding(conn, *image_id)?; + vector::delete_caption_embedding(conn, *image_id)?; + } let tx = conn.unchecked_transaction()?; for image_id in image_ids { - vector::delete_embedding(&tx, *image_id)?; tx.execute("DELETE FROM images WHERE id = ?1", [image_id])?; } tx.commit()?; @@ -539,6 +1034,58 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result Result Result Result> { - let mut stmt = conn.prepare( +fn get_pending_thumbnail_jobs_excluding( + conn: &Connection, + excluded_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let sql = format!( "SELECT j.image_id, i.folder_id, i.path, i.media_kind FROM thumbnail_jobs j JOIN images i ON i.id = j.image_id WHERE j.status = 'pending' + {} ORDER BY j.updated_at, j.image_id LIMIT ?1", - )?; + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; let rows = stmt.query_map([limit as i64], |row| { Ok(ThumbnailJob { image_id: row.get(0)?, @@ -585,18 +1145,20 @@ pub fn get_pending_thumbnail_jobs(conn: &Connection, limit: usize) -> Result, + paused_folder_ids: &std::collections::HashSet, fetch_limit: usize, claim_limit: usize, ) -> Result> { let tx = conn.transaction()?; - let candidates = get_pending_thumbnail_jobs(&tx, fetch_limit)?; + let excluded_folder_ids = active_folder_ids + .union(paused_folder_ids) + .copied() + .collect::>(); + let candidates = get_pending_thumbnail_jobs_excluding(&tx, &excluded_folder_ids, fetch_limit)?; let mut claimed = Vec::with_capacity(claim_limit); for job in candidates { - if active_folder_ids.contains(&job.folder_id) { - continue; - } - + debug_assert!(!excluded_folder_ids.contains(&job.folder_id)); let updated = tx.execute( "UPDATE thumbnail_jobs SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') @@ -617,15 +1179,22 @@ pub fn claim_thumbnail_jobs( Ok(claimed) } -pub fn get_pending_metadata_jobs(conn: &Connection, limit: usize) -> Result> { - let mut stmt = conn.prepare( +fn get_pending_metadata_jobs_excluding( + conn: &Connection, + excluded_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let sql = format!( "SELECT j.image_id, i.folder_id, i.path FROM metadata_jobs j JOIN images i ON i.id = j.image_id WHERE j.status = 'pending' AND i.media_kind = 'video' + {} ORDER BY j.updated_at, j.image_id LIMIT ?1", - )?; + folder_exclusion_clause("i", excluded_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; let rows = stmt.query_map([limit as i64], |row| { Ok(MetadataJob { image_id: row.get(0)?, @@ -639,18 +1208,20 @@ pub fn get_pending_metadata_jobs(conn: &Connection, limit: usize) -> Result, + paused_folder_ids: &std::collections::HashSet, fetch_limit: usize, claim_limit: usize, ) -> Result> { let tx = conn.transaction()?; - let candidates = get_pending_metadata_jobs(&tx, fetch_limit)?; + let excluded_folder_ids = active_folder_ids + .union(paused_folder_ids) + .copied() + .collect::>(); + let candidates = get_pending_metadata_jobs_excluding(&tx, &excluded_folder_ids, fetch_limit)?; let mut claimed = Vec::with_capacity(claim_limit); for job in candidates { - if active_folder_ids.contains(&job.folder_id) { - continue; - } - + debug_assert!(!excluded_folder_ids.contains(&job.folder_id)); let updated = tx.execute( "UPDATE metadata_jobs SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') @@ -765,7 +1336,9 @@ pub fn update_image_details( conn.query_row( "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, - favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error + favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, + generated_caption, caption_model, caption_updated_at, caption_error, + ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error FROM images WHERE id = ?1", [image_id], @@ -778,7 +1351,9 @@ pub fn get_image_by_id(conn: &Connection, image_id: i64) -> Result conn.query_row( "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, - favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error + favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, + generated_caption, caption_model, caption_updated_at, caption_error, + ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error FROM images WHERE id = ?1", [image_id], @@ -790,14 +1365,18 @@ pub fn get_image_by_id(conn: &Connection, image_id: i64) -> Result pub fn get_images_by_ids(conn: &Connection, image_ids: &[i64]) -> Result> { let mut images = Vec::with_capacity(image_ids.len()); for image_id in image_ids { - images.push(get_image_by_id(conn, *image_id)?); + match get_image_by_id(conn, *image_id) { + Ok(image) => images.push(image), + Err(error) if is_query_returned_no_rows(&error) => {} + Err(error) => return Err(error), + } } Ok(images) } pub fn get_folders(conn: &Connection) -> Result> { let mut stmt = - conn.prepare("SELECT id, path, name, image_count, indexed_at FROM folders ORDER BY name")?; + conn.prepare("SELECT id, path, name, image_count, indexed_at, scan_error FROM folders ORDER BY name")?; let rows = stmt.query_map([], |row| { Ok(Folder { id: row.get(0)?, @@ -805,17 +1384,63 @@ pub fn get_folders(conn: &Connection) -> Result> { name: row.get(2)?, image_count: row.get(3)?, indexed_at: row.get(4)?, + scan_error: row.get(5)?, }) })?; Ok(rows.collect::>>()?) } +pub fn rename_folder(conn: &Connection, folder_id: i64, new_name: &str) -> Result<()> { + conn.execute( + "UPDATE folders SET name = ?2 WHERE id = ?1", + params![folder_id, new_name], + )?; + Ok(()) +} + +pub fn update_folder_path(conn: &Connection, folder_id: i64, old_path: &str, new_path: &str, new_name: &str) -> Result<()> { + // Both updates must be atomic: if the image path rewrite fails (e.g. a + // uniqueness collision) the folder row must not remain at the new location. + let tx = conn.unchecked_transaction()?; + tx.execute( + "UPDATE folders SET path = ?2, name = ?3, scan_error = NULL WHERE id = ?1", + params![folder_id, new_path, new_name], + )?; + // Rewrite image paths so the indexer can match them by path and skip + // re-generating thumbnails and embeddings for unchanged files. + // Use SUBSTR to replace only the leading prefix; SQLite's replace() + // would corrupt paths where the old folder name also appears deeper in the tree. + tx.execute( + "UPDATE images SET path = ?2 || SUBSTR(path, LENGTH(?1) + 1) WHERE folder_id = ?3", + params![old_path, new_path, folder_id], + )?; + tx.commit()?; + Ok(()) +} + +pub fn set_folder_scan_error(conn: &Connection, folder_id: i64, error: &str) -> Result<()> { + conn.execute( + "UPDATE folders SET scan_error = ?2 WHERE id = ?1", + params![folder_id, error], + )?; + Ok(()) +} + +pub fn clear_folder_scan_error(conn: &Connection, folder_id: i64) -> Result<()> { + conn.execute( + "UPDATE folders SET scan_error = NULL WHERE id = ?1", + [folder_id], + )?; + Ok(()) +} + pub fn get_images( conn: &Connection, folder_id: Option, search: Option<&str>, media_kind: Option<&str>, favorites_only: bool, + rating_min: i64, embedding_failed_only: bool, sort: &str, offset: i64, @@ -828,6 +1453,8 @@ pub fn get_images( "date_desc" => "modified_at DESC NULLS LAST", "size_asc" => "file_size ASC", "size_desc" => "file_size DESC", + "rating_asc" => "rating ASC, modified_at DESC NULLS LAST", + "rating_desc" => "rating DESC, modified_at DESC NULLS LAST", "duration_asc" => "duration_ms ASC NULLS LAST", "duration_desc" => "duration_ms DESC NULLS LAST", _ => "modified_at DESC NULLS LAST", @@ -839,15 +1466,18 @@ pub fn get_images( let sql = format!( "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, - favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error + favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, + generated_caption, caption_model, caption_updated_at, caption_error, + ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error FROM images WHERE (?1 IS NULL OR folder_id = ?1) AND (?2 IS NULL OR filename LIKE ?2) AND (?3 IS NULL OR media_kind = ?3) AND (?4 = 0 OR favorite = 1) - AND (?5 = 0 OR embedding_status = 'failed') + AND rating >= ?5 + AND (?6 = 0 OR embedding_status = 'failed') ORDER BY {} - LIMIT ?6 OFFSET ?7", + LIMIT ?7 OFFSET ?8", order ); let mut stmt = conn.prepare(&sql)?; @@ -857,6 +1487,7 @@ pub fn get_images( search_pattern, media_kind, favorites_flag, + rating_min, embedding_failed_flag, limit, offset @@ -872,6 +1503,7 @@ pub fn count_images( search: Option<&str>, media_kind: Option<&str>, favorites_only: bool, + rating_min: i64, embedding_failed_only: bool, ) -> Result { let search_pattern = search.map(|value| format!("%{}%", value)); @@ -884,12 +1516,14 @@ pub fn count_images( AND (?2 IS NULL OR filename LIKE ?2) AND (?3 IS NULL OR media_kind = ?3) AND (?4 = 0 OR favorite = 1) - AND (?5 = 0 OR embedding_status = 'failed')", + AND rating >= ?5 + AND (?6 = 0 OR embedding_status = 'failed')", params![ folder_id, search_pattern, media_kind, favorites_flag, + rating_min, embedding_failed_flag ], |row| row.get(0), @@ -898,6 +1532,157 @@ pub fn count_images( Ok(count) } +pub fn search_images_by_tag( + conn: &Connection, + query: &str, + folder_id: Option, + media_kind: Option<&str>, + favorites_only: bool, + rating_min: i64, + limit: usize, + offset: usize, +) -> Result<(Vec, usize)> { + let normalized_query = query.trim().to_ascii_lowercase(); + if normalized_query.is_empty() { + return Ok((Vec::new(), 0)); + } + let favorites_flag = i64::from(favorites_only); + + // Total count (for pagination) + let total: usize = conn.query_row( + "SELECT COUNT(DISTINCT i.id) + FROM images i + JOIN image_tags t ON t.image_id = i.id + WHERE (?1 IS NULL OR i.folder_id = ?1) + AND (?2 IS NULL OR i.media_kind = ?2) + AND (?3 = 0 OR i.favorite = 1) + AND i.rating >= ?4 + AND LOWER(TRIM(t.tag)) = ?5", + params![folder_id, media_kind, favorites_flag, rating_min, normalized_query], + |row| row.get::<_, i64>(0), + )? as usize; + + let mut stmt = conn.prepare( + "SELECT DISTINCT i.id + FROM images i + JOIN image_tags t ON t.image_id = i.id + WHERE (?1 IS NULL OR i.folder_id = ?1) + AND (?2 IS NULL OR i.media_kind = ?2) + AND (?3 = 0 OR i.favorite = 1) + AND i.rating >= ?4 + AND LOWER(TRIM(t.tag)) = ?5 + ORDER BY i.rating DESC, i.modified_at DESC NULLS LAST, i.filename ASC + LIMIT ?6 OFFSET ?7", + )?; + + let image_ids = stmt + .query_map( + params![ + folder_id, + media_kind, + favorites_flag, + rating_min, + normalized_query, + limit as i64, + offset as i64 + ], + |row| row.get::<_, i64>(0), + )? + .collect::>>()?; + + Ok((get_images_by_ids(conn, &image_ids)?, total)) +} + +pub fn search_tags_autocomplete( + conn: &Connection, + query: &str, + folder_id: Option, + limit: usize, +) -> Result> { + let pattern = format!("%{}%", query.to_lowercase()); + let mut stmt = conn.prepare( + "SELECT t.tag, COUNT(DISTINCT t.image_id) AS tag_count, MIN(t.image_id) AS representative_image_id + FROM image_tags t + JOIN images i ON i.id = t.image_id + WHERE (?1 IS NULL OR i.folder_id = ?1) + AND LOWER(t.tag) LIKE ?2 + GROUP BY t.tag + ORDER BY tag_count DESC, t.tag ASC + LIMIT ?3", + )?; + let rows = stmt + .query_map(params![folder_id, pattern, limit as i64], |row| { + Ok(ExploreTagEntry { + tag: row.get(0)?, + count: row.get(1)?, + representative_image_id: row.get(2)?, + thumbnail_path: None, // skip per-suggestion thumbnail for speed + }) + })? + .collect::>>()?; + Ok(rows) +} + +pub struct ImagePathRecord { + pub id: i64, + pub path: String, + pub thumbnail_path: Option, +} + +pub fn get_all_image_paths( + conn: &Connection, + folder_id: Option, +) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, path, thumbnail_path FROM images WHERE (?1 IS NULL OR folder_id = ?1) ORDER BY id", + )?; + let rows = stmt + .query_map(params![folder_id], |row| { + Ok(ImagePathRecord { + id: row.get(0)?, + path: row.get(1)?, + thumbnail_path: row.get(2)?, + }) + })? + .collect::>>()?; + Ok(rows) +} + +pub fn get_explore_tags( + conn: &Connection, + folder_id: Option, + limit: usize, +) -> Result> { + let mut stmt = conn.prepare( + "SELECT t.tag, COUNT(DISTINCT t.image_id) AS tag_count, MIN(t.image_id) AS representative_image_id + FROM image_tags t + JOIN images i ON i.id = t.image_id + WHERE (?1 IS NULL OR i.folder_id = ?1) + GROUP BY t.tag + HAVING COUNT(DISTINCT t.image_id) >= 1 + ORDER BY tag_count DESC, t.tag ASC + LIMIT ?2", + )?; + + let rows = stmt + .query_map(params![folder_id, limit as i64], |row| { + let representative_image_id = row.get::<_, i64>(2)?; + let thumbnail_path = get_image_by_id(conn, representative_image_id) + .ok() + .and_then(|image| image.thumbnail_path); + + Ok(ExploreTagEntry { + tag: row.get(0)?, + count: row.get(1)?, + representative_image_id, + thumbnail_path, + }) + })? + .collect::>>()?; + + Ok(rows) +} + pub fn get_failed_embedding_images( conn: &Connection, folder_id: i64, @@ -918,11 +1703,419 @@ pub fn get_failed_embedding_images( Ok(rows.collect::>>()?) } -pub fn delete_folder(conn: &Connection, folder_id: i64) -> Result<()> { - conn.execute("DELETE FROM folders WHERE id = ?1", params![folder_id])?; +pub fn update_generated_caption( + conn: &Connection, + image_id: i64, + caption: &str, + model: &str, +) -> Result { + conn.execute( + "UPDATE images + SET generated_caption = ?2, + caption_model = ?3, + caption_updated_at = datetime('now'), + caption_error = NULL + WHERE id = ?1", + params![image_id, caption, model], + )?; + conn.execute("DELETE FROM caption_jobs WHERE image_id = ?1", [image_id])?; + get_image_by_id(conn, image_id) +} + +pub fn mark_caption_failed(conn: &Connection, image_id: i64, error: &str) -> Result<()> { + conn.execute( + "UPDATE images + SET caption_error = ?2 + WHERE id = ?1", + params![image_id, error], + )?; + conn.execute( + "UPDATE caption_jobs + SET status = 'failed', last_error = ?2, updated_at = datetime('now') + WHERE image_id = ?1", + params![image_id, error], + )?; Ok(()) } +// --------------------------------------------------------------------------- +// Tagging jobs +// --------------------------------------------------------------------------- + +pub fn enqueue_tagging_job(conn: &Connection, image_id: i64) -> Result<()> { + conn.execute( + "INSERT INTO tagging_jobs (image_id, status, attempts, last_error, created_at, updated_at) + VALUES (?1, 'pending', 0, NULL, datetime('now'), datetime('now')) + ON CONFLICT(image_id) DO UPDATE SET + status = CASE WHEN status != 'processing' THEN 'pending' ELSE status END, + last_error = CASE WHEN status != 'processing' THEN NULL ELSE last_error END, + updated_at = datetime('now')", + [image_id], + )?; + Ok(()) +} + +pub fn claim_tagging_jobs( + conn: &mut Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let tx = conn.transaction()?; + let candidates = get_pending_tagging_jobs_excluding(&tx, paused_folder_ids, limit * 2)?; + let mut claimed = Vec::new(); + for job in candidates { + let updated = tx.execute( + "UPDATE tagging_jobs + SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'pending'", + [job.image_id], + )?; + if updated > 0 { + claimed.push(job); + if claimed.len() >= limit { + break; + } + } + } + tx.commit()?; + Ok(claimed) +} + +fn get_pending_tagging_jobs_excluding( + conn: &Connection, + paused_folder_ids: &std::collections::HashSet, + limit: usize, +) -> Result> { + let sql = format!( + "SELECT j.image_id, i.folder_id, i.path + FROM tagging_jobs j + JOIN images i ON i.id = j.image_id + WHERE j.status = 'pending' + {} + ORDER BY j.created_at ASC + LIMIT ?1", + folder_exclusion_clause("i", paused_folder_ids) + ); + let mut stmt = conn.prepare(&sql)?; + let rows = stmt + .query_map([limit as i64], |row| { + Ok(TaggingJob { + image_id: row.get(0)?, + folder_id: row.get(1)?, + path: row.get(2)?, + }) + })? + .collect::>>()?; + Ok(rows) +} + +pub fn update_ai_tags( + conn: &Connection, + image_id: i64, + tags: &[(String, f64)], + rating: &str, + model: &str, +) -> Result<()> { + // NOTE: callers are responsible for wrapping this in a transaction. + // Do NOT open a nested transaction here — rusqlite/SQLite do not support + // nested transactions and will error with "cannot start a transaction + // within a transaction". + + // Remove previous AI tags for this image before inserting fresh ones + conn.execute( + "DELETE FROM image_tags WHERE image_id = ?1 AND source = 'ai'", + [image_id], + )?; + + let mut stmt = conn.prepare( + "INSERT INTO image_tags (image_id, tag, source, ai_model, confidence, created_at) + VALUES (?1, ?2, 'ai', ?3, ?4, datetime('now')) + ON CONFLICT(image_id, tag) DO UPDATE SET + source = 'ai', + ai_model = excluded.ai_model, + confidence = excluded.confidence + WHERE source != 'user'", + )?; + for (tag, confidence) in tags { + stmt.execute(params![image_id, tag, model, confidence])?; + } + drop(stmt); + + conn.execute( + "UPDATE images + SET ai_rating = ?2, + ai_tagger_model = ?3, + ai_tagged_at = datetime('now'), + ai_tagger_error = NULL + WHERE id = ?1", + params![image_id, rating, model], + )?; + conn.execute("DELETE FROM tagging_jobs WHERE image_id = ?1", [image_id])?; + Ok(()) +} + +pub fn mark_tagging_failed(conn: &Connection, image_id: i64, error: &str) -> Result<()> { + conn.execute( + "UPDATE images SET ai_tagger_error = ?2 WHERE id = ?1", + params![image_id, error], + )?; + conn.execute( + "UPDATE tagging_jobs + SET status = 'failed', last_error = ?2, updated_at = datetime('now') + WHERE image_id = ?1", + params![image_id, error], + )?; + Ok(()) +} + +pub fn clear_tagging_jobs(conn: &Connection, folder_id: Option) -> Result { + // Rows currently being processed by the worker must not be deleted mid-flight + // — the worker holds a reference and will write results back after inference. + // Mark them 'cancelled' so the worker discards the result instead of saving it, + // then delete every non-processing row. On the next poll the worker will find + // no pending work and the queue will appear empty. + let deleted = match folder_id { + Some(fid) => { + conn.execute( + "UPDATE tagging_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [fid], + )?; + let n = conn.execute( + "DELETE FROM tagging_jobs + WHERE status NOT IN ('processing', 'cancelled') + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [fid], + )?; + conn.execute( + "UPDATE images + SET ai_tagger_error = NULL + WHERE folder_id = ?1 AND ai_tagged_at IS NULL", + [fid], + )?; + n + } + None => { + conn.execute( + "UPDATE tagging_jobs + SET status = 'cancelled', last_error = NULL, updated_at = datetime('now') + WHERE status = 'processing'", + [], + )?; + let n = conn.execute("DELETE FROM tagging_jobs WHERE status NOT IN ('processing', 'cancelled')", [])?; + conn.execute( + "UPDATE images SET ai_tagger_error = NULL WHERE ai_tagged_at IS NULL", + [], + )?; + n + } + }; + Ok(deleted) +} + +pub fn get_image_tags(conn: &Connection, image_id: i64) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, image_id, tag, source, ai_model, confidence, created_at + FROM image_tags + WHERE image_id = ?1 + ORDER BY source DESC, confidence DESC NULLS LAST, tag ASC", + )?; + let rows = stmt + .query_map([image_id], |row| { + Ok(ImageTag { + id: row.get(0)?, + image_id: row.get(1)?, + tag: row.get(2)?, + source: row.get(3)?, + ai_model: row.get(4)?, + confidence: row.get(5)?, + created_at: row.get(6)?, + }) + })? + .collect::>>()?; + Ok(rows) +} + +pub fn add_user_tag(conn: &Connection, image_id: i64, tag: &str) -> Result { + conn.execute( + "INSERT INTO image_tags (image_id, tag, source, ai_model, confidence, created_at) + VALUES (?1, ?2, 'user', NULL, NULL, datetime('now')) + ON CONFLICT(image_id, tag) DO UPDATE SET + source = 'user', + ai_model = NULL, + confidence = NULL + WHERE source = 'ai'", + params![image_id, tag], + )?; + let row = conn.query_row( + "SELECT id, image_id, tag, source, ai_model, confidence, created_at + FROM image_tags WHERE image_id = ?1 AND tag = ?2", + params![image_id, tag], + |row| { + Ok(ImageTag { + id: row.get(0)?, + image_id: row.get(1)?, + tag: row.get(2)?, + source: row.get(3)?, + ai_model: row.get(4)?, + confidence: row.get(5)?, + created_at: row.get(6)?, + }) + }, + )?; + Ok(row) +} + +pub fn remove_tag(conn: &Connection, tag_id: i64) -> Result<()> { + conn.execute("DELETE FROM image_tags WHERE id = ?1", [tag_id])?; + Ok(()) +} + +pub fn enqueue_missing_tagging_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result { + let inserted = conn.execute( + "INSERT INTO tagging_jobs (image_id, status, attempts, last_error, created_at, updated_at) + SELECT id, 'pending', 0, NULL, datetime('now'), datetime('now') + FROM images + WHERE folder_id = ?1 + AND media_kind = 'image' + AND ai_tagged_at IS NULL + ON CONFLICT(image_id) DO UPDATE SET + -- Only reset to 'pending' if the row is not currently being + -- processed or cancelled; leave 'processing' rows untouched so + -- the worker can complete them and clean up normally. + status = CASE WHEN status NOT IN ('processing', 'cancelled') THEN 'pending' ELSE status END, + last_error = CASE WHEN status != 'processing' THEN NULL ELSE last_error END, + updated_at = datetime('now')", + [folder_id], + )?; + conn.execute( + "UPDATE images + SET ai_tagger_error = NULL + WHERE folder_id = ?1 + AND media_kind = 'image' + AND ai_tagged_at IS NULL", + [folder_id], + )?; + Ok(inserted) +} + +pub fn requeue_processing_tagging_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result<()> { + conn.execute( + "UPDATE tagging_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE status = 'processing' + AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)", + [folder_id], + )?; + Ok(()) +} + +/// Returns `true` when the job row for `image_id` currently has status = 'cancelled'. +/// Used by the worker to discard inference results for jobs that were cancelled +/// while inference was running. +pub fn is_tagging_job_cancelled(conn: &Connection, image_id: i64) -> Result { + let count: i64 = conn.query_row( + "SELECT COUNT(*) FROM tagging_jobs WHERE image_id = ?1 AND status = 'cancelled'", + [image_id], + |row| row.get(0), + )?; + Ok(count > 0) +} + +/// Returns `true` when the job row for `image_id` is still `processing`. +/// A `false` result means the row was reset to `pending` (pause) or `cancelled` +/// while inference was running — either way the result must be discarded. +pub fn is_tagging_job_processing(conn: &Connection, image_id: i64) -> Result { + let count: i64 = conn.query_row( + "SELECT COUNT(*) FROM tagging_jobs WHERE image_id = ?1 AND status = 'processing'", + [image_id], + |row| row.get(0), + )?; + Ok(count > 0) +} + +pub fn requeue_tagging_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()> { + for image_id in image_ids { + conn.execute( + "UPDATE tagging_jobs + SET status = 'pending', updated_at = datetime('now') + WHERE image_id = ?1 AND status = 'processing'", + [image_id], + )?; + } + Ok(()) +} + +pub fn suggest_tags_from_caption( + conn: &Connection, + image_id: i64, + limit: usize, +) -> Result> { + let caption = conn.query_row( + "SELECT generated_caption FROM images WHERE id = ?1", + [image_id], + |row| row.get::<_, Option>(0), + )?; + + Ok(caption + .as_deref() + .map(|caption| derive_caption_tags(caption, limit)) + .unwrap_or_default()) +} + +pub fn delete_folder(conn: &Connection, folder_id: i64) -> Result<()> { + let image_ids = { + let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?; + let rows = stmt + .query_map([folder_id], |row| row.get::<_, i64>(0))? + .collect::>>()?; + rows + }; + + // Delete sqlite-vec rows outside any transaction — DML on virtual tables is + // unreliable inside a transaction and will cause remove_folder to fail for + // folders that have generated embeddings. The folder cascade (images rows) + // is handled by the FK ON DELETE CASCADE when the folder row is deleted. + for image_id in &image_ids { + vector::delete_embedding(conn, *image_id)?; + vector::delete_caption_embedding(conn, *image_id)?; + } + + let tx = conn.unchecked_transaction()?; + tx.execute("DELETE FROM folders WHERE id = ?1", params![folder_id])?; + tx.commit()?; + Ok(()) +} + +fn derive_caption_tags(caption: &str, limit: usize) -> Vec { + const STOPWORDS: &[&str] = &[ + "a", "an", "and", "are", "as", "at", "be", "by", "for", "from", "in", "into", "is", "it", + "near", "of", "on", "or", "the", "to", "with", "without", "image", "photo", "picture", + "showing", "shows", "there", "this", "that", "over", "under", + ]; + + let stopwords = STOPWORDS + .iter() + .copied() + .collect::>(); + let mut tags = Vec::new(); + for word in caption + .split(|ch: char| !ch.is_alphanumeric()) + .map(|word| word.trim().to_lowercase()) + .filter(|word| word.len() >= 3 && !stopwords.contains(word.as_str())) + { + if !tags.contains(&word) { + tags.push(word); + } + if tags.len() >= limit { + break; + } + } + tags +} + fn map_image_row(row: &Row<'_>) -> rusqlite::Result { Ok(ImageRecord { id: row.get(0)?, @@ -948,6 +2141,14 @@ fn map_image_row(row: &Row<'_>) -> rusqlite::Result { embedding_model: row.get(20)?, embedding_updated_at: row.get(21)?, embedding_error: row.get(22)?, + generated_caption: row.get(23)?, + caption_model: row.get(24)?, + caption_updated_at: row.get(25)?, + caption_error: row.get(26)?, + ai_rating: row.get(27)?, + ai_tagger_model: row.get(28)?, + ai_tagged_at: row.get(29)?, + ai_tagger_error: row.get(30)?, }) } @@ -989,6 +2190,52 @@ pub fn set_tag_cloud_cache( Ok(()) } +/// Returns (groups_json, scanned_at_unix) for the given folder scope, if present. +pub fn get_duplicate_scan_cache( + conn: &Connection, + folder_scope: &str, +) -> Result> { + match conn.query_row( + "SELECT groups_json, scanned_at FROM duplicate_scan_cache WHERE folder_scope = ?1", + params![folder_scope], + |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + ) { + Ok(row) => Ok(Some(row)), + Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(e.into()), + } +} + +/// Deletes the duplicate scan cache for the given scope. +pub fn clear_duplicate_scan_cache(conn: &Connection, folder_scope: &str) -> Result<()> { + conn.execute( + "DELETE FROM duplicate_scan_cache WHERE folder_scope = ?1", + params![folder_scope], + )?; + Ok(()) +} + +/// Upserts the duplicate scan cache for the given scope. +pub fn set_duplicate_scan_cache( + conn: &Connection, + folder_scope: &str, + groups_json: &str, +) -> Result<()> { + let scanned_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0); + conn.execute( + "INSERT INTO duplicate_scan_cache (folder_scope, scanned_at, groups_json) + VALUES (?1, ?2, ?3) + ON CONFLICT(folder_scope) DO UPDATE SET + scanned_at = excluded.scanned_at, + groups_json = excluded.groups_json", + params![folder_scope, scanned_at, groups_json], + )?; + Ok(()) +} + fn ensure_column(conn: &Connection, table: &str, column: &str, definition: &str) -> Result<()> { let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", table))?; let mut rows = stmt.query([])?; @@ -1005,3 +2252,27 @@ fn ensure_column(conn: &Connection, table: &str, column: &str, definition: &str) )?; Ok(()) } + +fn is_query_returned_no_rows(error: &anyhow::Error) -> bool { + error + .downcast_ref::() + .is_some_and(|error| matches!(error, rusqlite::Error::QueryReturnedNoRows)) +} + +fn folder_exclusion_clause( + image_alias: &str, + excluded_folder_ids: &std::collections::HashSet, +) -> String { + if excluded_folder_ids.is_empty() { + return String::new(); + } + + let mut ids = excluded_folder_ids.iter().copied().collect::>(); + ids.sort_unstable(); + let id_list = ids + .into_iter() + .map(|id| id.to_string()) + .collect::>() + .join(","); + format!("AND {}.folder_id NOT IN ({})", image_alias, id_list) +} diff --git a/src-tauri/src/embedder.rs b/src-tauri/src/embedder.rs index 532417d..3a977c7 100644 --- a/src-tauri/src/embedder.rs +++ b/src-tauri/src/embedder.rs @@ -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> { + 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::()?) + } + pub fn embed_images(&self, paths: &[PathBuf]) -> Result>> { let images = load_images(paths, self.image_size)?.to_device(&self.device)?; let features = self.model.get_image_features(&images)?; diff --git a/src-tauri/src/hnsw_index.rs b/src-tauri/src/hnsw_index.rs new file mode 100644 index 0000000..0836a03 --- /dev/null +++ b/src-tauri/src/hnsw_index.rs @@ -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, + external_by_image_id: HashMap, + hnsw: Hnsw<'static, f32, DistCosine>, +} + +static IMAGE_HNSW_INDEX: OnceLock>> = OnceLock::new(); + +fn cache() -> &'static RwLock> { + IMAGE_HNSW_INDEX.get_or_init(|| RwLock::new(None)) +} + +fn build_index(conn: &Connection) -> Result { + // 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::::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::>(); + let external_by_image_id = image_ids_by_external + .iter() + .enumerate() + .map(|(external_id, image_id)| (*image_id, external_id)) + .collect::>(); + let data_with_id = embeddings + .iter() + .enumerate() + .map(|(external_id, (_, embedding))| (embedding, external_id)) + .collect::>(); + + 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, + threshold: f32, + offset: usize, + limit: usize, +) -> Result> { + 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> = 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 = 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::>(); + 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::>(); + + Ok(matches) +} diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index cd7e8a2..7b1426d 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -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>> = OnceLock::new(); static ACTIVE_INDEXING_FOLDERS: OnceLock>> = 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> = 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, + metadata: HashSet, + embedding: HashSet, + caption: HashSet, + tagging: HashSet, +} + +#[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 { + 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 { + 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>> = 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 = 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 = 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 = 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::>(); + // 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 = Vec::new(); let media_paths: Vec = 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::>(); @@ -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::>()); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), 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::>()); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); } Ok(()) @@ -518,14 +706,11 @@ fn process_embedding_batch( embedder: &mut Option, ) -> 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::>(); - emit_folder_job_progress(app, pool, &folder_ids.iter().copied().collect::>()); + emit_folder_job_progress( + app, + pool, + &folder_ids.iter().copied().collect::>(), + 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::>()); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), 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, +) -> 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::>(), + ) + })?; + return Err(error); + } + } + } + + let folder_ids = jobs.iter().map(|job| job.folder_id).collect::>(); + emit_folder_job_progress( + app, + pool, + &folder_ids.iter().copied().collect::>(), + 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::>(); + + 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::>(); + emit_media_updates( + app, + &MediaUpdateBatch { + images: updated_images, + }, + ); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); + } + + Ok(()) +} + +fn process_tagging_batch( + app: &AppHandle, + pool: &DbPool, + app_data_dir: &Path, + tagger_instance: &mut Option, +) -> 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::>(), + ) + })?; + return Err(error); + } + } + } + + let folder_ids = jobs.iter().map(|job| job.folder_id).collect::>(); + emit_folder_job_progress( + app, + pool, + &folder_ids.iter().copied().collect::>(), + 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::>(); + + 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 = 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::>(); + emit_media_updates( + app, + &MediaUpdateBatch { + images: updated_images, + }, + ); + emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::>(), true); + } + + Ok(()) +} + fn active_indexing_folders() -> HashSet { 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::>(); 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); } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 3592fbf..009851f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -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"); diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs new file mode 100644 index 0000000..94a284e --- /dev/null +++ b/src-tauri/src/tagger.rs @@ -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, +} + +#[derive(Clone, Serialize)] +pub struct TaggerModelProgress { + pub total_files: usize, + pub completed_files: usize, + pub current_file: Option, + 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, + pub outputs: Vec, +} + +// --------------------------------------------------------------------------- +// 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, + /// 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 { + 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::() + .unwrap_or(DEFAULT_THRESHOLD) + .clamp(0.01, 1.0) +} + +pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result { + 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::() + .unwrap_or(8) + .clamp(1, 100) +} + +pub fn set_tagger_batch_size(app_data_dir: &Path, batch_size: usize) -> Result { + 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::>(); + + 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 { + 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 { + 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 { + 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, + threshold: f32, + input_size: usize, +} + +impl WdTagger { + pub fn new(app_data_dir: &Path) -> Result { + 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 { + 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::() + .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 = 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 { + 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> { + 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> { + 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) +} diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 61217d5..98b5349 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -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> { - let embedding: Vec = 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, +) -> Result> { + let embedding: Vec = 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::>>()?); + } + + // 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::>>()?; 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, +// threshold: f32, +// offset: usize, +// limit: usize, +// ) -> Result> { +// let embedding: Vec = 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::>>()?), +// 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::>>()?), +// } +// } + +pub fn get_image_embedding(conn: &Connection, image_id: i64) -> Result>> { + let embedding: Result, 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 { + // 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> { +// 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::>>()?) +// } + /// 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, + limit: usize, +) -> Result> { + 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::>>()?) +} + +#[allow(dead_code)] +pub fn search_caption_ids_by_embedding( + conn: &Connection, + embedding: &[f32], + limit: usize, +) -> Result> { + 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 { + 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 { + conn.query_row("SELECT COUNT(*) FROM caption_vec", [], |row| row.get(0)) + .map_err(Into::into) +} + +pub fn delete_orphaned_embeddings(conn: &Connection) -> Result { + let image_ids = { + let mut stmt = conn.prepare("SELECT id FROM images")?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + 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::>>()?; + rows + }; + let orphaned_ids = vector_ids + .into_iter() + .filter(|image_id| !image_ids.contains(image_id)) + .collect::>(); + + 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 { + let image_ids = { + let mut stmt = conn.prepare("SELECT id FROM images")?; + let rows = stmt + .query_map([], |row| row.get::<_, i64>(0))? + .collect::>>()?; + 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::>>()?; + rows + }; + let orphaned_ids = vector_ids + .into_iter() + .filter(|image_id| !image_ids.contains(image_id)) + .collect::>(); + + 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 { + 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 { let mut out = Vec::with_capacity(values.len() * std::mem::size_of::()); diff --git a/src-tauri/tauri.conf.json b/src-tauri/tauri.conf.json index ea720ac..51aa343 100644 --- a/src-tauri/tauri.conf.json +++ b/src-tauri/tauri.conf.json @@ -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" ] } -} +} \ No newline at end of file diff --git a/src/App.tsx b/src/App.tsx index 9884233..ce2a339 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -6,18 +6,26 @@ import { Toolbar } from "./components/Toolbar"; import { Gallery } from "./components/Gallery"; import { Lightbox } from "./components/Lightbox"; import { TagCloud } from "./components/TagCloud"; +import { DuplicateFinder } from "./components/DuplicateFinder"; import { TitleBar } from "./components/TitleBar"; +import { SettingsModal } from "./components/SettingsModal"; +import { initializeNotifications } from "./notifications"; export default function App() { const loadFolders = useGalleryStore((state) => state.loadFolders); const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress); const loadImages = useGalleryStore((state) => state.loadImages); + const loadCaptionModelStatus = useGalleryStore((state) => state.loadCaptionModelStatus); + const loadDuplicateScanCache = useGalleryStore((state) => state.loadDuplicateScanCache); const subscribeToProgress = useGalleryStore((state) => state.subscribeToProgress); const activeView = useGalleryStore((state) => state.activeView); useEffect(() => { + void initializeNotifications(); loadFolders().then(() => { void loadBackgroundJobProgress(); + void loadCaptionModelStatus(); + void loadDuplicateScanCache(); return loadImages(true); }); let unlisten: (() => void) | undefined; @@ -43,6 +51,11 @@ export default function App() { + ) : activeView === "duplicates" ? ( + <> + + + ) : ( <> @@ -54,6 +67,7 @@ export default function App() { + ); } diff --git a/src/components/BackgroundTasks.tsx b/src/components/BackgroundTasks.tsx index 0433f63..d34ceb2 100644 --- a/src/components/BackgroundTasks.tsx +++ b/src/components/BackgroundTasks.tsx @@ -2,12 +2,13 @@ import { useEffect, useMemo, useState } from "react"; import { invoke } from "@tauri-apps/api/core"; import { useGalleryStore } from "../store"; -type WorkerKey = "thumbnail" | "metadata" | "embedding"; +type WorkerKey = "thumbnail" | "metadata" | "embedding" | "tagging"; const WORKER_FOR_STAGE: Record = { Thumbnails: "thumbnail", Metadata: "metadata", Embeddings: "embedding", + Tags: "tagging", }; interface TaskStage { @@ -22,6 +23,8 @@ interface Task { name: string; stages: TaskStage[]; hasFailedEmbeddings: boolean; + hasFailedTagging: boolean; + hasFailedCaptions: boolean; pendingMediaWork: number; embeddingProcessed: number; embeddingTotal: number; @@ -35,31 +38,57 @@ interface FailedEmbeddingItem { error: string | null; } +interface FolderWorkerStates { + folder_id: number; + thumbnail_paused: boolean; + metadata_paused: boolean; + embedding_paused: boolean; + tagging_paused: boolean; +} + +const DEFAULT_PAUSED_STATE: Record = { + thumbnail: false, + metadata: false, + embedding: false, + tagging: false, +}; + export function BackgroundTasks() { const folders = useGalleryStore((state) => state.folders); const indexingProgress = useGalleryStore((state) => state.indexingProgress); const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings); + const queueTaggingJobs = useGalleryStore((state) => state.queueTaggingJobs); + const duplicateScanning = useGalleryStore((state) => state.duplicateScanning); + const duplicateScanProgress = useGalleryStore((state) => state.duplicateScanProgress); const [expanded, setExpanded] = useState(false); const [dismissed, setDismissed] = useState>({}); - const [paused, setPaused] = useState>({ - thumbnail: false, - metadata: false, - embedding: false, - }); + const [paused, setPaused] = useState>>({}); const [failedItems, setFailedItems] = useState>({}); useEffect(() => { - invoke<{ thumbnail_paused: boolean; metadata_paused: boolean; embedding_paused: boolean }>( - "get_worker_states", - ).then((states) => { - setPaused({ - thumbnail: states.thumbnail_paused, - metadata: states.metadata_paused, - embedding: states.embedding_paused, - }); + const folderIds = folders.map((folder) => folder.id); + if (folderIds.length === 0) { + setPaused({}); + return; + } + + invoke("get_worker_states", { folderIds }).then((states) => { + setPaused( + Object.fromEntries( + states.map((state) => [ + state.folder_id, + { + thumbnail: state.thumbnail_paused, + metadata: state.metadata_paused, + embedding: state.embedding_paused, + tagging: state.tagging_paused, + }, + ]), + ), + ); }); - }, []); + }, [folders]); // Fetch failed embedding filenames whenever the expanded panel opens or failure counts change. const failedCounts = useMemo( @@ -83,13 +112,25 @@ export function BackgroundTasks() { } }, [expanded, failedCounts]); - const toggleWorker = (worker: WorkerKey) => { - const next = !paused[worker]; - setPaused((prev) => ({ ...prev, [worker]: next })); - void invoke("set_worker_paused", { worker, paused: next }); + const isWorkerPaused = (folderId: number, worker: WorkerKey) => { + return paused[folderId]?.[worker] ?? DEFAULT_PAUSED_STATE[worker]; + }; + + const toggleWorker = (folderId: number, worker: WorkerKey) => { + const next = !isWorkerPaused(folderId, worker); + setPaused((prev) => ({ + ...prev, + [folderId]: { + ...DEFAULT_PAUSED_STATE, + ...prev[folderId], + [worker]: next, + }, + })); + void invoke("set_worker_paused", { worker, folderId, paused: next }); }; const dismissTask = (id: number, snapshot: string) => { + if (id < 0) return; // system tasks (duplicate scan) cannot be dismissed setDismissed((prev) => ({ ...prev, [id]: snapshot })); setExpanded(false); }; @@ -105,13 +146,25 @@ export function BackgroundTasks() { const embeddingPending = jobs?.embedding_pending ?? 0; const embeddingReady = jobs?.embedding_ready ?? 0; const embeddingFailed = jobs?.embedding_failed ?? 0; + const taggingPending = jobs?.tagging_pending ?? 0; + const taggingReady = jobs?.tagging_ready ?? 0; + const taggingFailed = jobs?.tagging_failed ?? 0; + const captionPending = jobs?.caption_pending ?? 0; + const captionReady = jobs?.caption_ready ?? 0; + const captionFailed = jobs?.caption_failed ?? 0; - const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending; + const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending + taggingPending + captionPending; const embeddingProcessed = embeddingReady + embeddingFailed; const embeddingTotal = embeddingProcessed + embeddingPending; + const taggingProcessed = taggingReady + taggingFailed; + const taggingTotal = taggingProcessed + taggingPending; + const captionProcessed = captionReady + captionFailed; + const captionTotal = captionProcessed + captionPending; const hasFailedEmbeddings = embeddingFailed > 0; + const hasFailedTagging = taggingFailed > 0; + const hasFailedCaptions = captionFailed > 0; - if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings) return null; + if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings && !hasFailedTagging && !hasFailedCaptions) return null; const stages: TaskStage[] = []; @@ -153,6 +206,26 @@ export function BackgroundTasks() { }); } + if (taggingPending > 0) { + const pct = taggingTotal > 0 ? (taggingProcessed / taggingTotal) * 100 : 0; + stages.push({ + label: "Tags", + detail: `${taggingProcessed.toLocaleString()} / ${taggingTotal.toLocaleString()}`, + progress: pct, + failed: false, + }); + } + + if (captionPending > 0) { + const pct = captionTotal > 0 ? (captionProcessed / captionTotal) * 100 : 0; + stages.push({ + label: "Captions", + detail: `${captionProcessed.toLocaleString()} / ${captionTotal.toLocaleString()}`, + progress: pct, + failed: false, + }); + } + if (hasFailedEmbeddings && pendingMediaWork === 0) { stages.push({ label: "Failed", @@ -162,13 +235,33 @@ export function BackgroundTasks() { }); } - const snapshot = `${pendingMediaWork}:${embeddingFailed}`; + if (hasFailedTagging && pendingMediaWork === 0) { + stages.push({ + label: "Failed", + detail: `${taggingFailed.toLocaleString()} tags`, + progress: null, + failed: true, + }); + } + + if (hasFailedCaptions && pendingMediaWork === 0) { + stages.push({ + label: "Failed", + detail: `${captionFailed.toLocaleString()} captions`, + progress: null, + failed: true, + }); + } + + const snapshot = `${pendingMediaWork}:${embeddingFailed}:${taggingFailed}:${captionFailed}`; return { id: folder.id, name: folder.name, stages, hasFailedEmbeddings, + hasFailedTagging, + hasFailedCaptions, pendingMediaWork, embeddingProcessed, embeddingTotal, @@ -180,17 +273,44 @@ export function BackgroundTasks() { .filter((t) => dismissed[t.id] !== t.snapshot); }, [folders, indexingProgress, mediaJobProgress, dismissed]); - if (tasks.length === 0) return null; + // Synthetic task for duplicate scanning — negative id so dismiss/retry are suppressed + const duplicateScanTask: Task | null = duplicateScanning ? { + id: -1, + name: "Duplicate Scan", + stages: [{ + label: "Hashing", + detail: duplicateScanProgress + ? `${duplicateScanProgress.scanned.toLocaleString()} / ${duplicateScanProgress.total.toLocaleString()}` + : "Starting…", + progress: duplicateScanProgress && duplicateScanProgress.total > 0 + ? (duplicateScanProgress.scanned / duplicateScanProgress.total) * 100 + : null, + failed: false, + }], + hasFailedEmbeddings: false, + hasFailedTagging: false, + hasFailedCaptions: false, + pendingMediaWork: 1, + embeddingProcessed: 0, + embeddingTotal: 0, + currentFile: null, + snapshot: "", + } : null; - const primary = tasks[0]; - const extraCount = tasks.length - 1; - const hasFailed = tasks.some((t) => t.hasFailedEmbeddings && t.pendingMediaWork === 0); + const allTasks = duplicateScanTask ? [duplicateScanTask, ...tasks] : tasks; + + if (allTasks.length === 0) return null; + + const primary = allTasks[0]; + const extraCount = allTasks.length - 1; + const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedTagging || t.hasFailedCaptions) && t.pendingMediaWork === 0); // Best progress bar value: use embedding progress if available (most informative), - // otherwise fall back to scanning progress, otherwise indeterminate. + // otherwise tagging progress, otherwise fall back to scanning progress, otherwise indeterminate. const embeddingStage = primary.stages.find((s) => s.label === "Embeddings"); + const taggingStage = primary.stages.find((s) => s.label === "Tags"); const scanningStage = primary.stages.find((s) => s.label === "Scanning"); - const barProgress = embeddingStage?.progress ?? scanningStage?.progress ?? null; + const barProgress = embeddingStage?.progress ?? taggingStage?.progress ?? scanningStage?.progress ?? null; return (
@@ -214,7 +334,7 @@ export function BackgroundTasks() {
{primary.stages.map((stage) => { const workerKey = WORKER_FOR_STAGE[stage.label]; - const isPaused = workerKey ? paused[workerKey] : false; + const isPaused = workerKey ? isWorkerPaused(primary.id, workerKey) : false; return ( { e.stopPropagation(); toggleWorker(workerKey); }} + onClick={(e) => { e.stopPropagation(); toggleWorker(primary.id, workerKey); }} > {isPaused ? ( @@ -283,8 +403,8 @@ export function BackgroundTasks() { )} - {/* Expand chevron (only when multiple folders) */} - {tasks.length > 1 && ( + {/* Expand chevron (only when multiple tasks) */} + {allTasks.length > 1 && ( )} - {/* Dismiss */} - + {/* Dismiss — hidden for system tasks like duplicate scan */} + {primary.id >= 0 && ( + + )}
{/* Expanded panel — one row per folder */} {expanded && (
- {tasks.map((task) => { + {allTasks.map((task) => { const taskEmbeddingStage = task.stages.find((s) => s.label === "Embeddings"); + const taskTaggingStage = task.stages.find((s) => s.label === "Tags"); const taskScanningStage = task.stages.find((s) => s.label === "Scanning"); - const taskBarProgress = taskEmbeddingStage?.progress ?? taskScanningStage?.progress ?? null; - const taskHasFailed = task.hasFailedEmbeddings && task.pendingMediaWork === 0; + const taskBarProgress = taskEmbeddingStage?.progress ?? taskTaggingStage?.progress ?? taskScanningStage?.progress ?? null; + const taskHasFailed = (task.hasFailedEmbeddings || task.hasFailedTagging || task.hasFailedCaptions) && task.pendingMediaWork === 0; return (
@@ -322,7 +445,7 @@ export function BackgroundTasks() {
{task.stages.map((stage) => { const workerKey = WORKER_FOR_STAGE[stage.label]; - const isPaused = workerKey ? paused[workerKey] : false; + const isPaused = workerKey ? isWorkerPaused(task.id, workerKey) : false; return ( toggleWorker(workerKey)} + onClick={() => toggleWorker(task.id, workerKey)} > {isPaused ? ( @@ -381,21 +504,26 @@ export function BackgroundTasks() { {taskHasFailed && ( )} - + {task.id >= 0 && ( + + )}
{task.currentFile && ( diff --git a/src/components/DuplicateFinder.tsx b/src/components/DuplicateFinder.tsx new file mode 100644 index 0000000..cae277f --- /dev/null +++ b/src/components/DuplicateFinder.tsx @@ -0,0 +1,278 @@ +import { useState } from "react"; +import { convertFileSrc } from "@tauri-apps/api/core"; +import { DuplicateGroup, useGalleryStore } from "../store"; + +function formatBytes(bytes: number): string { + if (bytes >= 1_073_741_824) return `${(bytes / 1_073_741_824).toFixed(1)} GB`; + if (bytes >= 1_048_576) return `${(bytes / 1_048_576).toFixed(1)} MB`; + if (bytes >= 1024) return `${(bytes / 1024).toFixed(0)} KB`; + return `${bytes} B`; +} + +function DuplicateGroupCard({ group }: { group: DuplicateGroup }) { + const selectedIds = useGalleryStore((state) => state.duplicateSelectedIds); + const toggleDuplicateSelected = useGalleryStore((state) => state.toggleDuplicateSelected); + const selectAllDuplicates = useGalleryStore((state) => state.selectAllDuplicates); + const groupSelectedCount = group.images.filter((img) => selectedIds.has(img.id)).length; + const noneSelected = groupSelectedCount === 0; + + // "Keep all but the first" — a common quick action + const handleKeepFirst = () => { + const toDelete = group.images.slice(1).map((img) => img.id); + // Clear any selection for this group first, then add the ones to delete + for (const img of group.images) { + if (selectedIds.has(img.id)) toggleDuplicateSelected(img.id); + } + selectAllDuplicates(toDelete); + }; + + return ( +
+ {/* Group header */} +
+
+ + {group.images.length} copies + + {formatBytes(group.file_size)} each + + {formatBytes(group.file_size * (group.images.length - 1))} wasted + +
+
+ {noneSelected ? ( + + ) : ( + + )} +
+
+ + {/* Image grid */} +
+ {group.images.map((image) => { + const isSelected = selectedIds.has(image.id); + const src = image.thumbnail_path ? convertFileSrc(image.thumbnail_path) : null; + return ( + + ); + })} +
+
+ ); +} + +function formatRelativeTime(unixSecs: number): string { + const diff = Math.floor(Date.now() / 1000) - unixSecs; + if (diff < 60) return "just now"; + if (diff < 3600) return `${Math.floor(diff / 60)}m ago`; + if (diff < 86400) return `${Math.floor(diff / 3600)}h ago`; + return `${Math.floor(diff / 86400)}d ago`; +} + +export function DuplicateFinder() { + const duplicateGroups = useGalleryStore((state) => state.duplicateGroups); + const duplicateScanning = useGalleryStore((state) => state.duplicateScanning); + const duplicateScanProgress = useGalleryStore((state) => state.duplicateScanProgress); + const duplicateScanError = useGalleryStore((state) => state.duplicateScanError); + const duplicateSelectedIds = useGalleryStore((state) => state.duplicateSelectedIds); + const duplicateLastScanned = useGalleryStore((state) => state.duplicateLastScanned); + const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); + const scanDuplicates = useGalleryStore((state) => state.scanDuplicates); + const clearDuplicateSelection = useGalleryStore((state) => state.clearDuplicateSelection); + const selectKeepFirstAllGroups = useGalleryStore((state) => state.selectKeepFirstAllGroups); + const deleteSelectedDuplicates = useGalleryStore((state) => state.deleteSelectedDuplicates); + + const [deleting, setDeleting] = useState(false); + const [deleteResult, setDeleteResult] = useState(null); + + const selectedCount = duplicateSelectedIds.size; + const hasResults = duplicateGroups.length > 0; + const hasScanned = hasResults || duplicateLastScanned !== null || (!duplicateScanning && duplicateScanProgress !== null); + const totalWasted = duplicateGroups.reduce( + (sum, g) => sum + g.file_size * (g.images.length - 1), + 0, + ); + const totalDuplicateImages = duplicateGroups.reduce((sum, g) => sum + g.images.length - 1, 0); + + const handleDelete = async () => { + setDeleting(true); + setDeleteResult(null); + try { + const deleted = await deleteSelectedDuplicates(); + setDeleteResult(`Deleted ${deleted} file${deleted === 1 ? "" : "s"}.`); + } catch (e) { + setDeleteResult(String(e)); + } finally { + setDeleting(false); + } + }; + + const progressPercent = + duplicateScanProgress && duplicateScanProgress.total > 0 + ? Math.round((duplicateScanProgress.scanned / duplicateScanProgress.total) * 100) + : 0; + + return ( +
+ {/* Header */} +
+
+
+

Duplicate Finder

+

+ {duplicateScanning + ? duplicateScanProgress + ? `Scanning… ${duplicateScanProgress.scanned.toLocaleString()} / ${duplicateScanProgress.total.toLocaleString()}` + : "Starting scan…" + : hasResults + ? `${duplicateGroups.length} group${duplicateGroups.length === 1 ? "" : "s"} · ${formatBytes(totalWasted)} reclaimable` + : duplicateLastScanned !== null + ? "No duplicates found" + : "Scan your library for identical files"} +

+ {!duplicateScanning && duplicateLastScanned !== null && ( +

+ Last scanned {formatRelativeTime(duplicateLastScanned)} +

+ )} +
+
+ {/* Batch select — only shown when there are groups and nothing is selected yet */} + {hasResults && selectedCount === 0 && !deleting && ( + + )} + {selectedCount > 0 ? ( + <> + {selectedCount} marked for deletion + + + + ) : null} + +
+
+ + {/* Progress bar */} + {duplicateScanning && duplicateScanProgress ? ( +
+
+
+ ) : null} + + {duplicateScanError ? ( +

{duplicateScanError}

+ ) : null} + {deleteResult ? ( +

{deleteResult}

+ ) : null} +
+ + {/* Body */} + {duplicateScanning && !hasResults ? ( +
+
+ Hashing files… +
+ ) : !hasScanned ? ( +
+
+ + + +

+ Finds files with identical content regardless of filename or location. + Click Scan for duplicates to begin. +

+

+ Large libraries may take a minute — files are hashed from disk. +

+
+
+ ) : duplicateGroups.length === 0 ? ( +
+

No duplicate files found.

+
+ ) : ( +
+
+ {duplicateGroups.map((group) => ( + + ))} +
+
+ )} +
+ ); +} + diff --git a/src/components/Gallery.tsx b/src/components/Gallery.tsx index 95a5710..0f820d4 100644 --- a/src/components/Gallery.tsx +++ b/src/components/Gallery.tsx @@ -1,6 +1,6 @@ import { useEffect, useRef, useCallback, useState } from "react"; import { convertFileSrc } from "@tauri-apps/api/core"; -import { ImageRecord, tileSizeForZoom, useGalleryStore } from "../store"; +import { ImageRecord, parseSearchValue, tileSizeForZoom, useGalleryStore } from "../store"; const GAP = 6; @@ -30,6 +30,7 @@ function ContextMenu({ const openImage = useGalleryStore((state) => state.openImage); const updateImageDetails = useGalleryStore((state) => state.updateImageDetails); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); + const similarScope = useGalleryStore((state) => state.similarScope); const canFindSimilar = image.embedding_status === "ready"; return ( @@ -59,7 +60,7 @@ function ContextMenu({ }`} onClick={async () => { if (!canFindSimilar) return; - await loadSimilarImages(image.id); + await loadSimilarImages(image.id, similarScope === "current_folder" ? image.folder_id : null, true, image.folder_id); onClose(); }} disabled={!canFindSimilar} @@ -115,6 +116,7 @@ function ImageTile({ const [loaded, setLoaded] = useState(false); const [errored, setErrored] = useState(false); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); + const similarScope = useGalleryStore((state) => state.similarScope); const canFindSimilar = image.embedding_status === "ready"; const src = image.thumbnail_path @@ -181,6 +183,15 @@ function ImageTile({
)} + {image.rating > 0 && ( +
+ {Array.from({ length: image.rating }, (_, index) => ( + + + + ))} +
+ )} {image.media_kind === "video" && image.duration_ms && (
{formatDuration(image.duration_ms)} @@ -230,7 +241,7 @@ function ImageTile({ onClick={(event) => { event.stopPropagation(); if (!canFindSimilar) return; - void loadSimilarImages(image.id); + void loadSimilarImages(image.id, similarScope === "current_folder" ? image.folder_id : null, true, image.folder_id); }} disabled={!canFindSimilar} > @@ -250,7 +261,11 @@ export function Gallery() { const loadingImages = useGalleryStore((state) => state.loadingImages); const zoomPreset = useGalleryStore((state) => state.zoomPreset); const search = useGalleryStore((state) => state.search); - const searchMode = useGalleryStore((state) => state.searchMode); + const collectionTitle = useGalleryStore((state) => state.collectionTitle); + const imageLoadError = useGalleryStore((state) => state.imageLoadError); + const galleryScrollResetKey = useGalleryStore((state) => state.galleryScrollResetKey); + const isSimilarResults = collectionTitle === "Similar Images"; + const parsedSearch = parseSearchValue(search); const parentRef = useRef(null); const [contextMenu, setContextMenu] = useState<{ x: number; y: number; image: ImageRecord } | null>(null); @@ -258,6 +273,7 @@ export function Gallery() { const handleScroll = useCallback(() => { const element = parentRef.current; if (!element) return; + if (element.scrollTop < 24) return; const nearBottom = element.scrollTop + element.clientHeight >= element.scrollHeight - 600; if (nearBottom && !loadingImages && images.length < totalImages) { void loadMoreImages(); @@ -271,6 +287,10 @@ export function Gallery() { return () => element.removeEventListener("scroll", handleScroll); }, [handleScroll]); + useEffect(() => { + parentRef.current?.scrollTo({ top: 0, left: 0 }); + }, [galleryScrollResetKey]); + useEffect(() => { const close = (event: PointerEvent) => { if ((event.target as HTMLElement | null)?.closest("[data-gallery-context-menu]")) return; @@ -287,73 +307,87 @@ export function Gallery() { }; }, []); - if (images.length === 0 && loadingImages) { - return ( -
-
-
-

- {searchMode === "semantic" && search.trim().length > 0 - ? `Searching for matches to "${search}"` - : "Loading media"} -

-

- {searchMode === "semantic" && search.trim().length > 0 - ? "Semantic search can take a little longer than filename search" - : "Fetching results"} -

-
-
- ); - } - - if (images.length === 0 && !loadingImages) { - return ( -
-
- - - -

- {searchMode === "semantic" && search.trim().length > 0 - ? "No semantic matches found" - : "No media found"} -

-

- {searchMode === "semantic" && search.trim().length > 0 - ? "Try a broader phrase, or wait for more embeddings to finish processing" - : "Try adjusting your filters or add a new folder"} -

-
-
- ); - } - return (
-
- {images.map((image) => ( - openImage(image)} - onContextMenu={(event) => { - event.preventDefault(); - setContextMenu({ x: event.clientX, y: event.clientY, image }); - }} - /> - ))} -
+ {images.length === 0 && loadingImages ? ( +
+
+
+

+ {isSimilarResults + ? "Finding similar images" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? `Searching for matches to "${parsedSearch.query}"` + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? `Searching tags for "${parsedSearch.query}"` + : "Loading media"} +

+

+ {isSimilarResults + ? "Comparing visual embeddings" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? "Semantic search can take a little longer than filename search" + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? "Matching against AI and user tags" + : "Fetching results"} +

+
+
+ ) : images.length === 0 && !loadingImages ? ( +
+
+ + + +

+ {imageLoadError + ? "Could not load results" + : isSimilarResults + ? "No similar images found" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? "No semantic matches found" + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? "No tag matches found" + : "No media found"} +

+

+ {imageLoadError + ? imageLoadError + : isSimilarResults + ? "This item may be visually isolated, or more embeddings may need to finish processing" + : parsedSearch.mode === "semantic" && parsedSearch.query.length > 0 + ? "Try a broader phrase, or wait for more embeddings to finish processing" + : parsedSearch.mode === "tag" && parsedSearch.query.length > 0 + ? "Try a shorter tag, or wait for more tagging jobs to finish" + : "Try adjusting your filters or add a new folder"} +

+
+
+ ) : ( +
+ {images.map((image) => ( + openImage(image)} + onContextMenu={(event) => { + event.preventDefault(); + setContextMenu({ x: event.clientX, y: event.clientY, image }); + }} + /> + ))} +
+ )} - {loadingImages ? ( + {images.length > 0 && loadingImages ? (
diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index 1edfc37..dadec21 100644 --- a/src/components/Lightbox.tsx +++ b/src/components/Lightbox.tsx @@ -1,7 +1,7 @@ import { useEffect, useCallback, useRef, useState } from "react"; import { motion, AnimatePresence } from "framer-motion"; import { convertFileSrc } from "@tauri-apps/api/core"; -import { useGalleryStore } from "../store"; +import { useGalleryStore, ImageTag, AiRating } from "../store"; function formatBytes(bytes: number): string { if (bytes < 1024) return `${bytes} B`; @@ -45,18 +45,114 @@ function embeddingLabel(status: string, model: string | null): string { return "Queued"; } +function ratingPill(rating: AiRating): { label: string; className: string } { + switch (rating) { + case "general": + return { label: "General", className: "border-emerald-400/25 bg-emerald-500/10 text-emerald-300" }; + case "sensitive": + return { label: "Sensitive", className: "border-sky-400/25 bg-sky-500/10 text-sky-300" }; + case "questionable": + return { label: "Questionable", className: "border-amber-400/25 bg-amber-500/10 text-amber-300" }; + case "explicit": + return { label: "Explicit", className: "border-red-400/25 bg-red-500/10 text-red-300" }; + } +} + +interface DragRect { + startX: number; + startY: number; + endX: number; + endY: number; +} + +/** Compute a CSS-pixel rect (relative to the viewport container) from a DragRect. */ +function normaliseRect(r: DragRect): { left: number; top: number; width: number; height: number } { + return { + left: Math.min(r.startX, r.endX), + top: Math.min(r.startY, r.endY), + width: Math.abs(r.endX - r.startX), + height: Math.abs(r.endY - r.startY), + }; +} + +/** Convert a CSS-pixel drag rect (relative to viewport container) to normalised 0–1 crop coords + * relative to the actual rendered element bounds. */ +function rectToNormalisedCrop( + rect: DragRect, + imgEl: HTMLImageElement, +): { x: number; y: number; w: number; h: number } | null { + const imgBounds = imgEl.getBoundingClientRect(); + if (imgBounds.width === 0 || imgBounds.height === 0) return null; + + // rect coords are already in viewport space (client coords) + const rawX = Math.min(rect.startX, rect.endX); + const rawY = Math.min(rect.startY, rect.endY); + const rawW = Math.abs(rect.endX - rect.startX); + const rawH = Math.abs(rect.endY - rect.startY); + + // Clamp to image bounds + const clampedX = Math.max(rawX, imgBounds.left); + const clampedY = Math.max(rawY, imgBounds.top); + const clampedRight = Math.min(rawX + rawW, imgBounds.right); + const clampedBottom = Math.min(rawY + rawH, imgBounds.bottom); + + const croppedW = clampedRight - clampedX; + const croppedH = clampedBottom - clampedY; + + if (croppedW <= 0 || croppedH <= 0) return null; + + // Normalize by the CSS transform scale — getBoundingClientRect already returns + // the scaled (on-screen) size, so we normalize directly against that. + return { + x: (clampedX - imgBounds.left) / imgBounds.width, + y: (clampedY - imgBounds.top) / imgBounds.height, + w: croppedW / imgBounds.width, + h: croppedH / imgBounds.height, + }; +} + +/** Minimum selection size as a fraction of the viewport container dimension. */ +const MIN_SELECTION_FRACTION = 0.02; + export function Lightbox() { const selectedImage = useGalleryStore((state) => state.selectedImage); const closeImage = useGalleryStore((state) => state.closeImage); const images = useGalleryStore((state) => state.images); const openImage = useGalleryStore((state) => state.openImage); const loadSimilarImages = useGalleryStore((state) => state.loadSimilarImages); + const loadSimilarByRegion = useGalleryStore((state) => state.loadSimilarByRegion); + const similarScope = useGalleryStore((state) => state.similarScope); const updateImageDetails = useGalleryStore((state) => state.updateImageDetails); + const getImageTags = useGalleryStore((state) => state.getImageTags); + const addUserTag = useGalleryStore((state) => state.addUserTag); + const removeTag = useGalleryStore((state) => state.removeTag); + const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus); + const queueTaggingForImage = useGalleryStore((state) => state.queueTaggingForImage); + + // Tracks the image id that is currently displayed, used to discard async + // tag mutations that resolve after the user has navigated to another image. + const currentImageIdRef = useRef(null); + currentImageIdRef.current = selectedImage?.id ?? null; + const [zoom, setZoom] = useState(1); + const [imageTags, setImageTags] = useState([]); + const [tagInput, setTagInput] = useState(""); + const [tagAdding, setTagAdding] = useState(false); + const [tagsExpanded, setTagsExpanded] = useState(false); + const [taggingQueued, setTaggingQueued] = useState(false); + + // Region selection state + const [regionSelectMode, setRegionSelectMode] = useState(false); + const [isDragging, setIsDragging] = useState(false); + const [dragRect, setDragRect] = useState(null); + const [regionSearching, setRegionSearching] = useState(false); + const imageViewportRef = useRef(null); + const imgRef = useRef(null); const currentIndex = selectedImage ? images.findIndex((image) => image.id === selectedImage.id) : -1; const canFindSimilar = selectedImage?.embedding_status === "ready"; + const canSearchRegion = canFindSimilar && selectedImage?.media_kind === "image"; const goPrev = useCallback(() => { if (currentIndex > 0) openImage(images[currentIndex - 1]); @@ -66,15 +162,44 @@ export function Lightbox() { if (currentIndex >= 0 && currentIndex < images.length - 1) openImage(images[currentIndex + 1]); }, [currentIndex, images, openImage]); + const exitRegionMode = useCallback(() => { + setRegionSelectMode(false); + setIsDragging(false); + setDragRect(null); + }, []); + useEffect(() => { setZoom(1); - }, [selectedImage?.id]); + setImageTags([]); + setTagInput(""); + setTagsExpanded(false); + setTaggingQueued(false); + exitRegionMode(); + setRegionSearching(false); + }, [selectedImage?.id, exitRegionMode]); + + useEffect(() => { + if (!selectedImage) return; + // Capture the ID so a stale response for image A cannot overwrite B's tags + // when the user navigates before the request resolves. + let cancelled = false; + void getImageTags(selectedImage.id) + .then((tags) => { if (!cancelled) setImageTags(tags); }) + .catch(() => { if (!cancelled) setImageTags([]); }); + return () => { cancelled = true; }; + }, [selectedImage?.id, selectedImage?.ai_tagged_at, getImageTags]); + + // Reset the queued state once the worker finishes so the button is usable again + useEffect(() => { + if (selectedImage?.ai_tagged_at) setTaggingQueued(false); + }, [selectedImage?.ai_tagged_at]); useEffect(() => { const viewport = imageViewportRef.current; if (!viewport || !selectedImage || selectedImage.media_kind !== "image") return; const handleWheel = (event: WheelEvent) => { + if (regionSelectMode) return; // don't zoom during selection if (!event.ctrlKey && Math.abs(event.deltaY) < Math.abs(event.deltaX)) return; event.preventDefault(); setZoom((value) => { @@ -85,12 +210,19 @@ export function Lightbox() { viewport.addEventListener("wheel", handleWheel, { passive: false }); return () => viewport.removeEventListener("wheel", handleWheel); - }, [selectedImage]); + }, [selectedImage, regionSelectMode]); useEffect(() => { const handler = (event: KeyboardEvent) => { if (!selectedImage) return; - if (event.key === "Escape") closeImage(); + if (event.key === "Escape") { + if (regionSelectMode) { + exitRegionMode(); + } else { + closeImage(); + } + } + if (regionSelectMode) return; // block nav keys during selection if (event.key === "ArrowLeft") goPrev(); if (event.key === "ArrowRight") goNext(); if (event.key === "+" || event.key === "=") setZoom((value) => Math.min(3, value + 0.25)); @@ -99,7 +231,75 @@ export function Lightbox() { window.addEventListener("keydown", handler); return () => window.removeEventListener("keydown", handler); - }, [selectedImage, closeImage, goPrev, goNext]); + }, [selectedImage, closeImage, goPrev, goNext, regionSelectMode, exitRegionMode]); + + // ── Region selection pointer handlers ─────────────────────────────────────── + const handleRegionPointerDown = useCallback( + (event: React.PointerEvent) => { + if (!regionSelectMode) return; + event.preventDefault(); + event.currentTarget.setPointerCapture(event.pointerId); + setIsDragging(true); + setDragRect({ + startX: event.clientX, + startY: event.clientY, + endX: event.clientX, + endY: event.clientY, + }); + }, + [regionSelectMode], + ); + + const handleRegionPointerMove = useCallback( + (event: React.PointerEvent) => { + if (!isDragging) return; + setDragRect((prev) => + prev ? { ...prev, endX: event.clientX, endY: event.clientY } : null, + ); + }, + [isDragging], + ); + + const handleRegionPointerUp = useCallback( + (event: React.PointerEvent) => { + if (!isDragging || !dragRect || !selectedImage || !imgRef.current) { + setIsDragging(false); + return; + } + event.currentTarget.releasePointerCapture(event.pointerId); + + const finalRect: DragRect = { ...dragRect, endX: event.clientX, endY: event.clientY }; + const crop = rectToNormalisedCrop(finalRect, imgRef.current); + + setIsDragging(false); + setDragRect(null); + + // Ignore tiny accidental clicks + const containerBounds = imageViewportRef.current?.getBoundingClientRect(); + const containerSize = containerBounds + ? Math.min(containerBounds.width, containerBounds.height) + : 500; + const selW = Math.abs(finalRect.endX - finalRect.startX); + const selH = Math.abs(finalRect.endY - finalRect.startY); + if (!crop || selW < containerSize * MIN_SELECTION_FRACTION || selH < containerSize * MIN_SELECTION_FRACTION) { + exitRegionMode(); + return; + } + + exitRegionMode(); + setRegionSearching(true); + + const folderId = + similarScope === "current_folder" ? selectedImage.folder_id : null; + void loadSimilarByRegion(selectedImage.id, crop, folderId, selectedImage.folder_id) + .finally(() => setRegionSearching(false)); + }, + [isDragging, dragRect, selectedImage, similarScope, loadSimilarByRegion, exitRegionMode], + ); + + // Build the CSS rect for the selection overlay (viewport-relative) + const selectionOverlay = + isDragging && dragRect ? normaliseRect(dragRect) : null; return ( @@ -111,11 +311,11 @@ export function Lightbox() { animate={{ opacity: 1 }} exit={{ opacity: 0 }} transition={{ duration: 0.15 }} - onClick={closeImage} + onClick={regionSelectMode ? undefined : closeImage} > - {Math.round(zoom * 100)}% - + {!regionSelectMode && ( +
+
+ + {Math.round(zoom * 100)}% + +
-
+ )} )}
-
-
+
+

{selectedImage.filename}

Details

@@ -207,7 +442,7 @@ export function Lightbox() { }`} onClick={() => { if (!canFindSimilar) return; - void loadSimilarImages(selectedImage.id); + void loadSimilarImages(selectedImage.id, similarScope === "current_folder" ? selectedImage.folder_id : null, true, selectedImage.folder_id); }} disabled={!canFindSimilar} > @@ -221,7 +456,54 @@ export function Lightbox() {
-
+ {/* Search region button row */} + {canSearchRegion && ( +
+ +
+ )} + +

Rating

@@ -316,13 +598,117 @@ export function Lightbox() { ) : null}
+
+
+

Tags

+
+ {selectedImage.ai_rating ? ( + + {ratingPill(selectedImage.ai_rating).label} + + ) : null} + {selectedImage.media_kind === "image" ? ( + + ) : null} +
+
+ + {imageTags.length > 0 ? ( + <> +
+ {(tagsExpanded ? imageTags : imageTags.slice(0, 8)).map((t) => ( + + {t.tag} + + + ))} +
+ {imageTags.length > 8 && ( + + )} + + ) : ( +

No tags yet

+ )} + +
{ + e.preventDefault(); + const raw = tagInput.trim(); + if (!raw || tagAdding) return; + setTagAdding(true); + const taggedImageId = selectedImage.id; + void addUserTag(taggedImageId, raw) + .then((newTag) => { + // Discard if the user navigated away before the request resolved. + if (currentImageIdRef.current !== taggedImageId) return; + setImageTags((prev) => [...prev, newTag]); + setTagInput(""); + }) + .catch(() => undefined) + .finally(() => setTagAdding(false)); + }} + > + setTagInput(e.target.value)} + disabled={tagAdding} + /> + +
+
+

Path

{selectedImage.path}

-
+
{currentIndex + 1} / {images.length}
@@ -331,7 +717,7 @@ export function Lightbox() { + ); +} + +function TaggerAccelerationButton({ acceleration, current, onSelect, children }: { + acceleration: TaggerAcceleration; + current: TaggerAcceleration; + onSelect: (acceleration: TaggerAcceleration) => void; + children: React.ReactNode; +}) { + const active = acceleration === current; + return ( + + ); +} + +export function SettingsModal() { + const [activeSection, setActiveSection] = useState("workspace"); + const [taggerQueueStatus, setTaggerQueueStatus] = useState(null); + const [taggerQueueing, setTaggerQueueing] = useState(false); + const [taggerClearing, setTaggerClearing] = useState(false); + const [taggerAccelerationSaving, setTaggerAccelerationSaving] = useState(false); + const [taggerThresholdDraft, setTaggerThresholdDraft] = useState(null); + const [taggerThresholdSaving, setTaggerThresholdSaving] = useState(false); + const [taggerBatchSizeDraft, setTaggerBatchSizeDraft] = useState(null); + const [taggerBatchSizeSaving, setTaggerBatchSizeSaving] = useState(false); + const settingsOpen = useGalleryStore((state) => state.settingsOpen); + const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); + const folders = useGalleryStore((state) => state.folders); + const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); + const taggingQueueScope = useGalleryStore((state) => state.taggingQueueScope); + const taggingQueueFolderIds = useGalleryStore((state) => state.taggingQueueFolderIds); + const setTaggingQueueScope = useGalleryStore((state) => state.setTaggingQueueScope); + const toggleTaggingQueueFolder = useGalleryStore((state) => state.toggleTaggingQueueFolder); + const setTaggingQueueFolderIds = useGalleryStore((state) => state.setTaggingQueueFolderIds); + const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus); + const taggerModelPreparing = useGalleryStore((state) => state.taggerModelPreparing); + const taggerModelProgress = useGalleryStore((state) => state.taggerModelProgress); + const taggerModelError = useGalleryStore((state) => state.taggerModelError); + const taggerAcceleration = useGalleryStore((state) => state.taggerAcceleration); + const taggerThreshold = useGalleryStore((state) => state.taggerThreshold); + const taggerBatchSize = useGalleryStore((state) => state.taggerBatchSize); + const taggerRuntimeProbe = useGalleryStore((state) => state.taggerRuntimeProbe); + const taggerRuntimeChecking = useGalleryStore((state) => state.taggerRuntimeChecking); + const loadTaggerModelStatus = useGalleryStore((state) => state.loadTaggerModelStatus); + const prepareTaggerModel = useGalleryStore((state) => state.prepareTaggerModel); + const deleteTaggerModel = useGalleryStore((state) => state.deleteTaggerModel); + const loadTaggerAcceleration = useGalleryStore((state) => state.loadTaggerAcceleration); + const setTaggerAcceleration = useGalleryStore((state) => state.setTaggerAcceleration); + const loadTaggerThreshold = useGalleryStore((state) => state.loadTaggerThreshold); + const setTaggerThreshold = useGalleryStore((state) => state.setTaggerThreshold); + const loadTaggerBatchSize = useGalleryStore((state) => state.loadTaggerBatchSize); + const setTaggerBatchSize = useGalleryStore((state) => state.setTaggerBatchSize); + const probeTaggerRuntime = useGalleryStore((state) => state.probeTaggerRuntime); + const queueTaggingJobs = useGalleryStore((state) => state.queueTaggingJobs); + const queueTaggingJobsForFolders = useGalleryStore((state) => state.queueTaggingJobsForFolders); + const clearTaggingJobs = useGalleryStore((state) => state.clearTaggingJobs); + const clearTaggingJobsForFolders = useGalleryStore((state) => state.clearTaggingJobsForFolders); + + useEffect(() => { + if (!settingsOpen) return; + void loadTaggerModelStatus(); + void loadTaggerAcceleration(); + void loadTaggerThreshold(); + void loadTaggerBatchSize(); + + const handleKeyDown = (event: KeyboardEvent) => { + if (event.key === "Escape") setSettingsOpen(false); + }; + window.addEventListener("keydown", handleKeyDown); + return () => window.removeEventListener("keydown", handleKeyDown); + }, [settingsOpen, loadTaggerModelStatus, loadTaggerAcceleration, loadTaggerThreshold, loadTaggerBatchSize, setSettingsOpen]); + + const selectedFolders = useMemo( + () => folders.filter((folder) => taggingQueueFolderIds.includes(folder.id)), + [folders, taggingQueueFolderIds], + ); + + const totalQueuedJobs = useMemo( + () => Object.values(mediaJobProgress).reduce((sum, progress) => sum + (progress?.tagging_pending ?? 0), 0), + [mediaJobProgress], + ); + + if (!settingsOpen) return null; + + const taggerReady = taggerModelStatus?.ready ?? false; + const queueScopeLabel = + taggingQueueScope === "all" + ? "all media" + : selectedFolders.length > 0 + ? `${selectedFolders.length} selected folder${selectedFolders.length === 1 ? "" : "s"}` + : "no folders selected"; + const thresholdDisplay = taggerThresholdDraft ?? String(taggerThreshold); + const batchSizeDisplay = taggerBatchSizeDraft ?? String(taggerBatchSize); + const taggerDownloadLabel = taggerModelProgress + ? `Downloading ${taggerModelProgress.completed_files}/${taggerModelProgress.total_files}` + : taggerModelPreparing + ? "Preparing WD Tagger..." + : taggerReady + ? "Installed" + : "Install model"; + const taggerDownloadPercent = taggerModelProgress + ? Math.round((taggerModelProgress.completed_files / Math.max(taggerModelProgress.total_files, 1)) * 100) + : 0; + + const runQueueAction = (action: "queue" | "clear") => { + const selectedIds = taggingQueueFolderIds; + const perform = + taggingQueueScope === "all" + ? action === "queue" + ? queueTaggingJobs(null) + : clearTaggingJobs(null) + : selectedIds.length > 0 + ? action === "queue" + ? queueTaggingJobsForFolders(selectedIds) + : clearTaggingJobsForFolders(selectedIds) + : Promise.resolve(0); + + if (action === "queue") { + setTaggerQueueing(true); + } else { + setTaggerClearing(true); + } + setTaggerQueueStatus(null); + + void perform + .then((count) => { + if (taggingQueueScope === "selected" && selectedIds.length === 0) { + setTaggerQueueStatus("Choose at least one folder before running tagging jobs."); + return; + } + setTaggerQueueStatus( + count === 0 + ? action === "queue" + ? "No missing tags found for the current target." + : "No queued tagging jobs to clear for the current target." + : action === "queue" + ? `Queued ${count.toLocaleString()} image${count === 1 ? "" : "s"} for tagging.` + : `Cleared ${count.toLocaleString()} queued tagging job${count === 1 ? "" : "s"}.`, + ); + }) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => { + if (action === "queue") { + setTaggerQueueing(false); + } else { + setTaggerClearing(false); + } + }); + }; + + return ( +
setSettingsOpen(false)}> +
event.stopPropagation()} + > + + +
+
+
+

{activeSection === "workspace" ? "AI Workspace" : "Workers"}

+

{activeSection === "workspace" ? "Model setup and queue targets" : "Background processing status"}

+
+ +
+ +
+ {activeSection === "workspace" ? ( +
+ + +
+
+
+
+

WD SwinV2 Tagger v3

+ + {taggerReady ? "Installed" : taggerModelPreparing ? "Preparing" : "Not installed"} + +
+

+ Anime-focused vision model by SmilingWolf. Generates booru-style tags with configurable confidence thresholds. +

+
+
+ + {taggerReady ? ( + <> + + + + ) : null} +
+
+ +
+ +
+
+ {(["auto", "directml", "cpu"] as const).map((acceleration) => ( + { + setTaggerAccelerationSaving(true); + void setTaggerAcceleration(nextAcceleration) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => setTaggerAccelerationSaving(false)); + }} + > + {acceleration === "directml" ? "DirectML" : acceleration === "cpu" ? "CPU" : "Auto"} + + ))} +
+

{taggerAccelerationSaving ? "Saving..." : `Current: ${taggerAcceleration}`}

+
+ + + +
+ setTaggerThresholdDraft(event.target.value)} + onBlur={() => { + const value = parseFloat(thresholdDisplay); + if (!isNaN(value) && value >= 0.05 && value <= 0.99) { + setTaggerThresholdSaving(true); + void setTaggerThreshold(value) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => { + setTaggerThresholdDraft(null); + setTaggerThresholdSaving(false); + }); + } else { + setTaggerThresholdDraft(null); + } + }} + /> +

{taggerThresholdSaving ? "Saving..." : "Default: 0.35"}

+
+ + + +
+ setTaggerBatchSizeDraft(event.target.value)} + onBlur={() => { + const value = parseInt(batchSizeDisplay, 10); + if (!isNaN(value) && value >= 1 && value <= 100) { + setTaggerBatchSizeSaving(true); + void setTaggerBatchSize(value) + .catch((error: unknown) => setTaggerQueueStatus(String(error))) + .finally(() => { + setTaggerBatchSizeDraft(null); + setTaggerBatchSizeSaving(false); + }); + } else { + setTaggerBatchSizeDraft(null); + } + }} + /> +

{taggerBatchSizeSaving ? "Saving..." : "Default: 8"}

+
+ + +
+

Model location

+

+ {taggerReady ? taggerModelStatus?.local_dir : "Not downloaded"} +

+ {taggerModelProgress?.current_file ?

{taggerModelProgress.current_file}

: null} + {taggerModelError ?

{taggerModelError}

: null} + {taggerRuntimeProbe ? ( +
+
+

Runtime check

+ Ready +
+

Tagger acceleration: {taggerRuntimeProbe.acceleration}

+

{taggerRuntimeProbe.session.file}

+
+ ) : null} +
+
+
+ + + + +
+ All media + Selected folders +
+ + +
+
+
+

Folder selection

+

Current target: {queueScopeLabel}.

+
+
+ + +
+
+ +
+ {folders.map((folder) => { + const active = taggingQueueFolderIds.includes(folder.id); + const progress = mediaJobProgress[folder.id]; + return ( + + ); + })} + {folders.length === 0 ?

Add a folder first to enable targeted tagging queues.

: null} +
+
+ + +
+ + +
+ + + {taggerQueueStatus ?

{taggerQueueStatus}

: null} + + + +
+

Coming soon

+
+ + +
+ ) : ( +
+ + +
+
+

Tagging queued

+

{totalQueuedJobs.toLocaleString()}

+
+
+

Selected folders

+

{selectedFolders.length.toLocaleString()}

+
+
+

Library folders

+

{folders.length.toLocaleString()}

+
+
+ + + +
+

Workers are running in the background.

+ Live +
+ + +
+ )} +
+
+
+
+ ); +} diff --git a/src/components/Sidebar.tsx b/src/components/Sidebar.tsx index 4a4a33f..18f2fce 100644 --- a/src/components/Sidebar.tsx +++ b/src/components/Sidebar.tsx @@ -1,6 +1,73 @@ +import { useEffect, useRef, useState } from "react"; import { open } from "@tauri-apps/plugin-dialog"; import { useGalleryStore, Folder, IndexProgress } from "../store"; +interface ContextMenuState { + folderId: number; + x: number; + y: number; +} + +function FolderContextMenu({ + menu, + folder, + onClose, + onRename, + onReindex, + onLocate, + onRemove, +}: { + menu: ContextMenuState; + folder: Folder; + onClose: () => void; + onRename: () => void; + onReindex: () => void; + onLocate: () => void; + onRemove: () => void; +}) { + const ref = useRef(null); + + useEffect(() => { + const handleDown = (e: MouseEvent) => { + if (ref.current && !ref.current.contains(e.target as Node)) onClose(); + }; + const handleKey = (e: KeyboardEvent) => { if (e.key === "Escape") onClose(); }; + document.addEventListener("mousedown", handleDown); + document.addEventListener("keydown", handleKey); + return () => { + document.removeEventListener("mousedown", handleDown); + document.removeEventListener("keydown", handleKey); + }; + }, [onClose]); + + const item = (label: string, onClick: () => void, danger = false) => ( + + ); + + return ( +
+ {item("Reindex", onReindex)} + {item("Rename", onRename)} + {folder.scan_error && item("Locate Folder", onLocate)} +
+ {item("Remove from app", onRemove, true)} +
+ ); +} + function FolderItem({ folder, selected, @@ -10,64 +77,187 @@ function FolderItem({ selected: boolean; progress: IndexProgress | undefined; }) { - const { selectFolder, removeFolder, reindexFolder } = useGalleryStore(); + const { selectFolder, removeFolder, reindexFolder, renameFolder, updateFolderPath } = useGalleryStore(); const isIndexing = progress && !progress.done; + const isMissing = !!folder.scan_error && !isIndexing; + + const [contextMenu, setContextMenu] = useState(null); + const [renaming, setRenaming] = useState(false); + const [renameValue, setRenameValue] = useState(folder.name); + const [confirmingRemoval, setConfirmingRemoval] = useState(false); + const renameInputRef = useRef(null); + + useEffect(() => { + if (renaming) { + setRenameValue(folder.name); + setTimeout(() => renameInputRef.current?.select(), 0); + } + }, [renaming, folder.name]); + + const handleContextMenu = (e: React.MouseEvent) => { + e.preventDefault(); + e.stopPropagation(); + // Keep menu inside viewport + const x = Math.min(e.clientX, window.innerWidth - 180); + const y = Math.min(e.clientY, window.innerHeight - 160); + setContextMenu({ folderId: folder.id, x, y }); + }; + + const handleLocateFolder = async () => { + const picked = await open({ directory: true, multiple: false, title: `Locate "${folder.name}"` }); + if (picked && typeof picked === "string") { + await updateFolderPath(folder.id, picked); + } + }; + + const commitRename = async () => { + const trimmed = renameValue.trim(); + if (trimmed && trimmed !== folder.name) { + await renameFolder(folder.id, trimmed); + } + setRenaming(false); + }; + + const handleRenameKey = (e: React.KeyboardEvent) => { + if (e.key === "Enter") { e.preventDefault(); void commitRename(); } + if (e.key === "Escape") { setRenaming(false); } + }; return ( -
selectFolder(folder.id)} - > - - - - -
-
- {folder.name} -
- {isIndexing ? ( - <> -
{progress.indexed}/{progress.total}
-
-
0 ? (progress.indexed / progress.total) * 100 : 0}%` }} - /> -
- + <> +
!renaming && selectFolder(folder.id)} + onContextMenu={handleContextMenu} + > + {isMissing ? ( + + + + + ) : ( -
{folder.image_count.toLocaleString()}
+ + + + )} + +
+ {renaming ? ( + setRenameValue(e.target.value)} + onKeyDown={handleRenameKey} + onBlur={() => void commitRename()} + onClick={(e) => e.stopPropagation()} + /> + ) : ( +
+ {folder.name} +
+ )} + {isIndexing ? ( + <> +
{progress.indexed}/{progress.total}
+
+
0 ? (progress.indexed / progress.total) * 100 : 0}%` }} + /> +
+ + ) : ( +
{folder.image_count.toLocaleString()}
+ )} +
+ + {/* Hover action buttons */} + {!renaming && ( + confirmingRemoval ? ( +
e.stopPropagation()}> + + +
+ ) : ( +
+ + +
+ ) )}
-
- - -
-
+ {isMissing && ( +
+

Folder not found

+

+ This folder may have been moved or renamed. Locate it to resume, or remove it from the app. +

+
+ + +
+
+ )} + + {contextMenu && contextMenu.folderId === folder.id && ( + setContextMenu(null)} + onRename={() => setRenaming(true)} + onReindex={() => void reindexFolder(folder.id)} + onLocate={() => void handleLocateFolder()} + onRemove={() => setConfirmingRemoval(true)} + /> + )} + ); } @@ -138,6 +328,23 @@ export function Sidebar() { Explore
+ +
setView("duplicates")} + > + + + + + Duplicates + +
{/* Section label */} diff --git a/src/components/TagCloud.tsx b/src/components/TagCloud.tsx index 7c6d3f6..43d2fcc 100644 --- a/src/components/TagCloud.tsx +++ b/src/components/TagCloud.tsx @@ -1,243 +1,375 @@ -import { useEffect } from "react"; +import { useEffect, useLayoutEffect, useMemo, useRef, useState } from "react"; import { motion } from "framer-motion"; import { convertFileSrc } from "@tauri-apps/api/core"; -import { useGalleryStore, TagCloudEntry } from "../store"; +import { ExploreTagEntry, TagCloudEntry, useGalleryStore } from "../store"; -// Accent glow colours for the hover ring — cycled by index -const GLOWS: string[] = [ - "rgba(59,130,246,0.5)", - "rgba(168,85,247,0.5)", - "rgba(16,185,129,0.5)", - "rgba(245,158,11,0.5)", - "rgba(236,72,153,0.5)", - "rgba(6,182,212,0.5)", - "rgba(249,115,22,0.5)", - "rgba(34,197,94,0.5)", +const ACCENTS = [ + "#60a5fa", + "#c084fc", + "#4ade80", + "#fbbf24", + "#f472b4", + "#2dd4bf", + "#fb923c", + "#a78bfa", + "#34d399", + "#f87171", ]; -function pseudoRandom(seed: number): number { - const x = Math.sin(seed + 1) * 10000; +const GOLDEN_ANGLE = Math.PI * (3 - Math.sqrt(5)); + +function seeded(n: number): number { + const x = Math.sin(n * 9301 + 49297) * 233280; return x - Math.floor(x); } -// Map cluster size to a tile size bucket (px) -function getTileSize(count: number, maxCount: number): number { - if (maxCount === 0) return 72; - const ratio = count / maxCount; - if (ratio > 0.75) return 160; - if (ratio > 0.45) return 128; - if (ratio > 0.22) return 104; - if (ratio > 0.08) return 88; - return 72; -} - -function TagButton({ - entry, - index, - maxCount, - onSearch, -}: { +interface PlacedNode { entry: TagCloudEntry; index: number; - maxCount: number; - onSearch: (imageId: number) => void; -}) { - const size = getTileSize(entry.count, maxCount); - const glow = GLOWS[index % GLOWS.length]; + x: number; + y: number; + w: number; + h: number; + accent: string; + driftX: number; + driftY: number; + driftDuration: number; + rotateSeed: number; +} - // Small random rotation for organic feel — larger tiles stay flatter - const maxRot = size >= 128 ? 0 : size >= 104 ? 3 : size >= 88 ? 6 : 10; - const rotation = (pseudoRandom(index * 7) - 0.5) * 2 * maxRot; +function buildCloud(entries: TagCloudEntry[], containerW: number, containerH: number): PlacedNode[] { + if (!entries.length || containerW <= 0 || containerH <= 0) return []; - const mt = Math.floor(pseudoRandom(index * 3) * 10) + 4; - const mr = Math.floor(pseudoRandom(index * 5) * 12) + 4; - const mb = Math.floor(pseudoRandom(index * 11) * 10) + 4; - const ml = Math.floor(pseudoRandom(index * 13) * 12) + 4; + const maxCount = Math.max(...entries.map((e) => e.count)); + const cx = containerW / 2; + const cy = containerH / 2; + // Spread ellipse shrinks slightly to leave room for card half-widths at the edges + const spreadX = containerW * 0.42; + const spreadY = containerH * 0.36; + const n = entries.length; - const src = entry.thumbnail_path ? convertFileSrc(entry.thumbnail_path) : null; + // 1. Build initial positions using phyllotaxis spiral + const nodes: PlacedNode[] = entries.map((entry, i) => { + const ratio = Math.max(entry.count / maxCount, 0.08); + // Cards scale from 110px to 230px wide; height is 3/4 of width + const w = 110 + Math.sqrt(ratio) * 120; + const h = w * 0.75; + const radialRatio = Math.sqrt((i + 0.5) / n); + const angle = i * GOLDEN_ANGLE; + + return { + entry, + index: i, + x: cx + Math.cos(angle) * radialRatio * spreadX, + y: cy + Math.sin(angle) * radialRatio * spreadY, + w, + h, + accent: ACCENTS[i % ACCENTS.length], + driftX: (seeded(i + 11) - 0.5) * 18, + driftY: (seeded(i + 17) - 0.5) * 14, + driftDuration: 8 + seeded(i + 23) * 7, + rotateSeed: (seeded(i + 31) - 0.5) * 4, + }; + }); + + // 2. Iterative overlap resolution — no physics, just push apart + const PAD = 24; + for (let iter = 0; iter < 80; iter++) { + for (let a = 0; a < nodes.length; a++) { + const na = nodes[a]; + for (let b = a + 1; b < nodes.length; b++) { + const nb = nodes[b]; + const dx = nb.x - na.x; + const dy = nb.y - na.y; + const overlapX = (na.w + nb.w) / 2 + PAD - Math.abs(dx); + const overlapY = (na.h + nb.h) / 2 + PAD - Math.abs(dy); + if (overlapX <= 0 || overlapY <= 0) continue; + // Push along the smaller overlap axis + if (overlapX < overlapY) { + const push = overlapX * 0.5 * (dx >= 0 ? 1 : -1); + nb.x += push; + na.x -= push; + } else { + const push = overlapY * 0.5 * (dy >= 0 ? 1 : -1); + nb.y += push; + na.y -= push; + } + } + // Pull gently back toward anchor to prevent runaway drift + na.x += (cx + Math.cos(na.index * GOLDEN_ANGLE) * Math.sqrt((na.index + 0.5) / n) * spreadX - na.x) * 0.05; + na.y += (cy + Math.sin(na.index * GOLDEN_ANGLE) * Math.sqrt((na.index + 0.5) / n) * spreadY - na.y) * 0.05; + } + } + + // 3. Clamp so cards never poke outside the container + return nodes.map((node) => ({ + ...node, + x: Math.min(Math.max(node.x, node.w / 2 + 16), containerW - node.w / 2 - 16), + y: Math.min(Math.max(node.y, node.h / 2 + 16), containerH - node.h / 2 - 16), + })); +} + +function CloudCard({ node, onOpen }: { node: PlacedNode; onOpen: (imageIds: number[]) => void }) { + const src = node.entry.thumbnail_path ? convertFileSrc(node.entry.thumbnail_path) : null; + const { w, h, accent } = node; return ( onSearch(entry.representative_image_id)} - title={`${entry.count} similar ${entry.count === 1 ? "photo" : "photos"}`} - style={{ - width: size, - height: size, - margin: `${mt}px ${mr}px ${mb}px ${ml}px`, - borderRadius: 12, - border: "2px solid rgba(255,255,255,0.08)", - background: "rgba(255,255,255,0.04)", - cursor: "pointer", - padding: 0, - overflow: "hidden", - position: "relative", - flexShrink: 0, - boxShadow: "none", - transition: "border-color 0.15s, box-shadow 0.15s", - }} - onMouseEnter={(e) => { - const el = e.currentTarget; - el.style.borderColor = glow; - el.style.boxShadow = `0 0 16px ${glow}, 0 0 32px ${glow.replace("0.5", "0.25")}`; - }} - onMouseLeave={(e) => { - const el = e.currentTarget; - el.style.borderColor = "rgba(255,255,255,0.08)"; - el.style.boxShadow = "none"; + opacity: { duration: 0.3, delay: Math.min(node.index * 0.035, 0.7) }, + scale: { duration: 0.3, delay: Math.min(node.index * 0.035, 0.7) }, + x: { duration: node.driftDuration, repeat: Infinity, ease: "easeInOut", delay: seeded(node.index + 41) * 3 }, + y: { duration: node.driftDuration + 1.6, repeat: Infinity, ease: "easeInOut", delay: seeded(node.index + 51) * 3 }, + rotate: { duration: node.driftDuration + 0.9, repeat: Infinity, ease: "easeInOut" }, }} + whileHover={{ scale: 1.06, rotate: 0, transition: { duration: 0.18 } }} + onClick={() => onOpen(node.entry.image_ids)} + title={`Open cluster — ${node.entry.count.toLocaleString()} images`} > {src ? ( ) : ( - // Fallback placeholder when no thumbnail exists yet -
+
)} - - {/* Count badge — bottom-right corner */} +
+ {/* Accent glow on hover */}
- {entry.count} + className="absolute inset-0 opacity-0 transition-opacity duration-300 group-hover:opacity-100" + style={{ background: `radial-gradient(ellipse at bottom, ${accent}25, transparent 70%)` }} + /> +
+
+
+
+

Cluster

+

{node.entry.count.toLocaleString()}

+
+ + Open + +
); } +// Actual tag cloud — word size driven by log-scaled frequency +function TagWord({ + entry, + index, + logMin, + logRange, + onSearch, +}: { + entry: ExploreTagEntry; + index: number; + logMin: number; + logRange: number; + onSearch: (tag: string) => void; +}) { + const ratio = logRange > 0 ? (Math.log(Math.max(entry.count, 1)) - logMin) / logRange : 0.5; + const fontSize = 11 + ratio * 28; // 11px – 39px + const accent = ACCENTS[index % ACCENTS.length]; + const tilt = (seeded(index + 5) - 0.5) * 7; + + return ( + onSearch(entry.tag)} + title={`${entry.tag} — ${entry.count.toLocaleString()} images`} + > + 0.55 ? accent : "rgba(255,255,255,0.82)" }} + > + {entry.tag} + + + {entry.count.toLocaleString()} + + + ); +} + +function Spinner() { + return ( + + ); +} + +// Separate component so its useLayoutEffect fires when the canvas is actually +// mounted — not at TagCloud mount time when the container may still be hidden +// behind a loading state. +function ClusterCloud({ + entries, + onOpen, +}: { + entries: TagCloudEntry[]; + onOpen: (imageIds: number[]) => void; +}) { + const canvasRef = useRef(null); + const [canvasSize, setCanvasSize] = useState({ w: 0, h: 0 }); + + useLayoutEffect(() => { + const el = canvasRef.current; + if (!el) return; + const update = () => { + const r = el.getBoundingClientRect(); + setCanvasSize({ w: r.width, h: r.height }); + }; + update(); + const ro = new ResizeObserver(update); + ro.observe(el); + return () => ro.disconnect(); + }, []); + + const nodes = useMemo( + () => buildCloud(entries, canvasSize.w, canvasSize.h), + [entries, canvasSize.w, canvasSize.h], + ); + + return ( +
+
+ {nodes.map((node) => ( + + ))} +
+ ); +} + export function TagCloud() { + const exploreMode = useGalleryStore((state) => state.exploreMode); + const setExploreMode = useGalleryStore((state) => state.setExploreMode); const tagCloudEntries = useGalleryStore((state) => state.tagCloudEntries); const tagCloudLoading = useGalleryStore((state) => state.tagCloudLoading); const loadTagCloud = useGalleryStore((state) => state.loadTagCloud); - const searchByTag = useGalleryStore((state) => state.searchByTag); + const exploreTagEntries = useGalleryStore((state) => state.exploreTagEntries); + const exploreTagLoading = useGalleryStore((state) => state.exploreTagLoading); + const loadExploreTags = useGalleryStore((state) => state.loadExploreTags); + const showVisualCluster = useGalleryStore((state) => state.showVisualCluster); + const searchForTag = useGalleryStore((state) => state.searchForTag); const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); useEffect(() => { - void loadTagCloud(); - }, [selectedFolderId]); + if (exploreMode === "visual") void loadTagCloud(); + else void loadExploreTags(); + }, [exploreMode, selectedFolderId, loadTagCloud, loadExploreTags]); - const maxCount = - tagCloudEntries.length > 0 - ? Math.max(...tagCloudEntries.map((e) => e.count)) - : 1; + const { logMin, logRange } = useMemo(() => { + if (!exploreTagEntries.length) return { logMin: 0, logRange: 1 }; + const logs = exploreTagEntries.map((e) => Math.log(Math.max(e.count, 1))); + const lo = Math.min(...logs); + const hi = Math.max(...logs); + return { logMin: lo, logRange: hi - lo || 1 }; + }, [exploreTagEntries]); + + const loading = exploreMode === "visual" ? tagCloudLoading : exploreTagLoading; + const hasEntries = exploreMode === "visual" ? tagCloudEntries.length > 0 : exploreTagEntries.length > 0; + const entryCount = exploreMode === "visual" ? tagCloudEntries.length : exploreTagEntries.length; return ( -
+
{/* Header */} - -

- Explore your library -

-

- Visual clusters from your photos — sized by how many match -

- - - {/* Loading */} - {tagCloudLoading && ( -
- - - - -

Clustering your library…

+
+
+
+

Explore

+

+ {loading + ? exploreMode === "visual" ? "Computing visual clusters…" : "Loading tags…" + : hasEntries + ? exploreMode === "visual" + ? `${entryCount} cluster${entryCount !== 1 ? "s" : ""} — click any to open` + : `${entryCount} tag${entryCount !== 1 ? "s" : ""} — click any to search` + : exploreMode === "visual" + ? "No clusters — images need embeddings first" + : "No tags — run the AI tagger or add tags manually"} +

+
+
+ + +
- )} +
- {/* Empty state */} - {!tagCloudLoading && tagCloudEntries.length === 0 && ( -
-

- No embeddings yet. Add a folder and wait for the embedding worker to - finish, then come back here. + {loading ? ( +

+ + {exploreMode === "visual" ? "Computing clusters…" : "Loading tags…"} +
+ ) : !hasEntries ? ( +
+

+ {exploreMode === "visual" + ? "No visual clusters yet. Images need embeddings before they can be grouped. Check indexing progress in the sidebar." + : "No tags yet. Run the AI tagger from Settings, or add tags manually in the image preview."}

- )} - - {/* Cluster grid */} - {!tagCloudLoading && tagCloudEntries.length > 0 && ( -
- {tagCloudEntries.map((entry, index) => ( - - ))} + ) : exploreMode === "visual" ? ( + + ) : ( + /* Tag cloud — words sized by log-scaled frequency, wrapped freely */ +
+
+ {exploreTagEntries.map((entry, index) => ( + + ))} +
)} - - {!tagCloudLoading && tagCloudEntries.length > 0 && ( -

- Grouped by visual similarity · CLIP ViT-B/32 -

- )}
); } diff --git a/src/components/TitleBar.tsx b/src/components/TitleBar.tsx index 274367d..56a7f4e 100644 --- a/src/components/TitleBar.tsx +++ b/src/components/TitleBar.tsx @@ -1,5 +1,6 @@ import { useState, useEffect } from "react"; import { getCurrentWindow } from "@tauri-apps/api/window"; +import { useGalleryStore } from "../store"; // SVG icons for window controls function MinimizeIcon() { @@ -37,6 +38,7 @@ function CloseIcon() { export function TitleBar() { const [isMaximized, setIsMaximized] = useState(false); + const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); const appWindow = getCurrentWindow(); useEffect(() => { @@ -85,6 +87,17 @@ export function TitleBar() { className="flex items-stretch h-full" style={{ WebkitAppRegion: "no-drag" } as React.CSSProperties} > + + {/* Minimize */} - ))} -
-
- - - - { - userHasTyped.current = true; - setSearchValue(event.target.value); - }} - placeholder={searchMode === "semantic" ? "Search by meaning..." : "Search filenames..."} - className="w-64 bg-transparent py-1.5 pl-8 pr-9 text-sm text-white placeholder:text-gray-600 focus:outline-none transition-colors" - /> - {searchValue.trim().length > 0 && ( - - )} + onKeyDown={(event) => { + if (event.key === "Backspace" && searchQuery.length === 0 && searchCommand !== null) { + event.preventDefault(); + setSearchCommand(null); + } + }} + onFocus={() => setSearchPanelOpen(true)} + placeholder="Search files, or use /s /t" + className={`w-64 bg-transparent py-1.5 pr-9 text-sm text-white placeholder:text-gray-600 focus:outline-none transition-colors ${searchCommand !== null ? "pl-16" : "pl-8"}`} + /> + {searchCommand !== null ? ( +
+ +
+ ) : null} + {searchQuery.trim().length > 0 || searchCommand !== null ? ( + + ) : null} +
+ + {/* Tag autocomplete suggestions */} + {showTagSuggestions && tagSuggestions.length > 0 ? ( +
+ {tagSuggestions.map((entry) => ( + + ))} +
+ ) : null} + + {/* Tag mode with no suggestions yet — show a brief hint */} + {showTagSuggestions && tagSuggestions.length === 0 && searchQuery.trim().length > 0 ? ( +
+

No matching tags

+
+ ) : null} + + {/* Semantic mode hint */} + {searchCommand === "semantic" && searchPanelOpen ? ( +
+

Search by meaning and visual concepts

+
+ ) : null} + + {/* Command hints — only shown when no command is active */} + {showCommandHints ? ( +
+ {( + [ + { command: "tag" as SearchCommand, prefix: "/t", label: "Tags", description: "Search AI and user tags" }, + { command: "semantic" as SearchCommand, prefix: "/s", label: "Semantic", description: "Search by meaning" }, + ] as const + ).map((option) => ( + + ))} +
+ ) : null}
{/* Sort */} @@ -302,10 +435,14 @@ export function Toolbar() { {/* Filter row */}
- { setMediaFilter("all"); setFavoritesOnly(false); setFailedEmbeddingsOnly(false); }} /> + { setMediaFilter("all"); setFavoritesOnly(false); setMinimumRating(0); setFailedEmbeddingsOnly(false); }} /> { setMediaFilter("image"); setFavoritesOnly(false); setFailedEmbeddingsOnly(false); }} /> { setMediaFilter("video"); setFavoritesOnly(false); setFailedEmbeddingsOnly(false); }} /> { setFavoritesOnly(!favoritesOnly); setFailedEmbeddingsOnly(false); }} /> + { setMinimumRating(minimumRating === 1 ? 0 : 1); setFailedEmbeddingsOnly(false); }} /> + { setMinimumRating(minimumRating === 4 ? 0 : 4); setFailedEmbeddingsOnly(false); }} /> + setSimilarScope("current_folder")} /> + setSimilarScope("all_media")} /> {hasAnyFailedEmbeddings ? ( setFailedEmbeddingsOnly(!failedEmbeddingsOnly)} /> ) : null} + {isSimilarResults ? Current similar scope: {similarScope === "current_folder" ? "current folder" : "all media"} : null}
); diff --git a/src/notifications.ts b/src/notifications.ts new file mode 100644 index 0000000..ccfcac8 --- /dev/null +++ b/src/notifications.ts @@ -0,0 +1,31 @@ +import { + isPermissionGranted, + requestPermission, + sendNotification, +} from "@tauri-apps/plugin-notification"; + +let permissionPromise: Promise | null = null; + +export function initializeNotifications(): Promise { + permissionPromise ??= (async () => { + try { + if (await isPermissionGranted()) return true; + return (await requestPermission()) === "granted"; + } catch (error) { + console.warn("Windows notifications are unavailable:", error); + return false; + } + })(); + + return permissionPromise; +} + +export async function notifyTaskComplete(title: string, body: string): Promise { + if (!(await initializeNotifications())) return; + + try { + sendNotification({ title, body }); + } catch (error) { + console.warn("Could not send task completion notification:", error); + } +} diff --git a/src/store.ts b/src/store.ts index 47ba35d..82d456d 100644 --- a/src/store.ts +++ b/src/store.ts @@ -2,6 +2,7 @@ import { create } from "zustand"; import { invoke } from "@tauri-apps/api/core"; import { listen, UnlistenFn } from "@tauri-apps/api/event"; import { appDataDir, join } from "@tauri-apps/api/path"; +import { notifyTaskComplete } from "./notifications"; export interface Folder { id: number; @@ -9,12 +10,21 @@ export interface Folder { name: string; image_count: number; indexed_at: string | null; + scan_error: string | null; } export type MediaKind = "image" | "video"; export type MediaFilter = "all" | MediaKind; export type ZoomPreset = "compact" | "comfortable" | "detail"; export type SearchMode = "filename" | "semantic"; +export type SearchCommand = "filename" | "semantic" | "tag"; +export type CaptionAcceleration = "auto" | "cpu" | "directml"; +export type CaptionDetail = "short" | "detailed" | "paragraph"; +export type TaggerAcceleration = "auto" | "cpu" | "directml"; +export type AiRating = "general" | "sensitive" | "questionable" | "explicit"; +export type TaggingQueueScope = "all" | "selected"; +export type SimilarScope = "all_media" | "current_folder"; +export type ExploreMode = "visual" | "tags"; export interface ImageRecord { id: number; @@ -40,6 +50,39 @@ export interface ImageRecord { embedding_model: string | null; embedding_updated_at: string | null; embedding_error: string | null; + generated_caption: string | null; + caption_model: string | null; + caption_updated_at: string | null; + caption_error: string | null; + ai_rating: AiRating | null; + ai_tagger_model: string | null; + ai_tagged_at: string | null; + ai_tagger_error: string | null; +} + +export interface ImageTag { + id: number; + image_id: number; + tag: string; + source: "user" | "ai"; + ai_model: string | null; + confidence: number | null; + created_at: string; +} + +export interface TaggerModelStatus { + model_id: string; + model_name: string; + local_dir: string; + ready: boolean; + missing_files: string[]; +} + +export interface TaggerModelProgress { + total_files: number; + completed_files: number; + current_file: string | null; + done: boolean; } export interface IndexProgress { @@ -57,6 +100,12 @@ export interface FolderJobProgress { embedding_pending: number; embedding_ready: number; embedding_failed: number; + caption_pending: number; + caption_ready: number; + caption_failed: number; + tagging_pending: number; + tagging_ready: number; + tagging_failed: number; } export interface MediaJobProgressEvent { @@ -72,12 +121,87 @@ export interface ThumbnailBatch { images: ImageRecord[]; } -export type ActiveView = "gallery" | "explore"; +export type ActiveView = "gallery" | "explore" | "duplicates"; export interface TagCloudEntry { count: number; representative_image_id: number; thumbnail_path: string | null; + image_ids: number[]; +} + +export interface ExploreTagEntry { + tag: string; + count: number; + representative_image_id: number; + thumbnail_path: string | null; +} + +export interface DuplicateGroup { + file_hash: string; + file_size: number; + images: ImageRecord[]; +} + +export interface SimilarImagesPage { + images: ImageRecord[]; + offset: number; + limit: number; + has_more: boolean; +} + +export interface CaptionModelStatus { + model_id: string; + model_name: string; + local_dir: string; + ready: boolean; + missing_files: string[]; +} + +export interface CaptionModelProgress { + total_files: number; + completed_files: number; + current_file: string | null; + done: boolean; +} + +export interface CaptionRuntimeSessionProbe { + file: string; + inputs: string[]; + outputs: string[]; +} + +export interface CaptionRuntimeProbe { + ready: boolean; + acceleration: CaptionAcceleration; + detail: CaptionDetail; + tokenizer_vocab_size: number; + sessions: CaptionRuntimeSessionProbe[]; +} + +export interface CaptionVisionProbe { + input_shape: number[]; + output_shape: number[]; + output_values: number; + acceleration: CaptionAcceleration; +} + +export interface TaggerRuntimeSessionProbe { + file: string; + inputs: string[]; + outputs: string[]; +} + +export interface TaggerRuntimeProbe { + ready: boolean; + acceleration: TaggerAcceleration; + session: TaggerRuntimeSessionProbe; +} + +export interface ParsedSearch { + mode: SearchCommand; + query: string; + prefix: string | null; } export type SortOrder = @@ -87,6 +211,8 @@ export type SortOrder = | "name_desc" | "size_desc" | "size_asc" + | "rating_desc" + | "rating_asc" | "duration_desc" | "duration_asc"; @@ -97,28 +223,73 @@ interface GalleryState { totalImages: number; loadedCount: number; loadingImages: boolean; + imageLoadError: string | null; search: string; searchMode: SearchMode; sort: SortOrder; mediaFilter: MediaFilter; favoritesOnly: boolean; + minimumRating: number; failedEmbeddingsOnly: boolean; zoomPreset: ZoomPreset; selectedImage: ImageRecord | null; collectionTitle: string | null; + similarSourceImageId: number | null; + similarSourceFolderId: number | null; + similarHasMore: boolean; + similarScope: SimilarScope; + similarFolderId: number | null; + similarCrop: { x: number; y: number; w: number; h: number } | null; + galleryScrollResetKey: number; activeView: ActiveView; + exploreMode: ExploreMode; tagCloudEntries: TagCloudEntry[]; tagCloudLoading: boolean; tagCloudFolderId: number | null | undefined; // undefined = never loaded + exploreTagEntries: ExploreTagEntry[]; + exploreTagLoading: boolean; + exploreTagsFolderId: number | null | undefined; indexingProgress: Record; mediaJobProgress: Record; cacheDir: string; + captionModelStatus: CaptionModelStatus | null; + captionModelPreparing: boolean; + captionModelError: string | null; + captionModelProgress: CaptionModelProgress | null; + captionRuntimeProbe: CaptionRuntimeProbe | null; + captionRuntimeChecking: boolean; + captionAcceleration: CaptionAcceleration; + captionDetail: CaptionDetail; + aiCaptionsEnabled: boolean; + settingsOpen: boolean; + taggingQueueScope: TaggingQueueScope; + taggingQueueFolderIds: number[]; + + taggerModelStatus: TaggerModelStatus | null; + taggerModelPreparing: boolean; + taggerModelError: string | null; + taggerModelProgress: TaggerModelProgress | null; + taggerAcceleration: TaggerAcceleration; + taggerThreshold: number; + taggerBatchSize: number; + taggerRuntimeProbe: TaggerRuntimeProbe | null; + taggerRuntimeChecking: boolean; + + duplicateGroups: DuplicateGroup[]; + duplicateScanning: boolean; + duplicateScanProgress: { scanned: number; total: number } | null; + duplicateScanError: string | null; + duplicateSelectedIds: Set; + duplicateLastScanned: number | null; // Unix timestamp (seconds) + duplicateScanFolderId: number | null | undefined; // undefined = never scanned loadFolders: () => Promise; loadBackgroundJobProgress: () => Promise; addFolder: (path: string) => Promise; removeFolder: (folderId: number) => Promise; reindexFolder: (folderId: number) => Promise; + renameFolder: (folderId: number, newName: string) => Promise; + updateFolderPath: (folderId: number, newPath: string) => Promise; selectFolder: (folderId: number | null) => void; loadImages: (reset?: boolean) => Promise; loadMoreImages: () => Promise; @@ -129,21 +300,87 @@ interface GalleryState { setSort: (sort: SortOrder) => void; setMediaFilter: (filter: MediaFilter) => void; setFavoritesOnly: (favoritesOnly: boolean) => void; + setMinimumRating: (minimumRating: number) => void; setFailedEmbeddingsOnly: (failedEmbeddingsOnly: boolean) => void; setZoomPreset: (zoomPreset: ZoomPreset) => void; openImage: (image: ImageRecord) => void; closeImage: () => void; setView: (view: ActiveView) => void; + setExploreMode: (mode: ExploreMode) => void; loadTagCloud: () => Promise; - searchByTag: (imageId: number) => void; - loadSimilarImages: (imageId: number) => Promise; + loadExploreTags: () => Promise; + showVisualCluster: (imageIds: number[]) => Promise; + searchForTag: (tag: string) => void; + loadSimilarImages: (imageId: number, folderId?: number | null, reset?: boolean, sourceFolderId?: number | null) => Promise; + loadSimilarByRegion: (imageId: number, crop: { x: number; y: number; w: number; h: number }, folderId?: number | null, sourceFolderId?: number | null) => Promise; + setSimilarScope: (scope: SimilarScope) => void; + suggestImageTags: (imageId: number) => Promise; + loadCaptionModelStatus: () => Promise; + prepareCaptionModel: () => Promise; + deleteCaptionModel: () => Promise; + probeCaptionRuntime: () => Promise; + probeCaptionImage: (imageId: number) => Promise; + generateCaptionForImage: (imageId: number) => Promise; + queueCaptionJobs: (folderId?: number | null) => Promise; + queueCaptionForImage: (imageId: number) => Promise; + clearCaptionJobs: (folderId?: number | null) => Promise; + resetGeneratedCaptions: (folderId?: number | null) => Promise; + loadCaptionAcceleration: () => Promise; + setCaptionAcceleration: (acceleration: CaptionAcceleration) => Promise; + loadCaptionDetail: () => Promise; + setCaptionDetail: (detail: CaptionDetail) => Promise; + setAiCaptionsEnabled: (enabled: boolean) => void; + setSettingsOpen: (open: boolean) => void; + setTaggingQueueScope: (scope: TaggingQueueScope) => void; + toggleTaggingQueueFolder: (folderId: number) => void; + setTaggingQueueFolderIds: (folderIds: number[]) => void; retryFailedEmbeddings: (folderId: number) => Promise; updateImageDetails: (imageId: number, updates: { favorite?: boolean; rating?: number }) => Promise; setCacheDir: (dir: string) => void; subscribeToProgress: () => Promise; + + loadTaggerModelStatus: () => Promise; + prepareTaggerModel: () => Promise; + deleteTaggerModel: () => Promise; + loadTaggerAcceleration: () => Promise; + setTaggerAcceleration: (acceleration: TaggerAcceleration) => Promise; + loadTaggerThreshold: () => Promise; + setTaggerThreshold: (threshold: number) => Promise; + loadTaggerBatchSize: () => Promise; + setTaggerBatchSize: (batchSize: number) => Promise; + probeTaggerRuntime: () => Promise; + queueTaggingJobs: (folderId?: number | null) => Promise; + queueTaggingJobsForFolders: (folderIds: number[]) => Promise; + queueTaggingForImage: (imageId: number) => Promise; + clearTaggingJobs: (folderId?: number | null) => Promise; + clearTaggingJobsForFolders: (folderIds: number[]) => Promise; + loadDuplicateScanCache: (folderId?: number | null) => Promise; + scanDuplicates: (folderId?: number | null) => Promise; + toggleDuplicateSelected: (imageId: number) => void; + selectAllDuplicates: (imageIds: number[]) => void; + selectKeepFirstAllGroups: () => void; + clearDuplicateSelection: () => void; + deleteSelectedDuplicates: () => Promise; + getImageTags: (imageId: number) => Promise; + addUserTag: (imageId: number, tag: string) => Promise; + removeTag: (tagId: number) => Promise; } const PAGE_SIZE = 200; +const AI_CAPTIONS_ENABLED_KEY = "phokus.aiCaptionsEnabled"; +const SIMILAR_DISTANCE_THRESHOLD = 0.24; + +// Single token shared by all gallery-producing requests (folder loads, searches, +// similarity, region search). Any new request increments it so a stale response +// from a previous collection type cannot overwrite newer results. +let galleryRequestToken = 0; +let tagCloudRequestToken = 0; +let exploreTagRequestToken = 0; + +function initialAiCaptionsEnabled(): boolean { + if (typeof window === "undefined") return false; + return window.localStorage.getItem(AI_CAPTIONS_ENABLED_KEY) === "true"; +} function mergeIntoVisibleWindow( currentImages: ImageRecord[], @@ -160,19 +397,71 @@ function matchesSearch(image: ImageRecord, search: string): boolean { return image.filename.toLowerCase().includes(search.toLowerCase()); } +function isDerivedCollectionTitle(collectionTitle: string | null): boolean { + return collectionTitle !== null; +} + +export function parseSearchValue(search: string): ParsedSearch { + if (!search.trim()) { + return { mode: "filename", query: "", prefix: null }; + } + + const slashPrefix = search.match(/^\/([a-z])(?:\s|$)/i); + if (slashPrefix) { + const rawPrefix = slashPrefix[1].toLowerCase(); + const query = search.length > 3 ? search.slice(3) : ""; + if (rawPrefix === "s") { + return { mode: "semantic", query, prefix: "/s" }; + } + if (rawPrefix === "t") { + return { mode: "tag", query, prefix: "/t" }; + } + return { mode: "filename", query, prefix: rawPrefix === "f" ? "/f" : null }; + } + + const trimmed = search.trim(); + const match = trimmed.match(/^([a-z]):\s*(.*)$/i); + if (!match) { + return { mode: "filename", query: trimmed, prefix: null }; + } + + const rawPrefix = match[1].toLowerCase(); + const query = match[2].trim(); + if (rawPrefix === "s") { + return { mode: "semantic", query, prefix: "s:" }; + } + if (rawPrefix === "t") { + return { mode: "tag", query, prefix: "t:" }; + } + return { mode: "filename", query, prefix: rawPrefix === "f" ? "f:" : null }; +} + +export function searchModeLabel(mode: SearchCommand): string { + switch (mode) { + case "semantic": + return "Semantic Search"; + case "tag": + return "Tag Search"; + default: + return "Filename Search"; + } +} + function matchesFilters( image: ImageRecord, selectedFolderId: number | null, mediaFilter: MediaFilter, favoritesOnly: boolean, + minimumRating: number, failedEmbeddingsOnly: boolean, search: string, ): boolean { const matchesFolder = selectedFolderId === null || image.folder_id === selectedFolderId; const matchesMedia = mediaFilter === "all" || image.media_kind === mediaFilter; const matchesFavorite = !favoritesOnly || image.favorite; + const matchesRating = image.rating >= minimumRating; const matchesFailedEmbedding = !failedEmbeddingsOnly || image.embedding_status === "failed"; - return matchesFolder && matchesMedia && matchesFavorite && matchesFailedEmbedding && matchesSearch(image, search); + return matchesFolder && matchesMedia && matchesFavorite && matchesRating && matchesFailedEmbedding && matchesSearch(image, search); } function compareNullableNumber(a: number | null, b: number | null): number { @@ -197,6 +486,10 @@ function compareImages(a: ImageRecord, b: ImageRecord, sort: SortOrder): number return compareNullableNumber(a.file_size, b.file_size); case "size_desc": return compareNullableNumber(b.file_size, a.file_size); + case "rating_asc": + return compareNullableNumber(a.rating, b.rating); + case "rating_desc": + return compareNullableNumber(b.rating, a.rating); case "duration_asc": return compareNullableNumber(a.duration_ms, b.duration_ms); case "duration_desc": @@ -266,28 +559,83 @@ export const useGalleryStore = create((set, get) => ({ totalImages: 0, loadedCount: 0, loadingImages: false, + imageLoadError: null, search: "", searchMode: "filename", sort: "date_desc", mediaFilter: "all", favoritesOnly: false, + minimumRating: 0, failedEmbeddingsOnly: false, zoomPreset: "comfortable", selectedImage: null, collectionTitle: null, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarScope: "all_media", + similarFolderId: null, + similarCrop: null, + galleryScrollResetKey: 0, activeView: "gallery", + exploreMode: "visual", tagCloudEntries: [], tagCloudLoading: false, tagCloudFolderId: undefined, + exploreTagEntries: [], + exploreTagLoading: false, + exploreTagsFolderId: undefined, indexingProgress: {}, mediaJobProgress: {}, cacheDir: "", + captionModelStatus: null, + captionModelPreparing: false, + captionModelError: null, + captionModelProgress: null, + captionRuntimeProbe: null, + captionRuntimeChecking: false, + captionAcceleration: "auto", + captionDetail: "paragraph", + aiCaptionsEnabled: initialAiCaptionsEnabled(), + settingsOpen: false, + taggingQueueScope: "all", + taggingQueueFolderIds: [], + + taggerModelStatus: null, + taggerModelPreparing: false, + taggerModelError: null, + taggerModelProgress: null, + taggerAcceleration: "auto", + taggerThreshold: 0.35, + taggerBatchSize: 8, + taggerRuntimeProbe: null, + taggerRuntimeChecking: false, + + duplicateGroups: [], + duplicateScanning: false, + duplicateScanProgress: null, + duplicateScanError: null, + duplicateSelectedIds: new Set(), + duplicateLastScanned: null, + duplicateScanFolderId: undefined, setCacheDir: (cacheDir) => set({ cacheDir }), loadFolders: async () => { const folders = await invoke("get_folders"); - set({ folders }); + set((state) => { + const folderIds = new Set(folders.map((folder) => folder.id)); + const nextSelected = state.taggingQueueFolderIds.filter((folderId) => folderIds.has(folderId)); + return { + folders, + taggingQueueFolderIds: + nextSelected.length > 0 + ? nextSelected + : state.taggingQueueScope === "selected" && folders.length > 0 + ? [folders[0].id] + : nextSelected, + }; + }); }, loadBackgroundJobProgress: async () => { @@ -309,8 +657,8 @@ export const useGalleryStore = create((set, get) => ({ const { selectedFolderId, loadFolders, loadImages, loadBackgroundJobProgress } = get(); await loadFolders(); await loadBackgroundJobProgress(); - // Invalidate tag cloud cache since library content changed - set({ tagCloudFolderId: undefined, tagCloudEntries: [] }); + // Invalidate tag cloud and explore-tags cache since library content changed + set({ tagCloudFolderId: undefined, tagCloudEntries: [], exploreTagsFolderId: undefined }); if (selectedFolderId === folderId) { set({ selectedFolderId: null }); await loadImages(true); @@ -326,37 +674,97 @@ export const useGalleryStore = create((set, get) => ({ await loadBackgroundJobProgress(); }, + renameFolder: async (folderId, newName) => { + await invoke("rename_folder", { folderId, newName }); + await get().loadFolders(); + }, + + updateFolderPath: async (folderId, newPath) => { + const { loadFolders, loadBackgroundJobProgress } = get(); + await invoke("update_folder_path", { folderId, newPath }); + await loadFolders(); + await loadBackgroundJobProgress(); + }, + selectFolder: (folderId) => { - set({ selectedFolderId: folderId, images: [], loadedCount: 0, collectionTitle: null, activeView: "gallery", failedEmbeddingsOnly: false }); + set({ selectedFolderId: folderId, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, activeView: "gallery", failedEmbeddingsOnly: false, imageLoadError: null }); void get().loadImages(true); }, loadImages: async (reset = false) => { - const { selectedFolderId, search, searchMode, sort, loadedCount, mediaFilter, favoritesOnly, failedEmbeddingsOnly } = get(); - set({ loadingImages: true }); + const { selectedFolderId, search, sort, loadedCount, mediaFilter, favoritesOnly, minimumRating, failedEmbeddingsOnly } = get(); + const parsedSearch = parseSearchValue(search); + const requestToken = ++galleryRequestToken; + set({ loadingImages: true, imageLoadError: null }); try { - if (searchMode === "semantic" && search.trim()) { + if (parsedSearch.mode === "semantic" && parsedSearch.query) { const images = await invoke("semantic_search_images", { params: { - query: search, + query: parsedSearch.query, folder_id: selectedFolderId, media_kind: mediaFilter === "all" ? null : mediaFilter, favorites_only: favoritesOnly, + rating_min: minimumRating > 0 ? minimumRating : null, limit: PAGE_SIZE, }, }); + if (requestToken !== galleryRequestToken) return; set({ images, totalImages: images.length, loadedCount: images.length, loadingImages: false, - collectionTitle: `Semantic search: ${search}`, + collectionTitle: `Semantic search: ${parsedSearch.query}`, + selectedFolderId, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarFolderId: null, }); return; } + if (parsedSearch.mode === "tag" && parsedSearch.query) { + const offset = reset ? 0 : loadedCount; + const result = await invoke<{ images: ImageRecord[]; total: number; offset: number; limit: number }>("search_images_by_tag", { + params: { + query: parsedSearch.query, + folder_id: selectedFolderId, + media_kind: mediaFilter === "all" ? null : mediaFilter, + favorites_only: favoritesOnly, + rating_min: minimumRating > 0 ? minimumRating : null, + limit: PAGE_SIZE, + offset, + }, + }); + + if (requestToken !== galleryRequestToken) return; + if (reset) { + set({ + images: result.images, + totalImages: result.total, + loadedCount: result.images.length, + loadingImages: false, + collectionTitle: `Tag search: ${parsedSearch.query}`, + selectedFolderId, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarFolderId: null, + }); + } else { + set((state) => ({ + images: [...state.images, ...result.images], + loadedCount: state.loadedCount + result.images.length, + totalImages: result.total, + loadingImages: false, + })); + } + return; + } + const offset = reset ? 0 : loadedCount; const result = await invoke<{ images: ImageRecord[]; @@ -366,9 +774,10 @@ export const useGalleryStore = create((set, get) => ({ }>("get_images", { params: { folder_id: selectedFolderId, - search: search || null, + search: parsedSearch.query || null, media_kind: mediaFilter === "all" ? null : mediaFilter, favorites_only: favoritesOnly, + rating_min: minimumRating > 0 ? minimumRating : null, embedding_failed_only: failedEmbeddingsOnly, sort, offset, @@ -376,62 +785,110 @@ export const useGalleryStore = create((set, get) => ({ }, }); + if (requestToken !== galleryRequestToken) return; set((state) => ({ images: reset ? result.images : [...state.images, ...result.images], totalImages: result.total, loadedCount: reset ? result.images.length : state.loadedCount + result.images.length, loadingImages: false, collectionTitle: reset ? null : state.collectionTitle, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarFolderId: null, })); } catch (error) { + if (requestToken !== galleryRequestToken) return; console.error("Failed to load media:", error); - set({ loadingImages: false }); + set({ loadingImages: false, imageLoadError: String(error) }); } }, loadMoreImages: async () => { - const { loadedCount, totalImages, loadingImages } = get(); + const { loadedCount, totalImages, loadingImages, collectionTitle, similarSourceImageId, similarHasMore, similarFolderId, similarCrop } = get(); if (loadingImages || loadedCount >= totalImages) return; + if (collectionTitle === "Explore Cluster") return; + if (collectionTitle === "Similar Images" && similarSourceImageId !== null) { + if (!similarHasMore) return; + await get().loadSimilarImages(similarSourceImageId, similarFolderId, false, get().similarSourceFolderId ?? null); + return; + } + if (collectionTitle === "Region Search Results" && similarSourceImageId !== null && similarCrop !== null) { + if (!similarHasMore) return; + const requestToken = ++galleryRequestToken; + set({ loadingImages: true }); + try { + const result = await invoke("find_similar_by_region", { + params: { + image_id: similarSourceImageId, + crop_x: similarCrop.x, + crop_y: similarCrop.y, + crop_w: similarCrop.w, + crop_h: similarCrop.h, + folder_id: similarFolderId, + offset: loadedCount, + limit: PAGE_SIZE, + }, + }); + if (requestToken !== galleryRequestToken) return; + set((state) => ({ + images: [...state.images, ...result.images], + loadedCount: state.loadedCount + result.images.length, + totalImages: result.has_more ? state.loadedCount + result.images.length + 1 : state.loadedCount + result.images.length, + similarHasMore: result.has_more, + loadingImages: false, + })); + } catch { + if (requestToken !== galleryRequestToken) return; + set({ loadingImages: false }); + } + return; + } await get().loadImages(false); }, setSearch: (search) => { - set({ search, images: [], loadedCount: 0, collectionTitle: null }); + set({ search, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, clearSearch: () => { - set({ search: "", images: [], loadedCount: 0, collectionTitle: null }); + set({ search: "", images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, resetSearch: () => { - set({ search: "", searchMode: "filename", images: [], loadedCount: 0, collectionTitle: null }); + set({ search: "", searchMode: "filename", images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, setSearchMode: (searchMode) => { - set({ searchMode, images: [], loadedCount: 0, collectionTitle: null }); + set({ searchMode, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, setSort: (sort) => { - set({ sort, images: [], loadedCount: 0, collectionTitle: null }); + set({ sort, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, setMediaFilter: (mediaFilter) => { - set({ mediaFilter, images: [], loadedCount: 0, collectionTitle: null }); + set({ mediaFilter, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, setFavoritesOnly: (favoritesOnly) => { - set({ favoritesOnly, images: [], loadedCount: 0, collectionTitle: null }); + set({ favoritesOnly, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); + void get().loadImages(true); + }, + + setMinimumRating: (minimumRating) => { + set({ minimumRating, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, setFailedEmbeddingsOnly: (failedEmbeddingsOnly) => { - set({ failedEmbeddingsOnly, images: [], loadedCount: 0, collectionTitle: null }); + set({ failedEmbeddingsOnly, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, imageLoadError: null }); void get().loadImages(true); }, @@ -440,7 +897,19 @@ export const useGalleryStore = create((set, get) => ({ openImage: (image) => set({ selectedImage: image }), closeImage: () => set({ selectedImage: null }), - setView: (activeView) => set({ activeView }), + setView: (activeView) => { + if (activeView === "duplicates") { + const { selectedFolderId, duplicateScanFolderId } = get(); + if (duplicateScanFolderId !== selectedFolderId) { + set({ activeView, duplicateGroups: [], duplicateLastScanned: null, duplicateScanFolderId: undefined }); + void get().loadDuplicateScanCache(selectedFolderId); + return; + } + } + set({ activeView }); + }, + + setExploreMode: (exploreMode) => set({ exploreMode }), loadTagCloud: async () => { const { selectedFolderId, tagCloudFolderId, tagCloudLoading } = get(); @@ -448,39 +917,635 @@ export const useGalleryStore = create((set, get) => ({ if (!tagCloudLoading && tagCloudFolderId !== undefined && tagCloudFolderId === selectedFolderId) { return; } + const requestToken = ++tagCloudRequestToken; set({ tagCloudLoading: true, tagCloudFolderId: selectedFolderId }); try { const entries = await invoke("get_tag_cloud", { folderId: selectedFolderId, }); + if (requestToken !== tagCloudRequestToken) return; set({ tagCloudEntries: entries, tagCloudLoading: false }); } catch (error) { + if (requestToken !== tagCloudRequestToken) return; console.error("Failed to load tag cloud:", error); set({ tagCloudLoading: false }); } }, - searchByTag: (imageId) => { - set({ activeView: "gallery", images: [], loadedCount: 0, loadingImages: true, collectionTitle: "Similar Images" }); - void get().loadSimilarImages(imageId); + loadExploreTags: async () => { + const { selectedFolderId, exploreTagsFolderId, exploreTagLoading } = get(); + if (!exploreTagLoading && exploreTagsFolderId !== undefined && exploreTagsFolderId === selectedFolderId) { + return; + } + const requestToken = ++exploreTagRequestToken; + set({ exploreTagLoading: true, exploreTagsFolderId: selectedFolderId }); + try { + const entries = await invoke("get_explore_tags", { + params: { folder_id: selectedFolderId, limit: 48 }, + }); + if (requestToken !== exploreTagRequestToken) return; + set({ exploreTagEntries: entries, exploreTagLoading: false }); + } catch (error) { + if (requestToken !== exploreTagRequestToken) return; + console.error("Failed to load explore tags:", error); + set({ exploreTagLoading: false }); + } }, - loadSimilarImages: async (imageId) => { - set({ images: [], loadedCount: 0, loadingImages: true, collectionTitle: "Similar Images" }); - const images = await invoke("find_similar_images", { - params: { image_id: imageId, limit: PAGE_SIZE }, - }); - set({ - images, - totalImages: images.length, - loadedCount: images.length, - loadingImages: false, + showVisualCluster: async (imageIds) => { + const requestToken = ++galleryRequestToken; + set((state) => ({ + activeView: "gallery", + search: "", + images: [], + totalImages: imageIds.length, + loadedCount: 0, + loadingImages: true, + collectionTitle: "Explore Cluster", + imageLoadError: null, + similarSourceImageId: null, + similarSourceFolderId: null, + similarHasMore: false, + similarFolderId: null, + galleryScrollResetKey: state.galleryScrollResetKey + 1, + })); + + try { + const images = await invoke("get_images_by_ids", { + params: { image_ids: imageIds }, + }); + if (requestToken !== galleryRequestToken) return; + set({ + images, + totalImages: images.length, + loadedCount: images.length, + loadingImages: false, + imageLoadError: null, + collectionTitle: "Explore Cluster", + }); + } catch (error) { + if (requestToken !== galleryRequestToken) return; + set({ + images: [], + totalImages: 0, + loadedCount: 0, + loadingImages: false, + imageLoadError: String(error), + collectionTitle: "Explore Cluster", + }); + } + }, + + searchForTag: (tag) => { + set({ activeView: "gallery", search: `/t ${tag}`, images: [], loadedCount: 0, collectionTitle: null, similarSourceImageId: null, similarHasMore: false, similarFolderId: null, imageLoadError: null }); + void get().loadImages(true); + }, + + loadSimilarImages: async (imageId, folderId = get().selectedFolderId, reset = true, sourceFolderId = folderId ?? null) => { + const requestToken = ++galleryRequestToken; + const offset = reset ? 0 : get().loadedCount; + const similarScope = folderId === null ? "all_media" : "current_folder"; + set((state) => ({ + images: reset ? [] : state.images, + loadedCount: reset ? 0 : state.loadedCount, + loadingImages: true, collectionTitle: "Similar Images", - selectedFolderId: null, + imageLoadError: null, + similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarFolderId: folderId ?? null, + similarScope, + galleryScrollResetKey: reset ? state.galleryScrollResetKey + 1 : state.galleryScrollResetKey, + })); + + try { + const result = await invoke("find_similar_images", { + params: { + image_id: imageId, + folder_id: folderId ?? null, + offset, + limit: PAGE_SIZE, + threshold: SIMILAR_DISTANCE_THRESHOLD, + }, + }); + + if (requestToken !== galleryRequestToken) return; + + set((state) => { + const nextImages = reset ? result.images : [...state.images, ...result.images]; + const nextLoadedCount = nextImages.length; + return { + images: nextImages, + totalImages: result.has_more ? nextLoadedCount + 1 : nextLoadedCount, + loadedCount: nextLoadedCount, + loadingImages: false, + imageLoadError: null, + collectionTitle: "Similar Images", + similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarHasMore: result.has_more, + similarFolderId: folderId ?? null, + similarScope, + selectedImage: reset ? null : state.selectedImage, + }; + }); + } catch (error) { + if (requestToken !== galleryRequestToken) return; + console.error("Failed to load similar images:", error); + set({ + images: [], + totalImages: 0, + loadedCount: 0, + loadingImages: false, + imageLoadError: String(error), + collectionTitle: "Similar Images", + similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarHasMore: false, + similarFolderId: folderId ?? null, + similarScope, + selectedImage: null, + }); + } + }, + + loadSimilarByRegion: async (imageId, crop, folderId = get().selectedFolderId, sourceFolderId = folderId ?? null) => { + const requestToken = ++galleryRequestToken; + const similarScope = folderId === null ? "all_media" : "current_folder"; + set((state) => ({ + images: [], + loadedCount: 0, + loadingImages: true, + collectionTitle: "Region Search Results", + imageLoadError: null, + similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarFolderId: folderId ?? null, + similarCrop: crop, + similarScope, + galleryScrollResetKey: state.galleryScrollResetKey + 1, selectedImage: null, + })); + + try { + const result = await invoke("find_similar_by_region", { + params: { + image_id: imageId, + crop_x: crop.x, + crop_y: crop.y, + crop_w: crop.w, + crop_h: crop.h, + folder_id: folderId ?? null, + offset: 0, + limit: PAGE_SIZE, + }, + }); + + if (requestToken !== galleryRequestToken) return; + + set({ + images: result.images, + totalImages: result.has_more ? result.images.length + 1 : result.images.length, + loadedCount: result.images.length, + loadingImages: false, + imageLoadError: null, + collectionTitle: "Region Search Results", + similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarHasMore: result.has_more, + similarFolderId: folderId ?? null, + similarCrop: crop, + similarScope, + }); + } catch (error) { + if (requestToken !== galleryRequestToken) return; + console.error("Failed to load region search results:", error); + set({ + images: [], + totalImages: 0, + loadedCount: 0, + loadingImages: false, + imageLoadError: String(error), + collectionTitle: "Region Search Results", + similarSourceImageId: imageId, + similarSourceFolderId: sourceFolderId, + similarHasMore: false, + similarFolderId: folderId ?? null, + similarScope, + selectedImage: null, + }); + } + }, + + setSimilarScope: (similarScope) => { + set({ similarScope }); + const { similarSourceImageId, similarSourceFolderId, selectedFolderId, collectionTitle, similarCrop } = get(); + if (similarSourceImageId === null) return; + const folderId = similarScope === "current_folder" ? (similarSourceFolderId ?? selectedFolderId) : null; + if (collectionTitle === "Region Search Results" && similarCrop !== null) { + void get().loadSimilarByRegion(similarSourceImageId, similarCrop, folderId, similarSourceFolderId); + } else { + void get().loadSimilarImages(similarSourceImageId, folderId, true, similarSourceFolderId); + } + }, + + suggestImageTags: async (imageId) => { + return invoke("suggest_image_tags", { + params: { image_id: imageId, limit: 2 }, }); }, + loadCaptionModelStatus: async () => { + try { + const captionModelStatus = await invoke("get_caption_model_status"); + set({ captionModelStatus, captionModelError: null }); + } catch (error) { + set({ captionModelError: String(error) }); + } + }, + + loadCaptionAcceleration: async () => { + try { + const captionAcceleration = await invoke("get_caption_acceleration"); + set({ captionAcceleration }); + } catch (error) { + set({ captionModelError: String(error) }); + } + }, + + setCaptionAcceleration: async (acceleration) => { + const captionAcceleration = await invoke("set_caption_acceleration", { + params: { acceleration }, + }); + set({ captionAcceleration, captionRuntimeProbe: null }); + }, + + loadCaptionDetail: async () => { + try { + const captionDetail = await invoke("get_caption_detail"); + set({ captionDetail }); + } catch (error) { + set({ captionModelError: String(error) }); + } + }, + + setCaptionDetail: async (detail) => { + const captionDetail = await invoke("set_caption_detail", { + params: { detail }, + }); + set({ captionDetail, captionRuntimeProbe: null }); + }, + + prepareCaptionModel: async () => { + set({ captionModelPreparing: true, captionModelError: null, captionModelProgress: null }); + try { + const captionModelStatus = await invoke("prepare_caption_model"); + window.localStorage.setItem(AI_CAPTIONS_ENABLED_KEY, String(captionModelStatus.ready)); + set({ captionModelStatus, captionModelPreparing: false, captionModelError: null, captionModelProgress: null, aiCaptionsEnabled: captionModelStatus.ready }); + } catch (error) { + set({ captionModelPreparing: false, captionModelError: String(error), captionModelProgress: null }); + } + }, + + deleteCaptionModel: async () => { + set({ captionModelPreparing: true, captionModelError: null, captionModelProgress: null }); + try { + const captionModelStatus = await invoke("delete_caption_model"); + window.localStorage.setItem(AI_CAPTIONS_ENABLED_KEY, "false"); + set({ captionModelStatus, captionModelPreparing: false, captionModelError: null, captionModelProgress: null, captionRuntimeProbe: null, aiCaptionsEnabled: false }); + } catch (error) { + set({ captionModelPreparing: false, captionModelError: String(error), captionModelProgress: null }); + } + }, + + probeCaptionRuntime: async () => { + set({ captionRuntimeChecking: true, captionModelError: null }); + try { + const captionRuntimeProbe = await invoke("probe_caption_runtime"); + set({ captionRuntimeProbe, captionRuntimeChecking: false, captionModelError: null }); + } catch (error) { + set({ captionRuntimeChecking: false, captionModelError: String(error), captionRuntimeProbe: null }); + } + }, + + probeCaptionImage: async (imageId) => { + return invoke("probe_caption_image", { + params: { image_id: imageId }, + }); + }, + + generateCaptionForImage: async (imageId) => { + const updatedImage = await invoke("generate_caption_for_image", { + params: { image_id: imageId }, + }); + + set((state) => ({ + images: replaceImage(state.images, updatedImage, state.sort), + selectedImage: state.selectedImage?.id === updatedImage.id ? updatedImage : state.selectedImage, + })); + + return updatedImage; + }, + + queueCaptionJobs: async (folderId = get().selectedFolderId) => { + const queued = await invoke("queue_caption_jobs", { + params: { folder_id: folderId ?? null, image_id: null }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + + queueCaptionForImage: async (imageId) => { + const queued = await invoke("queue_caption_jobs", { + params: { folder_id: null, image_id: imageId }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + + clearCaptionJobs: async (folderId = get().selectedFolderId) => { + const cleared = await invoke("clear_caption_jobs", { + params: { folder_id: folderId ?? null }, + }); + await get().loadBackgroundJobProgress(); + return cleared; + }, + + resetGeneratedCaptions: async (folderId = get().selectedFolderId) => { + const reset = await invoke("reset_generated_captions", { + params: { folder_id: folderId ?? null }, + }); + await get().loadBackgroundJobProgress(); + await get().loadImages(true); + return reset; + }, + + setAiCaptionsEnabled: (aiCaptionsEnabled) => { + window.localStorage.setItem(AI_CAPTIONS_ENABLED_KEY, String(aiCaptionsEnabled)); + set({ aiCaptionsEnabled }); + }, + + setSettingsOpen: (settingsOpen) => set({ settingsOpen }), + + setTaggingQueueScope: (taggingQueueScope) => { + set((state) => ({ + taggingQueueScope, + taggingQueueFolderIds: + taggingQueueScope === "selected" && state.taggingQueueFolderIds.length === 0 && state.folders.length > 0 + ? [state.folders[0].id] + : state.taggingQueueFolderIds, + })); + }, + + toggleTaggingQueueFolder: (folderId) => { + set((state) => { + const next = state.taggingQueueFolderIds.includes(folderId) + ? state.taggingQueueFolderIds.filter((id) => id !== folderId) + : [...state.taggingQueueFolderIds, folderId].sort((a, b) => a - b); + return { taggingQueueFolderIds: next }; + }); + }, + + setTaggingQueueFolderIds: (taggingQueueFolderIds) => set({ taggingQueueFolderIds }), + + loadTaggerModelStatus: async () => { + try { + const taggerModelStatus = await invoke("get_tagger_model_status"); + set({ taggerModelStatus, taggerModelError: null }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + loadTaggerAcceleration: async () => { + try { + const taggerAcceleration = await invoke("get_tagger_acceleration"); + set({ taggerAcceleration }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + setTaggerAcceleration: async (acceleration) => { + const taggerAcceleration = await invoke("set_tagger_acceleration", { + params: { acceleration }, + }); + set({ taggerAcceleration, taggerRuntimeProbe: null }); + }, + + loadTaggerThreshold: async () => { + try { + const taggerThreshold = await invoke("get_tagger_threshold"); + set({ taggerThreshold }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + setTaggerThreshold: async (threshold) => { + const taggerThreshold = await invoke("set_tagger_threshold", { + params: { threshold }, + }); + set({ taggerThreshold }); + }, + + loadTaggerBatchSize: async () => { + try { + const taggerBatchSize = await invoke("get_tagger_batch_size"); + set({ taggerBatchSize }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + setTaggerBatchSize: async (batchSize) => { + const taggerBatchSize = await invoke("set_tagger_batch_size", { + params: { batch_size: batchSize }, + }); + set({ taggerBatchSize }); + }, + + prepareTaggerModel: async () => { + set({ taggerModelPreparing: true, taggerModelError: null, taggerModelProgress: null }); + try { + const taggerModelStatus = await invoke("prepare_tagger_model"); + set({ taggerModelStatus, taggerModelPreparing: false, taggerModelError: null, taggerModelProgress: null }); + } catch (error) { + set({ taggerModelPreparing: false, taggerModelError: String(error), taggerModelProgress: null }); + } + }, + + deleteTaggerModel: async () => { + set({ taggerModelPreparing: true, taggerModelError: null, taggerModelProgress: null }); + try { + const taggerModelStatus = await invoke("delete_tagger_model"); + set({ taggerModelStatus, taggerModelPreparing: false, taggerModelError: null, taggerModelProgress: null, taggerRuntimeProbe: null }); + } catch (error) { + set({ taggerModelPreparing: false, taggerModelError: String(error), taggerModelProgress: null }); + } + }, + + probeTaggerRuntime: async () => { + set({ taggerRuntimeChecking: true, taggerModelError: null }); + try { + const taggerRuntimeProbe = await invoke("probe_tagger_runtime"); + set({ taggerRuntimeProbe, taggerRuntimeChecking: false, taggerModelError: null }); + } catch (error) { + set({ taggerRuntimeChecking: false, taggerModelError: String(error), taggerRuntimeProbe: null }); + } + }, + + queueTaggingJobs: async (folderId = get().selectedFolderId) => { + const queued = await invoke("queue_tagging_jobs", { + params: { folder_id: folderId ?? null, image_id: null }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + + queueTaggingJobsForFolders: async (folderIds) => { + const queued = await invoke("queue_tagging_jobs", { + params: { folder_id: null, folder_ids: folderIds, image_id: null }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + + queueTaggingForImage: async (imageId) => { + const queued = await invoke("queue_tagging_jobs", { + params: { folder_id: null, image_id: imageId }, + }); + await get().loadBackgroundJobProgress(); + return queued; + }, + + clearTaggingJobs: async (folderId = get().selectedFolderId) => { + const cleared = await invoke("clear_tagging_jobs", { + params: { folder_id: folderId ?? null }, + }); + await get().loadBackgroundJobProgress(); + return cleared; + }, + + clearTaggingJobsForFolders: async (folderIds) => { + const cleared = await invoke("clear_tagging_jobs", { + params: { folder_id: null, folder_ids: folderIds }, + }); + await get().loadBackgroundJobProgress(); + return cleared; + }, + + getImageTags: async (imageId) => { + return invoke("get_image_tags", { + params: { image_id: imageId }, + }); + }, + + addUserTag: async (imageId, tag) => { + const result = await invoke("add_user_tag", { + params: { image_id: imageId, tag }, + }); + // Invalidate explore tags cache so new tag appears immediately + set({ exploreTagsFolderId: undefined }); + return result; + }, + + removeTag: async (tagId) => { + await invoke("remove_tag", { + params: { tag_id: tagId }, + }); + // Invalidate explore tags cache so removed tag disappears immediately + set({ exploreTagsFolderId: undefined }); + }, + + loadDuplicateScanCache: async (folderId = null) => { + interface CacheResult { groups: DuplicateGroup[]; scanned_at: number } + const cached = await invoke("load_duplicate_scan_cache", { folderId: folderId ?? null }); + if (cached) { + set({ duplicateGroups: cached.groups, duplicateLastScanned: cached.scanned_at, duplicateScanFolderId: folderId }); + } + }, + + scanDuplicates: async (folderId = null) => { + const { listen } = await import("@tauri-apps/api/event"); + set({ duplicateScanning: true, duplicateGroups: [], duplicateScanProgress: null, duplicateScanError: null, duplicateSelectedIds: new Set() }); + const unlisten = await listen<[number, number]>("duplicate_scan_progress", (event) => { + const [scanned, total] = event.payload; + set({ duplicateScanProgress: { scanned, total } }); + }); + try { + const groups = await invoke("find_duplicates", { folderId: folderId ?? null }); + set({ duplicateGroups: groups, duplicateLastScanned: Math.floor(Date.now() / 1000), duplicateScanFolderId: folderId }); + void notifyTaskComplete( + "Duplicate scan complete", + groups.length === 1 ? "Found 1 duplicate group." : `Found ${groups.length.toLocaleString()} duplicate groups.`, + ); + } catch (e) { + set({ duplicateScanError: String(e) }); + } finally { + unlisten(); + set({ duplicateScanning: false }); + } + }, + + toggleDuplicateSelected: (imageId) => { + set((state) => { + const next = new Set(state.duplicateSelectedIds); + if (next.has(imageId)) next.delete(imageId); + else next.add(imageId); + return { duplicateSelectedIds: next }; + }); + }, + + selectAllDuplicates: (imageIds) => { + set((state) => { + const next = new Set(state.duplicateSelectedIds); + for (const id of imageIds) next.add(id); + return { duplicateSelectedIds: next }; + }); + }, + + selectKeepFirstAllGroups: () => { + const { duplicateGroups } = get(); + const toMark = new Set(); + for (const group of duplicateGroups) { + for (const img of group.images.slice(1)) toMark.add(img.id); + } + set({ duplicateSelectedIds: toMark }); + }, + + clearDuplicateSelection: () => set({ duplicateSelectedIds: new Set() }), + + deleteSelectedDuplicates: async () => { + const { duplicateSelectedIds, duplicateGroups } = get(); + const ids = Array.from(duplicateSelectedIds); + if (ids.length === 0) return 0; + // Backend returns only the IDs that were actually removed from disk. + const succeededIds = await invoke("delete_images_from_disk", { params: { image_ids: ids } }); + const succeededSet = new Set(succeededIds); + // Only remove images confirmed deleted — failed files remain visible so the user can retry. + set((state) => ({ + duplicateSelectedIds: new Set(), + duplicateGroups: state.duplicateGroups + .map((g) => ({ ...g, images: g.images.filter((img) => !succeededSet.has(img.id)) })) + .filter((g) => g.images.length > 1), + })); + // Invalidate the persisted cache for every affected scope: + // - global "all" cache (always, since a folder-scoped deletion still makes the global result stale) + // - each folder that contained a deleted image (so a folder-scoped scan is also evicted) + const affectedFolderIds = new Set( + duplicateGroups + .flatMap((g) => g.images) + .filter((img) => succeededSet.has(img.id)) + .map((img) => img.folder_id), + ); + await invoke("invalidate_duplicate_scan_cache", { folderId: null }); // global + for (const folderId of affectedFolderIds) { + await invoke("invalidate_duplicate_scan_cache", { folderId }); + } + return succeededIds.length; + }, + retryFailedEmbeddings: async (folderId) => { await invoke("retry_failed_embeddings", { params: { folder_id: folderId } }); await get().loadBackgroundJobProgress(); @@ -504,6 +1569,7 @@ export const useGalleryStore = create((set, get) => ({ subscribeToProgress: async () => { const unlistenProgress = await listen("index-progress", (event) => { const progress = event.payload; + const previous = get().indexingProgress[progress.folder_id]; set((state) => ({ indexingProgress: { ...state.indexingProgress, @@ -512,9 +1578,23 @@ export const useGalleryStore = create((set, get) => ({ })); if (progress.done) { + if ( + previous && + !previous.done && + progress.total > 0 && + progress.indexed >= progress.total + ) { + const folderName = get().folders.find((folder) => folder.id === progress.folder_id)?.name; + void notifyTaskComplete( + "Folder scan complete", + folderName ? `${folderName} has finished scanning.` : "A folder has finished scanning.", + ); + } void get().loadFolders(); void get().loadBackgroundJobProgress(); - void get().loadImages(true); + if (get().activeView !== "explore" && !isDerivedCollectionTitle(get().collectionTitle)) { + void get().loadImages(true); + } setTimeout(() => { set((state) => { @@ -527,6 +1607,41 @@ export const useGalleryStore = create((set, get) => ({ }); const unlistenMediaJobs = await listen("media-job-progress", (event) => { + const previousProgress = get().mediaJobProgress; + + for (const progress of event.payload.progress) { + const previous = previousProgress[progress.folder_id]; + if (!previous) continue; + + const folderName = + get().folders.find((folder) => folder.id === progress.folder_id)?.name ?? "Folder"; + + if (previous.embedding_pending > 0 && progress.embedding_pending === 0) { + const failureDetail = + progress.embedding_failed > 0 + ? ` ${progress.embedding_failed.toLocaleString()} failed.` + : ""; + void notifyTaskComplete( + "Embeddings complete", + `${folderName} finished generating embeddings.${failureDetail}`, + ); + } + + if (previous.tagging_pending > 0 && progress.tagging_pending === 0) { + const failureDetail = + progress.tagging_failed > 0 + ? ` ${progress.tagging_failed.toLocaleString()} failed.` + : ""; + void notifyTaskComplete( + "AI tagging complete", + `${folderName} finished generating tags.${failureDetail}`, + ); + // New tags are now in the DB — invalidate the Explore tag cache so + // reopening Explore reflects the updated tag distribution. + set({ exploreTagsFolderId: undefined }); + } + } + set((state) => { const next = { ...state.mediaJobProgress }; for (const progress of event.payload.progress) { @@ -536,16 +1651,35 @@ export const useGalleryStore = create((set, get) => ({ }); }); + const unlistenCaptionModelProgress = await listen("caption-model-progress", (event) => { + set({ + captionModelProgress: event.payload.done ? null : event.payload, + captionModelPreparing: !event.payload.done, + }); + }); + + const unlistenTaggerModelProgress = await listen("tagger-model-progress", (event) => { + set({ + taggerModelProgress: event.payload.done ? null : event.payload, + taggerModelPreparing: !event.payload.done, + }); + }); + const unlistenImages = await listen("indexed-images", (event) => { const batch = event.payload; set((state) => { + if (isDerivedCollectionTitle(state.collectionTitle) || state.activeView === "explore") { + return state; + } + const visibleImages = batch.images.filter((image) => matchesFilters( image, state.selectedFolderId, state.mediaFilter, state.favoritesOnly, + state.minimumRating, state.failedEmbeddingsOnly, state.search, ), @@ -570,12 +1704,21 @@ export const useGalleryStore = create((set, get) => ({ const batch = event.payload; set((state) => { + if (isDerivedCollectionTitle(state.collectionTitle) || state.activeView === "explore") { + const selectedImage = + state.selectedImage && batch.images.some((image) => image.id === state.selectedImage?.id) + ? batch.images.find((image) => image.id === state.selectedImage?.id) ?? state.selectedImage + : state.selectedImage; + return { selectedImage }; + } + const visibleImages = batch.images.filter((image) => matchesFilters( image, state.selectedFolderId, state.mediaFilter, state.favoritesOnly, + state.minimumRating, state.failedEmbeddingsOnly, state.search, ), @@ -600,6 +1743,8 @@ export const useGalleryStore = create((set, get) => ({ return () => { unlistenProgress(); unlistenMediaJobs(); + unlistenCaptionModelProgress(); + unlistenTaggerModelProgress(); unlistenImages(); unlistenThumbnails(); };