From 400f47ca6f5b9b066cd163de5dab0c1cdfdb60ea Mon Sep 17 00:00:00 2001 From: Foxyz <58236817+Foxyz1302@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:01:38 +0300 Subject: [PATCH 1/5] feat(tagger): per-request per-tag thresholds backend Wire per_tag_thresholds/per_tag_column per request through TagImageBody/TagJobOptions into TaggingThresholds, resolve the per-tag dict per request, isolate cache slots per toggle+column with unknown-column collapse, validate column values (400), refresh thresholds on swap rollback and set_active_tagger. --- yadc/api/configuration.py | 6 ++ yadc/api/controllers/api_tagging.py | 56 ++++++++++-- yadc/api/services/tagging.py | 137 +++++++++++++++++++++++++++- yadc/taggers/onnx.py | 48 +++++++++- 4 files changed, 235 insertions(+), 12 deletions(-) diff --git a/yadc/api/configuration.py b/yadc/api/configuration.py index 926c636b..e617e16a 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 5b71dc20..61277d8e 100644 --- a/yadc/api/controllers/api_tagging.py +++ b/yadc/api/controllers/api_tagging.py @@ -101,6 +101,19 @@ 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 + from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS + + 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 +251,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 +440,24 @@ 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: + from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS + + if 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 +547,23 @@ 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: + from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS as _PER_TAG_COLS + + if per_tag_column not in _PER_TAG_COLS: + 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 +657,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 +806,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 cf5d5d12..b27d80e1 100644 --- a/yadc/api/services/tagging.py +++ b/yadc/api/services/tagging.py @@ -121,6 +121,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 +162,22 @@ 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 + from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS + + 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 +221,8 @@ class TaggerResultKey: rating_threshold: float general_threshold: float character_threshold: float + per_tag_enabled: bool + per_tag_column: str @dataclass @@ -363,6 +380,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 +392,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 +432,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 +551,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 +570,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 +580,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 +658,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 +764,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 +774,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 +794,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 +845,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 +855,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 +877,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 +912,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 +989,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 +1007,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 +1086,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 +1101,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 +1353,75 @@ 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 + from yadc.taggers.onnx import _load_per_tag_thresholds + + 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 +1989,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 2fc87f36..c93c03ac 100644 --- a/yadc/taggers/onnx.py +++ b/yadc/taggers/onnx.py @@ -96,6 +96,42 @@ def _load_labels_csv(path: Path) -> tuple[list[str], dict[str, list[str]]]: return names, categories +PER_TAG_THRESHOLD_COLUMNS = frozenset({"best_threshold", "best_recall"}) + + +def _load_per_tag_thresholds(path: Path) -> dict[str, dict[str, float]] | None: # pyright: ignore[reportUnusedFunction] + """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. + """ + SUPPORTED_COLUMNS = PER_TAG_THRESHOLD_COLUMNS + 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 = SUPPORTED_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 + + def _load_labels_txt(path: Path) -> tuple[list[str], dict[str, list[str]]]: """Load a flat one-label-per-line text file. No categorization.""" text = path.read_text(encoding="utf-8") @@ -366,6 +402,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 +420,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) From f05daa885fee17b0c04656d1a71ae699c8e7e6f7 Mon Sep 17 00:00:00 2001 From: Foxyz <58236817+Foxyz1302@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:01:42 +0300 Subject: [PATCH 2/5] feat(webui): wire per-tag thresholds through tagging stores and panels Forward per-tag options in tagImage/startTagJob/fetch/preview, pass through single/batch actions, dynamic column dropdown from per_tag_columns with reset on unsupported models, Tags tab fetch and preview carry per-tag, re-fetch on swap settled via SSE. --- .../lib/components/dataset/detail/Tags.svelte | 18 +- .../lib/components/tagging/TagSettings.svelte | 178 +++++++++++++++--- yadc/webui/src/lib/stores/tagging/actions.ts | 10 + yadc/webui/src/lib/stores/tagging/api.ts | 24 ++- yadc/webui/src/lib/stores/tagging/results.ts | 2 + yadc/webui/src/lib/stores/tagging/settings.ts | 18 +- yadc/webui/src/lib/stores/tagging/types.ts | 8 +- 7 files changed, 219 insertions(+), 39 deletions(-) diff --git a/yadc/webui/src/lib/components/dataset/detail/Tags.svelte b/yadc/webui/src/lib/components/dataset/detail/Tags.svelte index e25e38b7..af9ef3de 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 e97c4fff..abbed4bf 100644 --- a/yadc/webui/src/lib/components/tagging/TagSettings.svelte +++ b/yadc/webui/src/lib/components/tagging/TagSettings.svelte @@ -14,6 +14,9 @@ TAG_CATEGORIES } from '$lib/stores/tagging'; import type { TagSaveOptions, TaggerResult } from '$lib/stores/tagging'; + import { taggerStatus } from '$lib/stores/tagging'; + import { fetchActiveTagger } from '$lib/stores/tagging/api'; + import { TAGGER_SWAPPED_EVENT } from '$lib/stores/tagging/actions'; import Checkbox from '$lib/components/ui/Checkbox.svelte'; import Slider from '$lib/components/ui/Slider.svelte'; import PolicyList from './PolicyList.svelte'; @@ -27,6 +30,77 @@ 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 = { + best_threshold: 'Best F1 threshold', + best_recall: 'Best recall (higher)' + }; + + // --- Fetch active tagger to detect per-tag support --- + // Re-fetch when a model swap happens: the window event fires on 202 + // Accepted (immediate attempt, same-tab), and the ``tagger_status`` SSE + // stream settles on ready/failed (handles minutes-long first-run HF + // downloads, cross-tab). + + let swapInFlight = $state(false); + let sawSwapTransition = $state(false); + + async function fetchPerTagSupport(signal?: AbortSignal, retries = 3) { + for (let i = 0; i < retries; i++) { + try { + const res = await fetchActiveTagger(signal); + supportsPerTag = res.active?.has_per_tag_thresholds ?? false; + perTagColumns = res.active?.per_tag_columns ?? []; + if (!supportsPerTag) { + perTagThresholds = false; + perTagColumn = 'best_threshold'; + } else if (perTagColumns.length > 0 && !perTagColumns.includes(perTagColumn)) { + perTagColumn = perTagColumns[0]; + } + return; + } catch { + if (i < retries - 1) { + await new Promise((r) => setTimeout(r, 2000)); + } + } + } + } + + $effect(() => { + fetchPerTagSupport(); + const handler = () => { + swapInFlight = true; + sawSwapTransition = false; + fetchPerTagSupport(); + }; + window.addEventListener(TAGGER_SWAPPED_EVENT, handler); + return () => window.removeEventListener(TAGGER_SWAPPED_EVENT, handler); + }); + + $effect(() => { + const status = $taggerStatus; + if (!swapInFlight) { + return; + } + if ( + status.state === 'stopping' || + status.state === 'stopped' || + status.state === 'starting' + ) { + sawSwapTransition = true; + return; + } + if (!sawSwapTransition) { + return; + } + swapInFlight = false; + sawSwapTransition = false; + fetchPerTagSupport(); + }); // --- Save options --- @@ -61,6 +135,8 @@ draftName = saved.draftName || 'tags'; draftFormat = saved.draftFormat || 'comma'; overwrite = saved.overwrite; + perTagThresholds = saved.perTagThresholds; + perTagColumn = saved.perTagColumn || 'best_threshold'; } }); @@ -68,7 +144,7 @@ function _buildSettings() { return { - $version: 3, + $version: 4, ratingThreshold, generalThreshold, characterThreshold, @@ -76,7 +152,9 @@ saveMode, draftName, draftFormat, - overwrite + overwrite, + perTagThresholds, + perTagColumn }; } @@ -96,6 +174,8 @@ void draftName; void draftFormat; void overwrite; + void perTagThresholds; + void perTagColumn; persistSettings(); }); @@ -113,6 +193,8 @@ general_threshold: generalThreshold ?? undefined, character_threshold: characterThreshold ?? undefined, replace_underscores: replaceUnderscores, + per_tag_thresholds: perTagThresholds, + per_tag_column: perTagThresholds ? perTagColumn : undefined, save: assembledSave }); @@ -251,35 +333,73 @@

