From 2e60ce797ee264ca220b2e170fd11aaf9e068d52 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 26 Oct 2026 23:23:23 -0700 Subject: [PATCH] feat(autotag): one server-side tagging configuration --- py/api/annotation_suggest.py | 39 +++--- py/api/autotag_routes.py | 35 +++-- py/api/settings_routes.py | 85 +++++++++++- py/autotag.py | 41 ++---- py/config.py | 131 +++++++++++++++--- tests/browser/workspace_tagging_fixture.py | 13 +- .../js/workspace-image-tagging-state.test.js | 32 +++++ tests/test_api_annotations.py | 2 + tests/test_api_autotag.py | 98 +++++++++++++ tests/test_autotag_engine.py | 41 ++++-- tests/test_config.py | 96 ++++++++++++- tests/test_settings_sources.py | 126 +++++++++++++++++ web/js/workspace/image-tagging-state.mjs | 30 ++-- 13 files changed, 667 insertions(+), 102 deletions(-) diff --git a/py/api/annotation_suggest.py b/py/api/annotation_suggest.py index 563bdc6..2503546 100644 --- a/py/api/annotation_suggest.py +++ b/py/api/annotation_suggest.py @@ -10,12 +10,15 @@ filesystem or the database. import math -WD14_MODELS = frozenset({"wd14-swinv2", "wd14-vit"}) -LLM_MODELS = frozenset({"gguf", "hf"}) -MODEL_TYPES = WD14_MODELS | LLM_MODELS - -MAX_INSTRUCTION_CHARS = 4000 -MAX_INSTRUCTION_BYTES = 16 * 1024 +from ..config import ( # noqa: F401 - re-exported for the suggest route + LLM_MODEL_TYPES as LLM_MODELS, + MAX_INSTRUCTION_BYTES, + MAX_INSTRUCTION_CHARS, + MODEL_TYPES, + WD14_MODEL_TYPES as WD14_MODELS, + validate_template, + validate_threshold, +) SUGGEST_FIELDS = frozenset( {"target", "capture_token", "expected_epoch", "model_type", "options"} @@ -32,27 +35,17 @@ class OptionError(ValueError): def _threshold(name: str, value) -> float: - if type(value) not in (int, float) or not math.isfinite(value): - raise OptionError(f"{name} must be a finite number") - if not 0 <= value <= 1: - raise OptionError(f"{name} must be between 0 and 1") - return float(value) + try: + return validate_threshold(name, value) + except ValueError as error: + raise OptionError(str(error)) from None def _instructions(value) -> str: - if not isinstance(value, str) or not value.strip(): - raise OptionError("instructions must be non-empty text") - if len(value) > MAX_INSTRUCTION_CHARS: - raise OptionError( - f"instructions must be at most {MAX_INSTRUCTION_CHARS} characters" - ) try: - encoded = value.encode("utf-8") - except UnicodeEncodeError: - raise OptionError("instructions must be valid Unicode text") from None - if len(encoded) > MAX_INSTRUCTION_BYTES: - raise OptionError("instructions are too long") - return value + return validate_template(value) + except ValueError as error: + raise OptionError(str(error)) from None def validate_options(model_type, options, *, default_limit: int, max_limit: int): diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 6114bb6..5b77990 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -65,6 +65,17 @@ def path_is_within(candidate, root) -> bool: return child != parent and child.is_relative_to(parent) +def _thresholds_or_saved(general, character): + """The request's WD14 thresholds as floats; one it omits is the saved + default (``autotag.defaults``, OQ4). Raises ValueError/TypeError.""" + from ..config import AutotagConfig + + return ( + AutotagConfig.GENERAL_THRESHOLD if general is None else float(general), + AutotagConfig.CHARACTER_THRESHOLD if character is None else float(character), + ) + + async def _refusal(title): """The single error event of an auto-tag batch refused by another job.""" message = f"{title} already running" @@ -116,12 +127,15 @@ class AutotagRoutesMixin: { "success": True, "models": models_status, - "default_prompt": service.default_prompt, + # The saved tagging defaults (#328, System › Auto-tagging). + "default_prompt": AutotagConfig.TEMPLATE, "model_loaded": service.is_model_loaded(), "loaded_model_type": service.get_loaded_model_type(), - "wd14_general_threshold": service.wd14_general_threshold, - "wd14_character_threshold": service.wd14_character_threshold, + "wd14_general_threshold": AutotagConfig.GENERAL_THRESHOLD, + "wd14_character_threshold": AutotagConfig.CHARACTER_THRESHOLD, "suggestion_limit": AutotagConfig.suggestion_limit(), + "model_type": AutotagConfig.MODEL_TYPE, + "existing": AutotagConfig.EXISTING, } ) @@ -425,9 +439,11 @@ class AutotagRoutesMixin: "character_threshold": params.get("character_threshold"), } try: - for key in ("general_threshold", "character_threshold"): - if opts[key] is not None: - opts[key] = float(opts[key]) + opts["general_threshold"], opts["character_threshold"] = ( + _thresholds_or_saved( + opts["general_threshold"], opts["character_threshold"] + ) + ) except (ValueError, TypeError): return self._bad_request( "general_threshold and character_threshold must be numeric" @@ -607,10 +623,9 @@ class AutotagRoutesMixin: general_threshold = data.get("general_threshold") character_threshold = data.get("character_threshold") try: - if general_threshold is not None: - general_threshold = float(general_threshold) - if character_threshold is not None: - character_threshold = float(character_threshold) + general_threshold, character_threshold = _thresholds_or_saved( + general_threshold, character_threshold + ) except (ValueError, TypeError): return self._bad_request( "general_threshold and character_threshold must be numeric" diff --git a/py/api/settings_routes.py b/py/api/settings_routes.py index 466ba41..c38274c 100644 --- a/py/api/settings_routes.py +++ b/py/api/settings_routes.py @@ -47,11 +47,23 @@ _SETTINGS_FILE_KEYS = { "THUMBNAIL_FORMAT": ("thumbnails", "format"), "THUMBNAIL_QUALITY": ("thumbnails", "quality"), "THUMBNAIL_EDGE": ("thumbnails", "edge"), + # The one tagging configuration (#328), published to AutotagConfig. + "AUTOTAG_MODEL_TYPE": ("autotag", "defaults", "model_type"), + "AUTOTAG_EXISTING": ("autotag", "defaults", "existing"), + "AUTOTAG_GENERAL_THRESHOLD": ("autotag", "defaults", "general_threshold"), + "AUTOTAG_CHARACTER_THRESHOLD": ("autotag", "defaults", "character_threshold"), + "AUTOTAG_TEMPLATE": ("autotag", "defaults", "template"), + "AUTOTAG_SUGGESTION_LIMIT": ("autotag", "defaults", "suggestion_limit"), } # The thumbnail settings a running Convert outdated job fixes (#317). _THUMBNAIL_SETTINGS = ("THUMBNAIL_FORMAT", "THUMBNAIL_QUALITY", "THUMBNAIL_EDGE") # Attributes that live on GalleryConfig; the others on PromptManagerConfig. _GALLERY_SETTINGS = frozenset({"MONITORING_ENABLED"}) +# Attributes that live on AutotagConfig, named without the AUTOTAG_ prefix. +_AUTOTAG_SETTINGS = frozenset( + name for name in _SETTINGS_FILE_KEYS if name.startswith("AUTOTAG_") +) +_AUTOTAG_PREFIX = "AUTOTAG_" SOURCE_CHANGE_JOB_KIND = "source_change" SOURCE_CHANGE_JOB_TITLE = "Source change" @@ -151,6 +163,63 @@ def _set_nested(config, sections, key, value): config[sections[0]] = merged +def _validate_autotag_settings(data): + """The ``autotag_*`` settings as ``{AUTOTAG_ATTR: value}``. + + The first bad value raises ValueError with a message that names the + setting by its label in System › Auto-tagging. + """ + from ..config import ( + EXISTING_POLICIES, + MAX_SUGGESTION_LIMIT, + MODEL_TYPES, + validate_template, + validate_threshold, + ) + + updates = {} + if "autotag_model_type" in data: + value = data["autotag_model_type"] + if not isinstance(value, str) or value not in MODEL_TYPES: + raise ValueError( + "Tagging model must be one of " + ", ".join(sorted(MODEL_TYPES)) + "." + ) + updates["AUTOTAG_MODEL_TYPE"] = value + if "autotag_existing" in data: + value = data["autotag_existing"] + if not isinstance(value, str) or value not in EXISTING_POLICIES: + raise ValueError( + "Existing tags must be skip, merge or replace_after_review." + ) + updates["AUTOTAG_EXISTING"] = value + for key, label in ( + ("autotag_general_threshold", "General-tag confidence"), + ("autotag_character_threshold", "Character-tag confidence"), + ): + if key in data: + try: + value = validate_threshold(label, data[key]) + except ValueError as error: + raise ValueError(f"{error}.") from None + updates[key.upper()] = value + if "autotag_template" in data: + try: + updates["AUTOTAG_TEMPLATE"] = validate_template( + data["autotag_template"], "LLM tag template" + ) + except ValueError as error: + raise ValueError(f"{error}.") from None + if "autotag_suggestion_limit" in data: + value = data["autotag_suggestion_limit"] + if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT: + raise ValueError( + "Suggestions per image must be a whole number from 1 to " + f"{MAX_SUGGESTION_LIMIT}." + ) + updates["AUTOTAG_SUGGESTION_LIMIT"] = value + return updates + + async def _read_json(request): """Parse a JSON object body through the package's capped reader.""" from . import _read_json_body @@ -177,7 +246,7 @@ class SettingsRoutesMixin: async def get_settings(self, request): """Get current settings.""" try: - from ..config import PromptManagerConfig, GalleryConfig + from ..config import AutotagConfig, PromptManagerConfig, GalleryConfig root_paths = [ self._public_path(d) for d in GalleryConfig.MONITORING_DIRECTORIES @@ -203,6 +272,14 @@ class SettingsRoutesMixin: "thumbnail_format": PromptManagerConfig.THUMBNAIL_FORMAT, "thumbnail_quality": PromptManagerConfig.THUMBNAIL_QUALITY, "thumbnail_edge": PromptManagerConfig.THUMBNAIL_EDGE, + "autotag_model_type": AutotagConfig.MODEL_TYPE, + "autotag_existing": AutotagConfig.EXISTING, + "autotag_general_threshold": AutotagConfig.GENERAL_THRESHOLD, + "autotag_character_threshold": ( + AutotagConfig.CHARACTER_THRESHOLD + ), + "autotag_template": AutotagConfig.TEMPLATE, + "autotag_suggestion_limit": AutotagConfig.SUGGESTION_LIMIT, "cpu_count": PromptManagerConfig.max_worker_threads(), "gallery_root_paths": root_paths, "gallery_root_path": root_paths[0] if root_paths else "", @@ -425,7 +502,12 @@ class SettingsRoutesMixin: """ from ..config import GalleryConfig, PromptManagerConfig + from ..config import AutotagConfig + for name, value in updates.items(): + if name in _AUTOTAG_SETTINGS: + setattr(AutotagConfig, name[len(_AUTOTAG_PREFIX) :], value) + continue owner = GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig setattr(owner, name, value) restart_fields = [] @@ -564,6 +646,7 @@ class SettingsRoutesMixin: updates["THUMBNAIL_EDGE"] = thumbnail_settings.validate_edge( data["thumbnail_edge"] ) + updates.update(_validate_autotag_settings(data)) return updates def _parse_gallery_roots(self, data): diff --git a/py/autotag.py b/py/autotag.py index 68c72a5..8e33a03 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -25,6 +25,13 @@ except ImportError: from database.annotations import MAX_IMAGE_TAGS from database.vocabulary import validate_vocabulary_tag_name +# The tag template and thresholds live with the saved tagging defaults +# (#328), so config.py never imports PIL. +try: + from .config import DEFAULT_TAG_TEMPLATE, AutotagConfig +except ImportError: + from py.config import DEFAULT_TAG_TEMPLATE, AutotagConfig + # Try to import logging from utils, fallback to standard logging try: from ..utils.logging_config import get_logger @@ -85,17 +92,7 @@ MODELS = { } # Default prompts -DEFAULT_PROMPT = ( - "Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n" - "Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, " - "etc), age_approx (teen, young_woman), body_attributes, face_features, " - "hairstyle, appearance, clothing, accessories, pose, expression, action, " - "camera_angle, camera_shot, scene_composition (wide_shot, long_shot, " - "dutch_angle), environment, background_elements, lighting, mood.\n" - "Prioritize visual accuracy.\n" - "No sentences, no commentary, no captions. Only tags. Keep length short but " - "descriptive." -) +DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE # Upper bound of one suggestion request: the image tag set cap. MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS @@ -179,10 +176,6 @@ class AutoTagService: self._current_model_type: Optional[str] = None self._custom_prompt: str = DEFAULT_PROMPT - # WD14 threshold defaults - self._wd14_general_threshold: float = 0.35 - self._wd14_character_threshold: float = 0.85 - @property def models_config(self) -> Dict[str, Dict[str, Any]]: """Get the models configuration dictionary.""" @@ -205,19 +198,13 @@ class AutoTagService: @property def wd14_general_threshold(self) -> float: - return self._wd14_general_threshold - - @wd14_general_threshold.setter - def wd14_general_threshold(self, value: float): - self._wd14_general_threshold = max(0.0, min(1.0, value)) + """The saved General-tag confidence (``autotag.defaults``).""" + return AutotagConfig.GENERAL_THRESHOLD @property def wd14_character_threshold(self) -> float: - return self._wd14_character_threshold - - @wd14_character_threshold.setter - def wd14_character_threshold(self, value: float): - self._wd14_character_threshold = max(0.0, min(1.0, value)) + """The saved Character-tag confidence (``autotag.defaults``).""" + return AutotagConfig.CHARACTER_THRESHOLD def get_models_status(self) -> Dict[str, Dict[str, Any]]: """Get availability status for all model types. @@ -860,12 +847,12 @@ class AutoTagService: gen_thresh = ( general_threshold if general_threshold is not None - else self._wd14_general_threshold + else AutotagConfig.GENERAL_THRESHOLD ) char_thresh = ( character_threshold if character_threshold is not None - else self._wd14_character_threshold + else AutotagConfig.CHARACTER_THRESHOLD ) return self._generate_wd14( session, tags_list, image, gen_thresh, char_thresh diff --git a/py/config.py b/py/config.py index 46b4364..b71be4a 100644 --- a/py/config.py +++ b/py/config.py @@ -30,6 +30,7 @@ routes = server_instance.routes # Extension info extension_uri = None # Will be set in __init__.py +import math # noqa: E402 import os # noqa: E402 from typing import Any, Dict, List, Optional, Tuple # noqa: E402 @@ -563,18 +564,111 @@ class IntegrationConfig: DEFAULT_SUGGESTION_LIMIT = 10 MAX_SUGGESTION_LIMIT = 100 # pinned to MAX_IMAGE_TAGS by tests/test_config.py +# The tagging engines (``py/autotag.py`` MODELS) and the options they take. +WD14_MODEL_TYPES = frozenset({"wd14-swinv2", "wd14-vit"}) +LLM_MODEL_TYPES = frozenset({"gguf", "hf"}) +MODEL_TYPES = WD14_MODEL_TYPES | LLM_MODEL_TYPES +MAX_INSTRUCTION_CHARS = 4000 +MAX_INSTRUCTION_BYTES = 16 * 1024 +# What a batch does with a prompt that already has tags (#328). +EXISTING_POLICIES = ("skip", "merge", "replace_after_review") + +# The language-model tag template every page starts from until one is saved. +DEFAULT_TAG_TEMPLATE = ( + "Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n" + "Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, " + "etc), age_approx (teen, young_woman), body_attributes, face_features, " + "hairstyle, appearance, clothing, accessories, pose, expression, action, " + "camera_angle, camera_shot, scene_composition (wide_shot, long_shot, " + "dutch_angle), environment, background_elements, lighting, mood.\n" + "Prioritize visual accuracy.\n" + "No sentences, no commentary, no captions. Only tags. Keep length short but " + "descriptive." +) + + +def validate_threshold(name: str, value: Any) -> float: + """A WD14 confidence threshold in [0, 1], or ValueError naming ``name``.""" + if type(value) not in (int, float) or not math.isfinite(value): + raise ValueError(f"{name} must be a finite number") + if not 0 <= value <= 1: + raise ValueError(f"{name} must be between 0 and 1") + return float(value) + + +def validate_template(value: Any, name: str = "instructions") -> str: + """Non-empty language-model instructions within the character and byte + caps, or ValueError naming ``name``.""" + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{name} must be non-empty text") + if len(value) > MAX_INSTRUCTION_CHARS: + raise ValueError(f"{name} must be at most {MAX_INSTRUCTION_CHARS} characters") + try: + encoded = value.encode("utf-8") + except UnicodeEncodeError: + raise ValueError(f"{name} must be valid Unicode text") from None + if len(encoded) > MAX_INSTRUCTION_BYTES: + raise ValueError(f"{name} {'are' if name == 'instructions' else 'is'} too long") + return value + + +def _valid_model_type(value: Any) -> str: + if not isinstance(value, str) or value not in MODEL_TYPES: + raise ValueError("must be one of " + ", ".join(sorted(MODEL_TYPES))) + return value + + +def _valid_existing(value: Any) -> str: + if not isinstance(value, str) or value not in EXISTING_POLICIES: + raise ValueError("must be one of " + ", ".join(EXISTING_POLICIES)) + return value + + +def _valid_suggestion_limit(value: Any) -> int: + if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT: + raise ValueError(f"must be an integer from 1 to {MAX_SUGGESTION_LIMIT}") + return value + class AutotagConfig: """Persisted tagging defaults under ``autotag.defaults`` in config.json. - ``suggestion_limit`` is how many suggestions Generate returns (the top N - by score), an integer in ``1..MAX_SUGGESTION_LIMIT``. A hand-edited value - outside that range, or of another type, falls back to the default with a - warning rather than failing the load. + One tagging configuration (#328): the model, what a batch does with + prompts already tagged, the two WD14 thresholds, the language-model + template and how many suggestions Generate returns. Each key is checked + alone on load; a hand-edited value that fails its check falls back to + its default with one warning rather than failing the load. """ + MODEL_TYPE = "wd14-swinv2" + EXISTING = "skip" + GENERAL_THRESHOLD = 0.35 + CHARACTER_THRESHOLD = 0.85 + TEMPLATE = DEFAULT_TAG_TEMPLATE SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT + # config.json key -> (attribute, validator, default) + _KEYS = { + "model_type": ("MODEL_TYPE", _valid_model_type, "wd14-swinv2"), + "existing": ("EXISTING", _valid_existing, "skip"), + "general_threshold": ( + "GENERAL_THRESHOLD", + lambda v: validate_threshold("value", v), + 0.35, + ), + "character_threshold": ( + "CHARACTER_THRESHOLD", + lambda v: validate_threshold("value", v), + 0.85, + ), + "template": ("TEMPLATE", validate_template, DEFAULT_TAG_TEMPLATE), + "suggestion_limit": ( + "SUGGESTION_LIMIT", + _valid_suggestion_limit, + DEFAULT_SUGGESTION_LIMIT, + ), + } + @classmethod def suggestion_limit(cls) -> int: return cls.SUGGESTION_LIMIT @@ -587,7 +681,11 @@ class AutotagConfig: @classmethod def get_config(cls) -> Dict[str, Any]: - return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}} + return { + "defaults": { + key: getattr(cls, attr) for key, (attr, _v, _d) in cls._KEYS.items() + } + } @classmethod def update_config(cls, new_config: Any) -> None: @@ -598,17 +696,18 @@ class AutotagConfig: if not isinstance(defaults, dict): config_logger.warning("Ignoring autotag.defaults that is not an object") return - if "suggestion_limit" not in defaults: - return - value = defaults["suggestion_limit"] - if cls.valid_suggestion_limit(value): - cls.SUGGESTION_LIMIT = value - return - config_logger.warning( - "autotag.defaults.suggestion_limit must be an integer from 1 to " - f"{MAX_SUGGESTION_LIMIT}; using {DEFAULT_SUGGESTION_LIMIT}" - ) - cls.SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT + for key, (attr, validate, default) in cls._KEYS.items(): + if key not in defaults: + continue + try: + setattr(cls, attr, validate(defaults[key])) + except ValueError as error: + config_logger.warning( + f"autotag.defaults.{key} {error}; using {default!r}" + if key != "template" + else f"autotag.defaults.template {error}; using the built-in one" + ) + setattr(cls, attr, default) class PromptManagerConfig: diff --git a/tests/browser/workspace_tagging_fixture.py b/tests/browser/workspace_tagging_fixture.py index 6fa3219..621b0a2 100644 --- a/tests/browser/workspace_tagging_fixture.py +++ b/tests/browser/workspace_tagging_fixture.py @@ -182,6 +182,8 @@ def register_tagging_commands( ) -> None: """Register fixed seed, snapshot and fake-model commands; no paths or SQL input.""" + _set_defaults(model, DEFAULTS) # the server starts on the standard ones + @routes.post("/__test__/tagging/reset") async def reset(request: web.Request) -> web.Response: data = await request.json() if request.can_read_body else {} @@ -279,12 +281,19 @@ DEFAULTS = (0.35, 0.85, 10) def _set_defaults(model, values) -> None: + """Set the saved tagging defaults (``AutotagConfig``, #328); a reset also + puts back the saved model, batch policy and template (the fake model's + short instructions, which the concept comparisons show).""" from py.config import AutotagConfig general, character, limit = values - model.wd14_general_threshold = general - model.wd14_character_threshold = character + AutotagConfig.GENERAL_THRESHOLD = general + AutotagConfig.CHARACTER_THRESHOLD = character AutotagConfig.SUGGESTION_LIMIT = limit + if values == DEFAULTS: + AutotagConfig.MODEL_TYPE = "wd14-swinv2" + AutotagConfig.EXISTING = "skip" + AutotagConfig.TEMPLATE = model.default_prompt def _valid_defaults(data: object) -> bool: diff --git a/tests/js/workspace-image-tagging-state.test.js b/tests/js/workspace-image-tagging-state.test.js index 604716f..0540de9 100644 --- a/tests/js/workspace-image-tagging-state.test.js +++ b/tests/js/workspace-image-tagging-state.test.js @@ -197,6 +197,38 @@ test('instructions default to the server prompt', async () => { assert.equal(c.state.options.suggestion_limit, 7, 'the suggestion count survives a model change'); }); +// #328: one tagging configuration. The dialogs start on the model saved in System › Auto-tagging, unless a preset +// (a review session's choice) or the user already picked one; an unknown saved model keeps the first option. +test('dialog starts on the saved model unless a preset or the user chose one', async () => { + const saved = { [`GET ${MODELS}`]: () => response(200, models({ model_type: 'gguf' })) }; + const { c } = await opened(saved); + assert.equal(c.state.model, 'gguf'); + assert.equal(c.state.options.instructions, DEFAULT_PROMPT); + + const { c: preset } = await opened(saved, { preset: { model: 'wd14-vit', options: {} } }); + assert.equal(preset.state.model, 'wd14-vit'); + + const { createImageTaggingController, createTaggingOptionsController } = await load(); + const h = harness(defaultRoutes(saved)); + const picked = createImageTaggingController({ api: h.api, target: TARGET, source: 'images' }); + picked.setModel('hf'); + await picked.open(); + await settle(); + assert.equal(picked.state.model, 'hf', 'the user picked a model before the read landed'); + + const { c: unknown } = await opened({ [`GET ${MODELS}`]: () => response(200, models({ model_type: 'wd14-convnext' })) }); + assert.equal(unknown.state.model, 'wd14-swinv2'); + + const options = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api }); + await options.open(); + assert.equal(options.state.model, 'gguf'); + const chosen = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api }); + chosen.setModel('wd14-vit'); + await chosen.open(); + assert.equal(chosen.state.model, 'wd14-vit'); + for (const controller of [c, preset, picked, unknown, options, chosen]) controller.dispose(); +}); + test('an unreadable or malformed models read falls back to the documented defaults', async () => { for (const route of [() => fail(500, null), () => response(200, models({ default_prompt: 4, suggestion_limit: 0, wd14_general_threshold: 2, wd14_character_threshold: 'x' }))]) { const { c } = await opened({ [`GET ${MODELS}`]: route }); diff --git a/tests/test_api_annotations.py b/tests/test_api_annotations.py index 6476ea7..769afa8 100644 --- a/tests/test_api_annotations.py +++ b/tests/test_api_annotations.py @@ -1145,8 +1145,10 @@ class TestEnvelopes(AnnotationAPITestCase): sorted(await self.json_of(resp)), [ "default_prompt", + "existing", # #328 adds the saved batch policy and model "loaded_model_type", "model_loaded", + "model_type", "models", "success", "suggestion_limit", # B2 adds it (D68); nothing else changes diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index d2a9b15..d5e6baa 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -1288,3 +1288,101 @@ class TestGalleryReviewContract(AutotagAPITestCase): "/prompt_manager/autotag/single", {"path": "link.png", "root_index": 0} ) self.assertEqual(response.status, 403) + + +class TestSavedTaggingDefaults(AutotagAPITestCase): + """#328: /models reports the saved defaults; legacy requests without + thresholds use the saved ones (OQ4).""" + + def setUp(self): + super().setUp() + from py.config import AutotagConfig + + self.config = AutotagConfig + for attr in ( + "MODEL_TYPE", + "EXISTING", + "GENERAL_THRESHOLD", + "CHARACTER_THRESHOLD", + "TEMPLATE", + "SUGGESTION_LIMIT", + ): + self.addCleanup(setattr, AutotagConfig, attr, getattr(AutotagConfig, attr)) + + async def test_models_reports_saved_defaults(self): + from py.config import DEFAULT_TAG_TEMPLATE + + data = await (await self.client.get("/prompt_manager/autotag/models")).json() + self.assertEqual( + {k: data[k] for k in ("model_type", "existing", "default_prompt")}, + { + "model_type": "wd14-swinv2", + "existing": "skip", + "default_prompt": DEFAULT_TAG_TEMPLATE, + }, + ) + self.config.MODEL_TYPE = "gguf" + self.config.EXISTING = "replace_after_review" + self.config.GENERAL_THRESHOLD = 0.4 + self.config.CHARACTER_THRESHOLD = 0.7 + self.config.TEMPLATE = "Saved template" + self.config.SUGGESTION_LIMIT = 15 + data = await (await self.client.get("/prompt_manager/autotag/models")).json() + self.assertEqual( + { + k: data[k] + for k in ( + "model_type", + "existing", + "wd14_general_threshold", + "wd14_character_threshold", + "default_prompt", + "suggestion_limit", + ) + }, + { + "model_type": "gguf", + "existing": "replace_after_review", + "wd14_general_threshold": 0.4, + "wd14_character_threshold": 0.7, + "default_prompt": "Saved template", + "suggestion_limit": 15, + }, + ) + + async def test_single_without_thresholds_uses_saved_ones(self): + image = self._make_image(self.output_dir) + self.config.GENERAL_THRESHOLD = 0.4 + self.config.CHARACTER_THRESHOLD = 0.7 + resp = await self._post_json( + "/prompt_manager/autotag/single", + {"path": str(image), "model_type": "wd14-vit"}, + ) + self.assertEqual(resp.status, 200) + resp = await self._post_json( + "/prompt_manager/autotag/single", + { + "path": str(image), + "model_type": "wd14-vit", + "general_threshold": 0.2, + "character_threshold": 0.3, + }, + ) + self.assertEqual(resp.status, 200) + self.assertEqual( + [call[2:] for call in self.service.generate_calls], + [(0.4, 0.7), (0.2, 0.3)], + ) + + async def test_legacy_batch_without_thresholds_uses_saved_ones(self): + prompt = self._save_prompt("p") + self._link_image(prompt, self._make_image(self.output_dir, "a.png")) + self.config.GENERAL_THRESHOLD = 0.45 + self.config.CHARACTER_THRESHOLD = 0.75 + resp = await self.client.post( + "/prompt_manager/autotag/start?model_type=wd14-vit" + ) + self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete") + self.assertEqual( + [call[2:] for call in self.service.generate_calls], [(0.45, 0.75)] + ) diff --git a/tests/test_autotag_engine.py b/tests/test_autotag_engine.py index c76217a..53e0bca 100644 --- a/tests/test_autotag_engine.py +++ b/tests/test_autotag_engine.py @@ -19,8 +19,15 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from PIL import Image -from py import autotag -from py.autotag import MODELS, AutoTagService +# py.autotag reads its saved defaults from py.config, which imports ComfyUI's +# ``server`` at load: stand in for it when this module runs alone. +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules.setdefault("server", _mock_server) + +from py import autotag # noqa: E402 +from py.autotag import MODELS, AutoTagService # noqa: E402 +from py.config import AutotagConfig # noqa: E402 @contextlib.contextmanager @@ -254,12 +261,19 @@ class TestWd14Thresholding(EngineTestCase): self.assertEqual(list(feeds.keys()), ["input"]) self.assertEqual(tuple(feeds["input"].shape), (1, 8, 8, 3)) - def test_service_threshold_setters_clamp(self): - service = self._loaded([0.0] * 6) - service.wd14_general_threshold = 1.7 - service.wd14_character_threshold = -2 - self.assertEqual(service.wd14_general_threshold, 1.0) - self.assertEqual(service.wd14_character_threshold, 0.0) + def test_service_thresholds_are_the_saved_defaults(self): + # #328: one configuration; the service has no thresholds of its own. + service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5]) + with ( + patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.25), + patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.75), + ): + self.assertEqual( + (service.wd14_general_threshold, service.wd14_character_threshold), + (0.25, 0.75), + ) + tags = service.generate_tags(str(self.image_path)) + self.assertEqual(tags, ["1girl", "hatsune_miku", "long_hair", "smile"]) class TestGgufNormalisation(EngineTestCase): @@ -369,10 +383,13 @@ class TestTagSuggestions(EngineTestCase): def test_thresholds_default_to_the_service_values(self): service = self._wd14() - service.wd14_general_threshold = 0.95 - service.wd14_character_threshold = 0.95 - - suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10) + with ( + patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.95), + patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.95), + ): + suggestions = service.generate_tag_suggestions( + str(self.image_path), limit=10 + ) self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)]) diff --git a/tests/test_config.py b/tests/test_config.py index cf1da2f..5ba7efb 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -413,9 +413,21 @@ class TestAutotagConfig(unittest.TestCase): self.config_class = getattr(config_module, "AutotagConfig", None) self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented") self._orig = PromptManagerConfig.get_config() + self._defaults = { + attr: getattr(self.config_class, attr, None) + for attr in ( + "MODEL_TYPE", + "EXISTING", + "GENERAL_THRESHOLD", + "CHARACTER_THRESHOLD", + "TEMPLATE", + "SUGGESTION_LIMIT", + ) + } def tearDown(self): - self.config_class.SUGGESTION_LIMIT = 10 + for attr, value in self._defaults.items(): + setattr(self.config_class, attr, value) PromptManagerConfig.update_config(self._orig) def _load(self, value): @@ -446,7 +458,16 @@ class TestAutotagConfig(unittest.TestCase): self.assertEqual(self.config_class.suggestion_limit(), 10) self.assertEqual( PromptManagerConfig.get_config()["autotag"], - {"defaults": {"suggestion_limit": 10}}, + { + "defaults": { + "model_type": "wd14-swinv2", + "existing": "skip", + "general_threshold": 0.35, + "character_threshold": 0.85, + "template": config_module.DEFAULT_TAG_TEMPLATE, + "suggestion_limit": 10, + } + }, ) with tempfile.TemporaryDirectory() as tmpdir: path = os.path.join(tmpdir, "config.json") @@ -466,6 +487,77 @@ class TestAutotagConfig(unittest.TestCase): ) self.assertEqual(self.config_class.suggestion_limit(), value) + def test_autotag_defaults_fall_back_per_key_on_load(self): + """#328: each saved tagging default is checked alone on load; a bad + hand-edited value falls back to its default with one warning and + leaves the other keys as saved.""" + good = { + "model_type": "gguf", + "existing": "merge", + "general_threshold": 0.4, + "character_threshold": 0.7, + "template": "Tags only.", + "suggestion_limit": 15, + } + fallback = { + "model_type": "wd14-swinv2", + "existing": "skip", + "general_threshold": 0.35, + "character_threshold": 0.85, + "template": config_module.DEFAULT_TAG_TEMPLATE, + "suggestion_limit": 10, + } + bad_values = { + "model_type": ["wd14-convnext", 3, None], + "existing": ["replace", "", True], + "general_threshold": [1.5, -0.1, float("nan"), "0.4", True], + "character_threshold": [2, float("inf"), None], + # 4000 characters are at most 16,000 UTF-8 bytes, so the 16 KiB + # byte cap can only be met past the character cap. + "template": ["", " ", "x" * 4001, "\U0001f600" * 4097, "\ud800", 7], + } + for key, values in bad_values.items(): + for value in values: + with self.subTest(key=key, value=value): + saved = dict(good, **{key: value}) + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "config.json") + with open(path, "w") as f: + json.dump({"autotag": {"defaults": saved}}, f) + with self.assertLogs( + "prompt_manager.config", level="WARNING" + ) as logs: + PromptManagerConfig.load_from_file(path) + self.assertEqual(len(logs.records), 1, logs.output) + self.assertIn(f"autotag.defaults.{key}", logs.output[0]) + loaded = PromptManagerConfig.get_config()["autotag"]["defaults"] + self.assertEqual(loaded, dict(good, **{key: fallback[key]})) + + def test_tagging_validators_moved_with_their_messages(self): + validate_threshold = config_module.validate_threshold + validate_template = config_module.validate_template + self.assertEqual(validate_threshold("general_threshold", 1), 1.0) + for value, message in ( + (float("nan"), "general_threshold must be a finite number"), + ("0.4", "general_threshold must be a finite number"), + (1.01, "general_threshold must be between 0 and 1"), + ): + with self.assertRaises(ValueError) as caught: + validate_threshold("general_threshold", value) + self.assertEqual(str(caught.exception), message) + self.assertEqual(validate_template("Tags."), "Tags.") + for value, message in ( + (" ", "instructions must be non-empty text"), + ("x" * 4001, "instructions must be at most 4000 characters"), + ("\ud800", "instructions must be valid Unicode text"), + ): + with self.assertRaises(ValueError) as caught: + validate_template(value) + self.assertEqual(str(caught.exception), message) + self.assertEqual( + config_module.MODEL_TYPES, {"gguf", "hf", "wd14-swinv2", "wd14-vit"} + ) + def test_suggestion_limit_setting_validated_on_load(self): for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]): with self.subTest(value=value): diff --git a/tests/test_settings_sources.py b/tests/test_settings_sources.py index b964101..e85549f 100644 --- a/tests/test_settings_sources.py +++ b/tests/test_settings_sources.py @@ -976,3 +976,129 @@ class TestSourceFormError(unittest.TestCase): if __name__ == "__main__": unittest.main() + + +class TestAutotagDefaultsSettings(SourcesTestCase): + """#328: the tagging defaults save through POST /settings like the others.""" + + ATTRS = ( + "MODEL_TYPE", + "EXISTING", + "GENERAL_THRESHOLD", + "CHARACTER_THRESHOLD", + "TEMPLATE", + "SUGGESTION_LIMIT", + ) + SAVED = { + "autotag_model_type": "gguf", + "autotag_existing": "merge", + "autotag_general_threshold": 0.4, + "autotag_character_threshold": 0.8, + "autotag_template": "Tags only.", + "autotag_suggestion_limit": 15, + } + + async def asyncSetUp(self): + await super().asyncSetUp() + from py.config import AutotagConfig + + self.autotag = AutotagConfig + for attr in self.ATTRS: + self.addCleanup(setattr, AutotagConfig, attr, getattr(AutotagConfig, attr)) + + def _runtime(self): + return {attr: getattr(self.autotag, attr) for attr in self.ATTRS} + + async def test_autotag_defaults_save_round_trip(self): + self._write_config({"custom": {"keep": True}}) + revision = (await self._get())["revision"] + status, data = await self._post( + "/prompt_manager/settings", {"expected_revision": revision, **self.SAVED} + ) + self.assertEqual(status, 200, data) + self.assertEqual(data["restart_required_fields"], []) + saved = self._read_config() + self.assertEqual( + saved["autotag"]["defaults"], + {key[len("autotag_") :]: value for key, value in self.SAVED.items()}, + ) + self.assertEqual(saved["custom"], {"keep": True}) + self.assertEqual( + self._runtime(), + { + "MODEL_TYPE": "gguf", + "EXISTING": "merge", + "GENERAL_THRESHOLD": 0.4, + "CHARACTER_THRESHOLD": 0.8, + "TEMPLATE": "Tags only.", + "SUGGESTION_LIMIT": 15, + }, + ) + settings = (await self._get())["settings"] + self.assertEqual({key: settings[key] for key in self.SAVED}, self.SAVED) + + async def test_stale_revision_changes_no_default(self): + revision = (await self._get())["revision"] + self._write_config({"custom": {"edited": "elsewhere"}}) + before = (self._config_bytes(), self._runtime()) + status, data = await self._post( + "/prompt_manager/settings", {"expected_revision": revision, **self.SAVED} + ) + self.assertEqual(status, 409, data) + self.assertEqual(data["code"], "settings-changed") + self.assertEqual((self._config_bytes(), self._runtime()), before) + + async def test_bad_default_message_names_the_setting(self): + cases = ( + ( + "autotag_general_threshold", + 1.5, + "General-tag confidence must be between 0 and 1.", + ), + ( + "autotag_character_threshold", + "0.8", + "Character-tag confidence must be a finite number.", + ), + ( + "autotag_character_threshold", + True, + "Character-tag confidence must be a finite number.", + ), + ( + "autotag_model_type", + "wd14-convnext", + "Tagging model must be one of gguf, hf, wd14-swinv2, wd14-vit.", + ), + ( + "autotag_existing", + "replace", + "Existing tags must be skip, merge or replace_after_review.", + ), + ("autotag_template", " ", "LLM tag template must be non-empty text."), + ( + "autotag_template", + "x" * 4001, + "LLM tag template must be at most 4000 characters.", + ), + ( + "autotag_suggestion_limit", + 0, + "Suggestions per image must be a whole number from 1 to 100.", + ), + ( + "autotag_suggestion_limit", + 2.5, + "Suggestions per image must be a whole number from 1 to 100.", + ), + ) + self._write_config({"custom": {"keep": True}}) + before = (self._config_bytes(), self._runtime()) + for key, value, message in cases: + with self.subTest(key=key, value=value): + status, data = await self._post( + "/prompt_manager/settings", {**self.SAVED, key: value} + ) + self.assertEqual(status, 400, data) + self.assertEqual(data["error"], message) + self.assertEqual((self._config_bytes(), self._runtime()), before) diff --git a/web/js/workspace/image-tagging-state.mjs b/web/js/workspace/image-tagging-state.mjs index e85b959..b835f1f 100644 --- a/web/js/workspace/image-tagging-state.mjs +++ b/web/js/workspace/image-tagging-state.mjs @@ -179,14 +179,15 @@ function validateReceipt(value) { } /** - * Thresholds, instructions and suggestion count from the server, else the documented defaults, and the models' - * download status (`models`, null when unknown). + * Thresholds, instructions and suggestion count from the server, else the documented defaults, the models' + * download status (`models`, null when unknown) and the saved tagging model (`model`, null when unknown; #328). */ async function readDefaults(api, signal) { try { const value = await api.get(MODELS_URL, {signal, cache: 'no-store'}); - if (!object(value)) return {defaults: DEFAULTS, models: null}; + if (!object(value)) return {defaults: DEFAULTS, models: null, model: null}; return { + model: typeof value.model_type === 'string' && modelOf(value.model_type) ? value.model_type : null, defaults: { general_threshold: unit(value.wd14_general_threshold) ? value.wd14_general_threshold : DEFAULTS.general_threshold, character_threshold: unit(value.wd14_character_threshold) ? value.wd14_character_threshold : DEFAULTS.character_threshold, @@ -197,7 +198,7 @@ async function readDefaults(api, signal) { }; } catch { // The defaults keep the dialog usable; Generate reports a real failure itself. - return {defaults: DEFAULTS, models: null}; + return {defaults: DEFAULTS, models: null, model: null}; } } @@ -263,6 +264,7 @@ export function createImageTaggingController({api, target, source = 'images', on requireValue(SOURCES.has(source), 'Unknown tagging source'); requireValue(typeof onLateSave === 'function', 'onLateSave must be a function'); const chosen = validatePreset(preset); + let modelPicked = false; // the user chose a model: the saved one no longer applies let state = freeze(derive(initialState(source, chosen))); const listeners = new Set(); @@ -336,8 +338,11 @@ export function createImageTaggingController({api, target, source = 'images', on } capture = {token: value.capture.token, sha256: value.capture.sha256}; epoch = value.library_epoch; - ({defaults} = await defaultsRead); + const read = await defaultsRead; + ({defaults} = read); if (closed || serial !== openSerial) return; + // The saved model (System › Auto-tagging) unless a preset or the user already chose one (#328). + const savedModel = !chosen && !modelPicked && read.model ? {model: read.model} : {}; observedOnce = true; const {width, height, media_type: mediaType, preview_url: previewUrl, generate} = value.capture; // A kept draft stays as typed unless another writer changed the set meanwhile: then the set leads it (I2). @@ -350,7 +355,7 @@ export function createImageTaggingController({api, target, source = 'images', on savedTags: value.tags.slice(), tagRevision: value.tag_revision, manual: draft === null ? value.tags.join(', ') : changed ? mergeCurrent(value.tags, draft) : draft, draftRevision: state.draftRevision + 1, - options: {...defaults, ...chosen?.options}, optionsRevision: state.optionsRevision + 1, + ...savedModel, options: {...defaults, ...chosen?.options}, optionsRevision: state.optionsRevision + 1, generateSupported: generate.supported, generateReason: generate.reason, statusText: changed ? `${COPY.changedElsewhere} ${savedNow(value.tags)}${reason ? ` ${reason}` : ''}` : reason, }); @@ -362,7 +367,9 @@ export function createImageTaggingController({api, target, source = 'images', on function setModel(id) { requireValue(modelOf(id), 'Unknown tagging model'); - if (closed || id === state.model) return; + if (closed) return; + modelPicked = true; + if (id === state.model) return; stopGenerating(); const phase = state.phase === 'generating' || state.phase === 'generate-error' ? 'ready' : state.phase; set({ @@ -535,6 +542,7 @@ export function createTaggingOptionsController({api} = {}) { let defaults = DEFAULTS; let models = null; let abort = null; + let modelPicked = false; const deriveOptions = value => ({...value, canGenerate: false, unavailable: modelUnavailable(models, value.model), canStart: value.phase === 'ready' && !modelUnavailable(models, value.model)}); let state = freeze(deriveOptions({phase: 'loading', model: MODEL_OPTIONS[0].id, options: {...DEFAULTS}, optionsRevision: 0})); @@ -550,11 +558,15 @@ export function createTaggingOptionsController({api} = {}) { const read = await readDefaults(api, abort.signal); if (closed) return; ({defaults, models} = read); - set({phase: 'ready', options: {...defaults}, optionsRevision: state.optionsRevision + 1}); + // The saved model (System › Auto-tagging) unless the user already chose one (#328). + const savedModel = !modelPicked && read.model ? {model: read.model} : {}; + set({phase: 'ready', ...savedModel, options: {...defaults}, optionsRevision: state.optionsRevision + 1}); } function setModel(id) { requireValue(modelOf(id), 'Unknown tagging model'); - if (closed || id === state.model) return; + if (closed) return; + modelPicked = true; + if (id === state.model) return; // As in the Tag image dialog: the model's fields start from the defaults again; the count is kept (D91). set({model: id, optionsRevision: state.optionsRevision + 1, options: {...state.options, general_threshold: defaults.general_threshold,