diff --git a/tests/api/test_tagger_swap_endpoints.py b/tests/api/test_tagger_swap_endpoints.py
index 8c29890..bbd44ee 100644
--- a/tests/api/test_tagger_swap_endpoints.py
+++ b/tests/api/test_tagger_swap_endpoints.py
@@ -46,6 +46,8 @@ def mock_tagging() -> MagicMock:
tagging.effective_active_tagger = None
tagging.is_available = False
tagging.is_configured = False
+ tagging.has_per_tag_thresholds = False
+ tagging.per_tag_columns = []
tagging.swap_active_model = AsyncMock()
return tagging
@@ -142,6 +144,34 @@ async def test_reports_is_available_false_when_subprocess_down(self, client, moc
assert data["active"] is not None
assert data["is_available"] is False
+ @pytest.mark.asyncio
+ async def test_reports_per_tag_capability(self, client, mock_tagging: MagicMock):
+ mock_tagging.effective_active_tagger = ActiveTagger(kind="hf", repo_id="animetimm/convnext")
+ mock_tagging.is_available = True
+ mock_tagging.has_per_tag_thresholds = True
+ mock_tagging.per_tag_columns = ["best_threshold", "best_recall"]
+
+ resp = await client.get("/api/tagger/active")
+ data = await resp.get_json()
+ assert data["active"]["has_per_tag_thresholds"] is True
+ assert data["active"]["per_tag_columns"] == ["best_threshold", "best_recall"]
+
+
+class TestTagImageBodyPerTag:
+ def test_accepts_known_column(self):
+ from yadc.api.controllers.api_tagging import TagImageBody
+
+ body = TagImageBody.model_validate({"per_tag_thresholds": True, "per_tag_column": "best_recall"})
+ assert body.per_tag_column == "best_recall"
+
+ def test_rejects_unknown_column(self):
+ import pydantic
+
+ from yadc.api.controllers.api_tagging import TagImageBody
+
+ with pytest.raises(pydantic.ValidationError):
+ TagImageBody.model_validate({"per_tag_column": "best_f1"})
+
class TestSwapTagger:
"""``POST /api/tagger/swap`` — drive a swap; map service exceptions to HTTP."""
diff --git a/tests/taggers/test_onnx.py b/tests/taggers/test_onnx.py
index f155c3a..ec686ef 100644
--- a/tests/taggers/test_onnx.py
+++ b/tests/taggers/test_onnx.py
@@ -722,3 +722,91 @@ def _fake_download(repo_id, filename, **_):
filenames = sorted(c.kwargs["filename"] for c in mock_dl.call_args_list)
assert filenames == ["model.onnx", "model.onnx_data", "selected_tags.csv"]
+
+
+class TestLoadPerTagThresholds:
+ def _write_csv(self, tmp_path: Path, header: list[str], rows: list[list[str]]) -> Path:
+ path = tmp_path / "selected_tags.csv"
+ with open(path, "w", newline="", encoding="utf-8") as f:
+ writer = csv.writer(f)
+ writer.writerow(header)
+ writer.writerows(rows)
+ return path
+
+ def test_reads_supported_columns(self, tmp_path: Path):
+ from yadc.taggers.onnx_preprocess import _load_per_tag_thresholds
+
+ path = self._write_csv(
+ tmp_path,
+ ["name", "category", "best_threshold", "best_recall"],
+ [["1girl", "0", "0.35", "0.5"], ["solo", "0", "0.4", "0.6"]],
+ )
+ result = _load_per_tag_thresholds(path)
+ assert result is not None
+ assert result["best_threshold"] == {"1girl": 0.35, "solo": 0.4}
+ assert result["best_recall"] == {"1girl": 0.5, "solo": 0.6}
+
+ def test_returns_none_without_threshold_columns(self, tmp_path: Path):
+ from yadc.taggers.onnx_preprocess import _load_per_tag_thresholds
+
+ path = self._write_csv(tmp_path, ["name", "category"], [["1girl", "0"]])
+ assert _load_per_tag_thresholds(path) is None
+
+ def test_skips_invalid_values(self, tmp_path: Path):
+ from yadc.taggers.onnx_preprocess import _load_per_tag_thresholds
+
+ path = self._write_csv(
+ tmp_path,
+ ["name", "category", "best_threshold"],
+ [["good", "0", "0.5"], ["bad-float", "0", "abc"], ["bad-range", "0", "1.5"], ["", "0", "0.5"]],
+ )
+ result = _load_per_tag_thresholds(path)
+ assert result is not None
+ assert result["best_threshold"] == {"good": 0.5}
+
+ def test_ignores_best_f1_column(self, tmp_path: Path):
+ from yadc.taggers.onnx_preprocess import _load_per_tag_thresholds
+
+ path = self._write_csv(tmp_path, ["name", "category", "best_f1"], [["1girl", "0", "0.7"]])
+ assert _load_per_tag_thresholds(path) is None
+
+
+class TestApplyThresholdsPerTag:
+ def _result(self):
+ from yadc.taggers.base import TaggerResult
+
+ return TaggerResult(
+ tags={"1girl": 0.4, "solo": 0.3},
+ categories={"general": ["1girl", "solo"]},
+ )
+
+ def test_per_tag_overrides_category(self):
+ from yadc.taggers.onnx import apply_thresholds
+
+ out = apply_thresholds(
+ self._result(),
+ general_threshold=0.9,
+ per_tag_thresholds={"1girl": 0.35},
+ )
+ assert "1girl" in out.tags
+ assert "solo" not in out.tags
+
+ def test_missing_per_tag_falls_back_to_category(self):
+ from yadc.taggers.onnx import apply_thresholds
+
+ out = apply_thresholds(
+ self._result(),
+ general_threshold=0.35,
+ per_tag_thresholds={"1girl": 0.35},
+ )
+ assert "solo" not in out.tags
+
+ def test_zero_per_tag_keeps_tag(self):
+ from yadc.taggers.onnx import apply_thresholds
+
+ out = apply_thresholds(
+ self._result(),
+ general_threshold=0.9,
+ per_tag_thresholds={"solo": 0.0},
+ )
+ assert "solo" in out.tags
diff --git a/tests/taggers/test_service.py b/tests/taggers/test_service.py
index 8c9e5b5..d94cfd5 100644
--- a/tests/taggers/test_service.py
+++ b/tests/taggers/test_service.py
@@ -1639,6 +1639,8 @@ def _make_key(self, **overrides: Any) -> TaggerResultKey:
rating_threshold=0.0,
general_threshold=0.35,
character_threshold=0.85,
+ per_tag_enabled=False,
+ per_tag_column="",
)
defaults.update(overrides)
return TaggerResultKey(**defaults)
@@ -1665,6 +1667,8 @@ def test_distinct_field_makes_keys_unequal(self) -> None:
{"rating_threshold": 0.1},
{"general_threshold": 0.5},
{"character_threshold": 0.9},
+ {"per_tag_enabled": True},
+ {"per_tag_column": "best_recall"},
):
modified = self._make_key(**overrides)
assert modified != base, f"expected {overrides} to break equality"
@@ -1684,6 +1688,102 @@ def test_key_is_immutable(self) -> None:
key.image_id = 999 # pyright: ignore[reportAttributeAccessIssue]
+class TestPerTagThresholdResolution:
+ def test_thresholds_from_options_prefers_request(self, service: TaggingService) -> None:
+ from yadc.api.services.tagging import TagJobOptions
+
+ options = TagJobOptions(per_tag_thresholds=True, per_tag_column="best_recall")
+ thresholds = service._thresholds_from_options(options)
+ assert thresholds.per_tag_enabled is True
+ assert thresholds.per_tag_column == "best_recall"
+
+ def test_thresholds_from_options_falls_back_to_config(self, service: TaggingService, test_configuration: Configuration) -> None:
+ from yadc.api.services.tagging import TagJobOptions
+
+ test_configuration.tagger_per_tag_thresholds = True
+ test_configuration.tagger_per_tag_column = "best_recall"
+ thresholds = service._thresholds_from_options(TagJobOptions())
+ assert thresholds.per_tag_enabled is True
+ assert thresholds.per_tag_column == "best_recall"
+
+ def test_per_tag_for_thresholds_returns_none_when_disabled(self, service: TaggingService) -> None:
+ from yadc.api.services.tagging import TaggingThresholds
+
+ service._per_tag_thresholds_all = {"best_threshold": {"1girl": 0.5}}
+ out = service._per_tag_for_thresholds(TaggingThresholds(per_tag_enabled=False))
+ assert out is None
+
+ def test_per_tag_for_thresholds_resolves_column(self, service: TaggingService) -> None:
+ from yadc.api.services.tagging import TaggingThresholds
+
+ service._per_tag_thresholds_all = {"best_threshold": {"1girl": 0.5}, "best_recall": {"1girl": 0.7}}
+ out = service._per_tag_for_thresholds(TaggingThresholds(per_tag_enabled=True, per_tag_column="best_recall"))
+ assert out == {"1girl": 0.7}
+
+ def test_per_tag_for_thresholds_missing_column_returns_none(self, service: TaggingService) -> None:
+ from yadc.api.services.tagging import TaggingThresholds
+
+ service._per_tag_thresholds_all = {"best_threshold": {"1girl": 0.5}}
+ out = service._per_tag_for_thresholds(TaggingThresholds(per_tag_enabled=True, per_tag_column="best_recall"))
+ assert out is None
+
+ def test_tag_result_key_uses_request_per_tag(self, service: TaggingService) -> None:
+ from yadc.api.services.tagging import TaggingThresholds
+
+ service._per_tag_thresholds_all = {"best_recall": {"1girl": 0.7}}
+ on = service._tag_result_key("ds", 1, TaggingThresholds(per_tag_enabled=True, per_tag_column="best_recall"))
+ off = service._tag_result_key("ds", 1, TaggingThresholds(per_tag_enabled=False))
+ assert on != off
+ assert on.per_tag_column == "best_recall"
+ assert off.per_tag_column == ""
+
+ def test_tag_result_key_collapses_unknown_column(self, service: TaggingService) -> None:
+ from yadc.api.services.tagging import TaggingThresholds
+
+ service._per_tag_thresholds_all = {"best_threshold": {"1girl": 0.5}}
+ unknown = service._tag_result_key("ds", 1, TaggingThresholds(per_tag_enabled=True, per_tag_column="best_recall"))
+ off = service._tag_result_key("ds", 1, TaggingThresholds(per_tag_enabled=False))
+ assert unknown == off
+
+ def test_job_options_rejects_unknown_column(self) -> None:
+ import pydantic
+
+ from yadc.api.services.tagging import TagJobOptions
+
+ with pytest.raises(pydantic.ValidationError):
+ TagJobOptions(per_tag_column="best_f1")
+ assert TagJobOptions(per_tag_column="best_threshold").per_tag_column == "best_threshold"
+
+ def test_get_tag_result_applies_per_tag(self, service: TaggingService) -> None:
+ import asyncio
+
+ from yadc.api.services.tagging import TaggingThresholds
+ from yadc.taggers.base import TaggerResult
+
+ service._per_tag_thresholds_all = {"best_threshold": {"1girl": 0.5, "solo": 0.2}}
+ cached = TaggerResult(tags={"1girl": 0.4, "solo": 0.3}, categories={"general": ["1girl", "solo"]})
+ key = service._tag_result_key("ds", 1, TaggingThresholds(per_tag_enabled=True, per_tag_column="best_threshold"))
+ service._tag_results[key] = cached
+ out = asyncio.run(service.get_tag_result("ds", 1, thresholds=TaggingThresholds(per_tag_enabled=True, per_tag_column="best_threshold")))
+ assert out is not None
+ assert "1girl" not in out.tags
+ assert "solo" in out.tags
+
+ def test_get_tag_result_toggle_off_disables_per_tag(self, service: TaggingService) -> None:
+ import asyncio
+
+ from yadc.api.services.tagging import TaggingThresholds
+ from yadc.taggers.base import TaggerResult
+
+ service._per_tag_thresholds_all = {"best_threshold": {"1girl": 0.5}}
+ cached = TaggerResult(tags={"1girl": 0.4}, categories={"general": ["1girl"]})
+ key = service._tag_result_key("ds", 1, TaggingThresholds(per_tag_enabled=False))
+ service._tag_results[key] = cached
+ out = asyncio.run(service.get_tag_result("ds", 1, thresholds=TaggingThresholds(per_tag_enabled=False)))
+ assert out is not None
+ assert "1girl" in out.tags
+
+
class TestBucketThreshold:
"""``bucket_threshold`` — rounds a tagger threshold DOWN to the nearest 0.2 boundary."""
diff --git a/tests/taggers/test_swap_active_model.py b/tests/taggers/test_swap_active_model.py
index 4bd0576..b259d39 100644
--- a/tests/taggers/test_swap_active_model.py
+++ b/tests/taggers/test_swap_active_model.py
@@ -161,6 +161,24 @@ def test_respawn_failure_rolls_back_active_tagger(self, test_configuration: Conf
assert service.active_tagger == first_active
assert service.active_tagger.repo_id == "SmilingWolf/wd-eva02-large-tagger-v3"
+ def test_respawn_failure_refreshes_per_tag_thresholds(self, test_configuration: Configuration, service: TaggingService):
+ """Rollback re-resolves the label CSV so thresholds don't point at the failed model."""
+ test_configuration.tagger_repo_id = "SmilingWolf/wd-eva02-large-tagger-v3"
+ client = make_client_mock(alive=True)
+ with patch_client_factory(client):
+ run_swap(service, _selection("SmilingWolf/wd-eva02-large-tagger-v3"))
+
+ with patch_client_factory(client), pytest.MonkeyPatch.context() as mp:
+ mp.setattr(
+ service,
+ "_ensure_running_locked",
+ AsyncMock(side_effect=RuntimeError("boom")),
+ )
+ refresh = MagicMock()
+ mp.setattr(service, "_refresh_per_tag_thresholds", refresh)
+ run_swap(service, _selection("SmilingWolf/wd-vit-tagger-v3"))
+ assert refresh.call_count >= 1
+
class TestSwapPersistFailure:
"""Persist-after-respawn failure logs but doesn't roll back."""
diff --git a/yadc/api/configuration.py b/yadc/api/configuration.py
index 926c636..e617e16 100644
--- a/yadc/api/configuration.py
+++ b/yadc/api/configuration.py
@@ -175,6 +175,12 @@ class Configuration:
tagger_rating_threshold: float = 0.0
tagger_general_threshold: float = 0.35
tagger_character_threshold: float = 0.85
+ # Per-tag thresholds — when enabled, uses per-tag optimal thresholds
+ # from the CSV (e.g. best_threshold) instead of the global category
+ # thresholds above. Only available when the CSV contains threshold
+ # columns (e.g. animetimm models).
+ tagger_per_tag_thresholds: bool = False
+ tagger_per_tag_column: str = "best_threshold"
# Turn underscored tag names (``long_hair``) into spaces (``long hair``).
# Kaomojis are always preserved. Applied post-threshold so only
# surviving tags are touched. Off by default to preserve raw model
diff --git a/yadc/api/controllers/api_tagging.py b/yadc/api/controllers/api_tagging.py
index 5b71dc2..890536c 100644
--- a/yadc/api/controllers/api_tagging.py
+++ b/yadc/api/controllers/api_tagging.py
@@ -62,7 +62,7 @@
import pydantic
from quart import Response, jsonify, request
-from yadc.taggers.onnx_preprocess import list_profiles
+from yadc.taggers.onnx_preprocess import PER_TAG_THRESHOLD_COLUMNS, list_profiles
from ..configuration import Configuration
from ..modules.dataset_watcher import SELF_JOB_ID
@@ -101,6 +101,17 @@ class TagImageBody(pydantic.BaseModel):
general_threshold: float | None = None
character_threshold: float | None = None
replace_underscores: bool | None = None
+ per_tag_thresholds: bool | None = None
+ per_tag_column: str | None = None
+
+ @pydantic.field_validator("per_tag_column")
+ @classmethod
+ def _validate_per_tag_column(cls, v: str | None) -> str | None:
+ if v is None:
+ return None
+ if v not in PER_TAG_THRESHOLD_COLUMNS:
+ raise ValueError(f"Unknown per_tag_column: {v!r}")
+ return v
class TagCustomizationsBody(pydantic.BaseModel):
@@ -238,11 +249,22 @@ async def tag_image(name: str, image_id: int): # pyright: ignore[reportUnusedFu
source = request.args.get("source") or None
thresholds: TaggingThresholds | None = None
- if any(v is not None for v in (body.rating_threshold, body.general_threshold, body.character_threshold)):
+ if any(
+ v is not None
+ for v in (
+ body.rating_threshold,
+ body.general_threshold,
+ body.character_threshold,
+ body.per_tag_thresholds,
+ body.per_tag_column,
+ )
+ ):
thresholds = TaggingThresholds(
rating=body.rating_threshold if body.rating_threshold is not None else configuration.tagger_rating_threshold,
general=body.general_threshold if body.general_threshold is not None else configuration.tagger_general_threshold,
character=body.character_threshold if body.character_threshold is not None else configuration.tagger_character_threshold,
+ per_tag_enabled=body.per_tag_thresholds if body.per_tag_thresholds is not None else configuration.tagger_per_tag_thresholds,
+ per_tag_column=body.per_tag_column if body.per_tag_column is not None else configuration.tagger_per_tag_column,
)
try:
@@ -416,14 +438,21 @@ def _float_arg(name: str) -> float | None:
character = _float_arg("character_threshold")
replace_raw = request.args.get("replace_underscores")
replace = None if replace_raw is None else replace_raw.lower() in ("true", "1", "yes")
+ per_tag_raw = request.args.get("per_tag_thresholds")
+ per_tag_enabled = None if per_tag_raw is None else per_tag_raw.lower() in ("true", "1", "yes")
+ per_tag_column = request.args.get("per_tag_column")
+ if per_tag_column is not None and per_tag_column not in PER_TAG_THRESHOLD_COLUMNS:
+ return jsonify_error(f"Unknown per_tag_column: {per_tag_column!r}", status=400, code=ErrorCode.BAD_REQUEST)
thresholds = (
TaggingThresholds(
rating=rating if rating is not None else configuration.tagger_rating_threshold,
general=general if general is not None else configuration.tagger_general_threshold,
character=character if character is not None else configuration.tagger_character_threshold,
+ per_tag_enabled=per_tag_enabled if per_tag_enabled is not None else configuration.tagger_per_tag_thresholds,
+ per_tag_column=per_tag_column if per_tag_column is not None else configuration.tagger_per_tag_column,
)
- if any(v is not None for v in (rating, general, character))
+ if any(v is not None for v in (rating, general, character, per_tag_enabled, per_tag_column))
else None
)
@@ -513,13 +542,20 @@ def _float_arg(arg: str) -> float | None:
rating = _float_arg("rating_threshold")
general = _float_arg("general_threshold")
character = _float_arg("character_threshold")
+ per_tag_raw = request.args.get("per_tag_thresholds")
+ per_tag_enabled = None if per_tag_raw is None else per_tag_raw.lower() in ("true", "1", "yes")
+ per_tag_column = request.args.get("per_tag_column")
+ if per_tag_column is not None and per_tag_column not in PER_TAG_THRESHOLD_COLUMNS:
+ return jsonify_error(f"Unknown per_tag_column: {per_tag_column!r}", status=400, code=ErrorCode.BAD_REQUEST)
thresholds = (
TaggingThresholds(
rating=rating if rating is not None else configuration.tagger_rating_threshold,
general=general if general is not None else configuration.tagger_general_threshold,
character=character if character is not None else configuration.tagger_character_threshold,
+ per_tag_enabled=per_tag_enabled if per_tag_enabled is not None else configuration.tagger_per_tag_thresholds,
+ per_tag_column=per_tag_column if per_tag_column is not None else configuration.tagger_per_tag_column,
)
- if any(v is not None for v in (rating, general, character))
+ if any(v is not None for v in (rating, general, character, per_tag_enabled, per_tag_column))
else None
)
await tagging.set_tag_customizations(
@@ -613,9 +649,7 @@ async def suggest_tags(): # pyright: ignore[reportUnusedFunction]
return jsonify(
{
"query": raw_q,
- "suggestions": [
- {"name": name, "category": category} for name, category in suggestions
- ],
+ "suggestions": [{"name": name, "category": category} for name, category in suggestions],
}
)
@@ -764,6 +798,8 @@ async def get_active_tagger(): # pyright: ignore[reportUnusedFunction]
return jsonify({"active": None, "is_available": False})
payload = active.model_dump()
payload["source"] = active.source_label
+ payload["has_per_tag_thresholds"] = tagging.has_per_tag_thresholds
+ payload["per_tag_columns"] = tagging.per_tag_columns
return jsonify({"active": payload, "is_available": tagging.is_available})
@app.post("/tagger/swap")
diff --git a/yadc/api/services/tagging.py b/yadc/api/services/tagging.py
index cf5d5d1..254c8a7 100644
--- a/yadc/api/services/tagging.py
+++ b/yadc/api/services/tagging.py
@@ -55,6 +55,7 @@
from yadc.taggers import OnnxTagger, TaggerResult, apply_thresholds, extras_tags, format_draft
from yadc.taggers.base import TagCustomizations, tamer_result_size
from yadc.taggers.client import TaggerClient
+from yadc.taggers.onnx_preprocess import PER_TAG_THRESHOLD_COLUMNS, _load_per_tag_thresholds
from yadc.taggers.postprocessing import TagPolicy, apply_policy
from yadc.taggers.postprocessing import replace_underscores as replace_underscores_in
from yadc.utils import MemoryLRU, size_units
@@ -121,6 +122,8 @@ class TaggingThresholds:
rating: float = 0.0
general: float = 0.35
character: float = 0.85
+ per_tag_enabled: bool = False
+ per_tag_column: str = "best_threshold"
class TagSaveOptions(pydantic.BaseModel):
@@ -160,9 +163,20 @@ class TagJobOptions(pydantic.BaseModel):
general_threshold: float | None = None
character_threshold: float | None = None
replace_underscores: bool | None = None
+ per_tag_thresholds: bool | None = None
+ per_tag_column: str | None = None
save: TagSaveOptions = pydantic.Field(default_factory=TagSaveOptions)
source: str | None = None
+ @pydantic.field_validator("per_tag_column")
+ @classmethod
+ def _validate_per_tag_column(cls, v: str | None) -> str | None:
+ if v is None:
+ return None
+ if v not in PER_TAG_THRESHOLD_COLUMNS:
+ raise ValueError(f"Unknown per_tag_column: {v!r}")
+ return v
+
@dataclass(frozen=True)
class TaggerResultKey:
@@ -206,6 +220,8 @@ class TaggerResultKey:
rating_threshold: float
general_threshold: float
character_threshold: float
+ per_tag_enabled: bool
+ per_tag_column: str
@dataclass
@@ -363,6 +379,11 @@ def __init__(
size_fn=tamer_result_size,
)
+ # Per-tag threshold cache. Populated when the model loads;
+ # holds all columns (for the API response and per-request
+ # resolution via _per_tag_for_thresholds).
+ self._per_tag_thresholds_all: dict[str, dict[str, float]] | None = None
+
# Hydrate the persisted active-tagger selection last so the
# LRU and other fields above are ready, and so a malformed
# row can't take the service down. Warnings are logged; on
@@ -370,6 +391,7 @@ def __init__(
# service still operates against the flat ``Configuration``
# fallback once Phase 2 wires that path.
self._hydrate_active_tagger()
+ self._refresh_per_tag_thresholds()
# --- event handlers ---------------------------------------------------
@@ -409,6 +431,18 @@ def is_available(self) -> bool:
"""Whether the tagger subprocess is currently running."""
return self._tagger_client is not None and self._tagger_client.is_alive
+ @property
+ def has_per_tag_thresholds(self) -> bool:
+ """Whether the active model's CSV has per-tag threshold columns."""
+ return self._per_tag_thresholds_all is not None and len(self._per_tag_thresholds_all) > 0
+
+ @property
+ def per_tag_columns(self) -> list[str]:
+ """Available per-tag threshold column names from the CSV."""
+ if self._per_tag_thresholds_all is None:
+ return []
+ return list(self._per_tag_thresholds_all.keys())
+
# --- persisted active-tagger selection --------------------------------
async def swap_active_model(self, selection: ActiveTagger) -> ActiveTagger:
@@ -516,6 +550,7 @@ async def _swap_background(self, selection: ActiveTagger, old_active: ActiveTagg
self._active_tagger = old_active
self._swapped_at = None
self._logger.exception("Respawn with new model failed; rolling back.")
+ self._refresh_per_tag_thresholds()
if old_active is not None:
try:
await self._ensure_running_locked()
@@ -534,6 +569,7 @@ async def _swap_background(self, selection: ActiveTagger, old_active: ActiveTagg
selection.kind,
selection.source_label,
)
+ self._refresh_per_tag_thresholds()
finally:
self._lifecycle_lock.release()
except Exception:
@@ -543,6 +579,10 @@ async def _swap_background(self, selection: ActiveTagger, old_active: ActiveTagg
# the swap-in-progress lock held. Surface it as a ``failed`` SSE
# event and let the finally release the lock.
self._logger.exception("Unexpected error during tagger swap.")
+ try:
+ self._refresh_per_tag_thresholds()
+ except Exception:
+ pass
self._emit_status("failed", error="unexpected error during tagger swap")
finally:
self._swap_in_progress_lock.release()
@@ -617,6 +657,7 @@ def set_active_tagger(self, selection: ActiveTagger) -> None:
# usable, and log a warning so the cause is visible.
self._settings_service.set("tagger.active_model", selection.model_dump())
self._active_tagger = selection
+ self._refresh_per_tag_thresholds()
self._logger.info(
"Active tagger persisted. [kind=%s, source=%s]",
selection.kind,
@@ -722,6 +763,8 @@ def _tag_result_key(
# Floored so nearby request thresholds share a slot.
general_threshold=bucket_threshold(thresholds.general),
character_threshold=bucket_threshold(thresholds.character),
+ per_tag_enabled=thresholds.per_tag_enabled and self._per_tag_for_thresholds(thresholds) is not None,
+ per_tag_column=thresholds.per_tag_column if thresholds.per_tag_enabled and self._per_tag_for_thresholds(thresholds) is not None else "",
)
@staticmethod
@@ -730,6 +773,7 @@ def _refilter(
eff_thresholds: TaggingThresholds,
eff_replace: bool,
policy: TagPolicy | None = None,
+ per_tag_thresholds: dict[str, float] | None = None,
) -> TaggerResult:
"""Re-apply the request's effective thresholds + policy + ``replace_underscores`` to a cached value.
@@ -749,6 +793,7 @@ def _refilter(
rating_threshold=eff_thresholds.rating,
general_threshold=eff_thresholds.general,
character_threshold=eff_thresholds.character,
+ per_tag_thresholds=per_tag_thresholds,
)
if policy is not None:
re_filtered = apply_policy(re_filtered, policy)
@@ -799,6 +844,8 @@ async def get_tag_result(
rating=self._configuration.tagger_rating_threshold,
general=self._configuration.tagger_general_threshold,
character=self._configuration.tagger_character_threshold,
+ per_tag_enabled=self._configuration.tagger_per_tag_thresholds,
+ per_tag_column=self._configuration.tagger_per_tag_column,
)
eff_replace = replace_underscores if replace_underscores is not None else self._configuration.tagger_replace_underscores
eff_policy = self._resolve_policy(dataset_name, override=policy)
@@ -807,7 +854,8 @@ async def get_tag_result(
cached = self._tag_results.get(key)
if cached is None:
return None
- return self._refilter(cached, eff_thresholds, eff_replace, eff_policy)
+ per_tag = self._per_tag_for_thresholds(eff_thresholds)
+ return self._refilter(cached, eff_thresholds, eff_replace, eff_policy, per_tag)
async def evict_tag_result(
self,
@@ -828,6 +876,8 @@ async def evict_tag_result(
rating=self._configuration.tagger_rating_threshold,
general=self._configuration.tagger_general_threshold,
character=self._configuration.tagger_character_threshold,
+ per_tag_enabled=self._configuration.tagger_per_tag_thresholds,
+ per_tag_column=self._configuration.tagger_per_tag_column,
)
key = self._tag_result_key(dataset_name, image_id, eff_thresholds)
async with self._tag_lock:
@@ -861,6 +911,8 @@ async def set_tag_customizations(
rating=self._configuration.tagger_rating_threshold,
general=self._configuration.tagger_general_threshold,
character=self._configuration.tagger_character_threshold,
+ per_tag_enabled=self._configuration.tagger_per_tag_thresholds,
+ per_tag_column=self._configuration.tagger_per_tag_column,
)
key = self._tag_result_key(dataset_name, image_id, eff_thresholds)
async with self._tag_lock:
@@ -936,6 +988,8 @@ async def tag_image(
rating=self._configuration.tagger_rating_threshold,
general=self._configuration.tagger_general_threshold,
character=self._configuration.tagger_character_threshold,
+ per_tag_enabled=self._configuration.tagger_per_tag_thresholds,
+ per_tag_column=self._configuration.tagger_per_tag_column,
)
eff_replace = replace_underscores if replace_underscores is not None else self._configuration.tagger_replace_underscores
eff_policy = self._resolve_policy(dataset_name, override=policy)
@@ -952,7 +1006,8 @@ async def tag_image(
async with self._tag_lock:
cached = self._tag_results.get(cache_key)
if cached is not None:
- re_filtered = self._refilter(cached, eff_thresholds, eff_replace, eff_policy)
+ per_tag = self._per_tag_for_thresholds(eff_thresholds)
+ re_filtered = self._refilter(cached, eff_thresholds, eff_replace, eff_policy, per_tag)
# Yield once so the event loop can run other ready tasks
# (idle check, in-flight cancellations) between back-to-back
# cache hits in a batch. ``_refilter`` is the closest thing
@@ -1030,11 +1085,13 @@ async def tag_image(
# NOT applied to the cache value: both are read-time transforms
# applied post-hoc at retrieval so the slot can serve every
# variant of either without re-tagging.
+ per_tag = self._per_tag_for_thresholds(eff_thresholds)
bucketed = apply_thresholds(
result,
rating_threshold=cache_key.rating_threshold,
general_threshold=cache_key.general_threshold,
character_threshold=cache_key.character_threshold,
+ per_tag_thresholds=per_tag,
)
# Effective (request-specific) for caller + dispatched event.
@@ -1043,7 +1100,7 @@ async def tag_image(
# this only drops more tags, never adds any. The policy is
# applied here too (injects always-add, drops banned) so the
# returned result and SSE event match what the caller asked for.
- effective = self._refilter(bucketed, eff_thresholds, False, eff_policy) # replace applied below
+ effective = self._refilter(bucketed, eff_thresholds, False, eff_policy, per_tag) # replace applied below
duration_ms = int((time.monotonic() - start_t) * 1000)
self._event_dispatcher.dispatch(
@@ -1295,6 +1352,73 @@ def _resolve_label_path(self, model_path: str) -> str | None:
candidate = Path(model_path).parent / "selected_tags.csv"
return str(candidate) if candidate.exists() else None
+ def _resolve_label_csv_path(self) -> str | None:
+ """Resolve the path to the active model's selected_tags.csv."""
+ if self._active_tagger is not None:
+ if self._active_tagger.kind == "hf":
+ try:
+ from huggingface_hub import hf_hub_download
+
+ return hf_hub_download(
+ repo_id=self._active_tagger.repo_id,
+ filename=self._active_tagger.repo_label_filename,
+ )
+ except Exception:
+ return None
+ label_path = self._active_tagger.label_path
+ if label_path:
+ return label_path
+ model_path = self._active_tagger.model_path
+ if model_path:
+ return str(Path(model_path).parent / "selected_tags.csv")
+ return None
+ cfg = self._configuration
+ repo_id = cfg.tagger_repo_id.strip()
+ if repo_id:
+ try:
+ from huggingface_hub import hf_hub_download
+
+ return hf_hub_download(
+ repo_id=repo_id,
+ filename=cfg.tagger_repo_label_filename,
+ )
+ except Exception:
+ return None
+ label_path = cfg.tagger_label_path.strip()
+ if label_path:
+ return label_path
+ model_path = cfg.tagger_model_path.strip()
+ if model_path:
+ return str(Path(model_path).parent / "selected_tags.csv")
+ return None
+
+ def _refresh_per_tag_thresholds(self) -> None:
+ """Load and cache per-tag thresholds from the active model's CSV.
+
+ Called after model load/swap and at startup. Populates
+ ``_per_tag_thresholds_all`` (all columns, for the API response
+ and per-request resolution via :meth:`_per_tag_for_thresholds`).
+ """
+ self._per_tag_thresholds_all = None
+ csv_path = self._resolve_label_csv_path()
+ if csv_path is None:
+ return
+ try:
+ all_thresholds = _load_per_tag_thresholds(Path(csv_path))
+ except FileNotFoundError:
+ return
+ if all_thresholds is None:
+ return
+ self._per_tag_thresholds_all = all_thresholds
+
+ def _per_tag_for_thresholds(self, thresholds: TaggingThresholds) -> dict[str, float] | None:
+ """Resolve the per-tag dict for a request's effective thresholds."""
+ if not thresholds.per_tag_enabled:
+ return None
+ if self._per_tag_thresholds_all is None:
+ return None
+ return self._per_tag_thresholds_all.get(thresholds.per_tag_column)
+
def _catalog_sidecars(self, repo_id: str) -> list[str]:
"""Sidecar filenames to fetch alongside a known catalog repo's model.
@@ -1862,10 +1986,14 @@ async def _emit_tag_status(self, dataset_name: str, state: "_TagJobState") -> No
def _thresholds_from_options(self, options: TagJobOptions) -> TaggingThresholds:
"""Resolve per-request threshold overrides against the server config."""
cfg = self._configuration
+ per_tag_enabled = options.per_tag_thresholds if options.per_tag_thresholds is not None else cfg.tagger_per_tag_thresholds
+ per_tag_column = options.per_tag_column if options.per_tag_column is not None else cfg.tagger_per_tag_column
return TaggingThresholds(
rating=options.rating_threshold if options.rating_threshold is not None else cfg.tagger_rating_threshold,
general=options.general_threshold if options.general_threshold is not None else cfg.tagger_general_threshold,
character=options.character_threshold if options.character_threshold is not None else cfg.tagger_character_threshold,
+ per_tag_enabled=per_tag_enabled,
+ per_tag_column=per_tag_column,
)
def _resolve_job_images(self, dataset_name: str, image_ids: list[int] | None) -> list[ImageInfo]:
diff --git a/yadc/taggers/onnx.py b/yadc/taggers/onnx.py
index 2fc87f3..e6a36b2 100644
--- a/yadc/taggers/onnx.py
+++ b/yadc/taggers/onnx.py
@@ -366,6 +366,7 @@ def apply_thresholds(
rating_threshold: float = 0.0,
general_threshold: float = 0.35,
character_threshold: float = 0.85,
+ per_tag_thresholds: dict[str, float] | None = None,
) -> TaggerResult:
"""Return a new :class:`TaggerResult` with low-score tags dropped per category.
@@ -383,10 +384,15 @@ def apply_thresholds(
drop: set[str] = set()
for cat_name, cat_tags in result.categories.items():
- thr = thresholds.get(cat_name, 0.0)
- if thr <= 0:
- continue
+ cat_thr = thresholds.get(cat_name, 0.0)
for tag in cat_tags:
+ thr = 0.0
+ if per_tag_thresholds and tag in per_tag_thresholds:
+ thr = per_tag_thresholds[tag]
+ elif cat_thr > 0:
+ thr = cat_thr
+ if thr <= 0:
+ continue
score = result.tags.get(tag, 0.0)
if score < thr:
drop.add(tag)
diff --git a/yadc/taggers/onnx_preprocess.py b/yadc/taggers/onnx_preprocess.py
index 763b275..9226a16 100644
--- a/yadc/taggers/onnx_preprocess.py
+++ b/yadc/taggers/onnx_preprocess.py
@@ -32,9 +32,11 @@
from __future__ import annotations
+import csv
import io
import logging
from collections.abc import Sequence
+from pathlib import Path
from typing import Literal, NamedTuple
import numpy as np
@@ -383,3 +385,38 @@ def log_profile_info(
width,
profile.default_input_size,
)
+
+
+PER_TAG_THRESHOLD_COLUMNS = frozenset({"best_threshold", "best_recall"})
+
+
+def _load_per_tag_thresholds(path: Path) -> dict[str, dict[str, float]] | None:
+ """Extract per-tag threshold columns from a selected_tags.csv.
+
+ Returns {column_name: {tag_name: threshold_value}} when threshold
+ columns are present, None otherwise. Supported columns: best_threshold,
+ best_recall. Values that aren't valid floats between 0 and 1 are skipped.
+ """
+ result: dict[str, dict[str, float]] = {}
+ with open(path, newline="", encoding="utf-8") as f:
+ reader = csv.DictReader(f)
+ if reader.fieldnames is None:
+ return None
+ available = PER_TAG_THRESHOLD_COLUMNS.intersection(reader.fieldnames)
+ if not available:
+ return None
+ for col in available:
+ result[col] = {}
+ for row in reader:
+ name = row.get("name")
+ if not name:
+ continue
+ for col in available:
+ raw = row.get(col, "")
+ try:
+ val = float(raw)
+ except (ValueError, TypeError):
+ continue
+ if 0.0 <= val <= 1.0:
+ result[col][name] = val
+ return result if result else None
diff --git a/yadc/webui/src/lib/components/dataset/detail/Tags.svelte b/yadc/webui/src/lib/components/dataset/detail/Tags.svelte
index e25e38b..af9ef3d 100644
--- a/yadc/webui/src/lib/components/dataset/detail/Tags.svelte
+++ b/yadc/webui/src/lib/components/dataset/detail/Tags.svelte
@@ -71,6 +71,8 @@
let generalThreshold = $derived($tagSettings.generalThreshold);
let characterThreshold = $derived($tagSettings.characterThreshold);
let replaceUnderscores = $derived($tagSettings.replaceUnderscores);
+ let perTagThresholds = $derived($tagSettings.perTagThresholds);
+ let perTagColumn = $derived($tagSettings.perTagColumn);
// Policy lists are read reactively (not via ``get``) so toggling a
// tag in the always-add / banned Settings re-fetches the cached
// result — the policy is a backend read-time transform, so the
@@ -132,6 +134,12 @@
void alwaysAdd;
void banned;
void policyLoadedFor;
+ void ratingThreshold;
+ void generalThreshold;
+ void characterThreshold;
+ void replaceUnderscores;
+ void perTagThresholds;
+ void perTagColumn;
const imageId = item.id;
const controller = linkedController(parentSignal);
@@ -151,7 +159,9 @@
rating_threshold: ratingThreshold,
general_threshold: generalThreshold,
character_threshold: characterThreshold,
- replace_underscores: replaceUnderscores
+ replace_underscores: replaceUnderscores,
+ per_tag_thresholds: perTagThresholds,
+ per_tag_column: perTagThresholds ? perTagColumn : null
},
controller.signal
)) ?? undefined;
@@ -435,6 +445,8 @@
rating_threshold: number | null;
general_threshold: number | null;
character_threshold: number | null;
+ per_tag_thresholds: boolean | null;
+ per_tag_column: string | null;
};
customizations: TagCustomizations;
}) => {
@@ -483,7 +495,9 @@
thresholds: {
rating_threshold: ratingThreshold,
general_threshold: generalThreshold,
- character_threshold: characterThreshold
+ character_threshold: characterThreshold,
+ per_tag_thresholds: perTagThresholds,
+ per_tag_column: perTagThresholds ? perTagColumn : null
},
customizations: sel.customizations
});
diff --git a/yadc/webui/src/lib/components/tagging/TagSettings.svelte b/yadc/webui/src/lib/components/tagging/TagSettings.svelte
index e97c4ff..55e0caa 100644
--- a/yadc/webui/src/lib/components/tagging/TagSettings.svelte
+++ b/yadc/webui/src/lib/components/tagging/TagSettings.svelte
@@ -14,6 +14,7 @@
TAG_CATEGORIES
} from '$lib/stores/tagging';
import type { TagSaveOptions, TaggerResult } from '$lib/stores/tagging';
+ import { activeTagger, ensureActiveTaggerLoaded } from '$lib/stores/tagging';
import Checkbox from '$lib/components/ui/Checkbox.svelte';
import Slider from '$lib/components/ui/Slider.svelte';
import PolicyList from './PolicyList.svelte';
@@ -27,6 +28,37 @@
let generalThreshold: number | null = $state(null);
let characterThreshold: number | null = $state(null);
let replaceUnderscores = $state(false);
+ let perTagThresholds = $state(false);
+ let perTagColumn = $state('best_threshold');
+ let supportsPerTag = $state(false);
+ let perTagColumns: string[] = $state([]);
+
+ const PER_TAG_COLUMN_LABELS: Record
+ When the model's CSV contains per-tag optimal thresholds (e.g. + animetimm), use those instead of the three global category thresholds + below. +
+