per-tag thresholds for animetimm models

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."""

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

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."""
 

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."""

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