- {@render thresholdRow( - 'tag-rating-threshold', - 'Rating', - ratingThreshold, - CANONICAL_THRESHOLDS.rating, - ratingOverridden, - (v) => (ratingThreshold = v), - () => resetThreshold(TAG_CATEGORIES.rating) - )} - {@render thresholdRow( - 'tag-general-threshold', - 'General', - generalThreshold, - CANONICAL_THRESHOLDS.general, - generalOverridden, - (v) => (generalThreshold = v), - () => resetThreshold(TAG_CATEGORIES.general) - )} - {@render thresholdRow( - 'tag-character-threshold', - 'Character', - characterThreshold, - CANONICAL_THRESHOLDS.character, - characterOverridden, - (v) => (characterThreshold = v), - () => resetThreshold(TAG_CATEGORIES.character) - )} + {#if supportsPerTag} +
+ +
+ +

+ When the model's CSV contains per-tag optimal thresholds (e.g. + animetimm), use those instead of the three global category thresholds + below. +

+
+
+ {#if perTagThresholds} +
+ + +
+ {/if} + {/if}
+ {#if !perTagThresholds} +
+ {@render thresholdRow( + 'tag-rating-threshold', + 'Rating', + ratingThreshold, + CANONICAL_THRESHOLDS.rating, + ratingOverridden, + (v) => (ratingThreshold = v), + () => resetThreshold(TAG_CATEGORIES.rating) + )} + {@render thresholdRow( + 'tag-general-threshold', + 'General', + generalThreshold, + CANONICAL_THRESHOLDS.general, + generalOverridden, + (v) => (generalThreshold = v), + () => resetThreshold(TAG_CATEGORIES.general) + )} + {@render thresholdRow( + 'tag-character-threshold', + 'Character', + characterThreshold, + CANONICAL_THRESHOLDS.character, + characterOverridden, + (v) => (characterThreshold = v), + () => resetThreshold(TAG_CATEGORIES.character) + )} +
+ {/if} +
diff --git a/yadc/webui/src/lib/stores/tagging/actions.ts b/yadc/webui/src/lib/stores/tagging/actions.ts index 60ec244c..54d6d550 100644 --- a/yadc/webui/src/lib/stores/tagging/actions.ts +++ b/yadc/webui/src/lib/stores/tagging/actions.ts @@ -35,6 +35,8 @@ export interface TagOptions { general_threshold?: number; character_threshold?: number; replace_underscores?: boolean; + per_tag_thresholds?: boolean; + per_tag_column?: string; save?: TagSaveOptions; source?: string; } @@ -67,6 +69,8 @@ export async function tagSingleImage( general_threshold: options.general_threshold, character_threshold: options.character_threshold, replace_underscores: options.replace_underscores, + per_tag_thresholds: options.per_tag_thresholds, + per_tag_column: options.per_tag_column, source }); } finally { @@ -96,6 +100,8 @@ export async function startBatchTagging(datasetName: string): Promise { general_threshold: options.general_threshold, character_threshold: options.character_threshold, replace_underscores: options.replace_underscores, + per_tag_thresholds: options.per_tag_thresholds, + per_tag_column: options.per_tag_column, save: options.save, source: options.source }); @@ -205,12 +211,16 @@ export async function saveImageTagsAction( * ``GeneralSettings``). Returns ``null`` when the server refused the swap * (busy / in-progress / error); those refusals are toasted here since they * are immediate. */ +/** Dispatched on window after a successful model swap so sibling components can react. */ +export const TAGGER_SWAPPED_EVENT = 'yadc:tagger-swapped'; + export async function swapActiveModelAction(body: SwapTaggerBody): Promise { const result = await apiSwapTaggerModel(body); switch (result.status) { case 'ok': // 202 Accepted — don't toast success yet; the swap (including any // first-run model download) is still running in the background. + window.dispatchEvent(new CustomEvent(TAGGER_SWAPPED_EVENT)); return result.response.active ? { kind: result.response.active.kind, diff --git a/yadc/webui/src/lib/stores/tagging/api.ts b/yadc/webui/src/lib/stores/tagging/api.ts index 7c8af5af..2794aa09 100644 --- a/yadc/webui/src/lib/stores/tagging/api.ts +++ b/yadc/webui/src/lib/stores/tagging/api.ts @@ -34,6 +34,8 @@ export async function tagImage( general_threshold?: number; character_threshold?: number; replace_underscores?: boolean; + per_tag_thresholds?: boolean; + per_tag_column?: string; source?: string; } = {}, signal?: AbortSignal @@ -52,7 +54,9 @@ export async function tagImage( rating_threshold: options.rating_threshold, general_threshold: options.general_threshold, character_threshold: options.character_threshold, - replace_underscores: options.replace_underscores + replace_underscores: options.replace_underscores, + per_tag_thresholds: options.per_tag_thresholds, + per_tag_column: options.per_tag_column }), signal } @@ -76,6 +80,8 @@ export async function startTagJob( general_threshold?: number; character_threshold?: number; replace_underscores?: boolean; + per_tag_thresholds?: boolean; + per_tag_column?: string; save?: TagSaveOptions; source?: string; } = {}, @@ -180,6 +186,8 @@ export async function fetchTagResult( general_threshold?: number | null; character_threshold?: number | null; replace_underscores?: boolean | null; + per_tag_thresholds?: boolean | null; + per_tag_column?: string | null; } = {}, signal?: AbortSignal ): Promise { @@ -196,6 +204,12 @@ export async function fetchTagResult( if (options.replace_underscores != null) { params.set('replace_underscores', options.replace_underscores ? 'true' : 'false'); } + if (options.per_tag_thresholds != null) { + params.set('per_tag_thresholds', options.per_tag_thresholds ? 'true' : 'false'); + } + if (options.per_tag_column != null) { + params.set('per_tag_column', options.per_tag_column); + } const qs = params.toString(); const res = await fetch( `${API_BASE}/api/datasets/${encodeURIComponent(datasetName)}/images/${imageId}/tag${qs ? '?' + qs : ''}`, @@ -226,6 +240,8 @@ export async function previewImageTags( rating_threshold?: number | null; general_threshold?: number | null; character_threshold?: number | null; + per_tag_thresholds?: boolean | null; + per_tag_column?: string | null; customizations?: TagCustomizations | null; } = {}, signal?: AbortSignal @@ -240,6 +256,12 @@ export async function previewImageTags( if (options.character_threshold != null) { params.set('character_threshold', String(options.character_threshold)); } + if (options.per_tag_thresholds != null) { + params.set('per_tag_thresholds', options.per_tag_thresholds ? 'true' : 'false'); + } + if (options.per_tag_column != null) { + params.set('per_tag_column', options.per_tag_column); + } const qs = params.toString(); const res = await fetch( `${API_BASE}/api/datasets/${encodeURIComponent(datasetName)}/images/${imageId}/tags/preview${qs ? '?' + qs : ''}`, diff --git a/yadc/webui/src/lib/stores/tagging/results.ts b/yadc/webui/src/lib/stores/tagging/results.ts index fd9549d1..80c9a468 100644 --- a/yadc/webui/src/lib/stores/tagging/results.ts +++ b/yadc/webui/src/lib/stores/tagging/results.ts @@ -58,6 +58,8 @@ export async function fetchCachedTagResult( general_threshold?: number | null; character_threshold?: number | null; replace_underscores?: boolean | null; + per_tag_thresholds?: boolean | null; + per_tag_column?: string | null; } = {}, signal?: AbortSignal ): Promise { diff --git a/yadc/webui/src/lib/stores/tagging/settings.ts b/yadc/webui/src/lib/stores/tagging/settings.ts index 62e6c01f..378d8f08 100644 --- a/yadc/webui/src/lib/stores/tagging/settings.ts +++ b/yadc/webui/src/lib/stores/tagging/settings.ts @@ -17,6 +17,10 @@ export interface TagSettings { draftFormat: string; /** Batch-only: skip images that already have the target save artifact. */ overwrite: boolean; + /** Use per-tag optimal thresholds from the CSV instead of global category thresholds. */ + perTagThresholds: boolean; + /** Which per-tag threshold column to use (best_threshold or best_recall). */ + perTagColumn: string; } /** Canonical wd-tagger defaults — what the server falls back to when a @@ -38,13 +42,15 @@ const TagSettingsSchema = z.object({ saveMode: z.enum(['none', 'draft', 'extras']), draftName: z.string(), draftFormat: z.string(), - overwrite: z.boolean() + overwrite: z.boolean(), + perTagThresholds: z.boolean(), + perTagColumn: z.string() }); export const tagSettings = storable( 'yadc/tagSettings', { - $version: 3, + $version: 4, ratingThreshold: null, generalThreshold: null, characterThreshold: null, @@ -52,9 +58,11 @@ export const tagSettings = storable( saveMode: 'draft', draftName: 'tags', draftFormat: 'comma', - overwrite: false + overwrite: false, + perTagThresholds: false, + perTagColumn: 'best_threshold' }, - // v2 → v3: add ``overwrite`` (defaults off — preserve existing tags by default). - (data) => ({ ...(data as object), $version: 3, overwrite: false }) as TagSettings, + // v3 → v4: add per-tag threshold toggle and column selector. + (data) => ({ ...(data as object), $version: 4, perTagThresholds: false, perTagColumn: 'best_threshold' }) as TagSettings, TagSettingsSchema ); diff --git a/yadc/webui/src/lib/stores/tagging/types.ts b/yadc/webui/src/lib/stores/tagging/types.ts index 3f3d17ae..71e889e4 100644 --- a/yadc/webui/src/lib/stores/tagging/types.ts +++ b/yadc/webui/src/lib/stores/tagging/types.ts @@ -158,6 +158,10 @@ export interface ActiveTaggerSelection { default_size: number; /** Server-derived source label (``hf:`` / ``local:``). */ source: string; + /** Whether the active model's CSV contains per-tag threshold columns. */ + has_per_tag_thresholds?: boolean; + /** Available per-tag threshold column names (e.g. ["best_threshold", "best_recall"]). */ + per_tag_columns?: string[]; } /** ``GET /api/tagger/active`` response shape. ``active === null`` when nothing is configured. */ @@ -166,8 +170,8 @@ export interface ActiveTaggerResponse { is_available: boolean; } -/** Body for the swap request — ``ActiveTaggerSelection`` minus the server-derived ``source``. */ -export type SwapTaggerBody = Omit; +/** Body for the swap request — ``ActiveTaggerSelection`` minus server-derived fields. */ +export type SwapTaggerBody = Omit; /** Response from ``POST /api/tagger/swap`` — same shape as the GET on success. */ export type SwapTaggerResponse = ActiveTaggerResponse; From 42c095fcac4c3304c5d933246a6045c338065bae Mon Sep 17 00:00:00 2001 From: Foxyz <58236817+Foxyz1302@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:01:45 +0300 Subject: [PATCH 3/5] test(tagger): per-tag loader, apply, resolution, and endpoint coverage Cover CSV parsing edge cases, per-tag override/fallback semantics, options resolution, cache-key divergence and unknown-column collapse, get_tag_result filtering, active-endpoint capability, body validation, and rollback refresh. --- tests/api/test_tagger_swap_endpoints.py | 30 +++++++ tests/taggers/test_onnx.py | 88 +++++++++++++++++++++ tests/taggers/test_service.py | 100 ++++++++++++++++++++++++ tests/taggers/test_swap_active_model.py | 18 +++++ 4 files changed, 236 insertions(+) diff --git a/tests/api/test_tagger_swap_endpoints.py b/tests/api/test_tagger_swap_endpoints.py index 8c298902..bbd44ee2 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 f155c3a3..19640d8c 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 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 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 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 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 8c9e5b5f..d94cfd50 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 4bd05765..b259d396 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.""" From 25cc22648b211665503996b390e70c25fda1a559 Mon Sep 17 00:00:00 2001 From: Foxyz <58236817+Foxyz1302@users.noreply.github.com> Date: Sun, 13 Sep 2026 14:49:32 +0300 Subject: [PATCH 4/5] refactor(tagger): move per-tag columns and loader to onnx_preprocess onnx pulls numpy and the session machinery; the column allowlist and CSV loader need neither. Top-level imports everywhere, no more lazy dance in validators and refresh. --- tests/taggers/test_onnx.py | 8 +++---- yadc/api/controllers/api_tagging.py | 18 ++++---------- yadc/api/services/tagging.py | 5 +--- yadc/taggers/onnx.py | 36 ---------------------------- yadc/taggers/onnx_preprocess.py | 37 +++++++++++++++++++++++++++++ 5 files changed, 47 insertions(+), 57 deletions(-) diff --git a/tests/taggers/test_onnx.py b/tests/taggers/test_onnx.py index 19640d8c..ec686efe 100644 --- a/tests/taggers/test_onnx.py +++ b/tests/taggers/test_onnx.py @@ -734,7 +734,7 @@ def _write_csv(self, tmp_path: Path, header: list[str], rows: list[list[str]]) - return path def test_reads_supported_columns(self, tmp_path: Path): - from yadc.taggers.onnx import _load_per_tag_thresholds + from yadc.taggers.onnx_preprocess import _load_per_tag_thresholds path = self._write_csv( tmp_path, @@ -747,13 +747,13 @@ def test_reads_supported_columns(self, tmp_path: Path): 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 import _load_per_tag_thresholds + 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 import _load_per_tag_thresholds + from yadc.taggers.onnx_preprocess import _load_per_tag_thresholds path = self._write_csv( tmp_path, @@ -765,7 +765,7 @@ def test_skips_invalid_values(self, tmp_path: Path): assert result["best_threshold"] == {"good": 0.5} def test_ignores_best_f1_column(self, tmp_path: Path): - from yadc.taggers.onnx import _load_per_tag_thresholds + 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 diff --git a/yadc/api/controllers/api_tagging.py b/yadc/api/controllers/api_tagging.py index 61277d8e..890536c6 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 @@ -109,8 +109,6 @@ class TagImageBody(pydantic.BaseModel): def _validate_per_tag_column(cls, v: str | None) -> str | None: if v is None: return None - from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS - if v not in PER_TAG_THRESHOLD_COLUMNS: raise ValueError(f"Unknown per_tag_column: {v!r}") return v @@ -443,11 +441,8 @@ def _float_arg(name: str) -> float | None: 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: - from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS - - if 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) + 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( @@ -550,11 +545,8 @@ def _float_arg(arg: str) -> float | None: 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: - from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS as _PER_TAG_COLS - - if per_tag_column not in _PER_TAG_COLS: - return jsonify_error(f"Unknown per_tag_column: {per_tag_column!r}", status=400, code=ErrorCode.BAD_REQUEST) + 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, diff --git a/yadc/api/services/tagging.py b/yadc/api/services/tagging.py index b27d80e1..254c8a7d 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 @@ -172,8 +173,6 @@ class TagJobOptions(pydantic.BaseModel): def _validate_per_tag_column(cls, v: str | None) -> str | None: if v is None: return None - from yadc.taggers.onnx import PER_TAG_THRESHOLD_COLUMNS - if v not in PER_TAG_THRESHOLD_COLUMNS: raise ValueError(f"Unknown per_tag_column: {v!r}") return v @@ -1404,8 +1403,6 @@ def _refresh_per_tag_thresholds(self) -> None: csv_path = self._resolve_label_csv_path() if csv_path is None: return - from yadc.taggers.onnx import _load_per_tag_thresholds - try: all_thresholds = _load_per_tag_thresholds(Path(csv_path)) except FileNotFoundError: diff --git a/yadc/taggers/onnx.py b/yadc/taggers/onnx.py index c93c03ac..e6a36b26 100644 --- a/yadc/taggers/onnx.py +++ b/yadc/taggers/onnx.py @@ -96,42 +96,6 @@ def _load_labels_csv(path: Path) -> tuple[list[str], dict[str, list[str]]]: return names, categories -PER_TAG_THRESHOLD_COLUMNS = frozenset({"best_threshold", "best_recall"}) - - -def _load_per_tag_thresholds(path: Path) -> dict[str, dict[str, float]] | None: # pyright: ignore[reportUnusedFunction] - """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. - """ - SUPPORTED_COLUMNS = PER_TAG_THRESHOLD_COLUMNS - 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 = SUPPORTED_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 - - def _load_labels_txt(path: Path) -> tuple[list[str], dict[str, list[str]]]: """Load a flat one-label-per-line text file. No categorization.""" text = path.read_text(encoding="utf-8") diff --git a/yadc/taggers/onnx_preprocess.py b/yadc/taggers/onnx_preprocess.py index 763b275d..9226a166 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 From 0c2056d553dcd624d51b4356a532fbf24d5d4ce3 Mon Sep 17 00:00:00 2001 From: Foxyz <58236817+Foxyz1302@users.noreply.github.com> Date: Sun, 13 Sep 2026 14:49:32 +0300 Subject: [PATCH 5/5] refactor(webui): active-tagger store replaces swap event and retry fetch Capability derives from a shared store re-fetched on tagger_status settle (cross-tab, survives long HF downloads). Drops TAGGER_SWAPPED_EVENT and the abort-unaware retry loop. --- .../lib/components/tagging/TagSettings.svelte | 72 ++++--------------- yadc/webui/src/lib/stores/tagging/actions.ts | 6 +- .../src/lib/stores/tagging/activeTagger.ts | 62 ++++++++++++++++ yadc/webui/src/lib/stores/tagging/index.ts | 1 + 4 files changed, 80 insertions(+), 61 deletions(-) create mode 100644 yadc/webui/src/lib/stores/tagging/activeTagger.ts diff --git a/yadc/webui/src/lib/components/tagging/TagSettings.svelte b/yadc/webui/src/lib/components/tagging/TagSettings.svelte index abbed4bf..55e0caa5 100644 --- a/yadc/webui/src/lib/components/tagging/TagSettings.svelte +++ b/yadc/webui/src/lib/components/tagging/TagSettings.svelte @@ -14,9 +14,7 @@ TAG_CATEGORIES } from '$lib/stores/tagging'; import type { TagSaveOptions, TaggerResult } from '$lib/stores/tagging'; - import { taggerStatus } from '$lib/stores/tagging'; - import { fetchActiveTagger } from '$lib/stores/tagging/api'; - import { TAGGER_SWAPPED_EVENT } from '$lib/stores/tagging/actions'; + 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'; @@ -40,66 +38,26 @@ best_recall: 'Best recall (higher)' }; - // --- Fetch active tagger to detect per-tag support --- - // Re-fetch when a model swap happens: the window event fires on 202 - // Accepted (immediate attempt, same-tab), and the ``tagger_status`` SSE - // stream settles on ready/failed (handles minutes-long first-run HF - // downloads, cross-tab). - - let swapInFlight = $state(false); - let sawSwapTransition = $state(false); - - async function fetchPerTagSupport(signal?: AbortSignal, retries = 3) { - for (let i = 0; i < retries; i++) { - try { - const res = await fetchActiveTagger(signal); - supportsPerTag = res.active?.has_per_tag_thresholds ?? false; - perTagColumns = res.active?.per_tag_columns ?? []; - if (!supportsPerTag) { - perTagThresholds = false; - perTagColumn = 'best_threshold'; - } else if (perTagColumns.length > 0 && !perTagColumns.includes(perTagColumn)) { - perTagColumn = perTagColumns[0]; - } - return; - } catch { - if (i < retries - 1) { - await new Promise((r) => setTimeout(r, 2000)); - } - } - } - } + // --- Per-tag support (from the active-tagger store) --- + // Capability comes from ``GET /api/tagger/active`` via the shared + // store, which re-fetches whenever the ``tagger_status`` SSE stream + // settles — so a swap (including minutes-long first-run HF + // downloads) surfaces here in every tab without polling. $effect(() => { - fetchPerTagSupport(); - const handler = () => { - swapInFlight = true; - sawSwapTransition = false; - fetchPerTagSupport(); - }; - window.addEventListener(TAGGER_SWAPPED_EVENT, handler); - return () => window.removeEventListener(TAGGER_SWAPPED_EVENT, handler); + ensureActiveTaggerLoaded().catch(() => {}); }); $effect(() => { - const status = $taggerStatus; - if (!swapInFlight) { - return; - } - if ( - status.state === 'stopping' || - status.state === 'stopped' || - status.state === 'starting' - ) { - sawSwapTransition = true; - return; - } - if (!sawSwapTransition) { - return; + const active = $activeTagger?.active; + supportsPerTag = active?.has_per_tag_thresholds ?? false; + perTagColumns = active?.per_tag_columns ?? []; + if (!supportsPerTag) { + perTagThresholds = false; + perTagColumn = 'best_threshold'; + } else if (perTagColumns.length > 0 && !perTagColumns.includes(perTagColumn)) { + perTagColumn = perTagColumns[0]; } - swapInFlight = false; - sawSwapTransition = false; - fetchPerTagSupport(); }); // --- Save options --- diff --git a/yadc/webui/src/lib/stores/tagging/actions.ts b/yadc/webui/src/lib/stores/tagging/actions.ts index 54d6d550..a1efb79d 100644 --- a/yadc/webui/src/lib/stores/tagging/actions.ts +++ b/yadc/webui/src/lib/stores/tagging/actions.ts @@ -211,16 +211,14 @@ export async function saveImageTagsAction( * ``GeneralSettings``). Returns ``null`` when the server refused the swap * (busy / in-progress / error); those refusals are toasted here since they * are immediate. */ -/** Dispatched on window after a successful model swap so sibling components can react. */ -export const TAGGER_SWAPPED_EVENT = 'yadc:tagger-swapped'; - export async function swapActiveModelAction(body: SwapTaggerBody): Promise { const result = await apiSwapTaggerModel(body); switch (result.status) { case 'ok': // 202 Accepted — don't toast success yet; the swap (including any // first-run model download) is still running in the background. - window.dispatchEvent(new CustomEvent(TAGGER_SWAPPED_EVENT)); + // Progress (and the refreshed active-tagger store) arrives via + // the ``tagger_status`` SSE stream. return result.response.active ? { kind: result.response.active.kind, diff --git a/yadc/webui/src/lib/stores/tagging/activeTagger.ts b/yadc/webui/src/lib/stores/tagging/activeTagger.ts new file mode 100644 index 00000000..ff318680 --- /dev/null +++ b/yadc/webui/src/lib/stores/tagging/activeTagger.ts @@ -0,0 +1,62 @@ +import { get, writable } from 'svelte/store'; +import { fetchActiveTagger } from './api'; +import { taggerStatus } from './taggerStatus'; +import type { ActiveTaggerResponse } from './types'; + +/** Persisted active tagger selection + subprocess liveness + per-tag + * capability. Mirrors ``GET /api/tagger/active`` so components + * subscribe instead of fetching ad hoc. Fed by + * :func:`ensureActiveTaggerLoaded` on mount and re-fetched whenever + * the ``tagger_status`` SSE stream settles (swap done / failed) — + * cross-tab, unlike a window event. */ +export const activeTagger = writable(null); + +/** Module-level single-flight promise so parallel mounters share one + * in-flight fetch. */ +let _loadPromise: Promise | null = null; +let _loaded = false; + +/** Populate ``activeTagger`` from the server. Idempotent — the second + * call after the first resolves is a no-op. Throws on a non-200 so + * the caller can toast; the store keeps its current value on error. */ +export async function ensureActiveTaggerLoaded(): Promise { + if (_loaded) { + const current = get(activeTagger); + if (current !== null) { + return current; + } + } + if (_loadPromise === null) { + _loadPromise = (async () => { + try { + const response = await fetchActiveTagger(); + activeTagger.set(response); + _loaded = true; + return response; + } catch (e) { + _loadPromise = null; + throw e; + } + })(); + } + return _loadPromise; +} + +/** Force-refresh from the server. Used when the subprocess settles + * after a swap (or in tests). */ +export async function refreshActiveTagger(): Promise { + _loadPromise = null; + _loaded = false; + return ensureActiveTaggerLoaded(); +} + +// Re-fetch when the subprocess settles — a swap always ends in +// ``ready`` (new model live) or ``failed`` (rolled back), so the +// capability flags can't go stale after a minutes-long first-run HF +// download. Failures stay silent here (a stale value beats an +// unhandled rejection); the next mount retries. +taggerStatus.subscribe((status) => { + if (status.state === 'ready' || status.state === 'failed') { + void refreshActiveTagger().catch(() => {}); + } +}); diff --git a/yadc/webui/src/lib/stores/tagging/index.ts b/yadc/webui/src/lib/stores/tagging/index.ts index 12b35e69..9a1c9baa 100644 --- a/yadc/webui/src/lib/stores/tagging/index.ts +++ b/yadc/webui/src/lib/stores/tagging/index.ts @@ -1,4 +1,5 @@ export * from './actions'; +export * from './activeTagger'; export * from './api'; export * from './constants'; export * from './display';