diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index a10f80b..96ad3b0 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -2159,6 +2159,7 @@ pub struct SetTaggerModelParams { #[derive(Deserialize)] pub struct SetTaggerThresholdParams { pub threshold: f32, + pub model: Option, } #[derive(Deserialize)] @@ -2258,7 +2259,11 @@ pub async fn set_tagger_threshold( params: SetTaggerThresholdParams, ) -> Result { let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; - tagger::set_tagger_threshold(&app_dir, params.threshold).map_err(|e| e.to_string()) + match params.model { + Some(model) => tagger::set_tagger_threshold_for_model(&app_dir, model, params.threshold), + None => tagger::set_tagger_threshold(&app_dir, params.threshold), + } + .map_err(|e| e.to_string()) } #[tauri::command] diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index 9cac6a3..d2a7f03 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -16,6 +16,7 @@ pub const WD_TAGGER_MODEL_NAME: &str = "wd-swinv2-tagger-v3"; const TAGGER_ACCELERATION_FILE: &str = "settings/tagger_acceleration.txt"; const TAGGER_THRESHOLD_FILE: &str = "settings/tagger_threshold.txt"; +const JOYTAG_THRESHOLD_FILE: &str = "settings/joytag_threshold.txt"; const TAGGER_BATCH_SIZE_FILE: &str = "settings/tagger_batch_size.txt"; const TAGGER_MODEL_FILE: &str = "settings/tagger_model.txt"; @@ -27,8 +28,7 @@ pub const JOYTAG_MODEL_NAME: &str = "joytag"; // NCHW layout (not NHWC). These are the OpenAI CLIP normalization constants. const JOYTAG_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73]; const JOYTAG_STD: [f32; 3] = [0.268_629_54, 0.261_302_6, 0.275_777_1]; -// JoyTag's recommended detection threshold (used as the default; the user's -// tagger_threshold setting still overrides it). +// JoyTag's recommended detection threshold. const JOYTAG_DEFAULT_THRESHOLD: f32 = 0.4; // Shared ONNX Runtime DLLs (live in the caption model's `onnxruntime/` dir and @@ -64,8 +64,8 @@ pub const TAGGER_INFER_CHUNK: usize = 4; /// claim the GPU/CPU between dispatches. Negligible against per-chunk inference. pub const TAGGER_INFER_YIELD_MS: u64 = 40; -/// Set to `true` by `set_tagger_acceleration` so the tagging worker loop -/// knows to drop its cached `WdTagger` and reload with the new EP. +/// Set to `true` by tagger setting changes so the tagging worker loop knows to +/// drop its cached model session and reload with the current settings. pub static TAGGER_SESSION_DIRTY: AtomicBool = AtomicBool::new(false); // --------------------------------------------------------------------------- @@ -112,6 +112,20 @@ impl TaggerModel { } } + fn threshold_file(self) -> &'static str { + match self { + Self::Wd => TAGGER_THRESHOLD_FILE, + Self::JoyTag => JOYTAG_THRESHOLD_FILE, + } + } + + fn default_threshold(self) -> f32 { + match self { + Self::Wd => DEFAULT_THRESHOLD, + Self::JoyTag => JOYTAG_DEFAULT_THRESHOLD, + } + } + /// Hugging Face repo the model files are fetched from. fn repo_id(self) -> &'static str { match self { @@ -280,21 +294,33 @@ pub fn set_tagger_acceleration( Ok(acceleration) } -pub fn tagger_threshold(app_data_dir: &Path) -> f32 { - let path = app_data_dir.join(TAGGER_THRESHOLD_FILE); +fn tagger_threshold_for_model(app_data_dir: &Path, model: TaggerModel) -> f32 { + let path = app_data_dir.join(model.threshold_file()); let Ok(value) = std::fs::read_to_string(path) else { - return DEFAULT_THRESHOLD; + return model.default_threshold(); }; value .trim() .parse::() - .unwrap_or(DEFAULT_THRESHOLD) + .unwrap_or_else(|_| model.default_threshold()) .clamp(0.01, 1.0) } +pub fn tagger_threshold(app_data_dir: &Path) -> f32 { + tagger_threshold_for_model(app_data_dir, tagger_model(app_data_dir)) +} + pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result { + set_tagger_threshold_for_model(app_data_dir, tagger_model(app_data_dir), threshold) +} + +pub fn set_tagger_threshold_for_model( + app_data_dir: &Path, + model: TaggerModel, + threshold: f32, +) -> Result { let clamped = threshold.clamp(0.01, 1.0); - let path = app_data_dir.join(TAGGER_THRESHOLD_FILE); + let path = app_data_dir.join(model.threshold_file()); if let Some(parent) = path.parent() { std::fs::create_dir_all(parent)?; } @@ -1057,17 +1083,10 @@ fn load_joytag_labels(path: &Path) -> Result> { Ok(labels) } -/// JoyTag detection threshold. Honours the shared `tagger_threshold` setting if -/// the user has set it, otherwise uses JoyTag's recommended default (0.4). +/// JoyTag detection threshold. Uses JoyTag's own setting so WD tuning does not +/// leak into the general/photo-friendly model. fn joytag_threshold(app_data_dir: &Path) -> f32 { - match std::fs::read_to_string(app_data_dir.join(TAGGER_THRESHOLD_FILE)) { - Ok(value) => value - .trim() - .parse::() - .unwrap_or(JOYTAG_DEFAULT_THRESHOLD) - .clamp(0.01, 1.0), - Err(_) => JOYTAG_DEFAULT_THRESHOLD, - } + tagger_threshold_for_model(app_data_dir, TaggerModel::JoyTag) } fn sigmoid(x: f32) -> f32 { diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx index cee18b0..6ab464e 100644 --- a/src/components/SettingsModal.tsx +++ b/src/components/SettingsModal.tsx @@ -461,6 +461,9 @@ export function SettingsModal() { current={taggerModel} onSelect={(nextModel) => { if (nextModel === taggerModel) return; + setTaggerThresholdDraft(null); + setTaggerThresholdError(null); + if (thresholdErrorTimerRef.current) clearTimeout(thresholdErrorTimerRef.current); setTaggerModelSwitching(true); setTaggerModelSwitchError(null); void setTaggerModel(nextModel) @@ -562,12 +565,12 @@ export function SettingsModal() { className="w-20 rounded-md border border-white/10 bg-white/5 px-2 py-1.5 text-xs text-white focus:border-white/20 focus:outline-none" value={thresholdDisplay} onChange={(event) => setTaggerThresholdDraft(event.target.value)} - onBlur={() => { - const value = parseFloat(thresholdDisplay); + onBlur={(event) => { + const value = parseFloat(event.currentTarget.value); if (!isNaN(value) && value >= 0.05 && value <= 0.99) { setTaggerThresholdError(null); setTaggerThresholdSaving(true); - void setTaggerThreshold(value) + void setTaggerThreshold(value, taggerModel) .catch((error: unknown) => setTaggerThresholdError(String(error))) .finally(() => { setTaggerThresholdDraft(null); @@ -584,7 +587,7 @@ export function SettingsModal() { {taggerThresholdError ? (

{taggerThresholdError}

) : ( -

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

+

{taggerThresholdSaving ? "Saving..." : `Default: ${TAGGER_MODELS[taggerModel].defaultThreshold}`}

)} diff --git a/src/components/onboarding/StepAiFeatures.tsx b/src/components/onboarding/StepAiFeatures.tsx index c32b679..901a426 100644 --- a/src/components/onboarding/StepAiFeatures.tsx +++ b/src/components/onboarding/StepAiFeatures.tsx @@ -75,7 +75,7 @@ export function StepAiFeatures() { The AI tagger model labels images so you can search with{" "} /t — tags look like:

-
+
{(["wd", "joytag"] as const).map((model) => ( = { + wd: 0.35, + joytag: 0.4, +}; type AnyPayload = Record | undefined; type Rgb = [number, number, number]; @@ -42,6 +46,10 @@ function isRgb(value: unknown): value is Rgb { return Array.isArray(value) && value.length === 3 && value.every((entry) => typeof entry === "number"); } +function isTaggerModel(value: unknown): value is TaggerModel { + return value === "wd" || value === "joytag"; +} + function colorDistanceSq(left: Rgb, right: Rgb): number { return (left[0] - right[0]) ** 2 + (left[1] - right[1]) ** 2 + (left[2] - right[2]) ** 2; } @@ -528,8 +536,13 @@ export async function handleMockCommand(cmd: string, payload?: unknown): Promise case "set_tagger_acceleration": return p.acceleration ?? "auto"; case "get_tagger_threshold": + return mockTaggerThresholds[mockTaggerModel]; case "set_tagger_threshold": - return p.threshold ?? 0.35; + { + const model = isTaggerModel(p.model) ? p.model : mockTaggerModel; + mockTaggerThresholds[model] = p.threshold ?? mockTaggerThresholds[model]; + return mockTaggerThresholds[model]; + } case "get_tagger_batch_size": case "set_tagger_batch_size": return p.batch_size ?? 8; diff --git a/src/store.ts b/src/store.ts index 9f0933b..1cd1d50 100644 --- a/src/store.ts +++ b/src/store.ts @@ -589,7 +589,7 @@ interface GalleryState { loadTaggerModel: () => Promise; setTaggerModel: (model: TaggerModel) => Promise; loadTaggerThreshold: () => Promise; - setTaggerThreshold: (threshold: number) => Promise; + setTaggerThreshold: (threshold: number, model?: TaggerModel) => Promise; loadTaggerBatchSize: () => Promise; setTaggerBatchSize: (batchSize: number) => Promise; probeTaggerRuntime: () => Promise; @@ -2326,11 +2326,14 @@ export const useGalleryStore = create((set, get) => ({ 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. + // Switching models changes both readiness and the active threshold setting, + // so refresh them together for the selected model. try { - const taggerModelStatus = await invoke("get_tagger_model_status"); - set({ taggerModel, taggerModelStatus, taggerModelError: null, taggerRuntimeProbe: null }); + const [taggerModelStatus, taggerThreshold] = await Promise.all([ + invoke("get_tagger_model_status"), + invoke("get_tagger_threshold"), + ]); + set({ taggerModel, taggerModelStatus, taggerThreshold, taggerModelError: null, taggerRuntimeProbe: null }); } catch (error) { set({ taggerModel, taggerRuntimeProbe: null, taggerModelError: String(error) }); } @@ -2345,11 +2348,13 @@ export const useGalleryStore = create((set, get) => ({ } }, - setTaggerThreshold: async (threshold) => { + setTaggerThreshold: async (threshold, model) => { const taggerThreshold = await invoke("set_tagger_threshold", { - params: { threshold }, + params: { threshold, model }, }); - set({ taggerThreshold }); + if (!model || get().taggerModel === model) { + set({ taggerThreshold }); + } }, loadTaggerBatchSize: async () => { diff --git a/src/taggerModels.ts b/src/taggerModels.ts index e3ac289..53a3546 100644 --- a/src/taggerModels.ts +++ b/src/taggerModels.ts @@ -1,15 +1,17 @@ import type { TaggerModel } from "./store"; -export const TAGGER_MODELS: Record = { +export const TAGGER_MODELS: Record = { wd: { name: "WD SwinV2 Tagger v3", tab: "WD (anime)", + defaultThreshold: 0.35, description: "Anime-focused vision model by SmilingWolf. Generates booru-style tags with configurable confidence thresholds.", }, joytag: { name: "JoyTag", tab: "JoyTag (general)", + defaultThreshold: 0.4, description: "Booru-schema tagger that also handles photographic content and is strong on NSFW concepts. The explicitness rating is derived from its tags.", },