diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index bf4d4f6..217eabf 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -5,6 +5,7 @@ mod embedder; mod indexer; mod media; mod storage; +mod tagger; mod thumbnail; mod vector; @@ -67,7 +68,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()); - indexer::start_caption_worker(app.handle().clone(), pool.clone(), app_dir.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); @@ -87,18 +90,37 @@ pub fn run() { commands::retry_failed_embeddings, commands::semantic_search_images, 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_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::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, ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/src/components/BackgroundTasks.tsx b/src/components/BackgroundTasks.tsx index 6d4014c..bc007a5 100644 --- a/src/components/BackgroundTasks.tsx +++ b/src/components/BackgroundTasks.tsx @@ -2,13 +2,13 @@ import { useEffect, useMemo, useState } from "react"; import { invoke } from "@tauri-apps/api/core"; import { useGalleryStore } from "../store"; -type WorkerKey = "thumbnail" | "metadata" | "embedding" | "caption"; +type WorkerKey = "thumbnail" | "metadata" | "embedding" | "tagging"; const WORKER_FOR_STAGE: Record = { Thumbnails: "thumbnail", Metadata: "metadata", Embeddings: "embedding", - Captions: "caption", + Tags: "tagging", }; interface TaskStage { @@ -23,7 +23,7 @@ interface Task { name: string; stages: TaskStage[]; hasFailedEmbeddings: boolean; - hasFailedCaptions: boolean; + hasFailedTagging: boolean; pendingMediaWork: number; embeddingProcessed: number; embeddingTotal: number; @@ -42,14 +42,14 @@ interface FolderWorkerStates { thumbnail_paused: boolean; metadata_paused: boolean; embedding_paused: boolean; - caption_paused: boolean; + tagging_paused: boolean; } const DEFAULT_PAUSED_STATE: Record = { thumbnail: false, metadata: false, embedding: false, - caption: false, + tagging: false, }; export function BackgroundTasks() { @@ -57,6 +57,7 @@ export function BackgroundTasks() { const indexingProgress = useGalleryStore((state) => state.indexingProgress); const mediaJobProgress = useGalleryStore((state) => state.mediaJobProgress); const retryFailedEmbeddings = useGalleryStore((state) => state.retryFailedEmbeddings); + const clearTaggingJobs = useGalleryStore((state) => state.clearTaggingJobs); const [expanded, setExpanded] = useState(false); const [dismissed, setDismissed] = useState>({}); const [paused, setPaused] = useState>>({}); @@ -78,7 +79,7 @@ export function BackgroundTasks() { thumbnail: state.thumbnail_paused, metadata: state.metadata_paused, embedding: state.embedding_paused, - caption: state.caption_paused, + tagging: state.tagging_paused, }, ]), ), @@ -126,6 +127,7 @@ export function BackgroundTasks() { }; const dismissTask = (id: number, snapshot: string) => { + void clearTaggingJobs(id); setDismissed((prev) => ({ ...prev, [id]: snapshot })); setExpanded(false); }; @@ -141,16 +143,19 @@ export function BackgroundTasks() { const embeddingPending = jobs?.embedding_pending ?? 0; const embeddingReady = jobs?.embedding_ready ?? 0; const embeddingFailed = jobs?.embedding_failed ?? 0; - const captionPending = jobs?.caption_pending ?? 0; - const captionFailed = jobs?.caption_failed ?? 0; + const taggingPending = jobs?.tagging_pending ?? 0; + const taggingReady = jobs?.tagging_ready ?? 0; + const taggingFailed = jobs?.tagging_failed ?? 0; - const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending + captionPending; + const pendingMediaWork = thumbnailPending + metadataPending + embeddingPending + taggingPending; const embeddingProcessed = embeddingReady + embeddingFailed; const embeddingTotal = embeddingProcessed + embeddingPending; + const taggingProcessed = taggingReady + taggingFailed; + const taggingTotal = taggingProcessed + taggingPending; const hasFailedEmbeddings = embeddingFailed > 0; - const hasFailedCaptions = captionFailed > 0; + const hasFailedTagging = taggingFailed > 0; - if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings && !hasFailedCaptions) return null; + if (!index && pendingMediaWork === 0 && !hasFailedEmbeddings && !hasFailedTagging) return null; const stages: TaskStage[] = []; @@ -192,11 +197,12 @@ export function BackgroundTasks() { }); } - if (captionPending > 0) { + if (taggingPending > 0) { + const pct = taggingTotal > 0 ? (taggingProcessed / taggingTotal) * 100 : 0; stages.push({ - label: "Captions", - detail: captionPending.toLocaleString(), - progress: null, + label: "Tags", + detail: `${taggingProcessed.toLocaleString()} / ${taggingTotal.toLocaleString()}`, + progress: pct, failed: false, }); } @@ -210,23 +216,23 @@ export function BackgroundTasks() { }); } - if (hasFailedCaptions && pendingMediaWork === 0) { + if (hasFailedTagging && pendingMediaWork === 0) { stages.push({ label: "Failed", - detail: `${captionFailed.toLocaleString()} captions`, + detail: `${taggingFailed.toLocaleString()} tags`, progress: null, failed: true, }); } - const snapshot = `${pendingMediaWork}:${embeddingFailed}:${captionFailed}`; + const snapshot = `${pendingMediaWork}:${embeddingFailed}:${taggingFailed}`; return { id: folder.id, name: folder.name, stages, hasFailedEmbeddings, - hasFailedCaptions, + hasFailedTagging, pendingMediaWork, embeddingProcessed, embeddingTotal, @@ -242,13 +248,14 @@ export function BackgroundTasks() { const primary = tasks[0]; const extraCount = tasks.length - 1; - const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedCaptions) && t.pendingMediaWork === 0); + const hasFailed = tasks.some((t) => (t.hasFailedEmbeddings || t.hasFailedTagging) && 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 (
@@ -368,9 +375,10 @@ export function BackgroundTasks() {
{tasks.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.hasFailedCaptions) && task.pendingMediaWork === 0; + const taskBarProgress = taskEmbeddingStage?.progress ?? taskTaggingStage?.progress ?? taskScanningStage?.progress ?? null; + const taskHasFailed = (task.hasFailedEmbeddings || task.hasFailedTagging) && task.pendingMediaWork === 0; return (
diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx index 5460b16..cb3af78 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,29 +45,37 @@ 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" }; + } +} + export function Lightbox() { const selectedImage = useGalleryStore((state) => state.selectedImage); - const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); 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 updateImageDetails = useGalleryStore((state) => state.updateImageDetails); - const suggestImageTags = useGalleryStore((state) => state.suggestImageTags); - const captionModelStatus = useGalleryStore((state) => state.captionModelStatus); - const captionModelPreparing = useGalleryStore((state) => state.captionModelPreparing); - const captionModelError = useGalleryStore((state) => state.captionModelError); - const captionModelProgress = useGalleryStore((state) => state.captionModelProgress); - const aiCaptionsEnabled = useGalleryStore((state) => state.aiCaptionsEnabled); - const setAiCaptionsEnabled = useGalleryStore((state) => state.setAiCaptionsEnabled); - const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); - const prepareCaptionModel = useGalleryStore((state) => state.prepareCaptionModel); - const queueCaptionForImage = useGalleryStore((state) => state.queueCaptionForImage); - const queueCaptionJobs = useGalleryStore((state) => state.queueCaptionJobs); + const 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); const [zoom, setZoom] = useState(1); - const [suggestedTags, setSuggestedTags] = useState([]); - const [captionQueueing, setCaptionQueueing] = useState(false); - const [captionQueueStatus, setCaptionQueueStatus] = useState(null); + const [imageTags, setImageTags] = useState([]); + const [tagInput, setTagInput] = useState(""); + const [tagAdding, setTagAdding] = useState(false); + const [tagsExpanded, setTagsExpanded] = useState(false); + const [taggingQueued, setTaggingQueued] = useState(false); const imageViewportRef = useRef(null); const currentIndex = selectedImage ? images.findIndex((image) => image.id === selectedImage.id) : -1; @@ -83,16 +91,18 @@ export function Lightbox() { useEffect(() => { setZoom(1); - setSuggestedTags([]); - setCaptionQueueStatus(null); + setImageTags([]); + setTagInput(""); + setTagsExpanded(false); + setTaggingQueued(false); }, [selectedImage?.id]); useEffect(() => { - if (!selectedImage?.generated_caption) return; - void suggestImageTags(selectedImage.id) - .then(setSuggestedTags) - .catch(() => setSuggestedTags([])); - }, [selectedImage?.id, selectedImage?.generated_caption, suggestImageTags]); + if (!selectedImage) return; + void getImageTags(selectedImage.id) + .then(setImageTags) + .catch(() => setImageTags([])); + }, [selectedImage?.id, getImageTags]); useEffect(() => { const viewport = imageViewportRef.current; @@ -207,8 +217,8 @@ export function Lightbox() {
-
-
+
+

{selectedImage.filename}

Details

@@ -245,7 +255,7 @@ export function Lightbox() {
-
+

Rating

@@ -341,116 +351,102 @@ export function Lightbox() {
-

AI Caption

- {selectedImage.generated_caption ? ( - <> -

{selectedImage.generated_caption}

- {suggestedTags.length > 0 ? ( -
- {suggestedTags.map((tag) => ( - - {tag} - - ))} -
- ) : null} - {selectedImage.caption_model ? ( -

{selectedImage.caption_model}

- ) : null} - - ) : ( - <> -

- {selectedImage.caption_error ?? "Not generated"} -

- {captionModelStatus?.ready ? ( - <> -

- {aiCaptionsEnabled ? "Florence-2 enabled locally" : "Florence-2 downloaded but disabled"} -

-
- - {selectedImage.media_kind === "image" && aiCaptionsEnabled ? ( - - ) : null} - {aiCaptionsEnabled ? ( - - ) : null} -
- {captionQueueStatus ? ( -

{captionQueueStatus}

- ) : null} - - ) : ( - - )} - {captionModelProgress?.current_file ? ( -

{captionModelProgress.current_file}

- ) : null} - {captionModelError ? ( -

{captionModelError}

- ) : null} - {!captionModelStatus?.ready && !captionModelError ? ( -

Downloads Florence-2 on demand for offline captions.

+
+

Tags

+
+ {selectedImage.ai_rating ? ( + + {ratingPill(selectedImage.ai_rating).label} + ) : 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); + void addUserTag(selectedImage.id, raw) + .then((newTag) => { + setImageTags((prev) => [...prev, newTag]); + setTagInput(""); + }) + .catch(() => undefined) + .finally(() => setTagAdding(false)); + }} + > + setTagInput(e.target.value)} + disabled={tagAdding} + /> + +
@@ -459,7 +455,7 @@ export function Lightbox() {
-
+
{currentIndex + 1} / {images.length}
diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx index 3afb757..cdf9ddd 100644 --- a/src/components/SettingsModal.tsx +++ b/src/components/SettingsModal.tsx @@ -1,48 +1,15 @@ import { useEffect, useState } from "react"; -import { useGalleryStore } from "../store"; +import { TaggerAcceleration, useGalleryStore } from "../store"; -type SettingsSection = "ai" | "library" | "display" | "storage"; +type SettingsSection = "tagging" | "library" | "display" | "storage"; const SECTIONS: { id: SettingsSection; label: string; detail: string }[] = [ - { id: "ai", label: "Local AI", detail: "Captions and suggestions" }, + { id: "tagging", label: "AI Tagging", detail: "WD tagger model" }, { id: "library", label: "Library", detail: "Indexing and scanning" }, { id: "display", label: "Display", detail: "Gallery preferences" }, { id: "storage", label: "Storage", detail: "Cache and model files" }, ]; -function ToggleSwitch({ - checked, - disabled, - onChange, - label, -}: { - checked: boolean; - disabled?: boolean; - onChange: (checked: boolean) => void; - label: string; -}) { - return ( - - ); -} function StatusPill({ children, tone }: { children: React.ReactNode; tone: "ready" | "muted" | "busy" }) { const className = @@ -59,6 +26,34 @@ function StatusPill({ children, tone }: { children: React.ReactNode; tone: "read ); } +function TaggerAccelerationButton({ + acceleration, + current, + onSelect, + children, +}: { + acceleration: TaggerAcceleration; + current: TaggerAcceleration; + onSelect: (acceleration: TaggerAcceleration) => void; + children: React.ReactNode; +}) { + const active = acceleration === current; + return ( + + ); +} + + function SettingsRow({ title, description, @@ -98,50 +93,52 @@ function SectionShell({ } export function SettingsModal() { - const [activeSection, setActiveSection] = useState("ai"); - const [captionQueueStatus, setCaptionQueueStatus] = useState(null); - const [captionQueueing, setCaptionQueueing] = useState(false); + const [activeSection, setActiveSection] = useState("tagging"); + 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 settingsOpen = useGalleryStore((state) => state.settingsOpen); const setSettingsOpen = useGalleryStore((state) => state.setSettingsOpen); - const folders = useGalleryStore((state) => state.folders); const selectedFolderId = useGalleryStore((state) => state.selectedFolderId); - const captionModelStatus = useGalleryStore((state) => state.captionModelStatus); - const captionModelPreparing = useGalleryStore((state) => state.captionModelPreparing); - const captionModelProgress = useGalleryStore((state) => state.captionModelProgress); - const captionModelError = useGalleryStore((state) => state.captionModelError); - const captionRuntimeProbe = useGalleryStore((state) => state.captionRuntimeProbe); - const captionRuntimeChecking = useGalleryStore((state) => state.captionRuntimeChecking); - const aiCaptionsEnabled = useGalleryStore((state) => state.aiCaptionsEnabled); - const setAiCaptionsEnabled = useGalleryStore((state) => state.setAiCaptionsEnabled); - const prepareCaptionModel = useGalleryStore((state) => state.prepareCaptionModel); - const deleteCaptionModel = useGalleryStore((state) => state.deleteCaptionModel); - const probeCaptionRuntime = useGalleryStore((state) => state.probeCaptionRuntime); - const queueCaptionJobs = useGalleryStore((state) => state.queueCaptionJobs); + const folders = useGalleryStore((state) => state.folders); + 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 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 probeTaggerRuntime = useGalleryStore((state) => state.probeTaggerRuntime); + const queueTaggingJobs = useGalleryStore((state) => state.queueTaggingJobs); + const clearTaggingJobs = useGalleryStore((state) => state.clearTaggingJobs); useEffect(() => { if (!settingsOpen) return; + void loadTaggerModelStatus(); + void loadTaggerAcceleration(); + void loadTaggerThreshold(); const handleKeyDown = (event: KeyboardEvent) => { if (event.key === "Escape") setSettingsOpen(false); }; window.addEventListener("keydown", handleKeyDown); return () => window.removeEventListener("keydown", handleKeyDown); - }, [settingsOpen, setSettingsOpen]); + }, [settingsOpen, loadTaggerModelStatus, loadTaggerAcceleration, loadTaggerThreshold, setSettingsOpen]); if (!settingsOpen) return null; - const modelReady = captionModelStatus?.ready ?? false; - const downloadLabel = captionModelProgress - ? `Downloading ${captionModelProgress.completed_files}/${captionModelProgress.total_files}` - : captionModelPreparing - ? "Preparing Florence-2..." - : modelReady - ? "Downloaded" - : "Download Florence-2"; - const downloadPercent = captionModelProgress - ? Math.round((captionModelProgress.completed_files / Math.max(captionModelProgress.total_files, 1)) * 100) - : 0; const selectedFolder = folders.find((folder) => folder.id === selectedFolderId); - const captionScopeLabel = selectedFolder ? selectedFolder.name : "all libraries"; + const scopeLabel = selectedFolder ? selectedFolder.name : "all libraries"; return (
-
-
- {modelReady ? Florence-2 ready : Optional} - {captionModelPreparing ? Working : null} -
+
- {activeSection === "ai" ? ( - - - - - - -
- - {modelReady ? ( - <> - - - - ) : null} -
- - - - - +
+ + {taggerReady ? ( + <> + + + + ) : null} +
+ -
-

Model location

-

- {modelReady ? captionModelStatus?.local_dir : "Not downloaded"} -

- {captionModelProgress?.current_file ? ( -

{captionModelProgress.current_file}

- ) : null} - {captionModelError ? ( -

{captionModelError}

- ) : null} - {captionQueueStatus ? ( -

{captionQueueStatus}

- ) : null} - {captionRuntimeProbe ? ( -
-
-

Runtime check

- Ready + +
+
+ { + setTaggerAccelerationSaving(true); + void setTaggerAcceleration(acceleration) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => setTaggerAccelerationSaving(false)); + }} + > + Auto + + { + setTaggerAccelerationSaving(true); + void setTaggerAcceleration(acceleration) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => setTaggerAccelerationSaving(false)); + }} + > + DirectML + + { + setTaggerAccelerationSaving(true); + void setTaggerAcceleration(acceleration) + .catch((error) => setTaggerQueueStatus(String(error))) + .finally(() => setTaggerAccelerationSaving(false)); + }} + > + CPU +
-

- Tokenizer vocabulary: {captionRuntimeProbe.tokenizer_vocab_size.toLocaleString()} +

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

-
- {captionRuntimeProbe.sessions.map((session) => ( -
-

{session.file}

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

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

+
+ + + +
+ + +
+ + +
+

Model location

+

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

+ {taggerModelProgress?.current_file ? ( +

{taggerModelProgress.current_file}

+ ) : null} + {taggerModelError ? ( +

{taggerModelError}

+ ) : null} + {taggerQueueStatus ? ( +

{taggerQueueStatus}

+ ) : null} + {taggerRuntimeProbe ? ( +
+
+

Runtime check

+ Ready +
+

+ Tagger acceleration: {taggerRuntimeProbe.acceleration} +

+
+
+

{taggerRuntimeProbe.session.file}

- {session.inputs.length} input{session.inputs.length === 1 ? "" : "s"} · {session.outputs.length} output{session.outputs.length === 1 ? "" : "s"} + {taggerRuntimeProbe.session.inputs.length} input{taggerRuntimeProbe.session.inputs.length === 1 ? "" : "s"} · {taggerRuntimeProbe.session.outputs.join(", ")}

- ))} +
-
- ) : null} -
- - ) : null} + ) : null} +
+ + ); + })() : null} {activeSection === "library" ? ( @@ -335,11 +443,11 @@ export function SettingsModal() { {activeSection === "storage" ? ( - + diff --git a/src/store.ts b/src/store.ts index 5af3a44..32ae1fc 100644 --- a/src/store.ts +++ b/src/store.ts @@ -15,6 +15,10 @@ export type MediaKind = "image" | "video"; export type MediaFilter = "all" | MediaKind; export type ZoomPreset = "compact" | "comfortable" | "detail"; export type SearchMode = "filename" | "semantic"; +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 interface ImageRecord { id: number; @@ -44,6 +48,35 @@ export interface ImageRecord { 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 { @@ -64,6 +97,9 @@ export interface FolderJobProgress { caption_pending: number; caption_ready: number; caption_failed: number; + tagging_pending: number; + tagging_ready: number; + tagging_failed: number; } export interface MediaJobProgressEvent { @@ -110,6 +146,8 @@ export interface CaptionRuntimeSessionProbe { export interface CaptionRuntimeProbe { ready: boolean; + acceleration: CaptionAcceleration; + detail: CaptionDetail; tokenizer_vocab_size: number; sessions: CaptionRuntimeSessionProbe[]; } @@ -118,6 +156,19 @@ 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 type SortOrder = @@ -163,9 +214,20 @@ interface GalleryState { captionModelProgress: CaptionModelProgress | null; captionRuntimeProbe: CaptionRuntimeProbe | null; captionRuntimeChecking: boolean; + captionAcceleration: CaptionAcceleration; + captionDetail: CaptionDetail; aiCaptionsEnabled: boolean; settingsOpen: boolean; + taggerModelStatus: TaggerModelStatus | null; + taggerModelPreparing: boolean; + taggerModelError: string | null; + taggerModelProgress: TaggerModelProgress | null; + taggerAcceleration: TaggerAcceleration; + taggerThreshold: number; + taggerRuntimeProbe: TaggerRuntimeProbe | null; + taggerRuntimeChecking: boolean; + loadFolders: () => Promise; loadBackgroundJobProgress: () => Promise; addFolder: (path: string) => Promise; @@ -198,12 +260,33 @@ interface GalleryState { 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; 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; + probeTaggerRuntime: () => Promise; + queueTaggingJobs: (folderId?: number | null) => Promise; + queueTaggingForImage: (imageId: number) => Promise; + clearTaggingJobs: (folderId?: number | null) => Promise; + getImageTags: (imageId: number) => Promise; + addUserTag: (imageId: number, tag: string) => Promise; + removeTag: (tagId: number) => Promise; } const PAGE_SIZE = 200; @@ -361,9 +444,20 @@ export const useGalleryStore = create((set, get) => ({ captionModelProgress: null, captionRuntimeProbe: null, captionRuntimeChecking: false, + captionAcceleration: "auto", + captionDetail: "paragraph", aiCaptionsEnabled: initialAiCaptionsEnabled(), settingsOpen: false, + taggerModelStatus: null, + taggerModelPreparing: false, + taggerModelError: null, + taggerModelProgress: null, + taggerAcceleration: "auto", + taggerThreshold: 0.35, + taggerRuntimeProbe: null, + taggerRuntimeChecking: false, + setCacheDir: (cacheDir) => set({ cacheDir }), loadFolders: async () => { @@ -567,10 +661,10 @@ export const useGalleryStore = create((set, get) => ({ similarSourceImageId: imageId, galleryScrollResetKey: reset ? state.galleryScrollResetKey + 1 : state.galleryScrollResetKey, })); - try { - const images = await invoke("find_similar_images", { - params: { image_id: imageId, limit: requestedLimit }, - }); + try { + const images = await invoke("find_similar_images", { + params: { image_id: imageId, folder_id: folderId ?? null, limit: requestedLimit }, + }); const hasMore = images.length >= requestedLimit; set({ images, @@ -616,6 +710,38 @@ export const useGalleryStore = create((set, get) => ({ } }, + 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 { @@ -683,6 +809,23 @@ export const useGalleryStore = create((set, get) => ({ 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 }); @@ -690,6 +833,119 @@ export const useGalleryStore = create((set, get) => ({ setSettingsOpen: (settingsOpen) => set({ settingsOpen }), + 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 }); + }, + + 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; + }, + + 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; + }, + + getImageTags: async (imageId) => { + return invoke("get_image_tags", { + params: { image_id: imageId }, + }); + }, + + addUserTag: async (imageId, tag) => { + return invoke("add_user_tag", { + params: { image_id: imageId, tag }, + }); + }, + + removeTag: async (tagId) => { + await invoke("remove_tag", { + params: { tag_id: tagId }, + }); + }, + retryFailedEmbeddings: async (folderId) => { await invoke("retry_failed_embeddings", { params: { folder_id: folderId } }); await get().loadBackgroundJobProgress(); @@ -752,6 +1008,13 @@ export const useGalleryStore = create((set, get) => ({ }); }); + 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; @@ -817,6 +1080,7 @@ export const useGalleryStore = create((set, get) => ({ unlistenProgress(); unlistenMediaJobs(); unlistenCaptionModelProgress(); + unlistenTaggerModelProgress(); unlistenImages(); unlistenThumbnails(); };