From f1116c6c26192787673335884f45cfa84422dc8c Mon Sep 17 00:00:00 2001
From: LyAhn <admin@jezz.wtf>
Date: Fri, 3 Jul 2026 08:07:42 +0100
Subject: [PATCH] feat(onboarding): choose AI tagger model

Add WD and JoyTag selection to the Welcome Tour AI step so users can choose the model before downloading it.

Share tagger model metadata with Settings and keep the Settings close button anchored to the modal chrome.
---
 src/components/SettingsModal.tsx             | 39 +++++---------
 src/components/onboarding/StepAiFeatures.tsx | 53 +++++++++++++++++++-
 src/taggerModels.ts                          | 16 ++++++
 3 files changed, 80 insertions(+), 28 deletions(-)
 create mode 100644 src/taggerModels.ts

diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx
index f6ea576..cee18b0 100644
--- a/src/components/SettingsModal.tsx
+++ b/src/components/SettingsModal.tsx
@@ -4,6 +4,7 @@ import { FfmpegStatusRow } from "./onboarding/StepWelcome";
 import { ThemedDropdown } from "./ThemedDropdown";
 import { getChangelogForVersion } from "../changelog";
 import { Tooltip } from "./Tooltip";
+import { TAGGER_MODELS } from "../taggerModels";
 
 type SettingsSection = "workspace" | "general";
 
@@ -101,22 +102,6 @@ function ScopeButton({ scope, current, onSelect, children }: {
   );
 }
 
