From 04f2c921cfc51b6d16d60539fcb17b6d578b97d2 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 26 Oct 2026 23:00:48 -0700 Subject: [PATCH] feat(autotag): model status without paths, and a busy answer that names the operation --- py/api/autotag_routes.py | 87 +++++++++++++++++++-- py/api/jobs.py | 2 + py/autotag.py | 7 +- tests/browser/server.py | 5 +- tests/test_api_autotag.py | 156 +++++++++++++++++++++++++++++++++++++- 5 files changed, 246 insertions(+), 11 deletions(-) diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 6114bb6..45f4f68 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -3,6 +3,7 @@ import asyncio import json import os +import re import threading from pathlib import Path from typing import Optional @@ -22,14 +23,83 @@ MODEL_BUSY_ERROR = ( ) +# Who holds the model gate, as a 409 model-busy names it (``operation``). An +# owner is ``":"`` (``model_gate.new_owner``); a download names +# its own job, whose title says which model. +MODEL_DOWNLOAD_OWNER_KIND = "model-download" +OWNER_LABELS = { + "autotag-job": "Tag output images", + "legacy-batch": "Tag output images (legacy page)", + "legacy-single": "Tagging one image (legacy page)", + "legacy-unload": "Unloading the tagging model", + "image-tag-job": "Auto-tag selected images", + "image-annotations": "Generate suggestions", +} +ANOTHER_OPERATION = "Another tagging operation" +# The running download's gate owner and job title. Only one download runs at a +# time and each needs the gate, so a new start replaces the entry. +_download_titles = {} + + +def _holder_label(holder: Optional[str]) -> str: + if not holder: + return ANOTHER_OPERATION + kind = holder.split(":", 1)[0] + if kind == MODEL_DOWNLOAD_OWNER_KIND: + return _download_titles.get(holder, ANOTHER_OPERATION) + return OWNER_LABELS.get(kind, ANOTHER_OPERATION) + + def model_busy_response() -> web.Response: - """``409 model-busy``: another operation holds the shared model gate.""" + """``409 model-busy``: another operation holds the shared model gate. + + ``operation`` names the holder (``OWNER_LABELS``), or "Another tagging + operation" when the gate was freed meanwhile or the kind is unknown. + """ return web.json_response( - {"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR}, + { + "success": False, + "code": "model-busy", + "error": MODEL_BUSY_ERROR, + "operation": _holder_label(model_gate.holder()), + }, status=409, ) +_SIZE = re.compile(r"^(~?[0-9.]+)\s*([KMGT]?B)$") + + +def size_label(size: str) -> str: + """The engine table's size with a space before the unit (``~400 MB``).""" + match = _SIZE.match(str(size or "").strip()) + return f"{match.group(1)} {match.group(2)}" if match else str(size or "") + + +def public_models(status: dict, loaded_type: Optional[str]) -> dict: + """Each engine model's public entry: names and sizes from ``MODELS``, the + flags from ``status``; never a path (``model_path`` is dropped).""" + from ..autotag import MODELS + + models = {} + for model_type, config in MODELS.items(): + flags = status.get(model_type) or {} + entry = { + "name": config["name"], + "label": config["label"], + "description": config["description"], + "size": config["size"], + "size_label": size_label(config["size"]), + "downloaded": bool(flags.get("downloaded")), + "loaded": loaded_type == model_type, + } + for key in ("model_exists", "mmproj_exists"): + if key in flags: + entry[key] = bool(flags[key]) + models[model_type] = entry + return models + + # Only these may be handed to the tagging engine or listed by the scanner. _MEDIA_EXTENSIONS = frozenset( {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"} @@ -110,15 +180,18 @@ class AutotagRoutesMixin: from ..config import AutotagConfig service = get_autotag_service() - models_status = service.get_models_status() + # Off the loop: the HuggingFace check walks the whole HF cache. + status = await self._run_in_executor(service.get_models_status) + loaded = service.is_model_loaded() + loaded_type = service.get_loaded_model_type() if loaded else None return web.json_response( { "success": True, - "models": models_status, + "models": public_models(status, loaded_type), "default_prompt": service.default_prompt, - "model_loaded": service.is_model_loaded(), - "loaded_model_type": service.get_loaded_model_type(), + "model_loaded": loaded, + "loaded_model_type": loaded_type, "wd14_general_threshold": service.wd14_general_threshold, "wd14_character_threshold": service.wd14_character_threshold, "suggestion_limit": AutotagConfig.suggestion_limit(), @@ -467,7 +540,7 @@ class AutotagRoutesMixin: return await self._stream_sse(request, _refusal(busy.title)) async with self._jobs.hold( AUTOTAG_BATCH_JOB_KIND, - "Auto-tag images (legacy page)", + OWNER_LABELS["legacy-batch"], admission="records", ) as job: response = await self._start_autotag_batch(request, opts, events) diff --git a/py/api/jobs.py b/py/api/jobs.py index 9e5c463..a385f3b 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -109,6 +109,8 @@ DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup" THUMBNAILS_JOB_KIND = "thumbnails" LORA_IMPORT_JOB_KIND = "lora_import" AUTOTAG_BATCH_JOB_KIND = "autotag" +# A tagging model download (#328): read admission, not cancellable. +MODEL_DOWNLOAD_JOB_KIND = "model_download" # Duplicate image files (#249 System S6, py/api/duplicate_files.py): the # find is a cancellable ``read`` job; the reviewed removal is ``records``. DUPLICATES_JOB_KIND = "duplicates" diff --git a/py/autotag.py b/py/autotag.py index 68c72a5..4d6d4da 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -45,10 +45,12 @@ except ImportError: return logger -# Model configurations - same as standalone_tagger.py +# Model configurations - same as standalone_tagger.py. ``label`` is the short +# name System, Activity and a 409 model-busy show (#328). MODELS = { "gguf": { "name": "GGUF (Recommended)", + "label": "GGUF JoyCaption", "description": "Quantized model, ~4.5GB, faster loading", "size": "~4.5GB", "repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF", @@ -59,6 +61,7 @@ MODELS = { }, "hf": { "name": "HuggingFace", + "label": "HuggingFace JoyCaption", "description": "Full precision model, ~16GB, higher quality", "size": "~16GB", "repo": "fancyfeast/llama-joycaption-beta-one-hf-llava", @@ -66,6 +69,7 @@ MODELS = { }, "wd14-swinv2": { "name": "WD14 SwinV2", + "label": "WD14 SwinV2", "description": "Fast classifier tagger, ~400MB, no prompt needed", "size": "~400MB", "repo": "SmilingWolf/wd-swinv2-tagger-v3", @@ -75,6 +79,7 @@ MODELS = { }, "wd14-vit": { "name": "WD14 ViT", + "label": "WD14 ViT", "description": "Fast classifier tagger, ~400MB, ViT architecture", "size": "~400MB", "repo": "SmilingWolf/wd-vit-tagger-v3", diff --git a/tests/browser/server.py b/tests/browser/server.py index abccbfd..707634d 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -72,6 +72,8 @@ class TestModel: # Scored names the concept comparisons answer instead of the fixed list # (``/__test__/tagging/suggestions``); None keeps the fixed list. suggestions: list | None = None + # The engine's model table (``py/autotag.py`` MODELS), set at startup. + models_config: dict = {} def get_models_status(self) -> dict: return { @@ -79,7 +81,7 @@ class TestModel: "downloaded": name not in self.unavailable, "loaded": self.loaded == name, } - for name in ("gguf", "hf", "wd14-vit", "wd14-convnext", "wd14-swinv2") + for name in self.models_config } def is_model_loaded(self) -> bool: @@ -207,6 +209,7 @@ def create_app(root: Path) -> web.Application: from py import autotag from utils.hashing import generate_prompt_hash + TestModel.models_config = autotag.MODELS model = TestModel() autotag.get_autotag_service = lambda: model api = PromptManagerAPI() diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index d2a9b15..5132d8b 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -42,6 +42,10 @@ except ImportError: model_gate = None +# The engine's four model types (py/autotag.py MODELS), in its own order. +MODEL_TYPES = ("gguf", "hf", "wd14-swinv2", "wd14-vit") + + class FakeAutotagService: """Stand-in for AutoTagService that records calls and returns fixed tags.""" @@ -56,15 +60,24 @@ class FakeAutotagService: self.default_prompt = "default prompt" self.wd14_general_threshold = 0.35 self.wd14_character_threshold = 0.85 - self.models_config = {"gguf": {}, "wd14-vit": {}} + self.models_config = {name: {} for name in MODEL_TYPES} self._loaded_type = None self.generate_error = None self.download_result = True self.load_error = None + # Call counters and the thread each status read ran on. + self.status_calls = 0 + self.status_threads = [] + self.status_paths = {} def get_models_status(self): + self.status_calls += 1 + self.status_threads.append(threading.get_ident()) return { - name: {"downloaded": self.downloaded, "model_path": None} + name: { + "downloaded": self.downloaded, + "model_path": self.status_paths.get(name), + } for name in self.models_config } @@ -1288,3 +1301,142 @@ class TestGalleryReviewContract(AutotagAPITestCase): "/prompt_manager/autotag/single", {"path": "link.png", "root_index": 0} ) self.assertEqual(response.status, 403) + + +def _strings(value): + """Every string inside a JSON value (keys included).""" + if isinstance(value, str): + yield value + elif isinstance(value, dict): + for key, item in value.items(): + yield key + yield from _strings(item) + elif isinstance(value, list): + for item in value: + yield from _strings(item) + + +def _absolute(text): + import ntpath + import posixpath + import re + + return ( + posixpath.isabs(text) + or ntpath.isabs(text) + or re.match(r"^[A-Za-z]:[\\/]", text) is not None + ) + + +class TestModelStatus(AutotagAPITestCase): + """GET /autotag/models: labels, sizes and the loaded flag, never a path.""" + + async def _models(self): + resp = await self.client.get("/prompt_manager/autotag/models") + self.assertEqual(resp.status, 200) + return await resp.json() + + async def test_models_status_has_no_absolute_paths(self): + self.service.status_paths = { + "wd14-vit": "/srv/models/WD14-ViT/model.onnx", + "gguf": "C:\\models\\GGUF\\joycaption.gguf", + "hf": "D:/hf-cache/models--fancyfeast", + } + data = await self._models() + for model in data["models"].values(): + self.assertNotIn("model_path", model) + leaked = [text for text in _strings(data) if _absolute(text)] + self.assertEqual(leaked, []) + + async def test_models_status_adds_label_size_label_and_loaded(self): + self.service._loaded_type = "wd14-vit" + data = await self._models() + models = data["models"] + self.assertEqual(list(models), list(MODEL_TYPES)) + self.assertEqual( + {name: (m["label"], m["size_label"]) for name, m in models.items()}, + { + "gguf": ("GGUF JoyCaption", "~4.5 GB"), + "hf": ("HuggingFace JoyCaption", "~16 GB"), + "wd14-swinv2": ("WD14 SwinV2", "~400 MB"), + "wd14-vit": ("WD14 ViT", "~400 MB"), + }, + ) + self.assertEqual( + {name: m["loaded"] for name, m in models.items()}, + {"gguf": False, "hf": False, "wd14-swinv2": False, "wd14-vit": True}, + ) + # Kept fields legacy pages and the workspace read. + for model in models.values(): + self.assertTrue({"name", "description", "size", "downloaded"} <= set(model)) + self.assertTrue(data["model_loaded"]) + self.assertEqual(data["loaded_model_type"], "wd14-vit") + + async def test_opening_models_never_loads_or_downloads(self): + await self._models() + self.assertEqual(self.service.status_calls, 1) + self.assertEqual(self.service.load_calls, []) + self.assertEqual(self.service.download_calls, []) + self.assertEqual(self.service.unload_calls, 0) + self.assertEqual(self.service.generate_calls, []) + + async def test_models_status_reads_off_the_event_loop(self): + await self._models() + self.assertEqual(len(self.service.status_threads), 1) + self.assertNotEqual(self.service.status_threads[0], threading.get_ident()) + + +class TestModelBusyNamesTheOperation(AutotagAPITestCase): + """A 409 model-busy names who holds the gate (``operation``).""" + + LABELS = { + "autotag-job": "Tag output images", + "legacy-batch": "Tag output images (legacy page)", + "legacy-single": "Tagging one image (legacy page)", + "legacy-unload": "Unloading the tagging model", + "image-tag-job": "Auto-tag selected images", + "image-annotations": "Generate suggestions", + } + + async def test_model_busy_names_the_operation(self): + self.service._loaded_type = "wd14-vit" + for kind, label in self.LABELS.items(): + owner = model_gate.new_owner(kind) + self.assertTrue(model_gate.try_acquire(owner)) + try: + resp = await self.client.post("/prompt_manager/autotag/unload") + data = await resp.json() + finally: + model_gate.release(owner) + self.assertEqual(resp.status, 409, kind) + self.assertEqual(data["code"], "model-busy", kind) + self.assertEqual(data["operation"], label, kind) + self.assertIn("busy", data["error"].lower()) + self.assertEqual(self.service.unload_calls, 0) + + async def test_unknown_or_freed_holder_is_another_operation(self): + from py.api.autotag_routes import model_busy_response + + self.assertIsNone(model_gate.holder()) + freed = json.loads(model_busy_response().body) + self.assertEqual(freed["operation"], "Another tagging operation") + self.assertTrue(model_gate.try_acquire("test-holder")) + try: + unknown = json.loads(model_busy_response().body) + finally: + model_gate.release("test-holder") + self.assertEqual(unknown["operation"], "Another tagging operation") + + async def test_unload_during_generate_is_409_model_busy(self): + self.service._loaded_type = "wd14-vit" + owner = model_gate.new_owner("image-annotations") + self.assertTrue(model_gate.try_acquire(owner)) + try: + resp = await self.client.post("/prompt_manager/autotag/unload") + data = await resp.json() + finally: + model_gate.release(owner) + self.assertEqual(resp.status, 409) + self.assertEqual(data["operation"], "Generate suggestions") + self.assertEqual(self.service.unload_calls, 0) + self.assertEqual(self.service.get_loaded_model_type(), "wd14-vit")