fix(tagger): separate JoyTag confidence threshold
Store confidence thresholds per tagger model so JoyTag no longer inherits WD tuning. Refresh the active threshold when switching models, guard stale threshold saves, and keep UI Lab mocks in sync. Also tightens the onboarding model selector so the segmented control no longer stretches across the row.
This commit is contained in:
@@ -2159,6 +2159,7 @@ pub struct SetTaggerModelParams {
|
|||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
pub struct SetTaggerThresholdParams {
|
pub struct SetTaggerThresholdParams {
|
||||||
pub threshold: f32,
|
pub threshold: f32,
|
||||||
|
pub model: Option<TaggerModel>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
@@ -2258,7 +2259,11 @@ pub async fn set_tagger_threshold(
|
|||||||
params: SetTaggerThresholdParams,
|
params: SetTaggerThresholdParams,
|
||||||
) -> Result<f32, String> {
|
) -> Result<f32, String> {
|
||||||
let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
|
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]
|
#[tauri::command]
|
||||||
|
|||||||
+38
-19
@@ -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_ACCELERATION_FILE: &str = "settings/tagger_acceleration.txt";
|
||||||
const TAGGER_THRESHOLD_FILE: &str = "settings/tagger_threshold.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_BATCH_SIZE_FILE: &str = "settings/tagger_batch_size.txt";
|
||||||
const TAGGER_MODEL_FILE: &str = "settings/tagger_model.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.
|
// 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_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];
|
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
|
// JoyTag's recommended detection threshold.
|
||||||
// tagger_threshold setting still overrides it).
|
|
||||||
const JOYTAG_DEFAULT_THRESHOLD: f32 = 0.4;
|
const JOYTAG_DEFAULT_THRESHOLD: f32 = 0.4;
|
||||||
|
|
||||||
// Shared ONNX Runtime DLLs (live in the caption model's `onnxruntime/` dir and
|
// 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.
|
/// claim the GPU/CPU between dispatches. Negligible against per-chunk inference.
|
||||||
pub const TAGGER_INFER_YIELD_MS: u64 = 40;
|
pub const TAGGER_INFER_YIELD_MS: u64 = 40;
|
||||||
|
|
||||||
/// Set to `true` by `set_tagger_acceleration` so the tagging worker loop
|
/// Set to `true` by tagger setting changes so the tagging worker loop knows to
|
||||||
/// knows to drop its cached `WdTagger` and reload with the new EP.
|
/// drop its cached model session and reload with the current settings.
|
||||||
pub static TAGGER_SESSION_DIRTY: AtomicBool = AtomicBool::new(false);
|
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.
|
/// Hugging Face repo the model files are fetched from.
|
||||||
fn repo_id(self) -> &'static str {
|
fn repo_id(self) -> &'static str {
|
||||||
match self {
|
match self {
|
||||||
@@ -280,21 +294,33 @@ pub fn set_tagger_acceleration(
|
|||||||
Ok(acceleration)
|
Ok(acceleration)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn tagger_threshold(app_data_dir: &Path) -> f32 {
|
fn tagger_threshold_for_model(app_data_dir: &Path, model: TaggerModel) -> f32 {
|
||||||
let path = app_data_dir.join(TAGGER_THRESHOLD_FILE);
|
let path = app_data_dir.join(model.threshold_file());
|
||||||
let Ok(value) = std::fs::read_to_string(path) else {
|
let Ok(value) = std::fs::read_to_string(path) else {
|
||||||
return DEFAULT_THRESHOLD;
|
return model.default_threshold();
|
||||||
};
|
};
|
||||||
value
|
value
|
||||||
.trim()
|
.trim()
|
||||||
.parse::<f32>()
|
.parse::<f32>()
|
||||||
.unwrap_or(DEFAULT_THRESHOLD)
|
.unwrap_or_else(|_| model.default_threshold())
|
||||||
.clamp(0.01, 1.0)
|
.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<f32> {
|
pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result<f32> {
|
||||||
|
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<f32> {
|
||||||
let clamped = threshold.clamp(0.01, 1.0);
|
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() {
|
if let Some(parent) = path.parent() {
|
||||||
std::fs::create_dir_all(parent)?;
|
std::fs::create_dir_all(parent)?;
|
||||||
}
|
}
|
||||||
@@ -1057,17 +1083,10 @@ fn load_joytag_labels(path: &Path) -> Result<Vec<String>> {
|
|||||||
Ok(labels)
|
Ok(labels)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// JoyTag detection threshold. Honours the shared `tagger_threshold` setting if
|
/// JoyTag detection threshold. Uses JoyTag's own setting so WD tuning does not
|
||||||
/// the user has set it, otherwise uses JoyTag's recommended default (0.4).
|
/// leak into the general/photo-friendly model.
|
||||||
fn joytag_threshold(app_data_dir: &Path) -> f32 {
|
fn joytag_threshold(app_data_dir: &Path) -> f32 {
|
||||||
match std::fs::read_to_string(app_data_dir.join(TAGGER_THRESHOLD_FILE)) {
|
tagger_threshold_for_model(app_data_dir, TaggerModel::JoyTag)
|
||||||
Ok(value) => value
|
|
||||||
.trim()
|
|
||||||
.parse::<f32>()
|
|
||||||
.unwrap_or(JOYTAG_DEFAULT_THRESHOLD)
|
|
||||||
.clamp(0.01, 1.0),
|
|
||||||
Err(_) => JOYTAG_DEFAULT_THRESHOLD,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn sigmoid(x: f32) -> f32 {
|
fn sigmoid(x: f32) -> f32 {
|
||||||
|
|||||||
@@ -461,6 +461,9 @@ export function SettingsModal() {
|
|||||||
current={taggerModel}
|
current={taggerModel}
|
||||||
onSelect={(nextModel) => {
|
onSelect={(nextModel) => {
|
||||||
if (nextModel === taggerModel) return;
|
if (nextModel === taggerModel) return;
|
||||||
|
setTaggerThresholdDraft(null);
|
||||||
|
setTaggerThresholdError(null);
|
||||||
|
if (thresholdErrorTimerRef.current) clearTimeout(thresholdErrorTimerRef.current);
|
||||||
setTaggerModelSwitching(true);
|
setTaggerModelSwitching(true);
|
||||||
setTaggerModelSwitchError(null);
|
setTaggerModelSwitchError(null);
|
||||||
void setTaggerModel(nextModel)
|
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"
|
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}
|
value={thresholdDisplay}
|
||||||
onChange={(event) => setTaggerThresholdDraft(event.target.value)}
|
onChange={(event) => setTaggerThresholdDraft(event.target.value)}
|
||||||
onBlur={() => {
|
onBlur={(event) => {
|
||||||
const value = parseFloat(thresholdDisplay);
|
const value = parseFloat(event.currentTarget.value);
|
||||||
if (!isNaN(value) && value >= 0.05 && value <= 0.99) {
|
if (!isNaN(value) && value >= 0.05 && value <= 0.99) {
|
||||||
setTaggerThresholdError(null);
|
setTaggerThresholdError(null);
|
||||||
setTaggerThresholdSaving(true);
|
setTaggerThresholdSaving(true);
|
||||||
void setTaggerThreshold(value)
|
void setTaggerThreshold(value, taggerModel)
|
||||||
.catch((error: unknown) => setTaggerThresholdError(String(error)))
|
.catch((error: unknown) => setTaggerThresholdError(String(error)))
|
||||||
.finally(() => {
|
.finally(() => {
|
||||||
setTaggerThresholdDraft(null);
|
setTaggerThresholdDraft(null);
|
||||||
@@ -584,7 +587,7 @@ export function SettingsModal() {
|
|||||||
{taggerThresholdError ? (
|
{taggerThresholdError ? (
|
||||||
<p className="text-[11px] text-amber-300">{taggerThresholdError}</p>
|
<p className="text-[11px] text-amber-300">{taggerThresholdError}</p>
|
||||||
) : (
|
) : (
|
||||||
<p className="text-[11px] text-gray-600">{taggerThresholdSaving ? "Saving..." : `Default: ${taggerModel === "joytag" ? "0.4" : "0.35"}`}</p>
|
<p className="text-[11px] text-gray-600">{taggerThresholdSaving ? "Saving..." : `Default: ${TAGGER_MODELS[taggerModel].defaultThreshold}`}</p>
|
||||||
)}
|
)}
|
||||||
</div>
|
</div>
|
||||||
</SettingsItem>
|
</SettingsItem>
|
||||||
|
|||||||
@@ -75,7 +75,7 @@ export function StepAiFeatures() {
|
|||||||
The AI tagger model 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:
|
<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>
|
</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">
|
<div className="mt-3 inline-flex max-w-full 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) => (
|
{(["wd", "joytag"] as const).map((model) => (
|
||||||
<TaggerModelChoice
|
<TaggerModelChoice
|
||||||
key={model}
|
key={model}
|
||||||
|
|||||||
+14
-1
@@ -19,6 +19,10 @@ let nextAlbumId = 100;
|
|||||||
let nextTagId = 10_000;
|
let nextTagId = 10_000;
|
||||||
let mockTaggerModel: TaggerModel = db.scenario === "joytag-unready" ? "joytag" : "wd";
|
let mockTaggerModel: TaggerModel = db.scenario === "joytag-unready" ? "joytag" : "wd";
|
||||||
let mockTaggerReady = db.scenario !== "unready" && db.scenario !== "joytag-unready";
|
let mockTaggerReady = db.scenario !== "unready" && db.scenario !== "joytag-unready";
|
||||||
|
const mockTaggerThresholds: Record<TaggerModel, number> = {
|
||||||
|
wd: 0.35,
|
||||||
|
joytag: 0.4,
|
||||||
|
};
|
||||||
|
|
||||||
type AnyPayload = Record<string, any> | undefined;
|
type AnyPayload = Record<string, any> | undefined;
|
||||||
type Rgb = [number, number, number];
|
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");
|
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 {
|
function colorDistanceSq(left: Rgb, right: Rgb): number {
|
||||||
return (left[0] - right[0]) ** 2 + (left[1] - right[1]) ** 2 + (left[2] - right[2]) ** 2;
|
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":
|
case "set_tagger_acceleration":
|
||||||
return p.acceleration ?? "auto";
|
return p.acceleration ?? "auto";
|
||||||
case "get_tagger_threshold":
|
case "get_tagger_threshold":
|
||||||
|
return mockTaggerThresholds[mockTaggerModel];
|
||||||
case "set_tagger_threshold":
|
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 "get_tagger_batch_size":
|
||||||
case "set_tagger_batch_size":
|
case "set_tagger_batch_size":
|
||||||
return p.batch_size ?? 8;
|
return p.batch_size ?? 8;
|
||||||
|
|||||||
+13
-8
@@ -589,7 +589,7 @@ interface GalleryState {
|
|||||||
loadTaggerModel: () => Promise<void>;
|
loadTaggerModel: () => Promise<void>;
|
||||||
setTaggerModel: (model: TaggerModel) => Promise<void>;
|
setTaggerModel: (model: TaggerModel) => Promise<void>;
|
||||||
loadTaggerThreshold: () => Promise<void>;
|
loadTaggerThreshold: () => Promise<void>;
|
||||||
setTaggerThreshold: (threshold: number) => Promise<void>;
|
setTaggerThreshold: (threshold: number, model?: TaggerModel) => Promise<void>;
|
||||||
loadTaggerBatchSize: () => Promise<void>;
|
loadTaggerBatchSize: () => Promise<void>;
|
||||||
setTaggerBatchSize: (batchSize: number) => Promise<void>;
|
setTaggerBatchSize: (batchSize: number) => Promise<void>;
|
||||||
probeTaggerRuntime: () => Promise<void>;
|
probeTaggerRuntime: () => Promise<void>;
|
||||||
@@ -2326,11 +2326,14 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
const taggerModel = await invoke<TaggerModel>("set_tagger_model", {
|
const taggerModel = await invoke<TaggerModel>("set_tagger_model", {
|
||||||
params: { model },
|
params: { model },
|
||||||
});
|
});
|
||||||
// Switching models changes which files are required on disk, so refresh the
|
// Switching models changes both readiness and the active threshold setting,
|
||||||
// status too — the download/ready UI must reflect the newly-selected model.
|
// so refresh them together for the selected model.
|
||||||
try {
|
try {
|
||||||
const taggerModelStatus = await invoke<TaggerModelStatus>("get_tagger_model_status");
|
const [taggerModelStatus, taggerThreshold] = await Promise.all([
|
||||||
set({ taggerModel, taggerModelStatus, taggerModelError: null, taggerRuntimeProbe: null });
|
invoke<TaggerModelStatus>("get_tagger_model_status"),
|
||||||
|
invoke<number>("get_tagger_threshold"),
|
||||||
|
]);
|
||||||
|
set({ taggerModel, taggerModelStatus, taggerThreshold, taggerModelError: null, taggerRuntimeProbe: null });
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
set({ taggerModel, taggerRuntimeProbe: null, taggerModelError: String(error) });
|
set({ taggerModel, taggerRuntimeProbe: null, taggerModelError: String(error) });
|
||||||
}
|
}
|
||||||
@@ -2345,11 +2348,13 @@ export const useGalleryStore = create<GalleryState>((set, get) => ({
|
|||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
|
||||||
setTaggerThreshold: async (threshold) => {
|
setTaggerThreshold: async (threshold, model) => {
|
||||||
const taggerThreshold = await invoke<number>("set_tagger_threshold", {
|
const taggerThreshold = await invoke<number>("set_tagger_threshold", {
|
||||||
params: { threshold },
|
params: { threshold, model },
|
||||||
});
|
});
|
||||||
set({ taggerThreshold });
|
if (!model || get().taggerModel === model) {
|
||||||
|
set({ taggerThreshold });
|
||||||
|
}
|
||||||
},
|
},
|
||||||
|
|
||||||
loadTaggerBatchSize: async () => {
|
loadTaggerBatchSize: async () => {
|
||||||
|
|||||||
+3
-1
@@ -1,15 +1,17 @@
|
|||||||
import type { TaggerModel } from "./store";
|
import type { TaggerModel } from "./store";
|
||||||
|
|
||||||
export const TAGGER_MODELS: Record<TaggerModel, { name: string; tab: string; description: string }> = {
|
export const TAGGER_MODELS: Record<TaggerModel, { name: string; tab: string; description: string; defaultThreshold: number }> = {
|
||||||
wd: {
|
wd: {
|
||||||
name: "WD SwinV2 Tagger v3",
|
name: "WD SwinV2 Tagger v3",
|
||||||
tab: "WD (anime)",
|
tab: "WD (anime)",
|
||||||
|
defaultThreshold: 0.35,
|
||||||
description:
|
description:
|
||||||
"Anime-focused vision model by SmilingWolf. Generates booru-style tags with configurable confidence thresholds.",
|
"Anime-focused vision model by SmilingWolf. Generates booru-style tags with configurable confidence thresholds.",
|
||||||
},
|
},
|
||||||
joytag: {
|
joytag: {
|
||||||
name: "JoyTag",
|
name: "JoyTag",
|
||||||
tab: "JoyTag (general)",
|
tab: "JoyTag (general)",
|
||||||
|
defaultThreshold: 0.4,
|
||||||
description:
|
description:
|
||||||
"Booru-schema tagger that also handles photographic content and is strong on NSFW concepts. The explicitness rating is derived from its tags.",
|
"Booru-schema tagger that also handles photographic content and is strong on NSFW concepts. The explicitness rating is derived from its tags.",
|
||||||
},
|
},
|
||||||
|
|||||||
Reference in New Issue
Block a user