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);
|
||||
if !status.ready {
|
||||
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.join(", ")
|
||||
);
|
||||
|
||||
@@ -23,6 +23,8 @@ export default function App() {
|
||||
const loadBackgroundJobProgress = useGalleryStore((state) => state.loadBackgroundJobProgress);
|
||||
const loadImages = useGalleryStore((state) => state.loadImages);
|
||||
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 loadAlbums = useGalleryStore((state) => state.loadAlbums);
|
||||
const loadMutedFolderIds = useGalleryStore((state) => state.loadMutedFolderIds);
|
||||
@@ -53,6 +55,8 @@ export default function App() {
|
||||
loadFolders().then(async () => {
|
||||
void loadBackgroundJobProgress();
|
||||
void loadCaptionModelStatus();
|
||||
void loadTaggerModel();
|
||||
void loadTaggerModelStatus();
|
||||
void loadDuplicateScanCache();
|
||||
await loadAlbums();
|
||||
await loadImages(true);
|
||||
|
||||
@@ -156,6 +156,7 @@ export function Lightbox() {
|
||||
const addUserTag = useGalleryStore((state) => state.addUserTag);
|
||||
const removeTag = useGalleryStore((state) => state.removeTag);
|
||||
const taggerModelStatus = useGalleryStore((state) => state.taggerModelStatus);
|
||||
const loadTaggerModelStatus = useGalleryStore((state) => state.loadTaggerModelStatus);
|
||||
const queueTaggingForImage = useGalleryStore((state) => state.queueTaggingForImage);
|
||||
const albums = useGalleryStore((state) => state.albums);
|
||||
const addToAlbum = useGalleryStore((state) => state.addToAlbum);
|
||||
@@ -254,6 +255,13 @@ export function Lightbox() {
|
||||
const canStartSlideshow = slideshowImages.length > 0;
|
||||
const canFindSimilar = selectedImage?.embedding_status === "ready";
|
||||
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(() => {
|
||||
if (currentIndex > 0) openImage(images[currentIndex - 1]);
|
||||
@@ -421,6 +429,11 @@ export function Lightbox() {
|
||||
return () => { cancelled = true; };
|
||||
}, [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
|
||||
useEffect(() => {
|
||||
if (selectedImage?.ai_tagged_at) setTaggingQueued(false);
|
||||
@@ -1156,10 +1169,10 @@ export function Lightbox() {
|
||||
</span>
|
||||
) : null}
|
||||
{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
|
||||
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={() => {
|
||||
setTaggingQueued(true);
|
||||
void queueTaggingForImage(selectedImage.id).catch(() => undefined);
|
||||
|
||||
@@ -322,7 +322,7 @@ export function SettingsModal() {
|
||||
: taggerModelProgress
|
||||
? `Downloading ${taggerModelProgress.completed_files}/${taggerModelProgress.total_files}`
|
||||
: taggerModelPreparing
|
||||
? "Preparing WD Tagger..."
|
||||
? "Preparing AI tagger..."
|
||||
: taggerReady
|
||||
? "Installed"
|
||||
: "Install model";
|
||||
|
||||
@@ -147,6 +147,7 @@ export function TitleBar() {
|
||||
<Tooltip label="Settings (right-click to switch theme)" anchorToCursor>
|
||||
<button
|
||||
ref={settingsBtnRef}
|
||||
aria-label="Settings"
|
||||
onClick={() => setSettingsOpen(true)}
|
||||
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"
|
||||
|
||||
@@ -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"}`}
|
||||
/>
|
||||
{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>
|
||||
<button
|
||||
type="button"
|
||||
@@ -339,7 +339,7 @@ export function Toolbar() {
|
||||
</div>
|
||||
) : 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>
|
||||
<button
|
||||
aria-label="Clear search"
|
||||
|
||||
@@ -43,7 +43,7 @@ export function StepAiFeatures() {
|
||||
<div className="min-w-0">
|
||||
<p className="text-sm text-white">Automatic tags for every image</p>
|
||||
<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:
|
||||
</p>
|
||||
<span className="mt-2 flex flex-wrap gap-1.5">
|
||||
|
||||
+61
-4
@@ -1,5 +1,15 @@
|
||||
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 { getMockScenario } from "./mockScenarios";
|
||||
|
||||
@@ -7,6 +17,8 @@ const db = createMockDb(getMockScenario());
|
||||
let nextFolderId = 100;
|
||||
let nextAlbumId = 100;
|
||||
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 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[] {
|
||||
const p = params(payload);
|
||||
const query = String(p.search ?? p.query ?? "").trim().toLowerCase();
|
||||
@@ -463,10 +511,19 @@ export async function handleMockCommand(cmd: string, payload?: unknown): Promise
|
||||
return reset;
|
||||
}
|
||||
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":
|
||||
return mockTaggerModel;
|
||||
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 "set_tagger_acceleration":
|
||||
return p.acceleration ?? "auto";
|
||||
@@ -477,7 +534,7 @@ export async function handleMockCommand(cmd: string, payload?: unknown): Promise
|
||||
case "set_tagger_batch_size":
|
||||
return p.batch_size ?? 8;
|
||||
case "probe_tagger_runtime":
|
||||
return { ready: true, acceleration: "auto", session: { file: "model.onnx", inputs: ["input"], outputs: ["tags"] } };
|
||||
return mockTaggerRuntimeProbe();
|
||||
case "get_tagging_queue_scope":
|
||||
return "all";
|
||||
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>([
|
||||
"rich",
|
||||
@@ -9,6 +19,8 @@ const SCENARIOS = new Set<MockScenario>([
|
||||
"errors",
|
||||
"huge",
|
||||
"extreme",
|
||||
"unready",
|
||||
"joytag-unready",
|
||||
]);
|
||||
|
||||
export function getMockScenario(): MockScenario {
|
||||
|
||||
@@ -3010,6 +3010,9 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
||||
taggerModelProgress: event.payload.done ? null : event.payload,
|
||||
taggerModelPreparing: !event.payload.done,
|
||||
});
|
||||
if (event.payload.done) {
|
||||
void get().loadTaggerModelStatus();
|
||||
}
|
||||
});
|
||||
|
||||
const unlistenImages = await listen<IndexedImagesBatch>("indexed-images", (event) => {
|
||||
|
||||
Reference in New Issue
Block a user