From 8b2b0910986e28a40c897cab5cec2e35259f942e Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 26 Oct 2026 23:47:26 -0700 Subject: [PATCH] fix(autotag): server-free tagging defaults, cancel before load, and named legacy refusals --- py/api/annotation_suggest.py | 2 +- py/api/autotag_routes.py | 61 ++++++++------ py/autotag.py | 22 +++-- py/config.py | 83 +++++++------------ py/tagging_defaults.py | 55 ++++++++++++ .../js/workspace-image-tagging-state.test.js | 14 ++++ tests/test_api_autotag.py | 59 ++++++++++++- tests/test_config.py | 13 ++- 8 files changed, 222 insertions(+), 87 deletions(-) create mode 100644 py/tagging_defaults.py diff --git a/py/api/annotation_suggest.py b/py/api/annotation_suggest.py index 2503546..49e67a4 100644 --- a/py/api/annotation_suggest.py +++ b/py/api/annotation_suggest.py @@ -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, diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 8749339..45d4a51 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -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).""" diff --git a/py/autotag.py b/py/autotag.py index 821adca..9a6d55b 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -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 diff --git a/py/config.py b/py/config.py index b71be4a..2a14023 100644 --- a/py/config.py +++ b/py/config.py @@ -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: diff --git a/py/tagging_defaults.py b/py/tagging_defaults.py new file mode 100644 index 0000000..dedaf9b --- /dev/null +++ b/py/tagging_defaults.py @@ -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 diff --git a/tests/js/workspace-image-tagging-state.test.js b/tests/js/workspace-image-tagging-state.test.js index 0540de9..5607c0f 100644 --- a/tests/js/workspace-image-tagging-state.test.js +++ b/tests/js/workspace-image-tagging-state.test.js @@ -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'); diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 01fac1d..444594c 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -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()) diff --git a/tests/test_config.py b/tests/test_config.py index 5ba7efb..7a9e8ba 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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):