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

feat(autotag): scored suggestions with a configurable count behind a shared model gate

Browse Source 
- py/model_gate.py: one process-wide, non-blocking gate for loading,
  swapping, running, downloading and unloading a tagging model. Worker
  calls run shielded; a release requested while one still runs takes
  effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
  gate and answer 409 {code: "model-busy"} while another operation holds
  it; their bodies are otherwise unchanged. /autotag/models gains
  suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
  first, GGUF/HF (name, None) in output order, first `limit` distinct
  case-insensitively; generate_tags keeps its return; the image file is
  closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
  pinned capture, with strict per-engine options, never downloads
  (409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
  10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
  metadata over the inspection budgets no longer blocks it.
This commit is contained in:
Vito Sansevero
2026-10-22 09:12:01 -07:00
parent 06bd642611
commit 489d449ec8
15 changed files with 1772 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/annotation_suggest.py
+121
View File
@@ -0,0 +1,121 @@
"""Request validation and the worker of the Tag dialog's Generate (#249 B2).
``POST /prompt_manager/image-annotations/suggest`` asks one of legacy's four
engines for the top N tags of a captured image. Options are validated
strictly, per engine: WD14 classifiers take two thresholds, the language
models (GGUF and HuggingFace JoyCaption) take the instructions; every engine
takes the suggestion count and the GPU switch. Nothing here touches the
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
SUGGEST_FIELDS = frozenset(
{"target", "capture_token", "expected_epoch", "model_type", "options"}
)
_COMMON = frozenset({"suggestion_limit", "use_gpu"})
_WD14_REQUIRED = frozenset({"general_threshold", "character_threshold"})
_LLM_REQUIRED = frozenset({"instructions"})
UNAVAILABLE_HINT = "Download it from the legacy Auto-tag page."
class OptionError(ValueError):
"""A suggest request whose model type or options are not acceptable."""
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)
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
def validate_options(model_type, options, *, default_limit: int, max_limit: int):
"""Normalised ``(model_type, options)`` or :class:`OptionError`.
The answer echoes every option the request will run with, the defaults
(suggestion count from the settings, GPU on) filled in.
"""
if not isinstance(model_type, str) or model_type not in MODEL_TYPES:
raise OptionError("model_type must be one of " + ", ".join(sorted(MODEL_TYPES)))
if not isinstance(options, dict):
raise OptionError("options must be an object")
required = _WD14_REQUIRED if model_type in WD14_MODELS else _LLM_REQUIRED
keys = set(options)
if not required <= keys or not keys <= required | _COMMON:
raise OptionError(
f"options for {model_type} take exactly "
+ ", ".join(sorted(required))
+ " and optionally "
+ ", ".join(sorted(_COMMON))
)
limit = options.get("suggestion_limit", default_limit)
if type(limit) is not int or not 1 <= limit <= max_limit:
raise OptionError(f"suggestion_limit must be an integer from 1 to {max_limit}")
use_gpu = options.get("use_gpu", True)
if not isinstance(use_gpu, bool):
raise OptionError("use_gpu must be true or false")
normalised = {"use_gpu": use_gpu, "suggestion_limit": limit}
if model_type in WD14_MODELS:
for name in sorted(_WD14_REQUIRED):
normalised[name] = _threshold(name, options[name])
else:
normalised["instructions"] = _instructions(options["instructions"])
return model_type, normalised
def unavailable_message(model_name: str) -> str:
return f"{model_name} is not downloaded. {UNAVAILABLE_HINT}"
def run_suggestions(service, capture_path: str, model_type: str, options: dict):
"""Blocking: load (or swap to) *model_type*, then suggest for the capture.
The caller holds the model gate for the whole call; nothing is ever
downloaded (``load_model`` refuses a model that is not on disk).
"""
if not service.is_model_loaded() or service.get_loaded_model_type() != model_type:
service.load_model(model_type, options["use_gpu"])
return service.generate_tag_suggestions(
capture_path,
limit=options["suggestion_limit"],
prompt=options.get("instructions"),
general_threshold=options.get("general_threshold"),
character_threshold=options.get("character_threshold"),
)
def suggestion_rows(pairs) -> list:
"""``[{name, confidence}]``; a non-finite score is reported as unknown."""
rows = []
for name, score in pairs:
confidence = None
if score is not None and math.isfinite(score):
confidence = float(score)
rows.append({"name": str(name), "confidence": confidence})
return rows
py/api/annotations.py
+129
View File
@@ -16,8 +16,18 @@ from pathlib import Path
from aiohttp import web
from .annotation_suggest import (
SUGGEST_FIELDS,
OptionError,
run_suggestions,
suggestion_rows,
unavailable_message,
validate_options,
)
from .autotag_routes import model_busy_response
from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title
from .prompts import bad_request
from .. import model_gate
try:
from ...database.annotations import (
@@ -87,6 +97,7 @@ _STATUS = {
"storage-changed": 409,
"tag-conflict": 409,
"capture-expired": 409,
"model-unavailable": 409,
}
_MESSAGES = {
"source-changed": (
@@ -190,6 +201,9 @@ class AnnotationRoutesMixin:
routes.get(CAPTURE_URL + "{token}")(self.get_capture_preview)
routes.delete(CAPTURE_URL + "{token}")(self.release_capture)
routes.put("/prompt_manager/image-annotations/tags")(self.save_image_tags)
routes.post("/prompt_manager/image-annotations/suggest")(
self.suggest_image_tags
)
# -- capture store lifetime ------------------------------------------------
@@ -474,6 +488,121 @@ class AnnotationRoutesMixin:
{"success": True, **receipt}, headers={"Cache-Control": "no-store"}
)
# -- suggest ---------------------------------------------------------------
async def suggest_image_tags(self, request: web.Request) -> web.Response:
"""POST: the top N tags of the captured bytes from one legacy engine.
Runs only on the private capture (pinned for the whole inference),
under the shared model gate (``model-busy`` while anything else holds
it, including a legacy batch), never downloads (``model-unavailable``)
and never writes.
"""
from ..autotag import MODELS, get_autotag_service
from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig
data, error = await self._annotation_body(request, SUGGEST_FIELDS)
if error is not None:
return error
target, error = _target(data["target"])
if error is not None:
return error
try:
model_type, options = validate_options(
data["model_type"],
data["options"],
default_limit=AutotagConfig.suggestion_limit(),
max_limit=MAX_SUGGESTION_LIMIT,
)
except OptionError as exc:
return bad_request(str(exc))
token, epoch = data["capture_token"], data["expected_epoch"]
if not isinstance(token, str) or _TOKEN.fullmatch(token) is None:
return bad_request("Invalid capture token")
if not isinstance(epoch, str) or _UUID.fullmatch(epoch) is None:
return bad_request("expected_epoch must be a canonical UUID")
index, scope, relative = target
root, error = self._scoped_root(index, scope)
if error is not None:
return error
store = self._captures
capture = store.acquire(token) if store is not None else None
if capture is None:
return _coded("capture-expired")
started = False
try:
error = await self._suggest_guards(capture, target, root, epoch)
if error is not None:
return error
service = get_autotag_service()
status = await self._run_in_executor(service.get_models_status)
if not status.get(model_type, {}).get("downloaded"):
name = MODELS.get(model_type, {}).get("name", model_type)
return _coded("model-unavailable", unavailable_message(name))
owner = model_gate.new_owner("image-annotations")
if not model_gate.try_acquire(owner):
return model_busy_response()
try:
started = True
pairs = await model_gate.run(
owner,
_pinned_suggestions,
store,
token,
service,
capture,
model_type,
options,
)
except Exception:
self.logger.exception("Tag suggestion failed")
return _error(
500,
"suggest-failed",
"Auto-tagging failed. Check server logs for details.",
)
finally:
model_gate.release(owner)
finally:
if not started:
store.release_use(token)
return web.json_response(
{
"success": True,
"model_type": model_type,
"options": options,
"suggestions": suggestion_rows(pairs),
},
headers={"Cache-Control": "no-store"},
)
async def _suggest_guards(self, capture, target, root, epoch):
"""The capture must be this target's, of this epoch, and inferable."""
index, scope, relative = target
expected = os.path.normcase(
os.path.join(os.path.realpath(root), *relative.split("/"))
)
if (
capture.root_index != index
or capture.scope != scope
or capture.locator != expected
):
return _coded("capture-expired", _OTHER_FILE)
if epoch != capture.epoch:
current = await self._run_in_executor(self.db.get_execution_epoch)
return _coded("storage-changed", library_epoch=current)
if not capture.inferable:
return _coded("unsupported-media", capture.inferable_reason)
return None
def _pinned_suggestions(store, token, service, capture, model_type, options):
"""Worker thread: infer on the capture, then end its pin (even on error)."""
try:
return run_suggestions(service, capture.path, model_type, options)
finally:
store.release_use(token)
def _verify_captured_file(store: CaptureStore, capture, root, relative: str):
"""Re-hash the current file, or None when the token names another file.
py/api/autotag_routes.py
+76 -27
View File
@@ -9,11 +9,24 @@ from typing import Optional
from urllib.parse import quote
from .prompts import thumbnail_url_for
from .. import model_gate
from aiohttp import web
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
_BUSY_ERROR = "Auto-tagging is already running"
MODEL_BUSY_ERROR = (
"A tagging model is busy with another operation. Try again when it finishes."
)
def model_busy_response() -> web.Response:
"""``409 model-busy``: another operation holds the shared model gate."""
return web.json_response(
{"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR},
status=409,
)
# Only these may be handed to the tagging engine or listed by the scanner.
_MEDIA_EXTENSIONS = frozenset(
@@ -86,6 +99,7 @@ class AutotagRoutesMixin:
"""Get status of available AutoTag models."""
try:
from ..autotag import get_autotag_service
from ..config import AutotagConfig
service = get_autotag_service()
models_status = service.get_models_status()
@@ -99,6 +113,7 @@ class AutotagRoutesMixin:
"loaded_model_type": service.get_loaded_model_type(),
"wd14_general_threshold": service.wd14_general_threshold,
"wd14_character_threshold": service.wd14_character_threshold,
"suggestion_limit": AutotagConfig.suggestion_limit(),
}
)
@@ -138,6 +153,7 @@ class AutotagRoutesMixin:
async def download_autotag_model(self, request):
"""Download an AutoTag model with streaming progress (POST only)."""
model_type = request.match_info.get("model_type")
owner = model_gate.new_owner("legacy-download")
async def events():
try:
@@ -160,8 +176,8 @@ class AutotagRoutesMixin:
def progress_callback(status: str, progress: float):
pass
success = await self._run_in_executor(
service.download_model, model_type, progress_callback
success = await model_gate.run(
owner, service.download_model, model_type, progress_callback
)
if success:
yield {
@@ -178,7 +194,12 @@ class AutotagRoutesMixin:
"message": _INTERNAL_ERROR,
}
return await self._stream_sse(request, events())
if not model_gate.try_acquire(owner):
return model_busy_response()
try:
return await self._stream_sse(request, events())
finally:
model_gate.release(owner)
async def _read_start_params(self, request) -> Optional[dict]:
"""Merge query-string and optional JSON-body parameters for autotag/start.
@@ -239,8 +260,9 @@ class AutotagRoutesMixin:
return
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
run = self._gated_runner(opts.get("gate_owner"))
try:
await self._run_in_executor(service.load_model, model_type, True)
await run(service.load_model, model_type, True)
except Exception:
self.logger.exception("Failed to load model")
yield {
@@ -338,7 +360,7 @@ class AutotagRoutesMixin:
continue
try:
tags = await self._run_in_executor(
tags = await run(
service.generate_tags,
allowed_path,
general_threshold=opts["general_threshold"],
@@ -423,10 +445,26 @@ class AutotagRoutesMixin:
{"success": False, "error": _BUSY_ERROR}, status=409
)
try:
return await self._stream_sse(request, events())
opts["gate_owner"] = model_gate.new_owner("legacy-batch")
if not model_gate.try_acquire(opts["gate_owner"]):
return model_busy_response()
try:
return await self._stream_sse(request, events())
finally:
model_gate.release(opts["gate_owner"])
finally:
_autotag_batch_lock.release()
def _gated_runner(self, owner):
"""Executor calls under the model gate *owner* holds (plain when None)."""
if owner is None:
return self._run_in_executor
async def run(func, *args, **kwargs):
return await model_gate.run(owner, func, *args, **kwargs)
return run
def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]:
"""Return the real path when it is a media file inside an output directory."""
if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS:
@@ -549,22 +587,28 @@ class AutotagRoutesMixin:
from ..autotag import get_autotag_service
service = get_autotag_service()
owner = model_gate.new_owner("legacy-single")
if not model_gate.try_acquire(owner):
return model_busy_response()
try:
if (
not service.is_model_loaded()
or service.get_loaded_model_type() != model_type
):
await model_gate.run(owner, service.load_model, model_type, use_gpu)
if (
not service.is_model_loaded()
or service.get_loaded_model_type() != model_type
):
await self._run_in_executor(service.load_model, model_type, use_gpu)
if custom_prompt:
service.custom_prompt = custom_prompt
if custom_prompt:
service.custom_prompt = custom_prompt
tags = await self._run_in_executor(
service.generate_tags,
image_path,
general_threshold=general_threshold,
character_threshold=character_threshold,
)
tags = await model_gate.run(
owner,
service.generate_tags,
image_path,
general_threshold=general_threshold,
character_threshold=character_threshold,
)
finally:
model_gate.release(owner)
return web.json_response(
{
@@ -636,14 +680,19 @@ class AutotagRoutesMixin:
from ..autotag import get_autotag_service
service = get_autotag_service()
owner = model_gate.new_owner("legacy-unload")
if not model_gate.try_acquire(owner):
return model_busy_response()
try:
if not service.is_model_loaded():
return web.json_response(
{"success": True, "message": "No model was loaded"}
)
if not service.is_model_loaded():
return web.json_response(
{"success": True, "message": "No model was loaded"}
)
model_type = service.get_loaded_model_type()
service.unload_model()
model_type = service.get_loaded_model_type()
service.unload_model()
finally:
model_gate.release(owner)
return web.json_response(
{