diff --git a/TESTING.md b/TESTING.md index e0cf324..b7f56fa 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1389,6 +1389,12 @@ impact and command receipts report `image_count` / `images` / (spec amendment S1) use a byte record only while the file's current `(size, mtime_ns)` equals its verified fingerprint; an in-place rewrite that preserves both keeps the old set matching until the file is next observed. The +Images Tags filter (a file matches when every tag is on its linked positive or +its current image tag set) and `q` search stream the tagged records once and +stat only the listed files whose records carry a filter tag or a `q` hit; the +inspector's `image_tags` facet (`current` / `none` / `unavailable`) uses the +inspection's own stable stat (`tests/test_api_workspace_images.py` +`TestWorkspaceImageTags`, `tests/test_api_prompt_inspector.py`). The browser fixtures clear `media_image_tags`, `media_annotation_bytes` and `media_annotation_files` on reset and never the metadata row. @@ -1421,6 +1427,22 @@ window completes there with a digest of the copied bytes; POSIX and Windows on Vocabulary rename, merge and delete bump the `tag_revision` of every image set they change (and `image_index_revision`), so a Tag dialog opened before the change gets `tag-conflict` on Save (`tests/test_pair_tag_maintenance.py`). +A case or Unicode spelling of a listed name (`A.PNG` for `a.png`, an NFD name) +is an alias too: each path part must be spelled as its directory lists it, so a +case- or normalization-insensitive filesystem (macOS, Windows) cannot store a +second locator. The capture preview answers one byte range with 206 (416 for +an unsatisfiable or multi-range header) inside the capture's pin, and a +cancelled preview, observe or Generate still closes its handle and ends its pin +or capture (`tests/test_api_annotations.py` `TestCancellation`, +`tests/test_model_gate.py`). + +Follow-ups (tracked, not accepted): the legacy Auto-tag batch's cleanup +`unload_model()` in `py/api/autotag_routes.py` is not gated against the +batch's own in-flight inference when the batch coroutine is cancelled (other +owners are kept out; the thread may log an error); and `py/autotag.py`, +`py/api/images.py` and `database/operations.py` are over the 800-line ceiling +(all three were before #249 PR B), with a split of `py/autotag.py` into +backends and service planned first. Budgets: `test_filter_work_is_bounded` always runs. It asserts that a plain first page makes 1 link SELECT and a plain cursor page none, that every filter shape diff --git a/__init__.py b/__init__.py index f494589..140f7ea 100644 --- a/__init__.py +++ b/__init__.py @@ -90,6 +90,7 @@ try: host_app = getattr(config.server_instance, "app", None) if host_app is not None and hasattr(host_app, "on_cleanup"): host_app.on_cleanup.append(api.cleanup_history_reads) + host_app.on_cleanup.append(api.cleanup_annotation_captures) except Exception as e: # Log to internal logging system without console spam diff --git a/database/annotations.py b/database/annotations.py index 0a39be2..219a503 100644 --- a/database/annotations.py +++ b/database/annotations.py @@ -119,12 +119,16 @@ class AnnotationCommandError(LibraryCommandError): self.current = current +class ImageTagLimitError(ValueError): + """More than ``MAX_IMAGE_TAGS`` distinct names in one set.""" + + def normalize_image_tags(names: object) -> list[str]: """Trim and validate each name, fold case-insensitive repeats, cap at 100. The first occurrence keeps its spelling and position. Raises ``ValueError`` - for anything but a list of valid vocabulary names or for more than - ``MAX_IMAGE_TAGS`` distinct names. + for anything but a list of valid vocabulary names, and its subclass + ``ImageTagLimitError`` for more than ``MAX_IMAGE_TAGS`` distinct names. """ if not isinstance(names, list): raise ValueError("tags must be a list of strings") @@ -138,7 +142,7 @@ def normalize_image_tags(names: object) -> list[str]: seen.add(key) result.append(name) if len(result) > MAX_IMAGE_TAGS: - raise ValueError(f"An image can carry at most {MAX_IMAGE_TAGS} tags") + raise ImageTagLimitError(f"An image can carry at most {MAX_IMAGE_TAGS} tags") return result diff --git a/database/operations.py b/database/operations.py index 8fa24dd..28bf35d 100644 --- a/database/operations.py +++ b/database/operations.py @@ -232,6 +232,23 @@ JOIN prompts p ON p.id = COALESCE(pp.positive_prompt_id, gi.prompt_id) WHERE gi.file_path IS NOT NULL AND CASE WHEN {admit} THEN NOT EXISTS ( SELECT 1 FROM generated_images other WHERE other.file_path = gi.file_path AND other.id <> gi.id) ELSE 0 END""" +# The linked positive of each given path (one unambiguous link, pair scope as +# in _LINK_FILTER_SQL) with its category; {hit} as in link_filter_query. Used +# for the few files whose current image tags meet a Tags filter (#249 B3). +_LINK_FACTS_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id, p.category +FROM generated_images gi +LEFT JOIN prompt_pairs pp ON pp.id = ( + SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END + FROM image_prompt_pairs ip WHERE ip.image_id = gi.id) +JOIN prompts p ON p.id = COALESCE(pp.positive_prompt_id, gi.prompt_id) +WHERE gi.file_path IN ({marks}) AND NOT EXISTS ( + SELECT 1 FROM generated_images other + WHERE other.file_path = gi.file_path AND other.id <> gi.id)""" +_POSITIVE_TAG_NAMES_SQL = ( + "SELECT pt.prompt_id, t.name FROM prompt_tags pt JOIN tags t " + "ON t.id = pt.tag_id WHERE pt.prompt_id IN ({marks})" +) +_FACTS_CHUNK = 500 # The positive's text, its Library title or one of its tag names contains # :pattern (LIKE, so ASCII case is ignored as in Library search). A pair's own # title names its negative variant and is never searched (#249 D12). @@ -2450,6 +2467,46 @@ class PromptDatabase( matches[path] = (bool(matched), rating, positive) return matches + def get_image_link_facts( + self, paths: List[str], q: str = "" + ) -> Dict[str, Tuple[bool, Optional[int], int, Optional[str], Tuple[str, ...]]]: + """Map each normalized path with one unambiguous link to its positive's + (hit, rating, positive id, category, tag names). + + The link is resolved as in get_image_link_matches, and ``hit`` says + whether *q* occurs in the positive's text, Library title or a tag + name (always true without *q*). Paths go in chunks of 500 through the + ``file_path`` index; the tag names are read for those positives only. + """ + hit, params = "1", {} + if q: + hit = _LINK_FILTER_HIT + params["pattern"] = f"%{escape_like(q)}%" + rows: Dict[str, Tuple[bool, Optional[int], int, Optional[str]]] = {} + names: Dict[int, List[str]] = {} + unique = list(dict.fromkeys(paths)) + with self.model.get_connection() as conn: + for start in range(0, len(unique), _FACTS_CHUNK): + chunk = unique[start : start + _FACTS_CHUNK] + bound = {f"path{i}": path for i, path in enumerate(chunk)} + marks = ",".join(f":{key}" for key in bound) + sql = _LINK_FACTS_SQL.format(hit=hit, marks=marks) + for path, rating, matched, positive, category in conn.execute( + sql, {**params, **bound} + ): + rows[path] = (bool(matched), rating, positive, category) + positives = sorted({row[2] for row in rows.values()}) + for start in range(0, len(positives), _FACTS_CHUNK): + chunk = positives[start : start + _FACTS_CHUNK] + marks = ",".join("?" * len(chunk)) + for positive, name in conn.execute( + _POSITIVE_TAG_NAMES_SQL.format(marks=marks), chunk + ): + names.setdefault(positive, []).append(name) + return { + path: (*row, tuple(names.get(row[2], ()))) for path, row in rows.items() + } + def get_image_paths_matching( self, q: str, category: Optional[str], tags: List[str] ) -> Set[str]: diff --git a/py/api/__init__.py b/py/api/__init__.py index ce20699..b45b3d5 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -26,6 +26,7 @@ from .prompts import PromptRoutesMixin from .pairs import PairRoutesMixin from .library import LibraryRoutesMixin from .vocabulary import VocabularyRoutesMixin +from .annotations import AnnotationRoutesMixin from .saved_views import SavedViewRoutesMixin from .history import HistoryRoutesMixin from .images import ImageRoutesMixin @@ -252,6 +253,7 @@ class PromptManagerAPI( PairRoutesMixin, LibraryRoutesMixin, VocabularyRoutesMixin, + AnnotationRoutesMixin, SavedViewRoutesMixin, HistoryRoutesMixin, PromptRoutesMixin, @@ -285,6 +287,10 @@ class PromptManagerAPI( self.db = db if db is not None else PromptDatabase() self._history_read_fence = None self._history_reads_closed = False + # Private media captures of the annotation routes (py/api/annotations.py), + # created on first use with their restore subscription. + self._captures = None + self._captures_unsubscribe = None self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir() self._html_cache = {} # Cached HTML file contents keyed by path self._gallery_cache = {} # dict: path_str -> (files, timestamp) @@ -441,6 +447,7 @@ class PromptManagerAPI( self._register_pair_routes(routes) self._register_library_routes(routes) self._register_vocabulary_routes(routes) + self._register_annotation_routes(routes) self._register_saved_view_routes(routes) self._register_history_routes(routes) self._register_prompt_routes(routes) diff --git a/py/api/annotation_suggest.py b/py/api/annotation_suggest.py new file mode 100644 index 0000000..563bdc6 --- /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 new file mode 100644 index 0000000..1e08e8c --- /dev/null +++ b/py/api/annotations.py @@ -0,0 +1,759 @@ +"""Image annotation routes: observe, capture preview and byte-exact tag saves. + +Every route names its file through the listing identity ``{root_index, +root_scope, path}`` and checks the scope against the current gallery roots +before any file access. Observe copies the file into a private capture (the +token the preview and Save refer to), Save re-hashes the current file and +writes only when the bytes are still the captured ones. Responses carry no +server path and no locator. +""" + +import asyncio +import datetime +import functools +import os +import re +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 ( + MAX_IMAGE_TAGS, + AnnotationCommandError, + ImageTagLimitError, + normalize_image_tags, + ) + from ...database.execution_history import subscribe_history_restore + from ...database.library_commands import LibraryCommandError + from ...utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore +except ImportError: + from database.annotations import ( + MAX_IMAGE_TAGS, + AnnotationCommandError, + ImageTagLimitError, + normalize_image_tags, + ) + from database.execution_history import subscribe_history_restore + from database.library_commands import LibraryCommandError + from utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore + +# Request bodies of the annotation routes are small (a target plus a tag set). +MAX_BODY_BYTES = 64 * 1024 +MAX_TARGET_PATH = 1024 +CAPTURE_URL = "/prompt_manager/image-annotations/capture/" +_PREVIEW_CHUNK = 64 * 1024 + +_TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"}) +_OBSERVE_FIELDS = frozenset({"target"}) +_SAVE_FIELDS = frozenset( + { + "target", + "capture_token", + "sha256", + "expected_epoch", + "expected_tag_revision", + "tags", + } +) +_HEX64 = re.compile(r"[0-9a-f]{64}") +_UUID = re.compile(r"[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}") +# secrets.token_urlsafe(32) is 43 characters; the bound only rejects junk. +_TOKEN = re.compile(r"[A-Za-z0-9_-]{1,128}") + +_PREVIEW_TYPES = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".webp": "image/webp", + ".gif": "image/gif", + ".bmp": "image/bmp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".mp4": "video/mp4", + ".m4v": "video/mp4", + ".webm": "video/webm", + ".mov": "video/quicktime", + ".mkv": "video/x-matroska", + ".avi": "video/x-msvideo", + ".wmv": "video/x-ms-wmv", +} + +_STATUS = { + "review-capacity": 429, + "unsupported-media": 415, + "file-changed": 409, + "source-changed": 409, + "storage-changed": 409, + "tag-conflict": 409, + "capture-expired": 409, + "model-unavailable": 409, + "tag-limit": 400, + "metadata-conflict": 409, + "vocabulary-conflict": 409, +} +_MESSAGES = { + "source-changed": ( + "The source changed. Refresh the image collection before tagging this file." + ), + "file-changed": "The file changed since it was previewed. Reopen the Tag dialog.", + "review-capacity": "Too many previews are open. Close a Tag dialog and try again.", + "storage-changed": "Prompt storage was restored; reopen the Tag dialog.", + "tag-conflict": ( + "The tags of this image changed since the dialog opened; review them first." + ), + "capture-expired": "The preview expired. Open the Tag dialog again.", + "tag-limit": f"An image can carry at most {MAX_IMAGE_TAGS} tags.", +} +# Any other storage code (a conflict a later command adds) still answers a +# coded envelope rather than a bare 500. +_DEFAULT_STATUS = 409 +_DEFAULT_MESSAGE = "The tags could not be saved. Reopen the Tag dialog and try again." +_OTHER_FILE = "The preview belongs to another file. Open the Tag dialog again." + + +def preview_content_type(extension: str) -> str: + """The media type the preview is served as; never sniffed from bytes.""" + return _PREVIEW_TYPES.get(extension.lower(), "application/octet-stream") + + +def _error(status: int, code: str, message: str, **extra) -> web.Response: + return web.json_response( + {"success": False, "code": code, "error": message, **extra}, + status=status, + headers={"Cache-Control": "no-store"}, + ) + + +def _coded(code: str, message: str | None = None, **extra) -> web.Response: + status = _STATUS.get(code, _DEFAULT_STATUS) + return _error( + status, code, message or _MESSAGES.get(code, _DEFAULT_MESSAGE), **extra + ) + + +def _open_capture(path: str): + """Worker thread: the capture file opened for reading.""" + return open(path, "rb") + + +async def _settled(future, cleanup): + """Await an executor *future*; a cancelled caller still cleans up after it. + + The future is shielded, so the worker call always runs to its end. When the + caller is cancelled first, ``cleanup(result)`` runs once the call is done + (``cleanup(None)`` when it failed): a handle is closed, a capture released. + """ + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + future.add_done_callback(functools.partial(_cleanup_result, cleanup)) + raise + + +def _cleanup_result(cleanup, future) -> None: + if future.cancelled() or future.exception() is not None: + cleanup(None) + else: + cleanup(future.result()) + + +def _submit(func, *args, **kwargs): + return asyncio.get_running_loop().run_in_executor( + None, functools.partial(func, *args, **kwargs) + ) + + +def _capture_error(exc: CaptureError) -> web.Response: + """The envelope for a capture outcome; messages never name a path.""" + if exc.code not in _STATUS: + return _error(500, exc.code, "The file could not be captured.") + return _coded(exc.code, _MESSAGES.get(exc.code) or str(exc)) + + +def _target(data: object): + """``((index, scope, relative), None)`` for a target, else ``(None, error)``. + + Only the listing identity is accepted: an int root index, the root's + 64-hex scope and a bounded root-relative POSIX media path, never an + absolute, drive-qualified, backslash or ``..`` path. Nothing here touches + the filesystem or the database. + """ + if not isinstance(data, dict) or set(data) != _TARGET_FIELDS: + return None, bad_request("target must name root_index, root_scope and path") + index, scope, relative = data["root_index"], data["root_scope"], data["path"] + if type(index) is not int or index < 0: + return None, bad_request("Invalid gallery root") + if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None: + return None, bad_request("Invalid gallery root scope") + if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH: + return None, bad_request("path must be a relative media path") + parts = relative.split("/") + if ( + "\0" in relative + or "\\" in relative + or relative.startswith("/") + or _has_traversal(relative) + or any(part in ("", ".", "..") or ":" in part for part in parts) + ): + return None, _forbidden() + if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS: + return None, _forbidden("Only media files can be tagged") + return (index, scope, relative), None + + +def _expires_at(seconds: float | None) -> str: + moment = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta( + seconds=max(0.0, seconds or 0.0) + ) + return moment.replace(microsecond=0).isoformat().replace("+00:00", "Z") + + +def _capture_info(store: CaptureStore, capture) -> dict: + return { + "token": capture.token, + "expires_at": _expires_at(store.expires_in(capture.token)), + "preview_url": CAPTURE_URL + capture.token, + "sha256": capture.sha256, + "size": capture.size, + "width": capture.width, + "height": capture.height, + "media_type": capture.media_type, + "generate": { + "supported": capture.inferable, + "reason": capture.inferable_reason, + }, + } + + +class AnnotationRoutesMixin: + """Observe, capture preview and tag save routes over PR A's storage.""" + + def _register_annotation_routes(self, routes: web.RouteTableDef) -> None: + routes.post("/prompt_manager/image-annotations/observe")(self.observe_image) + 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 ------------------------------------------------ + + def _capture_store(self) -> CaptureStore: + """The process's capture store, created on first use. + + Created lazily so an API object that never observes a file (most tests, + a server that never opens the Tag dialog) owns no temp directory. A + restore of the database dooms every capture (``invalidate_all``). + """ + if self._captures is None: + self._captures = CaptureStore() + if self._captures_unsubscribe is None: + self._captures_unsubscribe = subscribe_history_restore( + self.db.model.db_path, self._invalidate_captures + ) + return self._captures + + def _invalidate_captures(self) -> None: + store = self._captures + if store is not None: + store.invalidate_all() + + def close_annotation_captures(self) -> None: + """Remove every capture file and the owned directory. Idempotent.""" + unsubscribe, self._captures_unsubscribe = self._captures_unsubscribe, None + if unsubscribe is not None: + unsubscribe() + store, self._captures = self._captures, None + if store is not None: + store.close() + + async def cleanup_annotation_captures(self, app: web.Application) -> None: + self.close_annotation_captures() + + # -- shared request steps -------------------------------------------------- + + async def _annotation_body(self, request: web.Request, fields: frozenset): + data, error = await self._read_json_body(request, max_bytes=MAX_BODY_BYTES) + if error is not None: + return None, error + if set(data) != fields: + return None, bad_request( + "Request must contain exactly " + ", ".join(sorted(fields)) + ) + return data, None + + def _scoped_root(self, index: int, scope: str): + """``(root, None)`` when the scope names a current root, else an error.""" + roots = list(self._get_all_output_dirs()) + error = self._gallery_scope_error( + {"root_scope": scope, "root": str(index)}, roots + ) + if error is not None: + if error.status == 409: + return None, _coded("source-changed") + return None, error + return roots[index], None + + # -- observe --------------------------------------------------------------- + + async def observe_image(self, request: web.Request) -> web.Response: + """POST ``{target}``: capture the file and answer its saved tag set. + + The capture is the byte-exact copy the preview and Save refer to; the + tags are those of exactly these bytes (an empty set at revision 0 for + unseen bytes). Observing never writes tags; equal bytes refresh the + record's verified fingerprint (S1). + """ + data, error = await self._annotation_body(request, _OBSERVE_FIELDS) + if error is not None: + return error + target, error = _target(data["target"]) + if error is not None: + return error + index, scope, relative = target + root, error = self._scoped_root(index, scope) + if error is not None: + return error + store = self._capture_store() + try: + epoch = await self._run_in_executor(self.db.get_execution_epoch) + capture = await _settled( + _submit( + store.capture, + root, + relative, + epoch=epoch, + scope=scope, + root_index=index, + ), + # A cancelled observe never hands its capture out: release it. + lambda made: made is not None and store.release(made.token), + ) + except CaptureError as exc: + return _capture_error(exc) + except Exception: + self.logger.exception("Image observation failed") + return _error(500, "observe-failed", "The image could not be observed.") + try: + stored, title = await self._run_in_executor( + self._observe_capture, store, capture, root, relative + ) + except asyncio.CancelledError: + store.release(capture.token) + raise + except Exception: + store.release(capture.token) + self.logger.exception("Image observation failed") + return _error(500, "observe-failed", "The image could not be observed.") + return web.json_response( + { + "success": True, + "library_epoch": capture.epoch, + "target": {"root_index": index, "root_scope": scope, "path": relative}, + "capture": _capture_info(store, capture), + "title": title, + "filename": relative.rsplit("/", 1)[-1], + "tags": stored["tags"], + "tag_revision": stored["tag_revision"], + }, + headers={"Cache-Control": "no-store"}, + ) + + def _observe_capture(self, store: CaptureStore, capture, root, relative: str): + """Blocking part of observe: fingerprint refresh, tag read, title.""" + self.db.refresh_image_fingerprint( + capture.locator, capture.sha256, capture.size, capture.mtime_ns + ) + stored = self.db.read_image_tags(capture.locator, capture.sha256) + return stored, self._observed_title(store, capture, root, relative) + + def _observed_title(self, store: CaptureStore, capture, root, relative: str): + """The dialog's title under the card and inspector rule (``_verified_title``). + + A stored Library title always shows; a first-line title only when the + captured bytes' own records confirm it. An ambiguous association has no + title. Videos read no card metadata, so only a stored title shows. + """ + path = str(Path(root) / relative) + summary = self._workspace_link_summaries([path]).get(path) + if summary is None: + summary = _unavailable_link() + positive_title = None + if ( + capture.media_type != "video" + and summary.get("title_source") == "text" + and store.acquire(capture.token) is not None + ): + try: + facts = self._read_card_facts(Path(capture.path), parse=True) + finally: + store.release_use(capture.token) + positive_title = facts.positive_title if facts is not None else None + return _verified_title(summary, positive_title) + + # -- capture preview ------------------------------------------------------- + + async def get_capture_preview(self, request: web.Request) -> web.StreamResponse: + """GET the captured bytes; the token is pinned while they stream.""" + token = request.match_info["token"] + store = self._captures + if store is None or _TOKEN.fullmatch(token) is None: + return _error(404, "capture-expired", _MESSAGES["capture-expired"]) + capture = store.acquire(token) + if capture is None: + return _error(404, "capture-expired", _MESSAGES["capture-expired"]) + pinned = True + try: + _root, error = self._scoped_root(capture.root_index, capture.scope) + if error is not None: + # The roots changed under the token: it is no longer usable. + store.release(token) + return error + try: + handle = await _settled( + _submit(_open_capture, capture.path), + functools.partial(_close_then_release, store, token), + ) + except asyncio.CancelledError: + # Cancelled while the file opens: the pin ends only after the + # handle closes (Windows cannot unlink an open file). + pinned = False + raise + except OSError: + store.release(token) + return _error(404, "capture-expired", _MESSAGES["capture-expired"]) + try: + return await self._stream_capture(request, capture, handle) + finally: + handle.close() + finally: + if pinned: + store.release_use(token) + + @staticmethod + async def _stream_capture(request, capture, handle) -> web.StreamResponse: + """The capture's bytes, or the one byte range a ``Range`` header asks for. + + The response is written here, inside the caller's pin, so the handle + is closed and the pin ended only after the last byte is sent. + """ + size = capture.size + headers = { + "Content-Type": preview_content_type(Path(capture.path).suffix), + "Accept-Ranges": "bytes", + "Cache-Control": "no-store", + "X-Content-Type-Options": "nosniff", + } + status, start, end = 200, 0, size + if "Range" in request.headers: + wanted = _byte_range(request, size) + if wanted is None: + return web.Response( + status=416, + headers={**headers, "Content-Range": f"bytes */{size}"}, + ) + status, (start, end) = 206, wanted + headers["Content-Range"] = f"bytes {start}-{end - 1}/{size}" + headers["Content-Length"] = str(end - start) + response = web.StreamResponse(status=status, headers=headers) + await response.prepare(request) + loop = asyncio.get_running_loop() + if start: + await loop.run_in_executor(None, handle.seek, start) + remaining = end - start + while remaining > 0 and request.method != "HEAD": + chunk = await loop.run_in_executor( + None, handle.read, min(_PREVIEW_CHUNK, remaining) + ) + if not chunk: + break + remaining -= len(chunk) + await response.write(chunk) + await response.write_eof() + return response + + async def release_capture(self, request: web.Request) -> web.Response: + """DELETE a capture token; releasing an unknown token is a success.""" + token = request.match_info["token"] + store = self._captures + if store is not None and _TOKEN.fullmatch(token) is not None: + store.release(token) + return web.json_response( + {"success": True}, headers={"Cache-Control": "no-store"} + ) + + # -- save ------------------------------------------------------------------ + + async def save_image_tags(self, request: web.Request) -> web.Response: + """PUT the tag set of the previewed bytes, re-hashing the current file. + + Writes only when the file still holds the captured bytes + (``file-changed`` with zero writes otherwise), under the storage epoch + and the tag revision the dialog read; an identical retry is + ``replayed``. The hash never runs inside a database transaction. + """ + data, error = await self._annotation_body(request, _SAVE_FIELDS) + if error is not None: + return error + target, error = _target(data["target"]) + if error is not None: + return error + guards, error = _save_guards(data) + if error is not None: + return error + token, sha256, expected_epoch, revision, tags = guards + index, scope, relative = target + root, error = self._scoped_root(index, scope) + if error is not None: + return error + store = self._captures + capture = store.get(token) if store is not None else None + if capture is None: + return _coded("capture-expired") + if capture.root_index != index or capture.scope != scope: + return _coded("capture-expired", _OTHER_FILE) + if sha256 != capture.sha256: + return bad_request("sha256 does not match the preview") + if expected_epoch != capture.epoch: + current = await self._run_in_executor(self.db.get_execution_epoch) + return _coded("storage-changed", library_epoch=current) + try: + facts = await self._run_in_executor( + _verify_captured_file, store, capture, root, relative + ) + except CaptureError as exc: + return _capture_error(exc) + if facts is None or facts.locator != capture.locator: + return _coded("capture-expired", _OTHER_FILE) + if facts.sha256 != capture.sha256: + return _coded("file-changed") + try: + receipt = await self._run_in_executor( + self.db.replace_image_tags, + locator=facts.locator, + sha256=facts.sha256, + size=facts.size, + mtime_ns=facts.mtime_ns, + expected_epoch=expected_epoch, + expected_tag_revision=revision, + tags=tags, + ) + except AnnotationCommandError as exc: + current = exc.current or {} + return _coded( + exc.code, + library_epoch=exc.epoch, + tags=current.get("tags", []), + tag_revision=current.get("tag_revision", 0), + ) + except LibraryCommandError as exc: + return _coded(exc.code, library_epoch=exc.epoch) + except ValueError as exc: + return bad_request(str(exc)) + except Exception: + self.logger.exception("Image tag save failed") + return _error(500, "save-failed", "The tags could not be saved.") + return web.json_response( + {"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: + # From here the pin is the gate's to end: the worker ends it, + # or ``on_skipped`` does when the worker never starts. + started = True + pairs = await model_gate.run( + owner, + _pinned_suggestions, + store, + token, + service, + capture, + model_type, + options, + on_skipped=functools.partial(store.release_use, token), + ) + 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 _close_then_release(store: CaptureStore, token: str, handle) -> None: + """Close a preview handle a cancelled request opened, then end its pin.""" + try: + if handle is not None: + handle.close() + finally: + store.release_use(token) + + +def _byte_range(request: web.Request, size: int): + """``(start, end)`` of the one satisfiable byte range asked for, else None. + + A malformed, multi-range or unsatisfiable header is None (416), as + aiohttp's own file responses answer it. + """ + try: + wanted = request.http_range + except ValueError: + return None + start, stop = wanted.start, wanted.stop + if start is None: + return None + if start < 0: # a suffix: the last -start bytes + start = max(0, size + start) + stop = size + stop = size if stop is None else min(stop, size) + if start >= size or stop <= start: + return None + return start, stop + + +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. + + The cheap locator comparison runs first, so a token sent for another + target never costs a whole-file hash. + """ + expected = os.path.normcase( + os.path.join(os.path.realpath(root), *relative.split("/")) + ) + if expected != capture.locator: + return None + return store.verify(root, relative) + + +def _save_guards(data: dict): + """Validated ``(token, sha256, epoch, revision, tags)`` or ``(None, error)``.""" + token = data["capture_token"] + if not isinstance(token, str) or _TOKEN.fullmatch(token) is None: + return None, bad_request("Invalid capture token") + sha256 = data["sha256"] + if not isinstance(sha256, str) or _HEX64.fullmatch(sha256) is None: + return None, bad_request("sha256 must be 64 lowercase hexadecimal characters") + epoch = data["expected_epoch"] + if not isinstance(epoch, str) or _UUID.fullmatch(epoch) is None: + return None, bad_request("expected_epoch must be a canonical UUID") + revision = data["expected_tag_revision"] + if type(revision) is not int or revision < 0: + return None, bad_request("expected_tag_revision must be a non-negative integer") + try: + tags = normalize_image_tags(data["tags"]) + except ImageTagLimitError as exc: + return None, _error(400, "tag-limit", str(exc)) + except ValueError as exc: + return None, bad_request(str(exc)) + return (token, sha256, epoch, revision, tags), None diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 50d62b2..999b941 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/api/images.py b/py/api/images.py index 415cde9..6edfdc1 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -14,7 +14,7 @@ import urllib.parse from bisect import bisect_left, bisect_right from collections.abc import Mapping from pathlib import Path -from typing import NamedTuple +from typing import Callable, NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError @@ -74,6 +74,7 @@ IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS # The header records the card parses (see utils.card_seed.header_params). _CARD_RECORDS = ("prompt", "workflow") _PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" +_LISTING_FAILED = "The image list could not be filtered. Check server logs for details." def _card_header(path: Path) -> tuple: @@ -418,6 +419,29 @@ def _listed(path: Path, mtime: float, root: Path, index: int) -> _Listed: ) +class _ImageTagMatch(NamedTuple): + """How one listed file's current image tags meet the filters (#249 B3).""" + + tags: bool # every filter tag is in its image tags or linked positive tags + query: bool # ``q`` is a case-folded substring of one image tag name + link: tuple | None # (hit, rating, positive id) of its unambiguous link + + +def _current_fingerprint(path: Path) -> tuple[int, int] | None: + """The file's present ``(st_size, st_mtime_ns)``; None when it is gone.""" + try: + observed = os.stat(path) + except OSError: + return None + return observed.st_size, observed.st_mtime_ns + + +def _image_tags_relevant(names: list, wanted: set, query: str) -> bool: + if wanted and not wanted.isdisjoint(names): + return True + return bool(query) and any(query in name.casefold() for name in names) + + def _workspace_key(entry: tuple, order: str) -> tuple: """The keyset sort key of *entry* (a _Listed or a 4-tuple entry).""" if not isinstance(entry, _Listed): @@ -752,6 +776,7 @@ def _filtered_page( filters: dict, links: dict | None, favorites: frozenset | None = None, + image_tags: dict | None = None, ) -> tuple[int, list, int | None]: """Count the *entries* the filters admit and take *size* from *start* on. @@ -761,6 +786,9 @@ def _filtered_page( or a linked-positive hit. Returns the total, the page and the number of distinct linked positives among all admitted entries (None without links). *favorites*, when given, admits only those listed paths (``favorites=1``). + *image_tags* maps normalized paths to their _ImageTagMatch: a tag filter + also admits a file whose current image tags (with its linked positive's + tags) carry every tag, and ``q`` also admits an image tag name hit. """ query, media, folder = filters["query"], filters["media"], filters["folder"] members_only = filters["category"] is not None or bool(filters["tags"]) @@ -780,9 +808,15 @@ def _filtered_page( if favorites is not None and entry.path not in favorites: continue match = None if links is None else links.get(entry.normalized) - if members_only and match is None: + image = None if image_tags is None else image_tags.get(entry.normalized) + if members_only and match is None and not (image and image.tags): continue - if query and query not in entry.folded and not (match and match[0]): + if ( + query + and query not in entry.folded + and not (match and match[0]) + and not (image and image.query) + ): continue total += 1 if match is not None: @@ -1140,7 +1174,7 @@ class ImageRoutesMixin: } ) except Exception as e: - self.logger.error(f"Output images error: {e}") + self.logger.error(f"Output images error: {e}", exc_info=True) return web.json_response( {"success": False, "error": _safe_error(e)}, status=500 ) @@ -1184,11 +1218,27 @@ class ImageRoutesMixin: for index, root in enumerate(output_dirs) if root_index is None or index == root_index ] - page, total, has_more, next_cursor, positive_count = ( - await self._run_in_executor( - self._select_workspace_page, listings, filters, after, limit, scope + try: + page, total, has_more, next_cursor, positive_count = ( + await self._run_in_executor( + self._select_workspace_page, listings, filters, after, limit, scope + ) + ) + except Exception: + # A filter that needs the prompt links or the image tags cannot be + # served without them. Never echo the exception: it can carry SQL + # or database text. The traceback goes to the server log. + self.logger.error( + "Workspace image list could not be filtered", exc_info=True + ) + return web.json_response( + { + "success": False, + "code": "listing-failed", + "error": _LISTING_FAILED, + }, + status=500, ) - ) images = await self._run_in_executor(self._format_workspace_page, page) result = { "success": True, @@ -1276,6 +1326,80 @@ class ImageRoutesMixin: ) return orders[order] + def _workspace_locators(self, listings: list[tuple]) -> dict: + """The selected listing's private locators -> _Listed (cached). + + A locator is ``normcase(join(realpath(root), relative))``, the key + image tags are stored under (database/annotations.py); the listing + never follows symlinks, so it equals the listed file's own locator. + """ + orders = self._workspace_orders(listings) + reals = tuple(os.path.realpath(root) for _index, root, _files in listings) + key = ("locators", reals) + if key not in orders: + located = {} + for (index, root, files), real in zip(listings, reals): + prefix = os.path.join(str(root), "") + for entry in self._workspace_listed(index, root, files): + text = str(entry.path) + relative = ( + text[len(prefix) :] + if text.startswith(prefix) + else str(entry.path.relative_to(root)) + ) + located[os.path.normcase(os.path.join(real, relative))] = entry + orders[key] = located + return orders[key] + + def _image_tag_matches(self, listings: list[tuple], filters: dict) -> dict: + """Normalized path -> _ImageTagMatch for files whose image tags count. + + Streams the tagged byte records once and keeps only listed locators + carrying a filter tag or a ``q`` hit; only those candidates are + stat'ed (never every listed file), and a candidate's tags count only + while a byte record is verified at its present ``(size, mtime_ns)`` + (spec amendment S1). Under a Tags filter the candidates' linked + positives are read too, so their tags and category combine. + """ + tags, query, category = filters["tags"], filters["query"], filters["category"] + wanted = set(tags) + located = self._workspace_locators(listings) + candidates = {} + records = self.db.iter_image_tag_records() + try: + for locator, _size, _mtime_ns, names in records: + entry = located.get(locator) + if entry is not None and locator not in candidates: + if _image_tags_relevant(names, wanted, query): + candidates[locator] = entry + finally: + records.close() + stats = {} + for locator, entry in candidates.items(): + fingerprint = _current_fingerprint(entry.path) + if fingerprint is not None: + stats[locator] = fingerprint + current = self.db.current_image_tags(stats) if stats else {} + facts = {} + if wanted and current: + facts = self.db.get_image_link_facts( + [candidates[locator].normalized for locator in current], + filters["text"], + ) + matches = {} + for locator, names in current.items(): + entry = candidates[locator] + link = facts.get(entry.normalized) + carried = set(names).union(link[4] if link else ()) + matches[entry.normalized] = _ImageTagMatch( + bool(wanted) + and wanted <= carried + and (category is None or (link is not None and link[3] == category)), + bool(query) and any(query in name.casefold() for name in names), + link[:3] if link else None, + ) + return matches + def _select_workspace_page( self, listings: list[tuple], @@ -1328,6 +1452,30 @@ class ImageRoutesMixin: ), exc_info=True, ) + image_tags = None + if filters["query"] or tags: + try: + image_tags = self._image_tag_matches(listings, filters) + except Exception: + # A Tags filter cannot be served without them (as without the + # links); a search then matches without image tags. + if tags: + raise + self.logger.warning( + "Workspace image tags could not be read; searching without " + "image tags", + exc_info=True, + ) + if image_tags and links is not None: + # A linked file admitted through its image tags keeps its rating + # and counts its positive. + admitted = { + path: match.link + for path, match in image_tags.items() + if match.tags and match.link is not None and path not in links + } + if admitted: + links = {**links, **admitted} if order == "rating": def rank(entry): @@ -1360,7 +1508,7 @@ class ImageRoutesMixin: self._workspace_favorite_paths(listings) if filters["favorites"] else None ) total, page, positive_count = _filtered_page( - entries, start, limit + 1, filters, links, favorites + entries, start, limit + 1, filters, links, favorites, image_tags ) if not needed and after is not None: positive_count = _COUNT_OMITTED @@ -2467,6 +2615,9 @@ class ImageRoutesMixin: library=self._inspection_library( saved, unavailable=unavailable ), + image_tags=self._inspection_image_tags( + roots[index], Path(relative).parts + ), ) except CurrentFileChanged: return None @@ -2505,6 +2656,30 @@ class ImageRoutesMixin: {"success": False, "error": "Prompt inspection failed"}, status=500 ) + def _inspection_image_tags(self, root: Path, parts: tuple) -> Callable: + """The inspector's ``image_tags`` reader for one rooted file (S1). + + The returned function takes the inspection's own stable stat and + answers ``current`` only for a byte record verified at exactly that + ``(size, mtime_ns)``; a failed read is ``unavailable``, never an error. + """ + locator = os.path.normcase(os.path.join(os.path.realpath(root), *parts)) + + def read(size: int, mtime_ns: int) -> dict: + try: + found = self.db.current_image_tags({locator: (size, mtime_ns)}) + except Exception: + self.logger.warning( + "Inspection image tags could not be read", exc_info=True + ) + return {"state": "unavailable", "tags": []} + tags = found.get(locator) + if tags is None: + return {"state": "none", "tags": []} + return {"state": "current", "tags": list(tags)} + + return read + def _inspection_library(self, saved: dict | None, *, unavailable: bool) -> dict: """The inspector's Library facts; the route sets the title afterwards.""" result = { diff --git a/py/api/inspection_service.py b/py/api/inspection_service.py index 787025c..ee8fa6c 100644 --- a/py/api/inspection_service.py +++ b/py/api/inspection_service.py @@ -3,6 +3,7 @@ import os import stat from pathlib import Path +from typing import Callable try: from ...utils.prompt_inspector import ( @@ -93,8 +94,19 @@ def _in_host_temp(path: Path) -> bool: def inspect_current_file( - path: Path, *, source_label: str, saved: dict | None, library: dict + path: Path, + *, + source_label: str, + saved: dict | None, + library: dict, + image_tags: Callable[[int, int], dict] | None = None, ) -> dict: + """Current-file evidence read between two equal stats of *path*. + + *image_tags*, when given, is called once with that stable + ``(st_size, st_mtime_ns)`` and its facet is returned as ``image_tags`` + (spec amendment S1: tags of a byte record verified at exactly that stat). + """ before = _current_file_facts(path) try: inspection = read_inspection_file(path) @@ -104,6 +116,9 @@ def inspect_current_file( fields = ("st_dev", "st_ino", "st_mode", "st_size", "st_mtime_ns", "st_ctime_ns") if any(getattr(before, key) != getattr(after, key) for key in fields): raise CurrentFileChanged + facets = {} + if image_tags is not None: + facets["image_tags"] = image_tags(before.st_size, before.st_mtime_ns) if inspection is None: evidence = unreadable_evidence() else: @@ -132,4 +147,5 @@ def inspect_current_file( **evidence, "inspection": inspection, "library": library, + **facets, } diff --git a/py/autotag.py b/py/autotag.py index 19f675d..68c72a5 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -16,6 +16,15 @@ from PIL import Image # A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend) ModelLoader = Callable[[str, bool], Tuple[str, Any]] +# Suggestions keep only names a Save accepts (the vocabulary tag rules), and +# never more than one image tag set holds. +try: + from ..database.annotations import MAX_IMAGE_TAGS + from ..database.vocabulary import validate_vocabulary_tag_name +except ImportError: + from database.annotations import MAX_IMAGE_TAGS + from database.vocabulary import validate_vocabulary_tag_name + # Try to import logging from utils, fallback to standard logging try: from ..utils.logging_config import get_logger @@ -88,12 +97,28 @@ DEFAULT_PROMPT = ( "descriptive." ) +# Upper bound of one suggestion request: the image tag set cap. +MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS + 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 +447,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 +495,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 +785,61 @@ 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). + A name Save would reject (over 50 characters, a control character + such as a newline, a comma) is skipped before the dedupe and the + ``limit`` count. ``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 raw, score in scored: + try: + name = validate_vocabulary_tag_name(raw) + except ValueError: + continue + 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 +847,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 +875,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 cc89108..a228524 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 # pinned to MAX_IMAGE_TAGS by tests/test_config.py + + +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 0000000..84f3053 --- /dev/null +++ b/py/model_gate.py @@ -0,0 +1,177 @@ +"""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, + on_skipped: Optional[Callable[[], None]] = None, + **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. + ``on_skipped`` runs once when ``func`` never starts (the gate is not + held, the executor refuses the call, or a loop shutdown drops it), so + a caller can free what ``func`` would have freed. + """ + try: + self.begin(owner) + except GateNotHeld: + if on_skipped is not None: + on_skipped() + raise + 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: + if state["finished"]: + # Already declared skipped (the future was cancelled as + # the thread picked the call up): never run it twice over. + return None + 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) + if skipped and on_skipped is not None: + on_skipped() + + try: + future = asyncio.get_running_loop().run_in_executor(None, work) + except BaseException: + finish() + if on_skipped is not None: + on_skipped() + 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, + on_skipped: Optional[Callable[[], None]] = None, + **kwargs, +): + return await _GATE.run(owner, func, *args, on_skipped=on_skipped, **kwargs) diff --git a/tests/browser/server.py b/tests/browser/server.py index c5ae192..ae67018 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/host_stubs.py b/tests/host_stubs.py new file mode 100644 index 0000000..5a61dc0 --- /dev/null +++ b/tests/host_stubs.py @@ -0,0 +1,29 @@ +"""One stand-in for ComfyUI's ``folder_paths`` module, shared by every test. + +Test modules that import the API standalone need a ``folder_paths`` module. +It must be a real module, not a ``MagicMock``: ``patch("folder_paths.x", +create=True)`` then adds and removes plain attributes, while a mock would +hand a later test an auto-created child (or lose the attribute on exit), so +the outcome depended on which test module discovery imported first. +""" + +import sys +import tempfile +import types + + +def install_folder_paths_stub() -> types.ModuleType: + """Install the shared ``folder_paths`` stub unless a real module is there. + + A test that installs its own stand-in for a while (and restores it) is + left alone; a ``MagicMock`` left behind by older code is replaced. + """ + current = sys.modules.get("folder_paths") + if isinstance(current, types.ModuleType): + return current + stub = types.ModuleType("folder_paths") + stub.models_dir = tempfile.gettempdir() + stub.base_path = tempfile.gettempdir() + stub.get_output_directory = tempfile.gettempdir + sys.modules["folder_paths"] = stub + return stub diff --git a/tests/test_api_annotation_suggest.py b/tests/test_api_annotation_suggest.py new file mode 100644 index 0000000..5f11ad8 --- /dev/null +++ b/tests/test_api_annotation_suggest.py @@ -0,0 +1,556 @@ +"""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) + + async def test_gate_work_that_never_starts_releases_the_capture_pin(self): + observed = await self.observed() + loop = asyncio.get_running_loop() + real = loop.run_in_executor + + def refuse_gate_work(executor, func, *args): + if getattr(func, "__qualname__", "").startswith("ModelGate.run"): + raise RuntimeError("cannot schedule new futures after shutdown") + return real(executor, func, *args) + + with patch.object(loop, "run_in_executor", refuse_gate_work): + status, body = await self.suggest(self.suggest_body(observed)) + + self.assertEqual((status, body["code"]), (500, "suggest-failed")) + self.assertIsNone(model_gate.holder()) + entry = self.api._captures._entries[observed["capture"]["token"]] + self.assertEqual(entry.uses, 0) + self.assertEqual(self.tagger.calls, []) + + +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 new file mode 100644 index 0000000..e59a65f --- /dev/null +++ b/tests/test_api_annotations.py @@ -0,0 +1,1222 @@ +"""Image annotation routes: observe, capture preview and tag save. + +Every request names its file through the listing identity +``{root_index, root_scope, path}``; the scope is checked before any file +access, tag reads and saves are byte-exact (the capture's digest, re-hashed +at Save) and no response ever carries a server path or a locator. +""" + +import asyncio +import hashlib +import json +import os +import sqlite3 +import tempfile +import threading +import unittest +import uuid +from pathlib import Path +from unittest.mock import AsyncMock, patch + +from aiohttp.test_utils import make_mocked_request + +try: + from test_api_images import ImageAPITestCase, make_png + from open_handles import assert_closed, open_handles_under +except ImportError: + from tests.test_api_images import ImageAPITestCase, make_png + from tests.open_handles import assert_closed, open_handles_under + +from database.execution_history_schema import rotate_history_epoch +from utils import media_capture +from utils.media_capture import CaptureStore + +try: + from py.api import annotations as annotations_module +except ImportError: # the module under test is not there yet (RED) + annotations_module = None + +OBSERVE = "/prompt_manager/image-annotations/observe" +TAGS = "/prompt_manager/image-annotations/tags" +CAPTURE = "/prompt_manager/image-annotations/capture/" +MAX_BODY = 64 * 1024 +WAIT = 30.0 # ceiling for cross-thread handoffs, never a short clock + +# Columns that an observe may legitimately refresh (S1) or that never matter +# for "nothing changed" comparisons. +_SNAPSHOT_TABLES = { + "prompts": "*", + "generated_images": "*", + "tags": "*", + "prompt_tags": "*", + "prompt_pair_tags": "*", + "image_favorites": "*", + "media_annotation_files": "id, locator", + "media_annotation_bytes": "id, file_id, sha256, first_seen_at, tag_revision", + "media_image_tags": "*", +} + + +class FakeClock: + def __init__(self, start=1000.0): + self.now = start + + def __call__(self): + return self.now + + +class AnnotationAPITestCase(ImageAPITestCase): + """The image API app plus helpers for the annotation routes.""" + + async def get_application(self): + app = await super().get_application() + self.assertIsNotNone( + annotations_module, "py.api.annotations is not implemented" + ) + self.epoch = self.api.db.get_execution_epoch() + return app + + async def tearDownAsync(self): + store = getattr(self.api, "_captures", None) + if store is not None: + # Every preview handle must already be closed (Windows unlink). + assert_closed(self, store.directory) + self.api.close_annotation_captures() + if store is not None: + self.assertFalse(os.path.exists(store.directory)) + await super().tearDownAsync() + + # -- helpers --------------------------------------------------------- + + def scope(self, index=0, roots=None): + roots = roots if roots is not None else self.api._get_all_output_dirs() + return self.api._gallery_root_scope(roots[index], index) + + def target(self, path="sub/a.png", **changes): + return {"root_index": 0, "root_scope": self.scope(), "path": path, **changes} + + async def json_of(self, response): + """The JSON envelope; an unrouted path (no JSON) fails as an assertion.""" + self.assertEqual( + response.content_type, + "application/json", + f"{response.method} {response.url.path} answered {response.status}", + ) + return await response.json() + + async def observe(self, body): + response = await self.client.post(OBSERVE, json=body) + return response.status, await self.json_of(response) + + async def observe_ok(self, path="sub/a.png"): + status, body = await self.observe({"target": self.target(path)}) + self.assertEqual(status, 200, body) + return body + + async def put_tags(self, body): + response = await self.client.put(TAGS, json=body) + return response.status, await self.json_of(response) + + def save_body(self, observed, tags, **changes): + return { + "target": dict(observed["target"]), + "capture_token": observed["capture"]["token"], + "sha256": observed["capture"]["sha256"], + "expected_epoch": observed["library_epoch"], + "expected_tag_revision": observed["tag_revision"], + "tags": list(tags), + **changes, + } + + def locator(self, path): + return os.path.normcase(os.path.realpath(str(self.output_dir / path))) + + def write_bytes(self, path, data, *, keep_mtime=False): + """Rewrite *path* in place, optionally restoring its ``mtime_ns``.""" + full = self.output_dir / path + before = full.stat() + full.write_bytes(data) + if keep_mtime: + os.utime(full, ns=(before.st_atime_ns, before.st_mtime_ns)) + return full + + def db_snapshot(self): + conn = sqlite3.connect(self._temp_db.name) + try: + snapshot = {} + for table, columns in _SNAPSHOT_TABLES.items(): + exists = conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", + (table,), + ).fetchone() + if exists is None: + continue + rows = conn.execute(f"SELECT {columns} FROM {table}").fetchall() + snapshot[table] = sorted(map(repr, rows)) + return snapshot + finally: + conn.close() + + def bytes_rows(self): + conn = sqlite3.connect(self._temp_db.name) + try: + return conn.execute( + "SELECT sha256, size, mtime_ns, tag_revision FROM " + "media_annotation_bytes ORDER BY first_seen_at" + ).fetchall() + finally: + conn.close() + + def saved_set(self, path, data, tags): + """Store *tags* for the exact *data* bytes of *path*; return the receipt.""" + full = self.output_dir / path + stat = full.stat() + return self.api.db.replace_image_tags( + locator=self.locator(path), + sha256=hashlib.sha256(data).hexdigest(), + size=stat.st_size, + mtime_ns=stat.st_mtime_ns, + expected_epoch=self.epoch, + expected_tag_revision=0, + tags=tags, + ) + + def use_fake_clock(self): + clock = FakeClock() + self.api.close_annotation_captures() + self.api._captures = CaptureStore(clock=clock) + return clock + + def assert_no_paths(self, body): + text = json.dumps(body) + self.assertNotIn(str(self.comfy_dir), text) + self.assertNotIn(str(self.output_dir), text) + self.assertNotIn(os.path.realpath(self.output_dir), text) + store = getattr(self.api, "_captures", None) + if store is not None: + self.assertNotIn(store.directory, text) + self.assertNotIn("\\\\", text.replace('\\"', "")) + self.assertNotIn(os.sep + "output" + os.sep, text) + + +class TestObserve(AnnotationAPITestCase): + async def test_observe_requires_valid_scope_before_any_file_access(self): + make_png(self.output_dir / "sub" / "a.png") + other = self.comfy_dir / "other" + other.mkdir() + make_png(other / "sub" / "a.png") + stale_scope = self.api._gallery_root_scope(other, 0) + reordered = self.api._gallery_root_scope(self.output_dir, 1) + with patch.object( + media_capture.CaptureStore, "capture", autospec=True + ) as capture: + capture.side_effect = AssertionError("capture must not run") + for name, changes, status in ( + ("stale root", {"root_scope": stale_scope}, 409), + ("reordered root", {"root_scope": reordered}, 409), + ("removed root", {"root_index": 1, "root_scope": reordered}, 409), + ("wrong scope", {"root_scope": "0" * 64}, 409), + ("uppercase scope", {"root_scope": "A" * 64}, 400), + ("short scope", {"root_scope": "abc"}, 400), + ("bool index", {"root_index": True}, 400), + ("negative index", {"root_index": -1}, 400), + ("string index", {"root_index": "0"}, 400), + ("traversal", {"path": "../a.png"}, 403), + ("dot segment", {"path": "./a.png"}, 403), + ("empty segment", {"path": "sub//a.png"}, 403), + ("absolute", {"path": "/etc/a.png"}, 403), + ("backslash", {"path": "sub\\a.png"}, 403), + ("drive", {"path": "C:/a.png"}, 403), + ("drive segment", {"path": "sub/C:a.png"}, 403), + ("nul", {"path": "a\x00.png"}, 403), + ("non-media", {"path": "notes.txt"}, 403), + ("empty path", {"path": ""}, 400), + ("long path", {"path": "a/" * 600 + "a.png"}, 400), + ("int path", {"path": 7}, 400), + ): + with self.subTest(name): + status_, body = await self.observe( + {"target": self.target(**changes)} + ) + self.assertEqual(status_, status, body) + self.assertFalse(body["success"]) + if status == 409: + self.assertEqual(body["code"], "source-changed") + self.assert_no_paths(body) + for name, body in ( + ("missing target", {}), + ("extra field", {"target": self.target(), "extra": 1}), + ("target not an object", {"target": "sub/a.png"}), + ("target missing path", {"target": {"root_index": 0}}), + ("target extra field", {"target": self.target(sha256="x")}), + ): + with self.subTest(name): + status_, body_ = await self.observe(body) + self.assertEqual(status_, 400, body_) + self.assertFalse(body_["success"]) + with self.subTest("body over 64 KiB"): + padding = "x" * MAX_BODY + response = await self.client.post( + OBSERVE, + data=json.dumps({"target": self.target(), "pad": padding}), + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(response.status, 413) + with self.subTest("not json"): + response = await self.client.post( + OBSERVE, data=b"{oops", headers={"Content-Type": "application/json"} + ) + self.assertEqual(response.status, 400) + capture.assert_not_called() + + async def test_observe_denies_case_spellings_of_the_listed_file(self): + make_png(self.output_dir / "sub" / "a.png") + for alias in ("sub/A.PNG", "SUB/a.png"): + with self.subTest(alias=alias): + status, body = await self.observe({"target": self.target(alias)}) + # Denied as an alias, like a symlink to the listed file. + self.assertEqual((status, body["code"]), (409, "source-changed")) + self.assert_no_paths(body) + store = self.api._captures + self.assertFalse(any(Path(store.directory).iterdir())) + await self.observe_ok("sub/a.png") + + async def test_observe_returns_saved_tags_for_exact_bytes_only(self): + path = make_png(self.output_dir / "sub" / "a.png") + original = path.read_bytes() + receipt = self.saved_set("sub/a.png", original, ["Forest", "mist"]) + self.assertEqual(receipt["tag_revision"], 1) + + body = await self.observe_ok() + self.assertEqual(body["tags"], ["Forest", "mist"]) + self.assertEqual(body["tag_revision"], 1) + self.assertEqual(body["library_epoch"], self.epoch) + self.assertEqual(body["filename"], "a.png") + self.assertIsNone(body["title"]) + self.assertEqual( + body["target"], + {"root_index": 0, "root_scope": self.scope(), "path": "sub/a.png"}, + ) + capture = body["capture"] + self.assertEqual( + sorted(capture), + sorted( + [ + "token", + "expires_at", + "preview_url", + "sha256", + "size", + "width", + "height", + "media_type", + "generate", + ] + ), + ) + self.assertEqual(capture["sha256"], hashlib.sha256(original).hexdigest()) + self.assertEqual(capture["size"], len(original)) + self.assertEqual((capture["width"], capture["height"]), (4, 4)) + self.assertEqual(capture["media_type"], "image") + self.assertEqual(capture["generate"], {"supported": True, "reason": None}) + self.assertEqual(capture["preview_url"], CAPTURE + capture["token"]) + self.assertTrue(capture["expires_at"].endswith("Z")) + + # Different bytes at the same path: an empty set at revision 0. + changed = make_png(path, color=(0, 255, 0)).read_bytes() + self.assertNotEqual(changed, original) + body = await self.observe_ok() + self.assertEqual((body["tags"], body["tag_revision"]), ([], 0)) + self.assertEqual(body["capture"]["sha256"], hashlib.sha256(changed).hexdigest()) + + # The original bytes again (new mtime): the saved set, and S1 refreshes + # the record's verified fingerprint so cheap reads see it as current. + self.write_bytes("sub/a.png", original) + stat = path.stat() + current = self.api.db.current_image_tags( + {self.locator("sub/a.png"): (stat.st_size, stat.st_mtime_ns)} + ) + body = await self.observe_ok() + self.assertEqual((body["tags"], body["tag_revision"]), (["Forest", "mist"], 1)) + refreshed = self.api.db.current_image_tags( + {self.locator("sub/a.png"): (stat.st_size, stat.st_mtime_ns)} + ) + if stat.st_mtime_ns != self.bytes_rows()[0][2]: + self.assertEqual(current, {}) + self.assertEqual(refreshed, {self.locator("sub/a.png"): ["Forest", "mist"]}) + + async def test_observe_changes_no_tags_revisions_prompts_or_links(self): + path = make_png(self.output_dir / "sub" / "a.png") + data = path.read_bytes() + self.saved_set("sub/a.png", data, ["forest"]) + make_png(self.output_dir / "sub" / "b.png", color=(1, 2, 3)) + self._link_image(self.output_dir / "sub" / "a.png", "a prompt") + prompt_id = self._save_prompt("tagged prompt", tags=["forest", "night"]) + self.assertTrue(prompt_id) + before = self.db_snapshot() + + for _ in range(2): + await self.observe_ok("sub/a.png") + await self.observe_ok("sub/b.png") + + self.assertEqual(self.db_snapshot(), before) + # Unseen bytes (b.png) never create a record. + self.assertEqual(len(self.bytes_rows()), 1) + self.assertEqual(self.bytes_rows()[0][3], 1) + + async def test_observe_title_follows_the_card_rule(self): + make_png(self.output_dir / "sub" / "a.png") + prompt_id = self._save_prompt("Same words\nrest") + self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png")) + body = await self.observe_ok() + # A first-line title is withheld unless the file's own records confirm + # it (the card rule); this PNG carries no prompt metadata. + self.assertIsNone(body["title"]) + with self.api.db.model.get_connection() as connection: + connection.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (prompt_id, "Same words"), + ) + # A stored title is user data and always shows. + body = await self.observe_ok() + self.assertEqual(body["title"], "Same words") + + async def test_observe_video_is_manual_only(self): + clip = self.output_dir / "clip.mp4" + clip.parent.mkdir(exist_ok=True) + clip.write_bytes(b"\x00\x00\x00\x18ftypmp42" + b"\x00" * 64) + body = await self.observe_ok("clip.mp4") + capture = body["capture"] + self.assertEqual(capture["media_type"], "video") + self.assertEqual((capture["width"], capture["height"]), (None, None)) + self.assertEqual( + capture["generate"], + {"supported": False, "reason": media_capture.VIDEO_ONLY_REASON}, + ) + self.assertIsNone(body["title"]) + response = await self.client.get(capture["preview_url"]) + self.assertEqual(response.status, 200) + self.assertEqual(response.headers["Content-Type"], "video/mp4") + self.assertEqual(await response.read(), clip.read_bytes()) + + async def test_observe_maps_capture_outcomes_to_codes(self): + make_png(self.output_dir / "sub" / "a.png") + outside = make_png(self.comfy_dir / "outside.png") + self._make_symlink(outside, self.output_dir / "sub" / "alias.png") + status, body = await self.observe({"target": self.target("sub/alias.png")}) + self.assertEqual((status, body["code"]), (409, "source-changed")) + self.assert_no_paths(body) + + status, body = await self.observe({"target": self.target("sub/missing.png")}) + self.assertEqual((status, body["code"]), (409, "file-changed")) + self.assert_no_paths(body) + + self.api.close_annotation_captures() + self.api._captures = CaptureStore(max_file=16) + status, body = await self.observe({"target": self.target("sub/a.png")}) + self.assertEqual((status, body["code"]), (415, "unsupported-media")) + self.assert_no_paths(body) + + self.api.close_annotation_captures() + self.api._captures = CaptureStore(max_captures=1) + await self.observe_ok("sub/a.png") + status, body = await self.observe({"target": self.target("sub/a.png")}) + self.assertEqual((status, body["code"]), (429, "review-capacity")) + self.assert_no_paths(body) + + +class TestCaptureRoutes(AnnotationAPITestCase): + async def test_capture_route_serves_exact_bytes_no_store_nosniff(self): + path = make_png(self.output_dir / "sub" / "a.png") + source_bytes = path.read_bytes() + body = await self.observe_ok() + token = body["capture"]["token"] + # The source may change after the capture; the preview never does. + make_png(path, color=(9, 9, 9)) + + resp = await self.client.get(body["capture"]["preview_url"]) + self.assertEqual(resp.status, 200) + self.assertEqual(resp.headers["Cache-Control"], "no-store") + self.assertEqual(resp.headers["X-Content-Type-Options"], "nosniff") + self.assertEqual(resp.headers["Content-Type"], "image/png") + self.assertEqual(resp.headers["Content-Length"], str(len(source_bytes))) + self.assertEqual(await resp.read(), source_bytes) + + store = self.api._captures + entry = store._entries[token] + self.assertEqual(entry.uses, 0) + self.assertEqual(open_handles_under(store.directory), []) + self.assertIsNotNone(store.get(token)) + + async def test_capture_route_rejects_unknown_expired_and_rescoped_tokens(self): + make_png(self.output_dir / "sub" / "a.png") + clock = self.use_fake_clock() + for token in ("nope", "A" * 300, "x.y", "%2e%2e%2fx", "a b"): + with self.subTest(token=token): + resp = await self.client.get(CAPTURE + token) + self.assertEqual(resp.status, 404) + data = await self.json_of(resp) + self.assertEqual(data["code"], "capture-expired") + self.assertEqual(resp.headers["Cache-Control"], "no-store") + + body = await self.observe_ok() + url = body["capture"]["preview_url"] + self.assertEqual((await self.client.get(url)).status, 200) + + # Released explicitly: gone, and releasing again stays successful. + resp = await self.client.delete(url) + self.assertEqual( + (resp.status, await self.json_of(resp)), (200, {"success": True}) + ) + self.assertEqual((await self.client.get(url)).status, 404) + resp = await self.client.delete(url) + self.assertEqual( + (resp.status, await self.json_of(resp)), (200, {"success": True}) + ) + + # Idle expiry. + body = await self.observe_ok() + url = body["capture"]["preview_url"] + clock.now += media_capture.DEFAULT_IDLE_SECONDS + 1 + resp = await self.client.get(url) + self.assertEqual(resp.status, 404) + self.assertEqual((await self.json_of(resp))["code"], "capture-expired") + + # The roots changed under a live token: denied and released. + body = await self.observe_ok() + url = body["capture"]["preview_url"] + moved = self.comfy_dir / "moved" + moved.mkdir() + original_roots = self.api._get_all_output_dirs + self.api._get_all_output_dirs = lambda: [moved] + try: + resp = await self.client.get(url) + self.assertEqual(resp.status, 409) + data = await self.json_of(resp) + self.assertEqual(data["code"], "source-changed") + self.assert_no_paths(data) + finally: + self.api._get_all_output_dirs = original_roots + resp = await self.client.get(url) + self.assertEqual(resp.status, 404) + self.assertFalse(any(Path(self.api._captures.directory).iterdir())) + + async def test_preview_pins_the_capture_while_it_streams(self): + path = make_png(self.output_dir / "sub" / "a.png", size=(64, 64)) + source_bytes = path.read_bytes() + body = await self.observe_ok() + token = body["capture"]["token"] + store = self.api._captures + capture_path = store.get(token).path + stream = annotations_module.AnnotationRoutesMixin._stream_capture + seen = {} + + async def release_mid_stream(request, capture, handle): + # A concurrent DELETE (or sweep) lands while the bytes stream. + self.assertEqual(store._entries[token].uses, 1) + self.assertTrue(store.release(token)) + seen["kept"] = os.path.exists(capture_path) + return await stream(request, capture, handle) + + with ( + patch.object( + annotations_module.AnnotationRoutesMixin, + "_stream_capture", + staticmethod(release_mid_stream), + ), + patch.object(annotations_module, "_PREVIEW_CHUNK", 16), + ): + resp = await self.client.get(body["capture"]["preview_url"]) + self.assertEqual(resp.status, 200) + self.assertEqual(await resp.read(), source_bytes) + self.assertTrue(seen["kept"]) + # Removed with the last use, after its handle closed. + self.assertFalse(os.path.exists(capture_path)) + self.assertEqual( + (await self.client.get(body["capture"]["preview_url"])).status, 404 + ) + + async def test_preview_answers_byte_ranges(self): + # Safari and WebKit play