-// Display metadata for each selectable tagging model.
-const TAGGER_MODELS: Record<TaggerModel, { name: string; tab: string; description: string }> = {
-  wd: {
-    name: "WD SwinV2 Tagger v3",
-    tab: "WD (anime)",
-    description:
-      "Anime-focused vision model by SmilingWolf. Generates booru-style tags with configurable confidence thresholds.",
-  },
-  joytag: {
-    name: "JoyTag",
-    tab: "JoyTag (general)",
-    description:
-      "Booru-schema tagger that also handles photographic content and is strong on NSFW concepts. The explicitness rating is derived from its tags.",
-  },
-};
-
 function TaggerModelButton({ model, current, onSelect, children }: {
   model: TaggerModel;
   current: TaggerModel;
@@ -445,16 +430,18 @@ export function SettingsModal() {
           </div>
         </aside>
 
-        <Tooltip label="Close settings" anchorToCursor className="absolute right-4 top-4 z-10">
-          <button
-            className="rounded-md p-1.5 text-gray-500 transition-colors hover:bg-white/[0.06] hover:text-white"
-            onClick={() => setSettingsOpen(false)}
-          >
-            <svg className="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
-              <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
-            </svg>
-          </button>
-        </Tooltip>
+        <div className="absolute right-4 top-4 z-10">
+          <Tooltip label="Close settings" anchorToCursor>
+            <button
+              className="rounded-md p-1.5 text-gray-500 transition-colors hover:bg-white/[0.06] hover:text-white"
+              onClick={() => setSettingsOpen(false)}
+            >
+              <svg className="h-4 w-4" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
+              </svg>
+            </button>
+          </Tooltip>
+        </div>
 
         <main className="min-w-0 flex-1 overflow-y-auto">
           <div className="px-10 py-8">
diff --git a/src/components/onboarding/StepAiFeatures.tsx b/src/components/onboarding/StepAiFeatures.tsx
index 049bf25..c32b679 100644
--- a/src/components/onboarding/StepAiFeatures.tsx
+++ b/src/components/onboarding/StepAiFeatures.tsx
@@ -1,5 +1,6 @@
 import { useEffect } from "react";
-import { TaggerModelProgress, useGalleryStore } from "../../store";
+import { TaggerModel, TaggerModelProgress, useGalleryStore } from "../../store";
+import { TAGGER_MODELS } from "../../taggerModels";
 import { FakeProgressBar, formatBytes } from "./fakes";
 
 const FAKE_TAGS = ["landscape", "sunset", "outdoors", "no_humans", "ocean", "cloudy_sky"];
@@ -15,19 +16,47 @@ function taggerDownloadFraction(progress: TaggerModelProgress | null): number |
   return null;
 }
 
+function TaggerModelChoice({ model, current, disabled, onSelect }: {
+  model: TaggerModel;
+  current: TaggerModel;
+  disabled: boolean;
+  onSelect: (model: TaggerModel) => void;
+}) {
+  const active = model === current;
+  return (
+    <button
+      type="button"
+      className={`rounded-md border px-3 py-1.5 text-xs transition-colors disabled:cursor-not-allowed disabled:opacity-45 ${
+        active
+          ? "border-emerald-400/35 bg-emerald-500/15 text-emerald-200 light-theme:border-emerald-600/50 light-theme:bg-emerald-100 light-theme:text-emerald-700"
+          : "border-transparent text-gray-500 hover:bg-white/[0.06] hover:text-gray-200 light-theme:text-gray-600 light-theme:hover:bg-gray-900 light-theme:hover:text-white"
+      }`}
+      disabled={disabled}
+      onClick={() => onSelect(model)}
+    >
+      {TAGGER_MODELS[model].tab}
+    </button>
+  );
+}
+
 export function StepAiFeatures() {
   const taggerModelStatus = useGalleryStore((s) => s.taggerModelStatus);
   const taggerModelPreparing = useGalleryStore((s) => s.taggerModelPreparing);
   const taggerModelProgress = useGalleryStore((s) => s.taggerModelProgress);
   const taggerModelError = useGalleryStore((s) => s.taggerModelError);
+  const taggerModel = useGalleryStore((s) => s.taggerModel);
   const prepareTaggerModel = useGalleryStore((s) => s.prepareTaggerModel);
+  const loadTaggerModel = useGalleryStore((s) => s.loadTaggerModel);
   const loadTaggerModelStatus = useGalleryStore((s) => s.loadTaggerModelStatus);
+  const setTaggerModel = useGalleryStore((s) => s.setTaggerModel);
 
   useEffect(() => {
+    void loadTaggerModel();
     void loadTaggerModelStatus();
-  }, [loadTaggerModelStatus]);
+  }, [loadTaggerModel, loadTaggerModelStatus]);
 
   const taggerReady = taggerModelStatus?.ready ?? false;
+  const taggerStatusLabel = taggerReady ? "Installed" : taggerModelPreparing ? "Preparing" : "Not installed";
 
   return (
     <div>
@@ -46,6 +75,26 @@ export function StepAiFeatures() {
                 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>
+              <div className="mt-3 flex flex-wrap items-center gap-1 rounded-lg border border-white/[0.07] p-0.5 light-theme:border-gray-300/80">
+                {(["wd", "joytag"] as const).map((model) => (
+                  <TaggerModelChoice
+                    key={model}
+                    model={model}
+                    current={taggerModel}
+                    disabled={taggerModelPreparing}
+                    onSelect={(nextModel) => {
+                      if (nextModel === taggerModel) return;
+                      void setTaggerModel(nextModel);
+                    }}
+                  />
+                ))}
+              </div>
+              <p className="mt-2 text-[11px] text-gray-600">
+                Current: {TAGGER_MODELS[taggerModel].name} · {taggerStatusLabel}
+              </p>
+              <p className="mt-1 text-xs leading-relaxed text-gray-500">
+                {TAGGER_MODELS[taggerModel].description}
+              </p>
               <span className="mt-2 flex flex-wrap gap-1.5">
                 {FAKE_TAGS.map((tag) => (
                   <span key={tag} className="rounded-md border border-white/10 bg-gray-900/50 px-2 py-0.5 text-[11px] text-gray-400 light-theme:border-gray-300/70 light-theme:bg-gray-900 light-theme:text-gray-600">
diff --git a/src/taggerModels.ts b/src/taggerModels.ts
new file mode 100644
index 0000000..e3ac289
--- /dev/null
+++ b/src/taggerModels.ts
@@ -0,0 +1,16 @@
+import type { TaggerModel } from "./store";
+
+export const TAGGER_MODELS: Record<TaggerModel, { name: string; tab: string; description: string }> = {
+  wd: {
+    name: "WD SwinV2 Tagger v3",
+    tab: "WD (anime)",
+    description:
+      "Anime-focused vision model by SmilingWolf. Generates booru-style tags with configurable confidence thresholds.",
+  },
+  joytag: {
+    name: "JoyTag",
+    tab: "JoyTag (general)",
+    description:
+      "Booru-schema tagger that also handles photographic content and is strong on NSFW concepts. The explicitness rating is derived from its tags.",
+  },
+};