diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx index 3d09b97..c36dc30 100644 --- a/src/components/SettingsModal.tsx +++ b/src/components/SettingsModal.tsx @@ -1,5 +1,5 @@ import { useEffect, useMemo, useRef, useState } from "react"; -import { AppTheme, CleanupOrphanedThumbnailsResult, DatabaseInfo, OrphanedThumbnailsInfo, TaggerAcceleration, TaggingQueueScope, VacuumResult, useGalleryStore } from "../store"; +import { AppTheme, CleanupOrphanedThumbnailsResult, DatabaseInfo, OrphanedThumbnailsInfo, TaggerAcceleration, TaggerModel, TaggingQueueScope, VacuumResult, useGalleryStore } from "../store"; import { FfmpegStatusRow } from "./onboarding/StepWelcome"; import { ThemedDropdown } from "./ThemedDropdown"; import { getChangelogForVersion } from "../changelog"; @@ -100,6 +100,44 @@ function ScopeButton({ scope, current, onSelect, children }: { ); } +// Display metadata for each selectable tagging model. +const TAGGER_MODELS: Record = { + 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; + onSelect: (model: TaggerModel) => void; + children: React.ReactNode; +}) { + const active = model === current; + return ( + + ); +} + function TaggerAccelerationButton({ acceleration, current, onSelect, children }: { acceleration: TaggerAcceleration; current: TaggerAcceleration; @@ -129,6 +167,8 @@ export function SettingsModal() { const [taggerClearing, setTaggerClearing] = useState(false); const [taggerAccelerationSaving, setTaggerAccelerationSaving] = useState(false); const [taggerAccelerationError, setTaggerAccelerationError] = useState(null); + const [taggerModelSwitching, setTaggerModelSwitching] = useState(false); + const [taggerModelSwitchError, setTaggerModelSwitchError] = useState(null); const [taggerThresholdDraft, setTaggerThresholdDraft] = useState(null); const [taggerThresholdSaving, setTaggerThresholdSaving] = useState(false); const [taggerThresholdError, setTaggerThresholdError] = useState(null); @@ -173,6 +213,9 @@ export function SettingsModal() { const deleteTaggerModel = useGalleryStore((state) => state.deleteTaggerModel); const loadTaggerAcceleration = useGalleryStore((state) => state.loadTaggerAcceleration); const setTaggerAcceleration = useGalleryStore((state) => state.setTaggerAcceleration); + const taggerModel = useGalleryStore((state) => state.taggerModel); + const loadTaggerModel = useGalleryStore((state) => state.loadTaggerModel); + const setTaggerModel = useGalleryStore((state) => state.setTaggerModel); const loadTaggerThreshold = useGalleryStore((state) => state.loadTaggerThreshold); const setTaggerThreshold = useGalleryStore((state) => state.setTaggerThreshold); const loadTaggerBatchSize = useGalleryStore((state) => state.loadTaggerBatchSize); @@ -210,6 +253,7 @@ export function SettingsModal() { useEffect(() => { if (!settingsOpen) return; void loadTaggerModelStatus(); + void loadTaggerModel(); void loadTaggerAcceleration(); void loadTaggerThreshold(); void loadTaggerBatchSize(); @@ -221,7 +265,7 @@ export function SettingsModal() { }; window.addEventListener("keydown", handleKeyDown); return () => window.removeEventListener("keydown", handleKeyDown); - }, [settingsOpen, loadTaggerModelStatus, loadTaggerAcceleration, loadTaggerThreshold, loadTaggerBatchSize, loadTaggingQueueScope, loadTaggingQueueFolderIds, setSettingsOpen]); + }, [settingsOpen, loadTaggerModelStatus, loadTaggerModel, loadTaggerAcceleration, loadTaggerThreshold, loadTaggerBatchSize, loadTaggingQueueScope, loadTaggingQueueFolderIds, setSettingsOpen]); useEffect(() => { if (!settingsOpen || activeSection !== "general") return; @@ -365,10 +409,39 @@ export function SettingsModal() { {activeSection === "workspace" ? (
+ +
+
+ {(["wd", "joytag"] as const).map((model) => ( + { + if (nextModel === taggerModel) return; + setTaggerModelSwitching(true); + setTaggerModelSwitchError(null); + void setTaggerModel(nextModel) + .catch((error: unknown) => setTaggerModelSwitchError(String(error))) + .finally(() => setTaggerModelSwitching(false)); + }} + > + {TAGGER_MODELS[model].tab} + + ))} +
+ {taggerModelSwitchError ? ( +

{taggerModelSwitchError}

+ ) : ( +

{taggerModelSwitching ? "Switching..." : `Current: ${TAGGER_MODELS[taggerModel].name}`}

+ )} +
+
+ - WD SwinV2 Tagger v3{" "} + {TAGGER_MODELS[taggerModel].name}{" "} {taggerReady ? "Installed" : taggerModelPreparing ? "Preparing" : "Not installed"} @@ -376,7 +449,7 @@ export function SettingsModal() { } - description="Anime-focused vision model by SmilingWolf. Generates booru-style tags with configurable confidence thresholds." + description={TAGGER_MODELS[taggerModel].description} >
{taggerReady ? ( @@ -469,7 +542,7 @@ export function SettingsModal() { {taggerThresholdError ? (

{taggerThresholdError}

) : ( -

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

+

{taggerThresholdSaving ? "Saving..." : `Default: ${taggerModel === "joytag" ? "0.4" : "0.35"}`}

)}
diff --git a/src/store.ts b/src/store.ts index 2270547..49f08d9 100644 --- a/src/store.ts +++ b/src/store.ts @@ -51,6 +51,7 @@ export type SearchCommand = "filename" | "semantic" | "tag"; export type CaptionAcceleration = "auto" | "cpu" | "directml"; export type CaptionDetail = "short" | "detailed" | "paragraph"; export type TaggerAcceleration = "auto" | "cpu" | "directml"; +export type TaggerModel = "wd" | "joytag"; export type AiRating = "general" | "sensitive" | "questionable" | "explicit"; export type TaggingQueueScope = "all" | "selected"; export type SimilarScope = "all_media" | "current_folder" | "current_album"; @@ -424,6 +425,7 @@ interface GalleryState { taggerModelPreparing: boolean; taggerModelError: string | null; taggerModelProgress: TaggerModelProgress | null; + taggerModel: TaggerModel; taggerAcceleration: TaggerAcceleration; taggerThreshold: number; taggerBatchSize: number; @@ -551,6 +553,8 @@ interface GalleryState { deleteTaggerModel: () => Promise; loadTaggerAcceleration: () => Promise; setTaggerAcceleration: (acceleration: TaggerAcceleration) => Promise; + loadTaggerModel: () => Promise; + setTaggerModel: (model: TaggerModel) => Promise; loadTaggerThreshold: () => Promise; setTaggerThreshold: (threshold: number) => Promise; loadTaggerBatchSize: () => Promise; @@ -902,6 +906,7 @@ export const useGalleryStore = create((set, get) => ({ taggerModelPreparing: false, taggerModelError: null, taggerModelProgress: null, + taggerModel: "wd", taggerAcceleration: "auto", taggerThreshold: 0.35, taggerBatchSize: 8, @@ -2158,6 +2163,29 @@ export const useGalleryStore = create((set, get) => ({ set({ taggerAcceleration, taggerRuntimeProbe: null }); }, + loadTaggerModel: async () => { + try { + const taggerModel = await invoke("get_tagger_model"); + set({ taggerModel }); + } catch (error) { + set({ taggerModelError: String(error) }); + } + }, + + setTaggerModel: async (model) => { + const taggerModel = await invoke("set_tagger_model", { + params: { model }, + }); + // Switching models changes which files are required on disk, so refresh the + // status too — the download/ready UI must reflect the newly-selected model. + try { + const taggerModelStatus = await invoke("get_tagger_model_status"); + set({ taggerModel, taggerModelStatus, taggerModelError: null, taggerRuntimeProbe: null }); + } catch (error) { + set({ taggerModel, taggerRuntimeProbe: null, taggerModelError: String(error) }); + } + }, + loadTaggerThreshold: async () => { try { const taggerThreshold = await invoke("get_tagger_threshold");