nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(autotag): one server-side tagging configuration

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-26 23:23:23 -07:00
parent cfe47d3cb0
commit 2e60ce797e
13 changed files with 667 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/annotation_suggest.py
+16 -23
View File
@@ -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):
py/api/autotag_routes.py
+25 -10
View File
@@ -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"
py/api/settings_routes.py
+84 -1
View File
@@ -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):
py/autotag.py
+14 -27
View File
@@ -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
py/config.py
+115 -16
View File
@@ -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:
tests/browser/workspace_tagging_fixture.py
+11 -2
View File
@@ -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:
tests/js/workspace-image-tagging-state.test.js
+32
View File
@@ -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 });
tests/test_api_annotations.py
+2
View File
@@ -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
tests/test_api_autotag.py
+98
View File
@@ -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)]
)
tests/test_autotag_engine.py
+29 -12
View File
@@ -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)])
tests/test_config.py
+94 -2
View File
@@ -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):
tests/test_settings_sources.py
+126
View File
@@ -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)
web/js/workspace/image-tagging-state.mjs
+21 -9
View File
@@ -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,