From b7e82dbf91b584e253d847986cf996c7a7bec802 Mon Sep 17 00:00:00 2001
From: LyAhn <admin@jezz.wtf>
Date: Thu, 2 Jul 2026 20:13:45 +0100
Subject: [PATCH] fix(ai-tags): refresh tagger readiness for lightbox

Load the selected tagger model and model status during app startup so the lightbox AI tags action does not stay disabled until Settings refreshes the state.

Refresh readiness after tagger downloads complete and replace stale WD-specific unavailable copy with generic AI tagger wording.
---
 src-tauri/src/tagger.rs                      |  3 ++-
 src/App.tsx                                  |  4 ++++
 src/components/Lightbox.tsx                  | 17 +++++++++++++++--
 src/components/SettingsModal.tsx             |  2 +-
 src/components/onboarding/StepAiFeatures.tsx |  2 +-
 src/store.ts                                 |  3 +++
 6 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs
index af1771e..9cac6a3 100644
--- a/src-tauri/src/tagger.rs
+++ b/src-tauri/src/tagger.rs
@@ -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(", ")
         );
diff --git a/src/App.tsx b/src/App.tsx
index 61945f4..1986c0c 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -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);
diff --git a/src/components/Lightbox.tsx b/src/components/Lightbox.tsx
index bedb93f..22c4730 100644
--- a/src/components/Lightbox.tsx
+++ b/src/components/Lightbox.tsx
@@ -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);
diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx
index 60df8f1..f6ea576 100644
--- a/src/components/SettingsModal.tsx
+++ b/src/components/SettingsModal.tsx
@@ -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";
diff --git a/src/components/onboarding/StepAiFeatures.tsx b/src/components/onboarding/StepAiFeatures.tsx
index f6adaa5..049bf25 100644
--- a/src/components/onboarding/StepAiFeatures.tsx
+++ b/src/components/onboarding/StepAiFeatures.tsx
@@ -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">
diff --git a/src/store.ts b/src/store.ts
index ea3b32e..9f0933b 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -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) => {