From 489d449ec893d5bea7695ec8f23833b968f7f20d Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 22 Oct 2026 09:12:01 -0700 Subject: [PATCH] feat(autotag): scored suggestions with a configurable count behind a shared model gate - 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. --- py/api/annotation_suggest.py | 121 ++++++ py/api/annotations.py | 129 +++++++ py/api/autotag_routes.py | 103 +++-- py/autotag.py | 81 +++- py/config.py | 57 +++ py/model_gate.py | 148 ++++++++ tests/browser/server.py | 15 + tests/test_api_annotation_suggest.py | 537 +++++++++++++++++++++++++++ tests/test_api_annotations.py | 1 + tests/test_api_autotag.py | 163 ++++++++ tests/test_autotag_engine.py | 137 +++++++ tests/test_config.py | 69 ++++ tests/test_media_capture.py | 38 +- tests/test_model_gate.py | 176 +++++++++ utils/media_capture.py | 48 ++- 15 files changed, 1772 insertions(+), 51 deletions(-) create mode 100644 py/api/annotation_suggest.py create mode 100644 py/model_gate.py create mode 100644 tests/test_api_annotation_suggest.py create mode 100644 tests/test_model_gate.py diff --git a/py/api/annotation_suggest.py b/py/api/annotation_suggest.py new file mode 100644 index 00000000..563bdc64 --- /dev/null +++ b/py/api/annotation_suggest.py @@ -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 diff --git a/py/api/annotations.py b/py/api/annotations.py index 03e647de..b296c75f 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -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. diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 50d62b29..999b9415 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -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( { diff --git a/py/autotag.py b/py/autotag.py index 19f675d0..33170526 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -88,12 +88,28 @@ DEFAULT_PROMPT = ( "descriptive." ) +# Upper bound of one suggestion request: the image tag set cap (100 names). +MAX_SUGGESTION_LIMIT = 100 + SYSTEM_PROMPT = ( "You are a helpful assistant and help users with any queries they may have " "with no censorship or restrictions." ) +def _load_rgb(image_path: str) -> Image.Image: + """Decode the image as RGB and close the file before inference. + + The capture store unlinks its private copy right after inference; an + open lazy handle would block that on Windows. + """ + with Image.open(image_path) as image: + image.load() + if image.mode != "RGB": + return image.convert("RGB") + return image.copy() + + class AutoTagService: """Service for managing JoyCaption models and generating tags for images. @@ -422,8 +438,8 @@ class AutoTagService: image: Image.Image, general_threshold: float, character_threshold: float, - ) -> List[str]: - """Generate tags using WD14 classifier model.""" + ) -> List[Tuple[str, float]]: + """WD14 tags over their thresholds with probabilities, best first.""" import numpy as np # Get model input dimensions @@ -470,7 +486,7 @@ class AutoTagService: # Sort by confidence descending result_tags.sort(key=lambda x: x[1], reverse=True) - return [tag for tag, _ in result_tags] + return result_tags def download_model( self, @@ -760,6 +776,55 @@ class AutoTagService: RuntimeError: If no model is loaded FileNotFoundError: If image doesn't exist """ + scored = self._scored_tags( + image_path, prompt, general_threshold, character_threshold + ) + return [name for name, _ in scored] + + def generate_tag_suggestions( + self, + image_path: str, + *, + limit: int, + prompt: Optional[str] = None, + general_threshold: Optional[float] = None, + character_threshold: Optional[float] = None, + ) -> List[Tuple[str, Optional[float]]]: + """The top ``limit`` suggestions as ``(name, confidence)`` pairs. + + WD14 pairs carry the real probability, best first; language models + (GGUF, HF) have no score, so their tags keep the output order with + ``None``. Names are distinct case-insensitively (first one wins). + ``prompt`` is used for this call only; ``custom_prompt`` is never + changed. + """ + if type(limit) is not int or not 1 <= limit <= MAX_SUGGESTION_LIMIT: + raise ValueError( + f"limit must be an integer from 1 to {MAX_SUGGESTION_LIMIT}" + ) + scored = self._scored_tags( + image_path, prompt, general_threshold, character_threshold + ) + suggestions: List[Tuple[str, Optional[float]]] = [] + seen = set() + for name, score in scored: + key = name.casefold() + if key in seen: + continue + seen.add(key) + suggestions.append((name, score)) + if len(suggestions) == limit: + break + return suggestions + + def _scored_tags( + self, + image_path: str, + prompt: Optional[str], + general_threshold: Optional[float], + character_threshold: Optional[float], + ) -> List[Tuple[str, Optional[float]]]: + """Run the loaded model: WD14 ``(name, probability)``, LLM ``(name, None)``.""" if self._tagger is None: raise RuntimeError("No model loaded. Call load_model() first.") @@ -767,18 +832,14 @@ class AutoTagService: raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path)) use_prompt = prompt or self._custom_prompt - - # Load image - image = Image.open(image_path) - if image.mode != "RGB": - image = image.convert("RGB") + image = _load_rgb(image_path) # Generate based on model type model_type, tagger_obj = self._tagger if model_type == "gguf": raw_tags = self._generate_gguf(tagger_obj, image, use_prompt) - return self._parse_tags(raw_tags) + return [(tag, None) for tag in self._parse_tags(raw_tags)] elif model_type == "wd14": session, tags_list = tagger_obj gen_thresh = ( @@ -799,7 +860,7 @@ class AutoTagService: raw_tags = self._generate_hf( model, processor, device, compute_dtype, image, use_prompt ) - return self._parse_tags(raw_tags) + return [(tag, None) for tag in self._parse_tags(raw_tags)] def _generate_gguf(self, model, image: Image.Image, prompt: str) -> str: """Generate tags using GGUF model.""" diff --git a/py/config.py b/py/config.py index cc89108a..5115bfe3 100644 --- a/py/config.py +++ b/py/config.py @@ -465,6 +465,59 @@ class IntegrationConfig: cls.CIVITAI_API_KEY = lora["civitai_api_key"] +# Number of tag suggestions the Tag dialogs request by default (D68). The cap +# is the image tag set limit: more suggestions could never all be kept. +DEFAULT_SUGGESTION_LIMIT = 10 +MAX_SUGGESTION_LIMIT = 100 + + +class AutotagConfig: + """Persisted tagging defaults under ``autotag.defaults`` in config.json. + + ``suggestion_limit`` is how many suggestions Generate returns (the top N + by score), an integer in ``1..MAX_SUGGESTION_LIMIT``. A hand-edited value + outside that range, or of another type, falls back to the default with a + warning rather than failing the load. + """ + + SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT + + @classmethod + def suggestion_limit(cls) -> int: + return cls.SUGGESTION_LIMIT + + @staticmethod + def valid_suggestion_limit(value: Any) -> bool: + return ( + type(value) is int and 1 <= value <= MAX_SUGGESTION_LIMIT + ) # bool, float and str are rejected + + @classmethod + def get_config(cls) -> Dict[str, Any]: + return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}} + + @classmethod + def update_config(cls, new_config: Any) -> None: + if not isinstance(new_config, dict): + config_logger.warning("Ignoring autotag settings that are not an object") + return + defaults = new_config.get("defaults", {}) + if not isinstance(defaults, dict): + config_logger.warning("Ignoring autotag.defaults that is not an object") + return + if "suggestion_limit" not in defaults: + return + value = defaults["suggestion_limit"] + if cls.valid_suggestion_limit(value): + cls.SUGGESTION_LIMIT = value + return + config_logger.warning( + "autotag.defaults.suggestion_limit must be an integer from 1 to " + f"{MAX_SUGGESTION_LIMIT}; using {DEFAULT_SUGGESTION_LIMIT}" + ) + cls.SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT + + class PromptManagerConfig: """Main configuration class for PromptManager core functionality. @@ -566,6 +619,7 @@ class PromptManagerConfig: }, "gallery": GalleryConfig.get_config(), "integrations": IntegrationConfig.get_config(), + "autotag": AutotagConfig.get_config(), } @classmethod @@ -706,6 +760,9 @@ class PromptManagerConfig: if "integrations" in new_config: IntegrationConfig.update_config(new_config["integrations"]) + if "autotag" in new_config: + AutotagConfig.update_config(new_config["autotag"]) + # Load configuration on import try: diff --git a/py/model_gate.py b/py/model_gate.py new file mode 100644 index 00000000..1ad7ec87 --- /dev/null +++ b/py/model_gate.py @@ -0,0 +1,148 @@ +"""One process-wide, non-blocking gate around the tagging models (#249 B2). + +Loading, swapping, running, downloading and unloading a tagging model all +happen under this gate, whether the caller is a legacy Auto-tag route, the +Tag dialog's Generate or (later) a bulk job. A caller that cannot take the +gate is refused at once (``409 model-busy``), never queued. + +Work that runs in a worker thread is admitted with :meth:`ModelGate.run`. A +release requested while such work still runs (a client disconnect cancels +the waiting coroutine, not the thread) takes effect only when the last +admitted call really exits, so nobody can swap the model under a running +inference. +""" + +from __future__ import annotations + +import asyncio +import threading +import uuid +from typing import Any, Callable, Optional + + +class GateNotHeld(RuntimeError): + """Raised when work is admitted for an owner that does not hold the gate.""" + + +class ModelGate: + """A single-owner gate whose release waits for admitted worker calls.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._holder: Optional[str] = None + self._inflight = 0 + self._releasing = False + + def try_acquire(self, owner: str) -> bool: + """Take the gate for *owner*; False at once when anyone holds it.""" + with self._lock: + if self._holder is not None: + return False + self._holder = owner + self._inflight = 0 + self._releasing = False + return True + + def release(self, owner: str) -> None: + """Let go; deferred until admitted work of *owner* has exited.""" + with self._lock: + if self._holder != owner: + return + if self._inflight: + self._releasing = True + else: + self._clear_locked() + + def holder(self) -> Optional[str]: + with self._lock: + return self._holder + + def begin(self, owner: str) -> None: + """Admit one worker call; the holder must still hold the gate.""" + with self._lock: + if self._holder != owner or self._releasing: + raise GateNotHeld(f"{owner} does not hold the model gate") + self._inflight += 1 + + def end(self, owner: str) -> None: + """One admitted call exited; completes a deferred release.""" + with self._lock: + self._end_locked(owner) + + async def run(self, owner: str, func: Callable[..., Any], *args, **kwargs): + """Run ``func`` in the default executor while *owner* holds the gate. + + The executor future is shielded: cancelling the caller never cancels + a call that has not started yet (which would leave it admitted + forever) and never ends the admission of a call that is running. + """ + self.begin(owner) + state = {"started": False, "finished": False} + + def finish() -> None: + with self._lock: + if not state["finished"]: + state["finished"] = True + self._end_locked(owner) + + def work(): + with self._lock: + state["started"] = True + try: + return func(*args, **kwargs) + finally: + finish() + + def never_ran(future) -> None: + # Only a loop shutdown cancels the shielded future; a call that + # never started must still give its admission back. + if future.cancelled(): + with self._lock: + skipped = not state["started"] and not state["finished"] + if skipped: + state["finished"] = True + self._end_locked(owner) + + try: + future = asyncio.get_running_loop().run_in_executor(None, work) + except BaseException: + finish() + raise + future.add_done_callback(never_ran) + return await asyncio.shield(future) + + def _end_locked(self, owner: str) -> None: + if self._holder != owner or self._inflight == 0: + return + self._inflight -= 1 + if self._releasing and self._inflight == 0: + self._clear_locked() + + def _clear_locked(self) -> None: + self._holder = None + self._inflight = 0 + self._releasing = False + + +_GATE = ModelGate() + + +def new_owner(kind: str) -> str: + """A unique owner name such as ``legacy-single:``.""" + return f"{kind}:{uuid.uuid4()}" + + +def try_acquire(owner: str) -> bool: + return _GATE.try_acquire(owner) + + +def release(owner: str) -> None: + _GATE.release(owner) + + +def holder() -> Optional[str]: + return _GATE.holder() + + +async def run(owner: str, func: Callable[..., Any], *args, **kwargs): + return await _GATE.run(owner, func, *args, **kwargs) diff --git a/tests/browser/server.py b/tests/browser/server.py index c5ae192d..ae67018c 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -80,6 +80,21 @@ class TestModel: def generate_tags(self, image_path: str, **kwargs: object) -> list[str]: return ["generated", Path(image_path).parent.parent.name] + def generate_tag_suggestions( + self, image_path: str, *, limit: int, **kwargs: object + ) -> list[tuple[str, float | None]]: + """Fixed scored suggestions (WD14) or unscored ones (language models).""" + scored = self.loaded is not None and self.loaded.startswith("wd14") + names = ["sandstone", "desert", "arch", "sky", "rock", "dune", "canyon"] + names += [f"tag_{index:02d}" for index in range(100)] + return [ + (name, round(0.98 - index * 0.008, 3) if scored else None) + for index, name in enumerate(names[:limit]) + ] + + def download_model(self, model_type: str, progress_callback=None) -> bool: + raise AssertionError("browser tests never download a model") + REVEAL_OUTCOMES = ("opened", "unavailable", "fail") diff --git a/tests/test_api_annotation_suggest.py b/tests/test_api_annotation_suggest.py new file mode 100644 index 00000000..7acac1f1 --- /dev/null +++ b/tests/test_api_annotation_suggest.py @@ -0,0 +1,537 @@ +"""Scored top-N tag suggestions for a captured image (#249 B2). + +``POST /prompt_manager/image-annotations/suggest`` runs the selected legacy +engine on the private capture (never the current file), behind the one +process-wide model gate shared with the legacy Auto-tag routes. It never +downloads a model and never writes. No real model is loaded: the engine is a +fake service. +""" + +import asyncio +import json +import math +import os +import threading +import unittest +import uuid +from unittest.mock import MagicMock, patch + +try: + from test_api_annotations import AnnotationAPITestCase + from test_api_images import make_png +except ImportError: + from tests.test_api_annotations import AnnotationAPITestCase + from tests.test_api_images import make_png + +try: # guarded so a missing module fails assertions, not test loading + from py import model_gate +except ImportError: + model_gate = None + +from py import config as config_module + +AutotagConfig = getattr(config_module, "AutotagConfig", None) + +SUGGEST = "/prompt_manager/image-annotations/suggest" +WAIT = 30.0 +WD14 = {"general_threshold": 0.35, "character_threshold": 0.85} +LLM = {"instructions": "Generate tags."} +UNAVAILABLE = "Download it from the legacy Auto-tag page." + +SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [ + (f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150) +] + + +class FakeTagger: + """Stand-in for ``AutoTagService`` with the suggestion entry point.""" + + default_prompt = "server default prompt" + custom_prompt = "legacy prompt" + wd14_general_threshold = 0.35 + wd14_character_threshold = 0.85 + + def __init__(self): + self.models_config = { + "wd14-swinv2": {}, + "wd14-vit": {}, + "gguf": {}, + "hf": {}, + } + self.downloaded = {name: True for name in self.models_config} + self.loaded = None + self.load_calls = [] + self.unload_calls = 0 + self.calls = [] + self.read_paths = [] + self.read_bytes = [] + self.download_model = MagicMock(side_effect=AssertionError("downloaded")) + self.load_error = None + self.block = None # (entered, proceed) events + + def get_models_status(self): + return { + name: {"downloaded": self.downloaded[name], "model_path": None} + for name in self.models_config + } + + def is_model_loaded(self): + return self.loaded is not None + + def get_loaded_model_type(self): + return self.loaded + + def load_model(self, model_type, use_gpu=True): + self.load_calls.append((model_type, use_gpu)) + if self.load_error is not None: + raise self.load_error + self.loaded = model_type + return True + + def unload_model(self): + self.unload_calls += 1 + self.loaded = None + + def generate_tags(self, image_path, **kwargs): + return ["legacy"] + + def generate_tag_suggestions( + self, + image_path, + *, + limit, + prompt=None, + general_threshold=None, + character_threshold=None, + ): + self.read_paths.append(image_path) + with open(image_path, "rb") as handle: + self.read_bytes.append(handle.read()) + self.calls.append( + { + "limit": limit, + "prompt": prompt, + "general_threshold": general_threshold, + "character_threshold": character_threshold, + "model": self.loaded, + } + ) + if self.block is not None: + entered, proceed = self.block + entered.set() + proceed.wait(WAIT) + if self.loaded.startswith("wd14"): + return list(SCORED[:limit]) + return [(name, None) for name, _ in SCORED[:limit]] + + +class SuggestTestCase(AnnotationAPITestCase): + async def get_application(self): + app = await super().get_application() + self.assertIsNotNone(model_gate, "py.model_gate is not implemented") + self.assertIsNotNone(AutotagConfig, "AutotagConfig is not implemented") + self.tagger = FakeTagger() + patcher = patch("py.autotag.get_autotag_service", return_value=self.tagger) + patcher.start() + self.addCleanup(patcher.stop) + AutotagConfig.SUGGESTION_LIMIT = 10 + self.addCleanup(setattr, AutotagConfig, "SUGGESTION_LIMIT", 10) + return app + + async def tearDownAsync(self): + leaked = model_gate.holder() if model_gate is not None else None + if leaked is not None: + model_gate.release(leaked) + await super().tearDownAsync() + self.assertIsNone(leaked, "the suggest route left the model gate held") + + async def observed(self, path="sub/a.png", size=(8, 6)): + make_png(self.output_dir / path, size=size) + return await self.observe_ok(path) + + def suggest_body(self, observed, model_type="wd14-vit", options=None, **changes): + if options is None: + options = dict(WD14 if model_type.startswith("wd14") else LLM) + return { + "target": dict(observed["target"]), + "capture_token": observed["capture"]["token"], + "expected_epoch": observed["library_epoch"], + "model_type": model_type, + "options": options, + **changes, + } + + async def suggest(self, body): + response = await self.client.post(SUGGEST, json=body) + return response.status, await self.json_of(response) + + async def suggest_ok(self, body): + status, data = await self.suggest(body) + self.assertEqual(status, 200, data) + return data + + def capture_path(self, observed): + return self.api._captures.get(observed["capture"]["token"]).path + + async def wait_until_free(self): + for _ in range(3000): + if model_gate.holder() is None: + return + await asyncio.sleep(0.01) + self.fail(f"model gate still held by {model_gate.holder()}") + + +class TestSuggestions(SuggestTestCase): + async def test_wd14_suggestions_are_scored_top_n_from_the_capture(self): + observed = await self.observed() + + body = await self.suggest_ok(self.suggest_body(observed)) + + self.assertEqual( + body["suggestions"][:2], + [ + {"name": "sandstone", "confidence": 0.98}, + {"name": "desert", "confidence": 0.96}, + ], + ) + self.assertEqual(len(body["suggestions"]), 10) + self.assertEqual(body["model_type"], "wd14-vit") + self.assertEqual( + body["options"], + { + "general_threshold": 0.35, + "character_threshold": 0.85, + "use_gpu": True, + "suggestion_limit": 10, + }, + ) + self.assertEqual(self.tagger.read_paths, [self.capture_path(observed)]) + self.assertEqual( + self.tagger.calls[0], + { + "limit": 10, + "prompt": None, + "general_threshold": 0.35, + "character_threshold": 0.85, + "model": "wd14-vit", + }, + ) + self.assertEqual(self.tagger.load_calls, [("wd14-vit", True)]) + self.tagger.download_model.assert_not_called() + self.assertIsNone(model_gate.holder()) + self.assert_no_paths(body) + + async def test_suggestions_keep_the_requested_limit(self): + observed = await self.observed() + for limit in (1, 10, 100): + with self.subTest(limit=limit): + options = {**WD14, "suggestion_limit": limit} + body = await self.suggest_ok( + self.suggest_body(observed, options=options) + ) + self.assertEqual(len(body["suggestions"]), limit) + self.assertEqual(body["options"]["suggestion_limit"], limit) + # Omitted: the persisted autotag.defaults.suggestion_limit. + AutotagConfig.SUGGESTION_LIMIT = 3 + body = await self.suggest_ok(self.suggest_body(observed)) + self.assertEqual(len(body["suggestions"]), 3) + self.assertEqual(body["options"]["suggestion_limit"], 3) + + async def test_llm_suggestions_send_instructions_and_null_confidence(self): + observed = await self.observed() + self.tagger.custom_prompt = "legacy prompt" + for model_type in ("gguf", "hf"): + with self.subTest(model_type=model_type): + options = {"instructions": f"tags for {model_type}", "use_gpu": False} + body = await self.suggest_ok( + self.suggest_body(observed, model_type, options) + ) + self.assertEqual(body["model_type"], model_type) + self.assertEqual(body["options"], {**options, "suggestion_limit": 10}) + self.assertEqual( + body["suggestions"][:2], + [ + {"name": "sandstone", "confidence": None}, + {"name": "desert", "confidence": None}, + ], + ) + call = self.tagger.calls[-1] + self.assertEqual(call["prompt"], f"tags for {model_type}") + self.assertIsNone(call["general_threshold"]) + self.assertEqual(self.tagger.load_calls[-1], (model_type, False)) + self.assertEqual(self.tagger.custom_prompt, "legacy prompt") + + async def test_loaded_model_is_reused_and_another_is_swapped_in(self): + observed = await self.observed() + self.tagger.loaded = "wd14-vit" + await self.suggest_ok(self.suggest_body(observed)) + self.assertEqual(self.tagger.load_calls, []) + await self.suggest_ok(self.suggest_body(observed, "wd14-swinv2")) + self.assertEqual(self.tagger.load_calls, [("wd14-swinv2", True)]) + self.assertEqual(self.tagger.calls[-1]["model"], "wd14-swinv2") + + async def test_suggest_uses_capture_bytes_not_current_file(self): + observed = await self.observed() + original = (self.output_dir / "sub" / "a.png").read_bytes() + make_png(self.output_dir / "sub" / "a.png", size=(9, 9), color=(0, 0, 255)) + self.assertNotEqual((self.output_dir / "sub" / "a.png").read_bytes(), original) + + await self.suggest_ok(self.suggest_body(observed)) + + capture_path = self.capture_path(observed) + self.assertEqual(self.tagger.read_paths, [capture_path]) + self.assertEqual(self.tagger.read_bytes, [original]) + self.assertNotEqual( + os.path.normcase(capture_path), + os.path.normcase(str(self.output_dir / "sub" / "a.png")), + ) + + async def test_suggest_writes_nothing(self): + observed = await self.observed() + before = self.db_snapshot() + rows = self.bytes_rows() + + await self.suggest_ok(self.suggest_body(observed)) + await self.suggest_ok(self.suggest_body(observed, "gguf")) + + self.assertEqual(self.db_snapshot(), before) + self.assertEqual(self.bytes_rows(), rows) + # The capture stays usable for Save after Generate. + self.assertIsNotNone(self.api._captures.get(observed["capture"]["token"])) + + +class TestModelAvailability(SuggestTestCase): + async def test_suggest_never_downloads_and_reports_model_unavailable(self): + observed = await self.observed() + self.tagger.downloaded["wd14-vit"] = False + + status, body = await self.suggest(self.suggest_body(observed)) + + self.assertEqual((status, body["code"]), (409, "model-unavailable")) + self.assertIn(UNAVAILABLE, body["error"]) + self.assertIn("WD14 ViT", body["error"]) + self.tagger.download_model.assert_not_called() + self.assertEqual(self.tagger.load_calls, []) + self.assertEqual(self.tagger.calls, []) + self.assertIsNone(model_gate.holder()) + + async def test_load_failure_is_a_generic_error_and_frees_everything(self): + observed = await self.observed() + self.tagger.load_error = RuntimeError(f"cannot load {self.output_dir}") + + status, body = await self.suggest(self.suggest_body(observed)) + + self.assertEqual((status, body["code"]), (500, "suggest-failed")) + self.assert_no_paths(body) + self.assertIsNone(model_gate.holder()) + entry = self.api._captures._entries[observed["capture"]["token"]] + self.assertEqual(entry.uses, 0) + + +class TestModelGate(SuggestTestCase): + async def _blocked_suggest(self, observed, model_type="wd14-vit"): + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + task = asyncio.ensure_future( + self.client.post(SUGGEST, json=self.suggest_body(observed, model_type)) + ) + await asyncio.get_running_loop().run_in_executor(None, entered.wait, WAIT) + self.assertTrue(entered.is_set()) + return task, proceed + + async def test_suggest_reports_busy_while_legacy_batch_runs(self): + observed = await self.observed() + prompt_id = self._save_prompt("batch prompt") + self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png")) + started, proceed = threading.Event(), threading.Event() + + def slow_generate(image_path, **kwargs): + started.set() + proceed.wait(WAIT) + return ["batch"] + + self.tagger.generate_tags = slow_generate + batch = asyncio.ensure_future( + self.client.post("/prompt_manager/autotag/start?model_type=wd14-vit") + ) + try: + await asyncio.get_running_loop().run_in_executor(None, started.wait, WAIT) + status, body = await self.suggest(self.suggest_body(observed, "gguf")) + self.assertEqual((status, body["code"]), (409, "model-busy")) + self.assertEqual(self.tagger.calls, []) + self.assertEqual(self.tagger.loaded, "wd14-vit") # not swapped + finally: + proceed.set() + response = await batch + await response.read() + await self.wait_until_free() + await self.suggest_ok(self.suggest_body(observed)) + + async def test_legacy_unload_returns_model_busy_while_native_generate_runs(self): + observed = await self.observed() + task, proceed = await self._blocked_suggest(observed) + try: + response = await self.client.post("/prompt_manager/autotag/unload") + self.assertEqual(response.status, 409) + self.assertEqual((await response.json())["code"], "model-busy") + response = await self.client.post( + "/prompt_manager/autotag/single", + json={"root_index": 0, "path": "sub/a.png", "model_type": "gguf"}, + ) + self.assertEqual(response.status, 409) + status, body = await self.suggest(self.suggest_body(observed, "gguf")) + self.assertEqual((status, body["code"]), (409, "model-busy")) + self.assertEqual(self.tagger.unload_calls, 0) + self.assertEqual(self.tagger.loaded, "wd14-vit") + finally: + proceed.set() + response = await task + self.assertEqual(response.status, 200) + await self.wait_until_free() + self.assertEqual([call["model"] for call in self.tagger.calls], ["wd14-vit"]) + + async def test_gate_held_after_client_disconnect_until_executor_finishes(self): + observed = await self.observed() + token = observed["capture"]["token"] + task, proceed = await self._blocked_suggest(observed) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + try: + # The client is gone; the inference still runs, so the gate and + # the capture pin stay until the worker thread exits. + self.assertIsNotNone(model_gate.holder()) + response = await self.client.post("/prompt_manager/autotag/unload") + self.assertEqual(response.status, 409) + self.assertEqual(self.api._captures._entries[token].uses, 1) + finally: + proceed.set() + await self.wait_until_free() + for _ in range(3000): + if self.api._captures._entries[token].uses == 0: + break + await asyncio.sleep(0.01) + self.assertEqual(self.api._captures._entries[token].uses, 0) + response = await self.client.post("/prompt_manager/autotag/unload") + self.assertEqual(response.status, 200) + + +class TestRequestValidation(SuggestTestCase): + async def test_strict_option_validation(self): + observed = await self.observed() + bad_options = { + "wd14-vit": [ + {**WD14, "general_threshold": True}, + {**WD14, "general_threshold": math.nan}, + {**WD14, "general_threshold": math.inf}, + {**WD14, "general_threshold": "0.3"}, + {**WD14, "general_threshold": 1.5}, + {**WD14, "character_threshold": -0.1}, + {"general_threshold": 0.35}, + {**WD14, "instructions": "x"}, + {**WD14, "unknown": 1}, + {**WD14, "use_gpu": "yes"}, + {**WD14, "suggestion_limit": 0}, + {**WD14, "suggestion_limit": 101}, + {**WD14, "suggestion_limit": 2.5}, + {**WD14, "suggestion_limit": "10"}, + {**WD14, "suggestion_limit": True}, + ], + "gguf": [ + {**LLM, "general_threshold": 0.35}, + {**LLM, "character_threshold": 0.85}, + {"instructions": "x" * 4001}, + {"instructions": "\ud800"}, # not encodable as UTF-8 + {"instructions": " "}, + {"instructions": 5}, + {}, + {**LLM, "suggestion_limit": 0}, + ], + } + for model_type, cases in bad_options.items(): + for options in cases: + with self.subTest(model_type=model_type, options=repr(options)[:80]): + status, body = await self.suggest( + self.suggest_body(observed, model_type, options) + ) + self.assertEqual(status, 400, body) + for body in ( + {**self.suggest_body(observed), "extra": 1}, + {k: v for k, v in self.suggest_body(observed).items() if k != "options"}, + self.suggest_body(observed, options=[]), + self.suggest_body(observed, model_type="ollama"), + self.suggest_body(observed, model_type=7, options=dict(WD14)), + self.suggest_body(observed, capture_token="../x"), + self.suggest_body(observed, expected_epoch="nope"), + self.suggest_body(observed, target={"root_index": 0}), + ): + with self.subTest(body=json.dumps(body)[:120]): + status, _ = await self.suggest(body) + self.assertEqual(status, 400) + # The edges that are allowed. + await self.suggest_ok( + self.suggest_body( + observed, + options={ + "general_threshold": 0, + "character_threshold": 1, + "suggestion_limit": 100, + }, + ) + ) + await self.suggest_ok( + self.suggest_body(observed, "hf", options={"instructions": "x" * 4000}) + ) + self.assertEqual(len(self.tagger.calls), 2) + self.tagger.download_model.assert_not_called() + + async def test_capture_token_must_match_target_and_epoch(self): + observed = await self.observed() + make_png(self.output_dir / "sub" / "b.png") + other = await self.observe_ok("sub/b.png") + + status, body = await self.suggest( + self.suggest_body(observed, capture_token="x" * 43) + ) + self.assertEqual((status, body["code"]), (409, "capture-expired")) + status, body = await self.suggest( + self.suggest_body(observed, capture_token=other["capture"]["token"]) + ) + self.assertEqual((status, body["code"]), (409, "capture-expired")) + status, body = await self.suggest( + self.suggest_body(observed, expected_epoch=str(uuid.uuid4())) + ) + self.assertEqual((status, body["code"]), (409, "storage-changed")) + self.assertEqual(body["library_epoch"], self.epoch) + status, body = await self.suggest( + self.suggest_body( + observed, target={**observed["target"], "root_scope": "0" * 64} + ) + ) + self.assertEqual((status, body["code"]), (409, "source-changed")) + self.api._captures.release(observed["capture"]["token"]) + status, body = await self.suggest(self.suggest_body(observed)) + self.assertEqual((status, body["code"]), (409, "capture-expired")) + self.assertEqual(self.tagger.calls, []) + for token in (observed, other): + entry = self.api._captures._entries.get(token["capture"]["token"]) + self.assertTrue(entry is None or entry.uses == 0) + + async def test_video_and_oversized_dimensions_unsupported_media(self): + (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 4) + video = await self.observe_ok("clip.mp4") + self.assertFalse(video["capture"]["generate"]["supported"]) + big = await self.observed("big.png", size=(4097, 1)) + self.assertFalse(big["capture"]["generate"]["supported"]) + + for observed in (video, big): + status, body = await self.suggest(self.suggest_body(observed)) + self.assertEqual((status, body["code"]), (415, "unsupported-media")) + self.assertEqual(body["error"], observed["capture"]["generate"]["reason"]) + self.assertEqual(self.tagger.calls, []) + self.assertEqual(self.tagger.load_calls, []) + self.assertIsNone(model_gate.holder()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_api_annotations.py b/tests/test_api_annotations.py index 2df315fa..624b41b2 100644 --- a/tests/test_api_annotations.py +++ b/tests/test_api_annotations.py @@ -986,6 +986,7 @@ class TestEnvelopes(AnnotationAPITestCase): "model_loaded", "models", "success", + "suggestion_limit", # B2 adds it (D68); nothing else changes "wd14_character_threshold", "wd14_general_threshold", ], diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 832221dd..6b9e2d7d 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -37,6 +37,11 @@ from database.operations import PromptDatabase from py.api import PromptManagerAPI from utils.hashing import generate_prompt_hash +try: # guarded so a missing module fails assertions, not test loading + from py import model_gate +except ImportError: + model_gate = None + class FakeAutotagService: """Stand-in for AutoTagService that records calls and returns fixed tags.""" @@ -128,6 +133,11 @@ class AutotagAPITestCase(AioHTTPTestCase): async def tearDownAsync(self): await super().tearDownAsync() # closes the aiohttp test client self._service_patch.stop() + if model_gate is not None: + leaked = model_gate.holder() + if leaked is not None: # never leak the process-wide gate + model_gate.release(leaked) + self.assertIsNone(leaked, "a route left the model gate held") self.api.db.close_all() for path in ( self._temp_db.name, @@ -947,6 +957,159 @@ class TestUnloadAutotag(AutotagAPITestCase): self.assertEqual(resp.status, 500) +class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase): + """While another model operation holds the gate, every legacy model route + answers ``409 model-busy`` and touches no model (the one allowed legacy + change).""" + + BUSY = {"success": False, "code": "model-busy"} + + def setUp(self): + super().setUp() + self.assertIsNotNone(model_gate, "py.model_gate is not implemented") + + async def _busy(self, resp): + self.assertEqual(resp.status, 409) + data = await resp.json() + self.assertEqual({k: data[k] for k in self.BUSY}, self.BUSY) + self.assertIn("busy", data["error"].lower()) + return data + + async def test_legacy_single_download_start_return_model_busy_while_gate_held( + self, + ): + image = self._make_image(self.output_dir) + self.service._loaded_type = "wd14-vit" + self.assertTrue(model_gate.try_acquire("test-holder")) + try: + await self._busy( + await self._post_json( + "/prompt_manager/autotag/single", + {"path": str(image), "model_type": "gguf"}, + ) + ) + await self._busy( + await self.client.request( + "POST", "/prompt_manager/autotag/download/gguf" + ) + ) + await self._busy( + await self.client.request( + "POST", "/prompt_manager/autotag/start?model_type=gguf" + ) + ) + await self._busy( + await self.client.request("POST", "/prompt_manager/autotag/unload") + ) + self.assertEqual(model_gate.holder(), "test-holder") + finally: + model_gate.release("test-holder") + + self.assertEqual(self.service.load_calls, []) + self.assertEqual(self.service.generate_calls, []) + self.assertEqual(self.service.download_calls, []) + self.assertEqual(self.service.unload_calls, 0) + self.assertEqual(self.service.get_loaded_model_type(), "wd14-vit") + + # A refused start released the batch lock: the next start runs. + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start?model_type=gguf" + ) + self.assertEqual(resp.status, 200) + self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete") + + async def test_legacy_routes_hold_the_gate_only_while_they_run(self): + image = self._make_image(self.output_dir) + seen = [] + + def generate(image_path, **kwargs): + seen.append(model_gate.holder()) + return ["x"] + + def download(model_type, progress_callback=None): + seen.append(model_gate.holder()) + return True + + def load(model_type, use_gpu=True): + seen.append(model_gate.holder()) + self.service._loaded_type = model_type + return True + + self.service.generate_tags = generate + self.service.download_model = download + self.service.load_model = load + + resp = await self._post_json( + "/prompt_manager/autotag/single", {"path": str(image)} + ) + self.assertEqual(resp.status, 200) + self.assertIsNone(model_gate.holder()) + resp = await self.client.request( + "POST", "/prompt_manager/autotag/download/gguf" + ) + self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete") + self.assertIsNone(model_gate.holder()) + resp = await self.client.request("POST", "/prompt_manager/autotag/unload") + self.assertEqual(resp.status, 200) + self.assertIsNone(model_gate.holder()) + + self.assertEqual(len(seen), 3) + self.assertTrue(all(owner and owner.startswith("legacy-") for owner in seen)) + + async def test_validation_errors_win_over_the_gate(self): + self.assertTrue(model_gate.try_acquire("test-holder")) + try: + resp = await self._post_json("/prompt_manager/autotag/single", {}) + self.assertEqual(resp.status, 400) + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start?general_threshold=x" + ) + self.assertEqual(resp.status, 400) + finally: + model_gate.release("test-holder") + + async def test_batch_holds_the_gate_until_it_completes(self): + prompt = self._save_prompt("p") + self._link_image(prompt, self._make_image(self.output_dir, "a.png")) + started = threading.Event() + proceed = threading.Event() + + def slow_generate(image_path, **kwargs): + started.set() + proceed.wait(30) + return ["slow"] + + self.service.generate_tags = slow_generate + run = asyncio.ensure_future( + self.client.request("POST", "/prompt_manager/autotag/start") + ) + try: + await asyncio.get_running_loop().run_in_executor(None, started.wait, 30) + self.assertTrue(model_gate.holder().startswith("legacy-batch:")) + await self._busy( + await self.client.request("POST", "/prompt_manager/autotag/unload") + ) + finally: + proceed.set() + resp = await run + events = self._sse_events(await resp.text()) + self.assertEqual(events[-1]["type"], "complete") + self.assertIsNone(model_gate.holder()) + + async def test_models_response_carries_the_suggestion_limit(self): + from py.config import AutotagConfig + + resp = await self.client.get("/prompt_manager/autotag/models") + data = await resp.json() + self.assertEqual(data["suggestion_limit"], 10) + AutotagConfig.SUGGESTION_LIMIT = 25 + try: + resp = await self.client.get("/prompt_manager/autotag/models") + self.assertEqual((await resp.json())["suggestion_limit"], 25) + finally: + AutotagConfig.SUGGESTION_LIMIT = 10 + + class TestScanOutputDir(AutotagAPITestCase): async def test_lists_images_with_thumbnails(self): diff --git a/tests/test_autotag_engine.py b/tests/test_autotag_engine.py index ddc2c957..a16f8ba2 100644 --- a/tests/test_autotag_engine.py +++ b/tests/test_autotag_engine.py @@ -301,6 +301,143 @@ class TestGgufNormalisation(EngineTestCase): ) +SCORED_TAGS = [ + {"name": "desert", "category": "0"}, + {"name": "sandstone", "category": "0"}, + {"name": "Sandstone", "category": "0"}, + {"name": "hatsune_miku", "category": "4"}, + {"name": "rating:safe", "category": "9"}, +] + [{"name": f"tag_{i:03d}", "category": "0"} for i in range(120)] +SCORED_PROBABILITIES = [0.96, 0.98, 0.5, 0.9, 0.99] + [ + 0.9 - i * 0.004 for i in range(120) +] + + +class TestTagSuggestions(EngineTestCase): + """``generate_tag_suggestions``: scored, ordered, bounded, distinct.""" + + def _wd14(self, tags=SCORED_TAGS, probabilities=SCORED_PROBABILITIES): + _touch_model_files(self.models_dir, "wd14-vit") + self.session = FakeOnnxSession(probabilities) + service = self._service(("wd14", (self.session, tags))) + service.load_model("wd14-vit") + return service + + def _gguf(self, text): + _touch_model_files(self.models_dir, "gguf") + self.model = FakeGgufModel(text) + service = self._service(("gguf", self.model)) + service.load_model("gguf") + return service + + def test_wd14_suggestions_carry_real_probabilities_sorted(self): + service = self._wd14() + + suggestions = service.generate_tag_suggestions( + str(self.image_path), + limit=10, + general_threshold=0.35, + character_threshold=0.85, + ) + + self.assertEqual(suggestions[:2], [("sandstone", 0.98), ("desert", 0.96)]) + # The character tag passes its own threshold; rating tags never show. + self.assertIn(("hatsune_miku", 0.9), suggestions) + self.assertNotIn("rating:safe", [name for name, _ in suggestions]) + scores = [score for _, score in suggestions] + self.assertEqual(scores, sorted(scores, reverse=True)) + self.assertTrue(all(isinstance(score, float) for score in scores)) + # "Sandstone" (0.5) repeats "sandstone" case-insensitively: dropped. + self.assertEqual( + [name.casefold() for name, _ in suggestions].count("sandstone"), 1 + ) + + def test_suggestions_keep_the_requested_limit(self): + service = self._wd14() + for limit in (1, 10, 100): + with self.subTest(limit=limit): + suggestions = service.generate_tag_suggestions( + str(self.image_path), + limit=limit, + general_threshold=0.0, + character_threshold=0.0, + ) + self.assertEqual(len(suggestions), limit) + for bad in (0, -1, 101, 2.5, "10", True, None): + with self.subTest(bad=bad), self.assertRaises(ValueError): + service.generate_tag_suggestions(str(self.image_path), limit=bad) + + def test_thresholds_default_to_the_service_values(self): + service = self._wd14() + service.wd14_general_threshold = 0.95 + service.wd14_character_threshold = 0.95 + + suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10) + + self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)]) + + def test_llm_suggestions_have_null_confidence_in_output_order(self): + service = self._gguf("Zebra, apple, ZEBRA, mango tree, kiwi, plum") + + suggestions = service.generate_tag_suggestions( + str(self.image_path), limit=3, prompt="tags please" + ) + + self.assertEqual( + suggestions, [("zebra", None), ("apple", None), ("mango_tree", None)] + ) + + def test_generate_tags_return_shape_unchanged(self): + service = self._wd14() + self.assertEqual( + service.generate_tags( + str(self.image_path), general_threshold=0.95, character_threshold=0.95 + ), + ["sandstone", "desert"], + ) + gguf = self._gguf("Smile, smile, long hair") + self.assertEqual( + gguf.generate_tags(str(self.image_path)), ["smile", "long_hair"] + ) + + def test_explicit_prompt_does_not_mutate_custom_prompt(self): + service = self._gguf("tag") + service.custom_prompt = "saved legacy prompt" + + service.generate_tag_suggestions( + str(self.image_path), limit=10, prompt="dialog instructions" + ) + + self.assertEqual( + self.model.calls[0]["messages"][1]["content"][0]["text"], + "dialog instructions", + ) + self.assertEqual(service.custom_prompt, "saved legacy prompt") + + def test_suggestions_need_a_loaded_model(self): + service = self._service(("gguf", FakeGgufModel("x"))) + with self.assertRaises(RuntimeError): + service.generate_tag_suggestions(str(self.image_path), limit=10) + + def test_image_file_is_closed_after_inference(self): + service = self._gguf("tag") + opened = [] + real_open = autotag.Image.open + + def tracking_open(*args, **kwargs): + image = real_open(*args, **kwargs) + opened.append(image) + return image + + with patch.object(autotag.Image, "open", tracking_open): + service.generate_tag_suggestions(str(self.image_path), limit=10) + service.generate_tags(str(self.image_path)) + + self.assertEqual(len(opened), 2) + for image in opened: + self.assertIsNone(getattr(image, "fp", None)) + + class TestGenerateTagsErrors(EngineTestCase): def test_no_model_loaded(self): diff --git a/tests/test_config.py b/tests/test_config.py index 5ddc9ba3..795de80c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -19,6 +19,7 @@ sys.modules.setdefault("server", _mock_server) sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from py import config as config_module from py.config import GalleryConfig, IntegrationConfig, PromptManagerConfig @@ -404,5 +405,73 @@ class TestConfigPathOverride(unittest.TestCase): self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 77) +class TestAutotagConfig(unittest.TestCase): + """``autotag.defaults.suggestion_limit``: integer 1..100, default 10 (D68).""" + + def setUp(self): + self.config_class = getattr(config_module, "AutotagConfig", None) + self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented") + self._orig = PromptManagerConfig.get_config() + + def tearDown(self): + self.config_class.SUGGESTION_LIMIT = 10 + PromptManagerConfig.update_config(self._orig) + + def _load(self, value): + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "config.json") + with open(path, "w") as f: + # json writes NaN for float("nan"), which json.load reads back. + json.dump({"autotag": {"defaults": {"suggestion_limit": value}}}, f) + with self.assertLogs("prompt_manager.config", level="WARNING") as logs: + PromptManagerConfig.load_from_file(path) + return logs + + def test_default_is_ten_and_persisted_under_autotag_defaults(self): + self.assertEqual(config_module.DEFAULT_SUGGESTION_LIMIT, 10) + self.assertEqual(config_module.MAX_SUGGESTION_LIMIT, 100) + self.assertEqual(self.config_class.suggestion_limit(), 10) + self.assertEqual( + PromptManagerConfig.get_config()["autotag"], + {"defaults": {"suggestion_limit": 10}}, + ) + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "config.json") + PromptManagerConfig.update_config( + {"autotag": {"defaults": {"suggestion_limit": 25}}} + ) + PromptManagerConfig.save_to_file(path) + self.config_class.SUGGESTION_LIMIT = 10 + PromptManagerConfig.load_from_file(path) + self.assertEqual(self.config_class.suggestion_limit(), 25) + + def test_valid_limits_are_kept(self): + for value in (1, 10, 100): + with self.subTest(value=value): + PromptManagerConfig.update_config( + {"autotag": {"defaults": {"suggestion_limit": value}}} + ) + self.assertEqual(self.config_class.suggestion_limit(), value) + + def test_suggestion_limit_setting_validated_on_load(self): + for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]): + with self.subTest(value=value): + self.config_class.SUGGESTION_LIMIT = 42 + logs = self._load(value) + self.assertEqual(self.config_class.suggestion_limit(), 10) + self.assertIn("suggestion_limit", "\n".join(logs.output)) + + def test_missing_or_malformed_sections_keep_the_current_value(self): + self.config_class.SUGGESTION_LIMIT = 7 + for config in ({}, {"autotag": {}}, {"autotag": {"defaults": {}}}): + PromptManagerConfig.update_config(config) + self.assertEqual(self.config_class.suggestion_limit(), 7) + with self.assertLogs("prompt_manager.config", level="WARNING"): + PromptManagerConfig.update_config({"autotag": "nope"}) + with self.assertLogs("prompt_manager.config", level="WARNING"): + PromptManagerConfig.update_config({"autotag": {"defaults": [1]}}) + self.assertEqual(self.config_class.suggestion_limit(), 7) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_media_capture.py b/tests/test_media_capture.py index e10a6470..69e43fa5 100644 --- a/tests/test_media_capture.py +++ b/tests/test_media_capture.py @@ -596,6 +596,42 @@ class CaptureStoreTestCase(unittest.TestCase): cap.sha256, hashlib.sha256(b"\xff\xd8not really a jpeg").hexdigest() ) + def test_metadata_over_inspection_budgets_does_not_block_generate(self): + # B2 carry-over (A M1): inferable is decided from the image header + # alone, so prompt metadata over the inspection budgets (here a 3 MiB + # text chunk) leaves a decodable image inferable. + base = png_bytes(64, 32) + huge = _chunk(b"tEXt", b"prompt\x00" + b"x" * (3 * 1024 * 1024)) + self.write("big-metadata.png", base[:-12] + huge + base[-12:]) + store = self.make_store() + + cap = self.capture(store, "big-metadata.png") + + self.assertEqual((cap.width, cap.height), (64, 32)) + self.assertTrue(cap.inferable) + self.assertIsNone(cap.inferable_reason) + assert_closed(self, self.parent) + + def test_decompression_bomb_header_is_not_inferable(self): + self.write("bomb.png", png_bytes(20_000, 20_000)) + store = self.make_store() + + cap = self.capture(store, "bomb.png") + + self.assertFalse(cap.inferable) + self.assertIn("4096", cap.inferable_reason) + assert_closed(self, self.parent) + + def test_header_dimensions_reads_only_the_header(self): + path = self.write("dims.png", png_bytes(7, 5)) + self.assertEqual(capture_module.header_dimensions(str(path)), (7, 5)) + assert_closed(self, self.parent) + with self.assertRaises(capture_module.HeaderUnreadable): + capture_module.header_dimensions( + str(self.write("bad.png", b"\x89PNG\r\n\x1a\nnope")) + ) + assert_closed(self, self.parent) + def test_video_capture_is_not_inferable(self): data = b"\x00\x00\x00\x18ftypmp42" + os.urandom(32) self.write("clip.mp4", data) @@ -604,7 +640,7 @@ class CaptureStoreTestCase(unittest.TestCase): def never(path): # pragma: no cover - guard only raise AssertionError("video captures are never probed") - with patch.object(capture_module, "read_inspection_file", never): + with patch.object(capture_module, "header_dimensions", never): cap = self.capture(store, "clip.mp4") self.assertEqual((cap.media_type, cap.width, cap.height), ("video", None, None)) self.assertFalse(cap.inferable) diff --git a/tests/test_model_gate.py b/tests/test_model_gate.py new file mode 100644 index 00000000..8b92d4b8 --- /dev/null +++ b/tests/test_model_gate.py @@ -0,0 +1,176 @@ +"""The shared, non-blocking model gate (#249 B2). + +One owner at a time may load, swap, run, download or unload a tagging model. +A release requested while gated work still runs in a worker thread takes +effect only when that work really exits, so a client disconnect never lets a +second owner swap the model under a running inference. +""" + +import asyncio +import os +import sys +import threading +import unittest +import unittest.mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +try: # guarded so a missing module fails assertions, not test loading + from py import model_gate +except ImportError: + model_gate = None + +WAIT = 30.0 # ceiling for cross-thread handoffs, never a short clock + + +class GateTestCase(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(model_gate, "py.model_gate is not implemented") + self.gate = model_gate.ModelGate() + + +class TestGate(GateTestCase): + def test_gate_is_nonblocking_and_single_owner(self): + self.assertIsNone(self.gate.holder()) + self.assertTrue(self.gate.try_acquire("legacy-batch")) + self.assertEqual(self.gate.holder(), "legacy-batch") + + # A second owner is refused at once, never queued. + self.assertFalse(self.gate.try_acquire("image-annotations")) + self.assertFalse(self.gate.try_acquire("legacy-batch")) + self.assertEqual(self.gate.holder(), "legacy-batch") + + # Only the holder can release. + self.gate.release("image-annotations") + self.assertEqual(self.gate.holder(), "legacy-batch") + self.gate.release("legacy-batch") + self.assertIsNone(self.gate.holder()) + self.gate.release("legacy-batch") # idempotent + + self.assertTrue(self.gate.try_acquire("image-annotations")) + self.assertEqual(self.gate.holder(), "image-annotations") + + def test_release_waits_for_running_work(self): + self.assertTrue(self.gate.try_acquire("owner")) + self.gate.begin("owner") + self.gate.release("owner") + self.assertEqual(self.gate.holder(), "owner") + self.assertFalse(self.gate.try_acquire("other")) + # No new work is admitted once the owner let go. + with self.assertRaises(model_gate.GateNotHeld): + self.gate.begin("owner") + self.gate.end("owner") + self.assertIsNone(self.gate.holder()) + + def test_begin_requires_the_holder(self): + with self.assertRaises(model_gate.GateNotHeld): + self.gate.begin("nobody") + self.assertTrue(self.gate.try_acquire("owner")) + with self.assertRaises(model_gate.GateNotHeld): + self.gate.begin("other") + self.gate.end("other") # a stray end never disturbs the holder + self.assertEqual(self.gate.holder(), "owner") + + def test_module_gate_is_one_process_wide_instance(self): + self.assertIsNone(model_gate.holder()) + self.assertTrue(model_gate.try_acquire("a")) + try: + self.assertFalse(model_gate.try_acquire("b")) + self.assertEqual(model_gate.holder(), "a") + finally: + model_gate.release("a") + self.assertIsNone(model_gate.holder()) + + def test_new_owner_names_are_unique(self): + first = model_gate.new_owner("legacy-single") + second = model_gate.new_owner("legacy-single") + self.assertNotEqual(first, second) + self.assertTrue(first.startswith("legacy-single:")) + + +class TestGatedExecutor(GateTestCase): + def test_cancelled_waiter_keeps_the_gate_until_the_thread_exits(self): + started = threading.Event() + proceed = threading.Event() + finished = threading.Event() + + def blocking_work(): + started.set() + proceed.wait(WAIT) + finished.set() + return "done" + + async def scenario(): + self.assertTrue(self.gate.try_acquire("native")) + task = asyncio.ensure_future(self.gate.run("native", blocking_work)) + try: + await asyncio.get_running_loop().run_in_executor( + None, started.wait, WAIT + ) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + finally: + self.gate.release("native") + # The waiter is gone but the inference still runs: still held. + self.assertEqual(self.gate.holder(), "native") + self.assertFalse(self.gate.try_acquire("legacy-unload")) + proceed.set() + await asyncio.get_running_loop().run_in_executor(None, finished.wait, WAIT) + for _ in range(1000): + if self.gate.holder() is None: + break + await asyncio.sleep(0.01) + self.assertIsNone(self.gate.holder()) + + asyncio.run(scenario()) + + def test_run_returns_the_result_and_propagates_errors(self): + async def scenario(): + self.assertTrue(self.gate.try_acquire("owner")) + try: + self.assertEqual(await self.gate.run("owner", lambda x: x * 2, 21), 42) + + def boom(): + raise ValueError("bad") + + with self.assertRaises(ValueError): + await self.gate.run("owner", boom) + self.assertEqual(self.gate.holder(), "owner") + finally: + self.gate.release("owner") + self.assertIsNone(self.gate.holder()) + with self.assertRaises(model_gate.GateNotHeld): + await self.gate.run("owner", lambda: None) + + asyncio.run(scenario()) + + def test_work_that_never_starts_gives_its_admission_back(self): + async def scenario(): + loop = asyncio.get_running_loop() + self.assertTrue(self.gate.try_acquire("owner")) + + def refused(executor, func): + raise RuntimeError("executor shut down") + + with unittest.mock.patch.object(loop, "run_in_executor", refused): + with self.assertRaises(RuntimeError): + await self.gate.run("owner", lambda: None) + + def cancelled(executor, func): + future = loop.create_future() + future.cancel() # a loop shutdown drops the queued call + return future + + with unittest.mock.patch.object(loop, "run_in_executor", cancelled): + with self.assertRaises(asyncio.CancelledError): + await self.gate.run("owner", lambda: None) + await asyncio.sleep(0) # let the done callback run + self.gate.release("owner") + self.assertIsNone(self.gate.holder()) + + asyncio.run(scenario()) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/media_capture.py b/utils/media_capture.py index 86988dbe..f9a54153 100644 --- a/utils/media_capture.py +++ b/utils/media_capture.py @@ -28,16 +28,13 @@ import threading import time import warnings from dataclasses import dataclass -from pathlib import Path from typing import Callable from PIL import Image try: - from .prompt_inspector import MetadataReadError, read_inspection_file from .video_metadata import VIDEO_EXTENSIONS, is_video_path except ImportError: # imported with the repo root on sys.path - from utils.prompt_inspector import MetadataReadError, read_inspection_file from utils.video_metadata import VIDEO_EXTENSIONS, is_video_path logger = logging.getLogger(__name__) @@ -578,20 +575,17 @@ class CaptureStore: def _probe(self, path: str, media_type: str): """Header-only dimensions of the capture; never decodes pixels. - Uses the inspection reader, so a PNG whose prompt metadata is over the - inspection budgets is reported as unreadable (not inferable) although - its pixels decode; B2 may switch to a dimensions-only header read. + Only the image header decides whether Generate is offered: prompt + metadata over the inspection budgets does not make a decodable image + unreadable (the metadata itself is read elsewhere, under its budgets). """ if media_type == "video": return None, None, False, VIDEO_ONLY_REASON try: - with warnings.catch_warnings(): - warnings.simplefilter("ignore", Image.DecompressionBombWarning) - read = read_inspection_file(Path(path)) - except (MetadataReadError, Image.DecompressionBombError): - return None, None, False, HEADER_REASON - width, height = read.width, read.height - if width is None or height is None: + width, height = header_dimensions(path) + except HeaderTooLarge: + return None, None, False, DIMENSION_REASON + except HeaderUnreadable: return None, None, False, HEADER_REASON if ( width > MAX_INFERENCE_SIDE @@ -602,6 +596,34 @@ class CaptureStore: return width, height, True, None +class HeaderUnreadable(Exception): + """The file has no image header Pillow can read.""" + + +class HeaderTooLarge(HeaderUnreadable): + """The header declares a decompression bomb (no pixel is decoded).""" + + +def header_dimensions(path: str) -> tuple[int, int]: + """``(width, height)`` from the image header alone; the file is closed. + + ``Image.open`` reads only the header; a decompression-bomb warning is an + error here, so an absurd size is refused before anything is decoded. + """ + try: + with warnings.catch_warnings(): + warnings.simplefilter("error", Image.DecompressionBombWarning) + with Image.open(path) as image: + width, height = image.size + except (Image.DecompressionBombError, Image.DecompressionBombWarning) as exc: + raise HeaderTooLarge(str(exc)) from None + except Exception as exc: # any decoder failure means "no usable header" + raise HeaderUnreadable(str(exc)) from None + if not width or not height: + raise HeaderUnreadable("The header declares no size") + return int(width), int(height) + + def _relative_parts(relative) -> tuple[list[str], str]: """Split a safe relative POSIX media path; anything else is denied.""" if not isinstance(relative, str) or not relative: