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

fix(autotag): server-free tagging defaults, cancel before load, and named legacy refusals

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-26 23:47:26 -07:00
parent ecc3189649
commit 8b2b091098
8 changed files with 222 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/annotation_suggest.py
+1 -1
View File
@@ -10,7 +10,7 @@ filesystem or the database.
import math
from ..config import ( # noqa: F401 - re-exported for the suggest route
from ..tagging_defaults import ( # noqa: F401 - re-exported for the route
LLM_MODEL_TYPES as LLM_MODELS,
MAX_INSTRUCTION_BYTES,
MAX_INSTRUCTION_CHARS,
py/api/autotag_routes.py
+37 -24
View File
@@ -507,7 +507,7 @@ class AutotagRoutesMixin:
``custom_prompt`` instead), ``cancelled`` (checked before each image:
a cancelled run returns without a ``complete`` frame) and
``on_item(index, total, group, image_path, detail)`` (one outcome per
image, ``BATCH_DETAILS``).
image, ``SKIP_DETAILS`` for the skips).
"""
model_type = opts["model_type"]
keep_in_memory = opts["keep_in_memory"]
@@ -518,6 +518,8 @@ class AutotagRoutesMixin:
yield {"type": "error", "message": f"Model {model_type} not downloaded"}
return
if cancelled():
return
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
run = self._gated_runner(opts.get("gate_owner"))
try:
@@ -725,7 +727,9 @@ class AutotagRoutesMixin:
# never runs beside a scan, restore or the workspace tagging job.
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
if busy is not None:
if busy.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload
# A second legacy batch keeps its JSON 409; the workspace's Tag
# output images job (cancellable) is named in a stream frame.
if busy.kind == AUTOTAG_BATCH_JOB_KIND and not busy.cancellable:
return web.json_response(
{"success": False, "error": _BUSY_ERROR}, status=409
)
@@ -971,6 +975,18 @@ class AutotagRoutesMixin:
ok, _backup = await self._safety_backup_step(job, self.db, backup_on)
if not ok:
return
try:
await self._tag_output_images(service, opts, job)
except Exception:
self.logger.exception("Tag output images failed")
job.update(status_line=_INTERNAL_ERROR)
job.log_line(f"Tagging failed · {_INTERNAL_ERROR}")
job.finish("failed")
finally:
if backup_on:
await self._prune_backup_pool(job, self.db)
async def _tag_output_images(self, service, opts, job):
roots = _scan_roots(self._get_all_output_dirs())
job.log_line(
f"Tagging started · {job.scope['label']} · "
@@ -999,28 +1015,25 @@ class AutotagRoutesMixin:
job.publish({"type": "progress", "progress": job.progress})
final = None
try:
async for event in self._autotag_batch_events(
service,
{**opts, "cancelled": lambda: job.cancelled, "on_item": on_item},
):
if event.get("type") in ("complete", "error"):
final = event
if final is None and job.cancelled:
job.update(status_line=BATCH_CANCELLED)
job.log_line(BATCH_CANCELLED)
job.finish("cancelled")
elif final is not None and final["type"] == "error":
job.update(status_line=final["message"])
job.log_line(f"Tagging failed · {final['message']}")
job.finish("failed")
else:
done = f"Tagged {_plural(job.counters.get('tagged', 0), 'prompt')}"
job.update(status_line=done, progress=100)
job.log_line(done)
job.finish("finished")
finally:
await self._prune_backup_pool(job, self.db)
async for event in self._autotag_batch_events(
service,
{**opts, "cancelled": lambda: job.cancelled, "on_item": on_item},
):
if event.get("type") in ("complete", "error"):
final = event
if final is not None and final["type"] == "error":
job.update(status_line=final["message"])
job.log_line(f"Tagging failed · {final['message']}")
job.finish("failed")
elif job.cancelled: # a cancel wins over a run that then completed
job.update(status_line=BATCH_CANCELLED)
job.log_line(BATCH_CANCELLED)
job.finish("cancelled")
else:
done = f"Tagged {_plural(job.counters.get('tagged', 0), 'prompt')}"
job.update(status_line=done, progress=100)
job.log_line(done)
job.finish("finished")
def _gated_runner(self, owner):
"""Executor calls under the model gate *owner* holds (plain when None)."""
py/autotag.py
+16 -6
View File
@@ -28,9 +28,19 @@ except ImportError:
# 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
from .tagging_defaults import DEFAULT_TAG_TEMPLATE
except ImportError:
from py.config import DEFAULT_TAG_TEMPLATE, AutotagConfig
from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
def _autotag_config():
"""The saved defaults; imported late, as config needs ComfyUI's server."""
try:
from .config import AutotagConfig
except ImportError:
from py.config import AutotagConfig
return AutotagConfig
# Try to import logging from utils, fallback to standard logging
try:
@@ -204,12 +214,12 @@ class AutoTagService:
@property
def wd14_general_threshold(self) -> float:
"""The saved General-tag confidence (``autotag.defaults``)."""
return AutotagConfig.GENERAL_THRESHOLD
return _autotag_config().GENERAL_THRESHOLD
@property
def wd14_character_threshold(self) -> float:
"""The saved Character-tag confidence (``autotag.defaults``)."""
return AutotagConfig.CHARACTER_THRESHOLD
return _autotag_config().CHARACTER_THRESHOLD
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
"""Get availability status for all model types.
@@ -852,12 +862,12 @@ class AutoTagService:
gen_thresh = (
general_threshold
if general_threshold is not None
else AutotagConfig.GENERAL_THRESHOLD
else _autotag_config().GENERAL_THRESHOLD
)
char_thresh = (
character_threshold
if character_threshold is not None
else AutotagConfig.CHARACTER_THRESHOLD
else _autotag_config().CHARACTER_THRESHOLD
)
return self._generate_wd14(
session, tags_list, image, gen_thresh, char_thresh
py/config.py
+32 -51
View File
@@ -30,7 +30,6 @@ 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
@@ -564,52 +563,31 @@ 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
# The tagging engines, options, template and validators (server-free, #328).
try:
from .tagging_defaults import ( # noqa: F401 - re-exported
DEFAULT_TAG_TEMPLATE,
EXISTING_POLICIES,
LLM_MODEL_TYPES,
MAX_INSTRUCTION_BYTES,
MAX_INSTRUCTION_CHARS,
MODEL_TYPES,
WD14_MODEL_TYPES,
validate_template,
validate_threshold,
)
except ImportError:
from py.tagging_defaults import ( # noqa: F401 - re-exported
DEFAULT_TAG_TEMPLATE,
EXISTING_POLICIES,
LLM_MODEL_TYPES,
MAX_INSTRUCTION_BYTES,
MAX_INSTRUCTION_CHARS,
MODEL_TYPES,
WD14_MODEL_TYPES,
validate_template,
validate_threshold,
)
def _valid_model_type(value: Any) -> str:
@@ -681,11 +659,14 @@ class AutotagConfig:
@classmethod
def get_config(cls) -> Dict[str, Any]:
return {
"defaults": {
key: getattr(cls, attr) for key, (attr, _v, _d) in cls._KEYS.items()
}
"""The saved defaults; the template only when it is not the built-in
one, so a later built-in template still reaches who never set one."""
defaults = {
key: getattr(cls, attr) for key, (attr, _v, _d) in cls._KEYS.items()
}
if defaults["template"] == DEFAULT_TAG_TEMPLATE:
del defaults["template"]
return {"defaults": defaults}
@classmethod
def update_config(cls, new_config: Any) -> None:
py/tagging_defaults.py
+55
View File
@@ -0,0 +1,55 @@
"""The tagging configuration's constants and validators (#328).
Server-free on purpose: ``py.api`` must import without ComfyUI's ``server``
module, which ``py/config.py`` needs; config re-exports all of this.
"""
import math
from typing import Any
# 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
tests/js/workspace-image-tagging-state.test.js
+14
View File
@@ -216,6 +216,20 @@ test('dialog starts on the saved model unless a preset or the user chose one', a
await settle();
assert.equal(picked.state.model, 'hf', 'the user picked a model before the read landed');
// A pick made while the models read is still in flight wins over the saved model that lands after it.
const gate = deferred();
const late = harness(defaultRoutes({ [`GET ${MODELS}`]: () => gate.promise }));
const racing = createImageTaggingController({ api: late.api, target: TARGET, source: 'images' });
const opening = racing.open();
await settle();
assert.equal(racing.state.phase, 'observing');
racing.setModel('wd14-vit');
gate.resolve(response(200, models({ model_type: 'gguf' })));
await opening;
await settle();
assert.equal(racing.state.model, 'wd14-vit');
racing.dispose();
const { c: unknown } = await opened({ [`GET ${MODELS}`]: () => response(200, models({ model_type: 'wd14-convnext' })) });
assert.equal(unknown.state.model, 'wd14-swinv2');
tests/test_api_autotag.py
+56 -3
View File
@@ -1885,8 +1885,10 @@ class TestTagOutputImagesPreview(BatchJobTestCase):
)
before = snapshot()
resp = await self._preview()
self.assertEqual(resp.status, 200, await resp.text())
with patch("PIL.Image.open") as image_open:
resp = await self._preview()
self.assertEqual(resp.status, 200, await resp.text())
image_open.assert_not_called() # opens no image file
self.assertEqual(snapshot(), before)
self.assertEqual(self.service.generate_calls, [])
self.assertEqual(self.service.load_calls, [])
@@ -2117,7 +2119,19 @@ class TestTagOutputImagesJob(BatchJobTestCase):
await asyncio.get_running_loop().run_in_executor(None, started.wait, 30)
self.assertTrue(model_gate.holder().startswith("autotag-job:"))
legacy = await self.client.post("/prompt_manager/autotag/start")
self.assertEqual(legacy.status, 409)
# Both legacy pages read a stream: the refusal is one error frame
# that names the running job.
self.assertEqual(legacy.status, 200)
self.assertEqual(
self._sse_events(await legacy.text()),
[
{
"type": "error",
"status": "error",
"message": "Tag output images already running",
}
],
)
unload = await self.client.post("/prompt_manager/autotag/unload")
self.assertEqual((await unload.json())["operation"], "Tag output images")
finally:
@@ -2298,3 +2312,42 @@ class TestLegacyBatchFrames(BatchJobTestCase):
},
],
)
class TestTagOutputImagesEdges(BatchJobTestCase):
async def test_cancel_during_backup_loads_no_model(self):
from py.config import PromptManagerConfig
self._seed()
PromptManagerConfig.SAFETY_BACKUP = True
started, proceed = threading.Event(), threading.Event()
def slow_backup(model, label="backup"):
started.set()
proceed.wait(30)
return "prompts.db.backup_x"
with patch("py.api.maintenance.take_safety_backup", side_effect=slow_backup):
resp = await self._start_job(existing="merge")
job = self.api._jobs.get((await resp.json())["job"]["id"])
await asyncio.get_running_loop().run_in_executor(None, started.wait, 30)
await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel")
proceed.set()
await asyncio.wait_for(asyncio.shield(job.task), 30)
self.assertEqual(job.state, "cancelled")
self.assertEqual(self.service.load_calls, [])
self.assertEqual(self.service.generate_calls, [])
async def test_unexpected_error_fails_the_job_with_a_reason(self):
self._seed()
with patch.object(
self.api.db, "autotag_batch_candidates", side_effect=[[], RuntimeError("x")]
):
job = await self._run_job(existing="merge")
self.assertEqual(job.state, "failed")
self.assertEqual(
job.status_line,
"An internal error occurred. Check server logs for details.",
)
self.assertIsNone(model_gate.holder())
tests/test_config.py
+11 -2
View File
@@ -464,11 +464,17 @@ class TestAutotagConfig(unittest.TestCase):
"existing": "skip",
"general_threshold": 0.35,
"character_threshold": 0.85,
"template": config_module.DEFAULT_TAG_TEMPLATE,
"suggestion_limit": 10,
}
},
)
# The built-in template is not written; a saved one is.
self.config_class.TEMPLATE = "Saved."
self.assertEqual(
PromptManagerConfig.get_config()["autotag"]["defaults"]["template"],
"Saved.",
)
self.config_class.TEMPLATE = config_module.DEFAULT_TAG_TEMPLATE
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "config.json")
PromptManagerConfig.update_config(
@@ -530,7 +536,10 @@ class TestAutotagConfig(unittest.TestCase):
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"]
loaded = {
k: getattr(self.config_class, attr)
for k, (attr, _v, _d) in self.config_class._KEYS.items()
}
self.assertEqual(loaded, dict(good, **{key: fallback[key]}))
def test_tagging_validators_moved_with_their_messages(self):