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

feat(autotag): model status without paths, and a busy answer that names the operation

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-26 23:00:48 -07:00
parent cfe47d3cb0
commit 04f2c921cf
5 changed files with 246 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/autotag_routes.py
+80 -7
View File
@@ -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 ``"<kind>:<uuid>"`` (``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)
py/api/jobs.py
+2
View File
@@ -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"
py/autotag.py
+6 -1
View File
@@ -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",
tests/browser/server.py
+4 -1
View File
@@ -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()
tests/test_api_autotag.py
+154 -2
View File
@@ -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")