Merge AI tagger readiness fixes
Refresh selected tagger readiness at startup so the lightbox AI tags action reflects installed model state without needing a Settings refresh. Add UI Lab coverage for uninstalled WD and JoyTag scenarios, and include the toolbar clear-button alignment follow-up.
This commit is contained in:
@@ -466,7 +466,8 @@ pub fn probe_tagger_runtime(app_data_dir: &Path) -> Result<TaggerRuntimeProbe> {
|
|||||||
let status = tagger_model_status(app_data_dir);
|
let status = tagger_model_status(app_data_dir);
|
||||||
if !status.ready {
|
if !status.ready {
|
||||||
anyhow::bail!(
|
anyhow::bail!(
|
||||||
"WD Tagger model is missing {} required file(s): {}",
|
"{} is missing {} required file(s): {}",
|
||||||
|
status.model_name,
|
||||||
status.missing_files.len(),
|
status.missing_files.len(),
|
||||||
status.missing_files.join(", ")
|
status.missing_files.join(", ")
|
||||||
);
|
);
|
||||||
|
|||||||
@@ -23,6 +23,8 @@ export default function App() {
|
|||||||
const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress);
|
const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress);
|
||||||
const loadImages = useGalleryStore((state) => state.loadImages);
|
const loadImages = useGalleryStore((state) => state.loadImages);
|
||||||
const loadCaptionModelStatus = useGalleryStore((state) => state.loadCaptionModelStatus);
|
const loadCaptionModelStatus = useGalleryStore((state) => state.loadCaptionModelStatus);
|
||||||
|
const loadTaggerModelStatus = useGalleryStore((state) => state.loadTaggerModelStatus);
|
||||||
|
const loadTaggerModel = useGalleryStore((state) => state.loadTaggerModel);
|
||||||
const loadDuplicateScanCache = useGalleryStore((state) => state.loadDuplicateScanCache);
|
const loadDuplicateScanCache = useGalleryStore((state) => state.loadDuplicateScanCache);
|
||||||
const loadAlbums = useGalleryStore((state) => state.loadAlbums);
|
const loadAlbums = useGalleryStore((state) => state.loadAlbums);
|
||||||
const loadMutedFolderIds = useGalleryStore((state) => state.loadMutedFolderIds);
|
const loadMutedFolderIds = useGalleryStore((state) => state.loadMutedFolderIds);
|
||||||
@@ -53,6 +55,8 @@ export default function App() {
|
|||||||
loadFolders().then(async () => {
|
loadFolders().then(async () => {
|
||||||
void loadBackgroundJobProgress();
|
void loadBackgroundJobProgress();
|
||||||
void loadCaptionModelStatus();
|
void loadCaptionModelStatus();
|
||||||
|
void loadTaggerModel();
|
||||||
|
void loadTaggerModelStatus();
|
||||||
void loadDuplicateScanCache();
|
void loadDuplicateScanCache();
|
||||||
await loadAlbums();
|
await loadAlbums();
|
||||||
await loadImages(true);
|
await loadImages(true);
|
||||||
|
|||||||
@@ -156,6 +156,7 @@ export function Lightbox() {
|
|||||||
const addUserTag = useGalleryStore((state) => state.addUserTag);
|
const addUserTag = useGalleryStore((state) => state.addUserTag);
|
||||||
const removeTag = useGalleryStore((state) => state.removeTag);
|
const removeTag = useGalleryStore((state) => state.removeTag);
|
||||||
const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus);
|
const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus);
|
||||||
|
const loadTaggerModelStatus = useGalleryStore((state) => state.loadTaggerModelStatus);
|
||||||
const queueTaggingForImage = useGalleryStore((state) => state.queueTaggingForImage);
|
const queueTaggingForImage = useGalleryStore((state) => state.queueTaggingForImage);
|
||||||
const albums = useGalleryStore((state) => state.albums);
|
const albums = useGalleryStore((state) => state.albums);
|
||||||
const addToAlbum = useGalleryStore((state) => state.addToAlbum);
|
const addToAlbum = useGalleryStore((state) => state.addToAlbum);
|
||||||
@@ -254,6 +255,13 @@ export function Lightbox() {
|
|||||||
const canStartSlideshow = slideshowImages.length > 0;
|
const canStartSlideshow = slideshowImages.length > 0;
|
||||||
const canFindSimilar = selectedImage?.embedding_status === "ready";
|
const canFindSimilar = selectedImage?.embedding_status === "ready";
|
||||||
const canSearchRegion = canFindSimilar && selectedImage?.media_kind === "image";
|
const canSearchRegion = canFindSimilar && selectedImage?.media_kind === "image";
|
||||||
|
const taggerReady = taggerModelStatus?.ready ?? false;
|
||||||
|
const taggerStatusKnown = taggerModelStatus !== null;
|
||||||
|
const taggerButtonTooltip = !taggerStatusKnown
|
||||||
|
? "Checking AI tagger model..."
|
||||||
|
: taggerReady
|
||||||
|
? "Queue AI tagging for this image"
|
||||||
|
: "AI tagger model not installed";
|
||||||
|
|
||||||
const goPrev = useCallback(() => {
|
const goPrev = useCallback(() => {
|
||||||
if (currentIndex > 0) openImage(images[currentIndex - 1]);
|
if (currentIndex > 0) openImage(images[currentIndex - 1]);
|
||||||
@@ -421,6 +429,11 @@ export function Lightbox() {
|
|||||||
return () => { cancelled = true; };
|
return () => { cancelled = true; };
|
||||||
}, [selectedImage?.id, selectedImage?.media_kind, getImageExif]);
|
}, [selectedImage?.id, selectedImage?.media_kind, getImageExif]);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
if (selectedImage?.media_kind !== "image" || taggerStatusKnown) return;
|
||||||
|
void loadTaggerModelStatus();
|
||||||
|
}, [loadTaggerModelStatus, selectedImage?.media_kind, taggerStatusKnown]);
|
||||||
|
|
||||||
// Reset the queued state once the worker finishes so the button is usable again
|
// Reset the queued state once the worker finishes so the button is usable again
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
if (selectedImage?.ai_tagged_at) setTaggingQueued(false);
|
if (selectedImage?.ai_tagged_at) setTaggingQueued(false);
|
||||||
@@ -1156,10 +1169,10 @@ export function Lightbox() {
|
|||||||
</span>
|
</span>
|
||||||
) : null}
|
) : null}
|
||||||
{selectedImage.media_kind === "image" ? (
|
{selectedImage.media_kind === "image" ? (
|
||||||
<Tooltip label={!(taggerModelStatus?.ready ?? false) ? "WD Tagger model not downloaded" : "Queue AI tagging for this image"} followCursor>
|
<Tooltip label={taggerButtonTooltip} followCursor>
|
||||||
<button
|
<button
|
||||||
className="rounded-md border border-white/10 bg-white/5 px-2 py-0.5 text-[10px] text-gray-400 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-40"
|
className="rounded-md border border-white/10 bg-white/5 px-2 py-0.5 text-[10px] text-gray-400 transition-colors hover:bg-white/10 hover:text-white disabled:cursor-not-allowed disabled:opacity-40"
|
||||||
disabled={!(taggerModelStatus?.ready ?? false) || taggingQueued}
|
disabled={!taggerReady || taggingQueued}
|
||||||
onClick={() => {
|
onClick={() => {
|
||||||
setTaggingQueued(true);
|
setTaggingQueued(true);
|
||||||
void queueTaggingForImage(selectedImage.id).catch(() => undefined);
|
void queueTaggingForImage(selectedImage.id).catch(() => undefined);
|
||||||
|
|||||||
@@ -322,7 +322,7 @@ export function SettingsModal() {
|
|||||||
: taggerModelProgress
|
: taggerModelProgress
|
||||||
? `Downloading ${taggerModelProgress.completed_files}/${taggerModelProgress.total_files}`
|
? `Downloading ${taggerModelProgress.completed_files}/${taggerModelProgress.total_files}`
|
||||||
: taggerModelPreparing
|
: taggerModelPreparing
|
||||||
? "Preparing WD Tagger..."
|
? "Preparing AI tagger..."
|
||||||
: taggerReady
|
: taggerReady
|
||||||
? "Installed"
|
? "Installed"
|
||||||
: "Install model";
|
: "Install model";
|
||||||
|
|||||||
@@ -147,6 +147,7 @@ export function TitleBar() {
|
|||||||
<Tooltip label="Settings (right-click to switch theme)" anchorToCursor>
|
<Tooltip label="Settings (right-click to switch theme)" anchorToCursor>
|
||||||
<button
|
<button
|
||||||
ref={settingsBtnRef}
|
ref={settingsBtnRef}
|
||||||
|
aria-label="Settings"
|
||||||
onClick={() => setSettingsOpen(true)}
|
onClick={() => setSettingsOpen(true)}
|
||||||
onContextMenu={handleSettingsContextMenu}
|
onContextMenu={handleSettingsContextMenu}
|
||||||
className="group flex h-full w-10 items-center justify-center text-gray-600 transition-colors duration-100 hover:bg-white/6 hover:text-gray-300"
|
className="group flex h-full w-10 items-center justify-center text-gray-600 transition-colors duration-100 hover:bg-white/6 hover:text-gray-300"
|
||||||
|
|||||||
@@ -326,7 +326,7 @@ export function Toolbar() {
|
|||||||
className={`w-40 bg-transparent py-1.5 pr-9 text-sm text-white placeholder:text-gray-600 focus:outline-none transition-colors lg:w-52 xl:w-64 ${searchCommand !== null ? "pl-16" : "pl-8"}`}
|
className={`w-40 bg-transparent py-1.5 pr-9 text-sm text-white placeholder:text-gray-600 focus:outline-none transition-colors lg:w-52 xl:w-64 ${searchCommand !== null ? "pl-16" : "pl-8"}`}
|
||||||
/>
|
/>
|
||||||
{searchCommand !== null ? (
|
{searchCommand !== null ? (
|
||||||
<div className="absolute left-8 top-1/2 -translate-y-1/2">
|
<div className="absolute left-8 top-1/2 flex -translate-y-1/2">
|
||||||
<Tooltip label="Remove search command" anchorToCursor>
|
<Tooltip label="Remove search command" anchorToCursor>
|
||||||
<button
|
<button
|
||||||
type="button"
|
type="button"
|
||||||
@@ -339,7 +339,7 @@ export function Toolbar() {
|
|||||||
</div>
|
</div>
|
||||||
) : null}
|
) : null}
|
||||||
{searchQuery.trim().length > 0 || searchCommand !== null ? (
|
{searchQuery.trim().length > 0 || searchCommand !== null ? (
|
||||||
<div className="absolute right-2 top-1/2 -translate-y-1/2">
|
<div className="absolute right-2 top-1/2 flex -translate-y-1/2">
|
||||||
<Tooltip label="Clear search" anchorToCursor>
|
<Tooltip label="Clear search" anchorToCursor>
|
||||||
<button
|
<button
|
||||||
aria-label="Clear search"
|
aria-label="Clear search"
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ export function StepAiFeatures() {
|
|||||||
<div className="min-w-0">
|
<div className="min-w-0">
|
||||||
<p className="text-sm text-white">Automatic tags for every image</p>
|
<p className="text-sm text-white">Automatic tags for every image</p>
|
||||||
<p className="mt-1 text-xs leading-relaxed text-gray-500">
|
<p className="mt-1 text-xs leading-relaxed text-gray-500">
|
||||||
The WD tagger model (~1.3 GB download) labels images so you can search with{" "}
|
The AI tagger model labels images so you can search with{" "}
|
||||||
<code className="rounded bg-gray-900 px-1 py-0.5 text-[11px] text-gray-200 light-theme:bg-gray-800 light-theme:text-gray-100">/t</code> — tags look like:
|
<code className="rounded bg-gray-900 px-1 py-0.5 text-[11px] text-gray-200 light-theme:bg-gray-800 light-theme:text-gray-100">/t</code> — tags look like:
|
||||||
</p>
|
</p>
|
||||||
<span className="mt-2 flex flex-wrap gap-1.5">
|
<span className="mt-2 flex flex-wrap gap-1.5">
|
||||||
|
|||||||
+61
-4
@@ -1,5 +1,15 @@
|
|||||||
import { emit } from "@tauri-apps/api/event";
|
import { emit } from "@tauri-apps/api/event";
|
||||||
import type { Album, ExploreTagEntry, Folder, ImageRecord, ImageTag, SortOrder } from "../store";
|
import type {
|
||||||
|
Album,
|
||||||
|
ExploreTagEntry,
|
||||||
|
Folder,
|
||||||
|
ImageRecord,
|
||||||
|
ImageTag,
|
||||||
|
SortOrder,
|
||||||
|
TaggerModel,
|
||||||
|
TaggerModelStatus,
|
||||||
|
TaggerRuntimeProbe,
|
||||||
|
} from "../store";
|
||||||
import { compareImages, createMockDb, mockExif } from "./mockFixtures";
|
import { compareImages, createMockDb, mockExif } from "./mockFixtures";
|
||||||
import { getMockScenario } from "./mockScenarios";
|
import { getMockScenario } from "./mockScenarios";
|
||||||
|
|
||||||
@@ -7,6 +17,8 @@ const db = createMockDb(getMockScenario());
|
|||||||
let nextFolderId = 100;
|
let nextFolderId = 100;
|
||||||
let nextAlbumId = 100;
|
let nextAlbumId = 100;
|
||||||
let nextTagId = 10_000;
|
let nextTagId = 10_000;
|
||||||
|
let mockTaggerModel: TaggerModel = db.scenario === "joytag-unready" ? "joytag" : "wd";
|
||||||
|
let mockTaggerReady = db.scenario !== "unready" && db.scenario !== "joytag-unready";
|
||||||
|
|
||||||
type AnyPayload = Record<string, any> | undefined;
|
type AnyPayload = Record<string, any> | undefined;
|
||||||
type Rgb = [number, number, number];
|
type Rgb = [number, number, number];
|
||||||
@@ -55,6 +67,42 @@ function page<T>(items: T[], offset = 0, limit = 200) {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function mockTaggerStatus(): TaggerModelStatus {
|
||||||
|
const modelInfo =
|
||||||
|
mockTaggerModel === "joytag"
|
||||||
|
? {
|
||||||
|
model_id: "fancyfeast/joytag",
|
||||||
|
model_name: "JoyTag",
|
||||||
|
local_dir: "mock://models/joytag",
|
||||||
|
label_file: "top_tags.txt",
|
||||||
|
}
|
||||||
|
: {
|
||||||
|
model_id: "SmilingWolf/wd-swinv2-tagger-v3",
|
||||||
|
model_name: "wd-swinv2-tagger-v3",
|
||||||
|
local_dir: "mock://models/wd-swinv2-tagger-v3",
|
||||||
|
label_file: "selected_tags.csv",
|
||||||
|
};
|
||||||
|
|
||||||
|
return {
|
||||||
|
model_id: modelInfo.model_id,
|
||||||
|
model_name: modelInfo.model_name,
|
||||||
|
local_dir: modelInfo.local_dir,
|
||||||
|
ready: mockTaggerReady,
|
||||||
|
missing_files: mockTaggerReady ? [] : ["model.onnx", modelInfo.label_file],
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function mockTaggerRuntimeProbe(): TaggerRuntimeProbe {
|
||||||
|
if (!mockTaggerReady) {
|
||||||
|
throw new Error(`${mockTaggerStatus().model_name} is missing required model files`);
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
ready: true,
|
||||||
|
acceleration: "auto",
|
||||||
|
session: { file: "model.onnx", inputs: ["input"], outputs: ["tags"] },
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
function filterImages(input: ImageRecord[], payload: unknown): ImageRecord[] {
|
function filterImages(input: ImageRecord[], payload: unknown): ImageRecord[] {
|
||||||
const p = params(payload);
|
const p = params(payload);
|
||||||
const query = String(p.search ?? p.query ?? "").trim().toLowerCase();
|
const query = String(p.search ?? p.query ?? "").trim().toLowerCase();
|
||||||
@@ -463,10 +511,19 @@ export async function handleMockCommand(cmd: string, payload?: unknown): Promise
|
|||||||
return reset;
|
return reset;
|
||||||
}
|
}
|
||||||
case "get_tagger_model_status":
|
case "get_tagger_model_status":
|
||||||
return { model_id: "SmilingWolf/wd-vit-tagger-v3", model_name: "WD ViT Tagger", local_dir: "mock://models/tagger", ready: true, missing_files: [] };
|
return mockTaggerStatus();
|
||||||
case "get_tagger_model":
|
case "get_tagger_model":
|
||||||
|
return mockTaggerModel;
|
||||||
case "set_tagger_model":
|
case "set_tagger_model":
|
||||||
return p.model ?? "wd";
|
mockTaggerModel = p.model ?? "wd";
|
||||||
|
mockTaggerReady = db.scenario !== "unready" && db.scenario !== "joytag-unready";
|
||||||
|
return mockTaggerModel;
|
||||||
|
case "prepare_tagger_model":
|
||||||
|
mockTaggerReady = true;
|
||||||
|
return mockTaggerStatus();
|
||||||
|
case "delete_tagger_model":
|
||||||
|
mockTaggerReady = false;
|
||||||
|
return mockTaggerStatus();
|
||||||
case "get_tagger_acceleration":
|
case "get_tagger_acceleration":
|
||||||
case "set_tagger_acceleration":
|
case "set_tagger_acceleration":
|
||||||
return p.acceleration ?? "auto";
|
return p.acceleration ?? "auto";
|
||||||
@@ -477,7 +534,7 @@ export async function handleMockCommand(cmd: string, payload?: unknown): Promise
|
|||||||
case "set_tagger_batch_size":
|
case "set_tagger_batch_size":
|
||||||
return p.batch_size ?? 8;
|
return p.batch_size ?? 8;
|
||||||
case "probe_tagger_runtime":
|
case "probe_tagger_runtime":
|
||||||
return { ready: true, acceleration: "auto", session: { file: "model.onnx", inputs: ["input"], outputs: ["tags"] } };
|
return mockTaggerRuntimeProbe();
|
||||||
case "get_tagging_queue_scope":
|
case "get_tagging_queue_scope":
|
||||||
return "all";
|
return "all";
|
||||||
case "get_tagging_queue_folder_ids":
|
case "get_tagging_queue_folder_ids":
|
||||||
|
|||||||
@@ -1,4 +1,14 @@
|
|||||||
export type MockScenario = "rich" | "empty" | "busy" | "duplicates" | "album" | "errors" | "huge" | "extreme";
|
export type MockScenario =
|
||||||
|
| "rich"
|
||||||
|
| "empty"
|
||||||
|
| "busy"
|
||||||
|
| "duplicates"
|
||||||
|
| "album"
|
||||||
|
| "errors"
|
||||||
|
| "huge"
|
||||||
|
| "extreme"
|
||||||
|
| "unready"
|
||||||
|
| "joytag-unready";
|
||||||
|
|
||||||
const SCENARIOS = new Set<MockScenario>([
|
const SCENARIOS = new Set<MockScenario>([
|
||||||
"rich",
|
"rich",
|
||||||
@@ -9,6 +19,8 @@ const SCENARIOS = new Set<MockScenario>([
|
|||||||
"errors",
|
"errors",
|
||||||
"huge",
|
"huge",
|
||||||
"extreme",
|
"extreme",
|
||||||
|
"unready",
|
||||||
|
"joytag-unready",
|
||||||
]);
|
]);
|
||||||
|
|
||||||
export function getMockScenario(): MockScenario {
|
export function getMockScenario(): MockScenario {
|
||||||
|
|||||||
@@ -3010,6 +3010,9 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
taggerModelProgress: event.payload.done ? null : event.payload,
|
taggerModelProgress: event.payload.done ? null : event.payload,
|
||||||
taggerModelPreparing: !event.payload.done,
|
taggerModelPreparing: !event.payload.done,
|
||||||
});
|
});
|
||||||
|
if (event.payload.done) {
|
||||||
|
void get().loadTaggerModelStatus();
|
||||||
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
const unlistenImages = await listen<IndexedImagesBatch>("indexed-images", (event) => {
|
const unlistenImages = await listen<IndexedImagesBatch>("indexed-images", (event) => {
|
||||||
|
|||||||
Reference in New Issue
Block a user