Vito Sansevero
489d449ec8
- 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.
122 lines
4.7 KiB
Python
122 lines
4.7 KiB
Python
"""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
|