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

feat(annotations): image tag routes, model gate and scored suggestions (#278)

Browse Source 
* feat(annotations): add scoped observe, capture preview and tag save routes

POST /prompt_manager/image-annotations/observe captures the file named by
the listing identity {root_index, root_scope, path} (scope checked against
the current roots before any file access) and answers the saved tag set of
exactly those bytes, refreshing a matching record's verified fingerprint.
GET/DELETE .../capture/{token} serve the private copy (no-store, nosniff,
pinned with acquire/release_use while it streams) and release it.
PUT .../tags re-hashes the current file and writes only when it still
holds the captured bytes, under the storage epoch and tag revision
(file-changed, storage-changed, tag-conflict, tag-limit, replayed).
Bodies are capped at 64 KiB with unknown fields refused; envelopes never
carry server paths or locators. Restores invalidate every capture and the
host's cleanup removes the capture directory.

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

- py/model_gate.py: one process-wide, non-blocking gate for loading,
  swapping, running, downloading and unloading a tagging model. Worker
  calls run shielded; a release requested while one still runs takes
  effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
  gate and answer 409 {code: "model-busy"} while another operation holds
  it; their bodies are otherwise unchanged. /autotag/models gains
  suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
  first, GGUF/HF (name, None) in output order, first `limit` distinct
  case-insensitively; generate_tags keeps its return; the image file is
  closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
  pinned capture, with strict per-engine options, never downloads
  (409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
  10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
  metadata over the inspection budgets no longer blocks it.

* feat(images): match image tags in the Tags filter, search and inspector

The workspace Images listing now admits a file under a Tags filter when
every requested tag is carried by its linked positive or its current image
tag set, and q also matches image tag names (case-folded substring). Image
tags count only while a byte record is verified at the file's present
(size, mtime_ns) (spec amendment S1): the tagged records are streamed
once, only listed files carrying a filter tag or q hit are stat'ed, and
their linked positives are read through the file_path index so tags,
category, rating and the positive count combine.

The rooted inspector envelope gains a validated image_tags facet
(current / none / unavailable) read at the inspection's own stable stat.
A failed image-tag read fails a Tags filter like a failed link read and
degrades a q search to no image tags; plain listings read none. Legacy
gallery and search responses are unchanged.

* fix(annotations): answer a fixed coded message when the filtered image list fails

A failed prompt-link or image-tag read under a Tags, category or rating
filter returned the raw exception text (SQL or database messages) in the
500 body. The workspace listing now answers code listing-failed with a
fixed sentence and logs the traceback. The legacy output listing and the
inspector image-tag read now log with exc_info as well.

* fix(annotations): drop suggestion names that Save would reject

A JoyCaption part over 50 characters or holding a newline came back as a
preselected chip, so the first Save of the dialog failed with a 400 and
nothing was stored. generate_tag_suggestions now applies the vocabulary
tag-name rules before the case-insensitive dedupe and the limit count, so
the requested number of valid chips still comes back.

* fix(annotations): stop changing the process warnings filters during capture

header_dimensions wrapped Image.open in warnings.catch_warnings with an
error filter for DecompressionBombWarning. catch_warnings is not
thread-safe, so two overlapping captures could leave that error filter
installed for every later Pillow open in the ComfyUI process. Only
Pillow's DecompressionBombError is now mapped to HeaderTooLarge; a size
that only warns is refused by the existing inference pixel cap.

* fix(annotations): free pins on cancellation, serve preview ranges, code every storage error

- A preview cancelled while its file opens closes the handle and only
  then ends the pin, so a doomed capture is never unlinked while open.
- An observe cancelled while capturing or reading releases the capture
  instead of leaving it until the idle expiry.
- The model gate takes an on_skipped callback that runs once when the
  work never starts; Generate passes the capture's release_use, so an
  executor refusal no longer leaks a pin against the capture budget.
- The preview answers a single byte range with 206 and Content-Range
  (416 for an unsatisfiable or multi-range header), still streamed inside
  the pin, so Safari and WebKit can play video captures.
- _coded maps tag-limit (400), metadata-conflict and vocabulary-conflict,
  and falls back to 409 with a fixed sentence for any other code instead
  of raising KeyError.

* fix(annotations): deny case and Unicode spellings of a listed file

On macOS (case- and normalization-insensitive APFS) a crafted target such
as A.PNG or an NFD name opened the listed a.png or NFC file while
realpath kept the requested spelling, so the tags were stored under a
second locator the Images filter never shows. Capture and verify now
require each path part to be spelled as its directory lists it and deny
any other case or Unicode spelling as an alias (409 source-changed, like
a symlink alias). A missing file with no other spelling keeps its usual
outcome.

* test(annotations): one shared folder_paths stub for the API test modules

test_api_images and test_api_prompts installed a MagicMock folder_paths,
so whenever one of them was imported before test_api_autotag the later
patch("folder_paths.get_output_directory") in test_api_history_outputs
left a mock attribute behind and two test_diagnostics_config tests
failed. tests/host_stubs.py now installs one real-module stub, used by
test_api_images, test_api_prompts and test_api_autotag (and through
test_api_images by the annotation modules, whose copy is gone), so the
modules pass in either import order.

* fix(annotations): derive the suggestion cap from the image tag set cap

py/autotag.py now takes MAX_SUGGESTION_LIMIT from database.annotations
MAX_IMAGE_TAGS. The two copies that cannot import it (py/config.py loads
before the database, utils/inspection_response.py stays storage-free) are
pinned equal by a test, so the 100 limits cannot drift apart.

* docs(testing): record capture aliases, preview ranges and the PR B follow-ups

TESTING.md now states that case and Unicode spellings are aliases, that
the preview serves byte ranges inside the capture pin, where the
cancellation tests live, and two tracked follow-ups: the legacy batch's
ungated cleanup unload under its own last inference, and the three files
over the 800-line ceiling (split py/autotag.py first).
This commit is contained in:
Vito
2026-10-22 10:36:58 -07:00
committed by GitHub
parent 0095db8cbc
commit 1b8fbdbb75
30 changed files with 4786 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
TESTING.md
+22
View File
@@ -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
__init__.py
+1
View File
@@ -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
database/annotations.py
+7 -3
View File
@@ -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
database/operations.py
+57
View File
@@ -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]:
py/api/__init__.py
+7
View File
@@ -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)
py/api/annotation_suggest.py
+121
View File
@@ -0,0 +1,121 @@
"""Request validation and the worker of the Tag dialog's Generate (#249 B2).
``POST /prompt_manager/image-annotations/suggest`` asks one of legacy's four
engines for the top N tags of a captured image. Options are validated
strictly, per engine: WD14 classifiers take two thresholds, the language
models (GGUF and HuggingFace JoyCaption) take the instructions; every engine
takes the suggestion count and the GPU switch. Nothing here touches the
filesystem or the database.
"""
import math
WD14_MODELS = frozenset({"wd14-swinv2", "wd14-vit"})
LLM_MODELS = frozenset({"gguf", "hf"})
MODEL_TYPES = WD14_MODELS | LLM_MODELS
MAX_INSTRUCTION_CHARS = 4000
MAX_INSTRUCTION_BYTES = 16 * 1024
SUGGEST_FIELDS = frozenset(
{"target", "capture_token", "expected_epoch", "model_type", "options"}
)
_COMMON = frozenset({"suggestion_limit", "use_gpu"})
_WD14_REQUIRED = frozenset({"general_threshold", "character_threshold"})
_LLM_REQUIRED = frozenset({"instructions"})
UNAVAILABLE_HINT = "Download it from the legacy Auto-tag page."
class OptionError(ValueError):
"""A suggest request whose model type or options are not acceptable."""
def _threshold(name: str, value) -> float:
if type(value) not in (int, float) or not math.isfinite(value):
raise OptionError(f"{name} must be a finite number")
if not 0 <= value <= 1:
raise OptionError(f"{name} must be between 0 and 1")
return float(value)
def _instructions(value) -> str:
if not isinstance(value, str) or not value.strip():
raise OptionError("instructions must be non-empty text")
if len(value) > MAX_INSTRUCTION_CHARS:
raise OptionError(
f"instructions must be at most {MAX_INSTRUCTION_CHARS} characters"
)
try:
encoded = value.encode("utf-8")
except UnicodeEncodeError:
raise OptionError("instructions must be valid Unicode text") from None
if len(encoded) > MAX_INSTRUCTION_BYTES:
raise OptionError("instructions are too long")
return value
def validate_options(model_type, options, *, default_limit: int, max_limit: int):
"""Normalised ``(model_type, options)`` or :class:`OptionError`.
The answer echoes every option the request will run with, the defaults
(suggestion count from the settings, GPU on) filled in.
"""
if not isinstance(model_type, str) or model_type not in MODEL_TYPES:
raise OptionError("model_type must be one of " + ", ".join(sorted(MODEL_TYPES)))
if not isinstance(options, dict):
raise OptionError("options must be an object")
required = _WD14_REQUIRED if model_type in WD14_MODELS else _LLM_REQUIRED
keys = set(options)
if not required <= keys or not keys <= required | _COMMON:
raise OptionError(
f"options for {model_type} take exactly "
+ ", ".join(sorted(required))
+ " and optionally "
+ ", ".join(sorted(_COMMON))
)
limit = options.get("suggestion_limit", default_limit)
if type(limit) is not int or not 1 <= limit <= max_limit:
raise OptionError(f"suggestion_limit must be an integer from 1 to {max_limit}")
use_gpu = options.get("use_gpu", True)
if not isinstance(use_gpu, bool):
raise OptionError("use_gpu must be true or false")
normalised = {"use_gpu": use_gpu, "suggestion_limit": limit}
if model_type in WD14_MODELS:
for name in sorted(_WD14_REQUIRED):
normalised[name] = _threshold(name, options[name])
else:
normalised["instructions"] = _instructions(options["instructions"])
return model_type, normalised
def unavailable_message(model_name: str) -> str:
return f"{model_name} is not downloaded. {UNAVAILABLE_HINT}"
def run_suggestions(service, capture_path: str, model_type: str, options: dict):
"""Blocking: load (or swap to) *model_type*, then suggest for the capture.
The caller holds the model gate for the whole call; nothing is ever
downloaded (``load_model`` refuses a model that is not on disk).
"""
if not service.is_model_loaded() or service.get_loaded_model_type() != model_type:
service.load_model(model_type, options["use_gpu"])
return service.generate_tag_suggestions(
capture_path,
limit=options["suggestion_limit"],
prompt=options.get("instructions"),
general_threshold=options.get("general_threshold"),
character_threshold=options.get("character_threshold"),
)
def suggestion_rows(pairs) -> list:
"""``[{name, confidence}]``; a non-finite score is reported as unknown."""
rows = []
for name, score in pairs:
confidence = None
if score is not None and math.isfinite(score):
confidence = float(score)
rows.append({"name": str(name), "confidence": confidence})
return rows
py/api/annotations.py
+759
View File
@@ -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
py/api/autotag_routes.py
+76 -27
View File
@@ -9,11 +9,24 @@ from typing import Optional
from urllib.parse import quote
from .prompts import thumbnail_url_for
from .. import model_gate
from aiohttp import web
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
_BUSY_ERROR = "Auto-tagging is already running"
MODEL_BUSY_ERROR = (
"A tagging model is busy with another operation. Try again when it finishes."
)
def model_busy_response() -> web.Response:
"""``409 model-busy``: another operation holds the shared model gate."""
return web.json_response(
{"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR},
status=409,
)
# Only these may be handed to the tagging engine or listed by the scanner.
_MEDIA_EXTENSIONS = frozenset(
@@ -86,6 +99,7 @@ class AutotagRoutesMixin:
"""Get status of available AutoTag models."""
try:
from ..autotag import get_autotag_service
from ..config import AutotagConfig
service = get_autotag_service()
models_status = service.get_models_status()
@@ -99,6 +113,7 @@ class AutotagRoutesMixin:
"loaded_model_type": service.get_loaded_model_type(),
"wd14_general_threshold": service.wd14_general_threshold,
"wd14_character_threshold": service.wd14_character_threshold,
"suggestion_limit": AutotagConfig.suggestion_limit(),
}
)
@@ -138,6 +153,7 @@ class AutotagRoutesMixin:
async def download_autotag_model(self, request):
"""Download an AutoTag model with streaming progress (POST only)."""
model_type = request.match_info.get("model_type")
owner = model_gate.new_owner("legacy-download")
async def events():
try:
@@ -160,8 +176,8 @@ class AutotagRoutesMixin:
def progress_callback(status: str, progress: float):
pass
success = await self._run_in_executor(
service.download_model, model_type, progress_callback
success = await model_gate.run(
owner, service.download_model, model_type, progress_callback
)
if success:
yield {
@@ -178,7 +194,12 @@ class AutotagRoutesMixin:
"message": _INTERNAL_ERROR,
}
return await self._stream_sse(request, events())
if not model_gate.try_acquire(owner):
return model_busy_response()
try:
return await self._stream_sse(request, events())
finally:
model_gate.release(owner)
async def _read_start_params(self, request) -> Optional[dict]:
"""Merge query-string and optional JSON-body parameters for autotag/start.
@@ -239,8 +260,9 @@ class AutotagRoutesMixin:
return
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
run = self._gated_runner(opts.get("gate_owner"))
try:
await self._run_in_executor(service.load_model, model_type, True)
await run(service.load_model, model_type, True)
except Exception:
self.logger.exception("Failed to load model")
yield {
@@ -338,7 +360,7 @@ class AutotagRoutesMixin:
continue
try:
tags = await self._run_in_executor(
tags = await run(
service.generate_tags,
allowed_path,
general_threshold=opts["general_threshold"],
@@ -423,10 +445,26 @@ class AutotagRoutesMixin:
{"success": False, "error": _BUSY_ERROR}, status=409
)
try:
return await self._stream_sse(request, events())
opts["gate_owner"] = model_gate.new_owner("legacy-batch")
if not model_gate.try_acquire(opts["gate_owner"]):
return model_busy_response()
try:
return await self._stream_sse(request, events())
finally:
model_gate.release(opts["gate_owner"])
finally:
_autotag_batch_lock.release()
def _gated_runner(self, owner):
"""Executor calls under the model gate *owner* holds (plain when None)."""
if owner is None:
return self._run_in_executor
async def run(func, *args, **kwargs):
return await model_gate.run(owner, func, *args, **kwargs)
return run
def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]:
"""Return the real path when it is a media file inside an output directory."""
if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS:
@@ -549,22 +587,28 @@ class AutotagRoutesMixin:
from ..autotag import get_autotag_service
service = get_autotag_service()
owner = model_gate.new_owner("legacy-single")
if not model_gate.try_acquire(owner):
return model_busy_response()
try:
if (
not service.is_model_loaded()
or service.get_loaded_model_type() != model_type
):
await model_gate.run(owner, service.load_model, model_type, use_gpu)
if (
not service.is_model_loaded()
or service.get_loaded_model_type() != model_type
):
await self._run_in_executor(service.load_model, model_type, use_gpu)
if custom_prompt:
service.custom_prompt = custom_prompt
if custom_prompt:
service.custom_prompt = custom_prompt
tags = await self._run_in_executor(
service.generate_tags,
image_path,
general_threshold=general_threshold,
character_threshold=character_threshold,
)
tags = await model_gate.run(
owner,
service.generate_tags,
image_path,
general_threshold=general_threshold,
character_threshold=character_threshold,
)
finally:
model_gate.release(owner)
return web.json_response(
{
@@ -636,14 +680,19 @@ class AutotagRoutesMixin:
from ..autotag import get_autotag_service
service = get_autotag_service()
owner = model_gate.new_owner("legacy-unload")
if not model_gate.try_acquire(owner):
return model_busy_response()
try:
if not service.is_model_loaded():
return web.json_response(
{"success": True, "message": "No model was loaded"}
)
if not service.is_model_loaded():
return web.json_response(
{"success": True, "message": "No model was loaded"}
)
model_type = service.get_loaded_model_type()
service.unload_model()
model_type = service.get_loaded_model_type()
service.unload_model()
finally:
model_gate.release(owner)
return web.json_response(
{
py/api/images.py
+184 -9
View File
@@ -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 = {
py/api/inspection_service.py
+17 -1
View File
@@ -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,
}
py/autotag.py
+86 -10
View File
@@ -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."""
py/config.py
+57
View File
@@ -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:
py/model_gate.py
+177
View File
@@ -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:<uuid>``."""
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)
tests/browser/server.py
+15
View File
@@ -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")
tests/host_stubs.py
+29
View File
@@ -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
tests/test_api_annotation_suggest.py
+556
View File
@@ -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()
tests/test_api_annotations.py
+1222
View File
File diff suppressed because it is too large Load Diff
tests/test_api_autotag.py
+169 -7
View File
@@ -11,7 +11,6 @@ import shutil
import sys
import tempfile
import threading
import types
import unittest
from unittest.mock import MagicMock
from pathlib import Path
@@ -23,12 +22,12 @@ sys.modules.setdefault("server", _mock_server) # py.config imports it at module
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if "folder_paths" not in sys.modules:
_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
try:
from host_stubs import install_folder_paths_stub
except ImportError: # imported as tests.test_api_autotag
from tests.host_stubs import install_folder_paths_stub
install_folder_paths_stub()
from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase
@@ -37,6 +36,11 @@ from database.operations import PromptDatabase
from py.api import PromptManagerAPI
from utils.hashing import generate_prompt_hash
try: # guarded so a missing module fails assertions, not test loading
from py import model_gate
except ImportError:
model_gate = None
class FakeAutotagService:
"""Stand-in for AutoTagService that records calls and returns fixed tags."""
@@ -128,6 +132,11 @@ class AutotagAPITestCase(AioHTTPTestCase):
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
self._service_patch.stop()
if model_gate is not None:
leaked = model_gate.holder()
if leaked is not None: # never leak the process-wide gate
model_gate.release(leaked)
self.assertIsNone(leaked, "a route left the model gate held")
self.api.db.close_all()
for path in (
self._temp_db.name,
@@ -947,6 +956,159 @@ class TestUnloadAutotag(AutotagAPITestCase):
self.assertEqual(resp.status, 500)
class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase):
"""While another model operation holds the gate, every legacy model route
answers ``409 model-busy`` and touches no model (the one allowed legacy
change)."""
BUSY = {"success": False, "code": "model-busy"}
def setUp(self):
super().setUp()
self.assertIsNotNone(model_gate, "py.model_gate is not implemented")
async def _busy(self, resp):
self.assertEqual(resp.status, 409)
data = await resp.json()
self.assertEqual({k: data[k] for k in self.BUSY}, self.BUSY)
self.assertIn("busy", data["error"].lower())
return data
async def test_legacy_single_download_start_return_model_busy_while_gate_held(
self,
):
image = self._make_image(self.output_dir)
self.service._loaded_type = "wd14-vit"
self.assertTrue(model_gate.try_acquire("test-holder"))
try:
await self._busy(
await self._post_json(
"/prompt_manager/autotag/single",
{"path": str(image), "model_type": "gguf"},
)
)
await self._busy(
await self.client.request(
"POST", "/prompt_manager/autotag/download/gguf"
)
)
await self._busy(
await self.client.request(
"POST", "/prompt_manager/autotag/start?model_type=gguf"
)
)
await self._busy(
await self.client.request("POST", "/prompt_manager/autotag/unload")
)
self.assertEqual(model_gate.holder(), "test-holder")
finally:
model_gate.release("test-holder")
self.assertEqual(self.service.load_calls, [])
self.assertEqual(self.service.generate_calls, [])
self.assertEqual(self.service.download_calls, [])
self.assertEqual(self.service.unload_calls, 0)
self.assertEqual(self.service.get_loaded_model_type(), "wd14-vit")
# A refused start released the batch lock: the next start runs.
resp = await self.client.request(
"POST", "/prompt_manager/autotag/start?model_type=gguf"
)
self.assertEqual(resp.status, 200)
self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete")
async def test_legacy_routes_hold_the_gate_only_while_they_run(self):
image = self._make_image(self.output_dir)
seen = []
def generate(image_path, **kwargs):
seen.append(model_gate.holder())
return ["x"]
def download(model_type, progress_callback=None):
seen.append(model_gate.holder())
return True
def load(model_type, use_gpu=True):
seen.append(model_gate.holder())
self.service._loaded_type = model_type
return True
self.service.generate_tags = generate
self.service.download_model = download
self.service.load_model = load
resp = await self._post_json(
"/prompt_manager/autotag/single", {"path": str(image)}
)
self.assertEqual(resp.status, 200)
self.assertIsNone(model_gate.holder())
resp = await self.client.request(
"POST", "/prompt_manager/autotag/download/gguf"
)
self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete")
self.assertIsNone(model_gate.holder())
resp = await self.client.request("POST", "/prompt_manager/autotag/unload")
self.assertEqual(resp.status, 200)
self.assertIsNone(model_gate.holder())
self.assertEqual(len(seen), 3)
self.assertTrue(all(owner and owner.startswith("legacy-") for owner in seen))
async def test_validation_errors_win_over_the_gate(self):
self.assertTrue(model_gate.try_acquire("test-holder"))
try:
resp = await self._post_json("/prompt_manager/autotag/single", {})
self.assertEqual(resp.status, 400)
resp = await self.client.request(
"POST", "/prompt_manager/autotag/start?general_threshold=x"
)
self.assertEqual(resp.status, 400)
finally:
model_gate.release("test-holder")
async def test_batch_holds_the_gate_until_it_completes(self):
prompt = self._save_prompt("p")
self._link_image(prompt, self._make_image(self.output_dir, "a.png"))
started = threading.Event()
proceed = threading.Event()
def slow_generate(image_path, **kwargs):
started.set()
proceed.wait(30)
return ["slow"]
self.service.generate_tags = slow_generate
run = asyncio.ensure_future(
self.client.request("POST", "/prompt_manager/autotag/start")
)
try:
await asyncio.get_running_loop().run_in_executor(None, started.wait, 30)
self.assertTrue(model_gate.holder().startswith("legacy-batch:"))
await self._busy(
await self.client.request("POST", "/prompt_manager/autotag/unload")
)
finally:
proceed.set()
resp = await run
events = self._sse_events(await resp.text())
self.assertEqual(events[-1]["type"], "complete")
self.assertIsNone(model_gate.holder())
async def test_models_response_carries_the_suggestion_limit(self):
from py.config import AutotagConfig
resp = await self.client.get("/prompt_manager/autotag/models")
data = await resp.json()
self.assertEqual(data["suggestion_limit"], 10)
AutotagConfig.SUGGESTION_LIMIT = 25
try:
resp = await self.client.get("/prompt_manager/autotag/models")
self.assertEqual((await resp.json())["suggestion_limit"], 25)
finally:
AutotagConfig.SUGGESTION_LIMIT = 10
class TestScanOutputDir(AutotagAPITestCase):
async def test_lists_images_with_thumbnails(self):
tests/test_api_core.py
+14
View File
@@ -448,6 +448,20 @@ class TestUiRoutes(CoreRoutesTestCase):
resp = await self.client.request("GET", "/prompt_manager/test")
self.assertTrue((await resp.json())["success"])
async def test_image_annotation_routes_are_registered(self):
table = {
(route.method, route.resource.canonical)
for route in self.app.router.routes()
if route.resource is not None
}
for entry in (
("POST", "/prompt_manager/image-annotations/observe"),
("GET", "/prompt_manager/image-annotations/capture/{token}"),
("DELETE", "/prompt_manager/image-annotations/capture/{token}"),
("PUT", "/prompt_manager/image-annotations/tags"),
):
self.assertIn(entry, table)
class TestStaticRoutes(CoreRoutesTestCase):
tests/test_api_images.py
+6 -1
View File
@@ -18,7 +18,12 @@ from unittest.mock import AsyncMock, MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# ComfyUI-only modules are stubbed so the API can be imported standalone.
sys.modules.setdefault("folder_paths", MagicMock())
try:
from host_stubs import install_folder_paths_stub
except ImportError: # imported as tests.test_api_images
from tests.host_stubs import install_folder_paths_stub
install_folder_paths_stub()
sys.modules.setdefault("server", MagicMock())
from aiohttp import web # noqa: E402
tests/test_api_prompt_inspector.py
+83
View File
@@ -1642,3 +1642,86 @@ class TestPromptInspector(ImageAPITestCase):
lookup.assert_called_once_with(
str(path), False, report_ambiguity=True
)
class TestInspectorImageTags(ImageAPITestCase):
"""B3: the inspector envelope's ``image_tags`` facet (spec amendment S1)."""
inspect = TestPromptInspector.inspect
def tag(self, path, tags):
import hashlib
import os
stat = path.stat()
return self.api.db.replace_image_tags(
locator=os.path.normcase(os.path.realpath(path)),
sha256=hashlib.sha256(path.read_bytes()).hexdigest(),
size=stat.st_size,
mtime_ns=stat.st_mtime_ns,
expected_epoch=self.api.db.get_execution_epoch(),
expected_tag_revision=0,
tags=tags,
)
async def test_inspector_reports_current_none_and_unavailable_image_tags(self):
import os
tagged = png(self.output_dir / "sub" / "tagged.png", prompt=graph())
changed = png(self.output_dir / "changed.png", prompt=graph())
png(self.output_dir / "untagged.png", prompt=graph())
empty = png(self.output_dir / "empty.png", prompt=graph())
self.tag(tagged, ["dunes", "warm_light"])
self.tag(changed, ["dunes"])
self.tag(empty, [])
stat = changed.stat()
os.utime(changed, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000))
response, envelope = await self.inspect("sub/tagged.png")
self.assertEqual(response.status, 200, envelope)
self.assertEqual(
envelope.get("image_tags"),
{"state": "current", "tags": ["dunes", "warm_light"]},
)
# The tags never change the Library facts of the prompt link.
self.assertEqual(envelope["library"]["tags"], [])
_, changed_body = await self.inspect("changed.png")
self.assertEqual(changed_body.get("image_tags"), {"state": "none", "tags": []})
_, untagged = await self.inspect("untagged.png")
self.assertEqual(untagged.get("image_tags"), {"state": "none", "tags": []})
_, saved_empty = await self.inspect("empty.png")
self.assertEqual(
saved_empty.get("image_tags"), {"state": "current", "tags": []}
)
with patch.object(
self.api.db, "current_image_tags", side_effect=RuntimeError("boom")
):
with self.assertLogs("prompt_manager", level="WARNING") as logged:
response, failed = await self.inspect("sub/tagged.png")
self.assertTrue(
any(r.exc_info and "image tags" in r.getMessage() for r in logged.records),
logged.output,
)
self.assertEqual(response.status, 200, failed)
self.assertEqual(failed.get("image_tags"), {"state": "unavailable", "tags": []})
self.assertEqual(failed["prompts"], envelope["prompts"])
self.assertNotIn(str(self.comfy_dir), json.dumps(envelope))
async def test_inspector_image_tags_use_the_inspection_stable_stat(self):
tagged = png(self.output_dir / "tagged.png", prompt=graph())
self.tag(tagged, ["dunes"])
stat = tagged.stat()
seen = []
real = self.api.db.current_image_tags
def spy(stats):
seen.append(dict(stats))
return real(stats)
with patch.object(self.api.db, "current_image_tags", side_effect=spy):
_, body = await self.inspect("tagged.png")
self.assertEqual(
body.get("image_tags"), {"state": "current", "tags": ["dunes"]}
)
self.assertEqual(list(seen[0].values()), [(stat.st_size, stat.st_mtime_ns)])
self.assertEqual(len(seen), 1)
tests/test_api_prompts.py
+6 -1
View File
@@ -11,7 +11,12 @@ from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.modules.setdefault("folder_paths", MagicMock())
try:
from host_stubs import install_folder_paths_stub
except ImportError: # imported as tests.test_api_prompts
from tests.host_stubs import install_folder_paths_stub
install_folder_paths_stub()
sys.modules.setdefault("server", MagicMock())
from aiohttp import web # noqa: E402
tests/test_api_workspace_images.py
+359 -3
View File
@@ -4,6 +4,7 @@ import asyncio
import json
import os
import shutil
import sqlite3
import subprocess # nosec B404
import sys
import unittest
@@ -1232,15 +1233,21 @@ class TestWorkspaceImages(ImageAPITestCase):
{"q": "zebra", "category": "Night"},
{"sort": "rating"},
):
leak = 'near "gi": syntax error FROM generation_images'
with (
self.subTest(failing=params),
patch.object(
db, "get_image_link_matches", side_effect=RuntimeError("gone")
db, "get_image_link_matches", side_effect=RuntimeError(leak)
),
self.assertLogs("prompt_manager", level="ERROR") as logged,
):
response, data = await self.get(**params)
self.assertEqual(response.status, 500, data)
self.assertNotIn("images", data)
self.assertEqual(data["code"], "listing-failed")
self.assertNotIn(leak, json.dumps(data))
self.assertNotIn("syntax error", json.dumps(data))
self.assertTrue(any(r.exc_info for r in logged.records))
async def test_filename_search_survives_a_database_failure(self):
for name, mtime in (("zebra-1.png", 300), ("zebra-2.png", 200)):
@@ -1361,8 +1368,10 @@ class TestWorkspaceImages(ImageAPITestCase):
for query in statements
if query.lstrip().upper().startswith(("SELECT", "WITH"))
]
self.assertEqual(len(selects), 3, selects)
# Plus the one stream of tagged image records q and tags read (B3).
self.assertEqual(len(selects), 4, selects)
self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects)
self.assertEqual(sum("media_image_tags" in q for q in selects), 1, selects)
async def test_rating_sort_orders_linked_then_unlinked(self):
self._filter_fixtures()
@@ -1626,7 +1635,15 @@ class TestWorkspaceImages(ImageAPITestCase):
response, data = await self.get(limit=10, **params)
self.assertEqual(response.status, 200, data)
self.assertEqual(built, [])
self.assertEqual(len(selects()), 1)
found = selects()
# q and tags also stream the tagged image records once
# (#249 B3); none are tagged here, so nothing is stat'ed.
streamed = "q" in params or "tags" in params
self.assertEqual(len(found), 2 if streamed else 1, found)
self.assertEqual(
sum("media_image_tags" in query for query in found),
int(streamed),
)
self.assertIs(
self.api._workspace_order_cache[((0, str(self.output_dir)),)][1][
"newest"
@@ -1773,3 +1790,342 @@ class TestWorkspaceCursorEncoding(unittest.TestCase):
token = _workspace_cursor(key, "scope")
self.assertLessEqual(len(token), 4096)
self.assertEqual(_workspace_cursor_key(token, "scope", order), key)
class TestWorkspaceImageTags(ImageAPITestCase):
"""B3: image tags in the Tags filter and search (spec amendment S1)."""
media = TestWorkspaceImages.media
get = TestWorkspaceImages.get
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def tag(self, path, tags):
"""Save *tags* for the current bytes of *path*, as the Tag dialog does."""
import hashlib
stat = path.stat()
return self.api.db.replace_image_tags(
locator=os.path.normcase(os.path.realpath(path)),
sha256=hashlib.sha256(path.read_bytes()).hexdigest(),
size=stat.st_size,
mtime_ns=stat.st_mtime_ns,
expected_epoch=self.api.db.get_execution_epoch(),
expected_tag_revision=0,
tags=tags,
)
def names(self, data):
return [item["path"] for item in data["images"]]
async def test_tags_filter_admits_unlinked_image_with_current_image_tags(self):
tagged = self.media("unlinked_tagged.png", mtime=103)
self.media("plain.png", mtime=102)
linked = self.media("linked.png", mtime=101)
self.api.db.link_image_to_prompt(
self._save_prompt("Linked", tags=["other"]), str(linked)
)
self.tag(tagged, ["dunes", "warm_light"])
for tags in ("dunes", "dunes,warm_light", "warm_light"):
with self.subTest(tags=tags):
response, body = await self.get(tags=tags)
self.assertEqual(response.status, 200, body)
self.assertEqual(
[i["path"] for i in body["images"]], ["unlinked_tagged.png"]
)
self.assertEqual(body["total"], 1)
self.assertEqual(body["positive_count"], 0)
# Exact names, as for prompt tags; a missing tag excludes the file.
for tags in ("Dunes", "dunes,missing", "dune"):
with self.subTest(tags=tags):
_, body = await self.get(tags=tags)
self.assertEqual(self.names(body), [])
# A category filter reads only the linked positive.
_, body = await self.get(tags="dunes", category="Night")
self.assertEqual(self.names(body), [])
async def test_tags_filter_combines_image_and_positive_tags(self):
both = self._save_prompt("Both", category="Night", tags=["a", "b"], rating=2)
only_a = self._save_prompt("Only a", category="Night", tags=["a"], rating=5)
day_a = self._save_prompt("Day a", category="Day", tags=["a"], rating=4)
files = {
name: self.media(f"{name}.png", mtime=100 + index)
for index, name in enumerate(
("positive_ab", "positive_a_image_b", "image_ab", "image_a", "day")
)
}
self.api.db.link_image_to_prompt(both, str(files["positive_ab"]))
self.api.db.link_image_to_prompt(only_a, str(files["positive_a_image_b"]))
self.api.db.link_image_to_prompt(day_a, str(files["day"]))
self.tag(files["positive_a_image_b"], ["b"])
self.tag(files["image_ab"], ["a", "b"])
self.tag(files["image_a"], ["a"])
self.tag(files["day"], ["b"])
_, body = await self.get(tags="a,b", sort="name")
self.assertEqual(
self.names(body),
["day.png", "image_ab.png", "positive_a_image_b.png", "positive_ab.png"],
)
self.assertEqual(body["total"], 4)
self.assertEqual(body["positive_count"], 3)
_, body = await self.get(tags="a,b", category="Night", sort="name")
self.assertEqual(
self.names(body), ["positive_a_image_b.png", "positive_ab.png"]
)
# The rating sort ranks an image-admitted linked file by its positive.
_, body = await self.get(tags="a,b", sort="rating")
self.assertEqual(
self.names(body),
["positive_a_image_b.png", "day.png", "positive_ab.png", "image_ab.png"],
)
# With q, a linked positive hit admits as before.
_, body = await self.get(tags="b", q="only", sort="name")
self.assertEqual(self.names(body), ["positive_a_image_b.png"])
async def test_q_matches_image_tag_names(self):
tagged = self.media("unlinked_tagged.png", mtime=102)
self.media("plain.png", mtime=101)
self.tag(tagged, ["Sand Dunes", "warm_light"])
for query in ("dune", "DUNES", "sand d", "warm_"):
with self.subTest(q=query):
response, body = await self.get(q=query)
self.assertEqual(response.status, 200, body)
self.assertEqual(self.names(body), ["unlinked_tagged.png"])
_, body = await self.get(q="plain")
self.assertEqual(self.names(body), ["plain.png"])
_, body = await self.get(q="nothing")
self.assertEqual(self.names(body), [])
async def test_filter_ignores_tags_of_a_changed_fingerprint(self):
moved = self.media("touched.png", mtime=102)
grown = self.media("grown.png", mtime=101)
kept = self.media("kept.png", mtime=100)
for path in (moved, grown, kept):
self.tag(path, ["dunes"])
_, body = await self.get(tags="dunes")
self.assertEqual(self.names(body), ["touched.png", "grown.png", "kept.png"])
os.utime(moved, (500, 500))
with grown.open("ab") as handle:
handle.write(b"\0")
os.utime(grown, (101, 101))
self.api.invalidate_gallery_cache()
_, body = await self.get(tags="dunes")
self.assertEqual(self.names(body), ["kept.png"])
_, body = await self.get(q="dunes")
self.assertEqual(self.names(body), ["kept.png"])
async def test_stat_calls_equal_tagged_candidates_not_listing_size(self):
from py.api import images as images_module
listed = [
(self.output_dir / f"day-{i % 50:02}" / f"img_{i:05}.png", 1_000.0 + i)
for i in range(1_997)
]
tagged = [self.media(f"tagged_{i}.png", mtime=5_000 + i) for i in range(3)]
listed += [(path, path.stat().st_mtime) for path in tagged]
# Tagged but outside the listing, and listed with a non-matching tag.
outside = self.media("elsewhere.png", root=self.comfy_dir)
unrelated = self.media("unrelated.png", mtime=4_000)
listed.append((unrelated, unrelated.stat().st_mtime))
for path in tagged:
self.tag(path, ["dunes"])
self.tag(outside, ["dunes"])
self.tag(unrelated, ["other"])
self.assertEqual(len(listed), 2_001)
stat_calls = 0
real = getattr(images_module, "_current_fingerprint", None)
self.assertIsNotNone(real, "the candidate stat helper is not implemented")
def counting(path):
nonlocal stat_calls
stat_calls += 1
return real(path)
with (
patch.object(
self.api, "_get_gallery_files", AsyncMock(return_value=listed)
),
patch.object(
self.api, "_format_workspace_page", TestWorkspaceImages._stub_page
),
patch.object(images_module, "_current_fingerprint", counting),
):
response, body = await self.get(tags="dunes", limit=10)
self.assertEqual(response.status, 200, body)
self.assertEqual(
sorted(i["filename"] for i in body["images"]),
["tagged_0.png", "tagged_1.png", "tagged_2.png"],
)
self.assertEqual(stat_calls, 3)
async def test_plain_listing_reads_no_image_tags(self):
prompt = self._save_prompt("Linked", category="Night", tags=["a"])
for index in range(4):
path = self.media(f"{index}.png", mtime=100 + index)
self.tag(path, ["dunes"])
self.api.db.link_image_to_prompt(prompt, str(path))
statements = []
original = self.api.db.model.get_connection
def trace_connection():
connection = original()
connection.set_trace_callback(statements.append)
return connection
with patch.object(
self.api.db.model, "get_connection", side_effect=trace_connection
):
_, first = await self.get(limit=2)
_, rest = await self.get(limit=2, cursor=first["next_cursor"])
for params in (
{"sort": "rating"},
{"category": "Night"},
{"media": "image"},
):
response, _ = await self.get(**params)
self.assertEqual(response.status, 200)
self.assertEqual(len(first["images"] + rest["images"]), 4)
self.assertFalse(
[q for q in statements if "media_" in q or "annotation" in q], statements
)
async def test_image_tag_read_failure_fails_tags_and_degrades_search(self):
tagged = self.media("unlinked_tagged.png", mtime=102)
self.media("dunes_by_name.png", mtime=101)
self.tag(tagged, ["dunes"])
leak = "no such table: media_annotation_files"
with patch.object(
self.api.db,
"iter_image_tag_records",
side_effect=sqlite3.OperationalError(leak),
):
# As a failed link read does under a Tags filter: a fixed coded
# sentence, never the exception or SQL text; the traceback is logged.
with self.assertLogs("prompt_manager", level="ERROR") as failed:
response, body = await self.get(tags="dunes")
self.assertEqual((response.status, body["success"]), (500, False))
self.assertEqual(body["code"], "listing-failed")
self.assertNotIn(leak, body["error"])
self.assertNotIn("media_annotation_files", json.dumps(body))
self.assertTrue(any(r.exc_info for r in failed.records))
with self.assertLogs("prompt_manager", level="WARNING") as logged:
response, body = await self.get(q="dunes")
self.assertEqual(response.status, 200, body)
self.assertEqual(self.names(body), ["dunes_by_name.png"])
self.assertIn("image tags", "\n".join(logged.output))
async def test_legacy_output_failure_logs_its_traceback(self):
self.media("legacy.png", mtime=100)
with (
patch.object(
self.api, "_collect_output_images", side_effect=RuntimeError("x")
),
self.assertLogs("prompt_manager", level="ERROR") as logged,
):
response = await self.client.get("/prompt_manager/images/output")
self.assertEqual(response.status, 500)
self.assertTrue(any(r.exc_info for r in logged.records), logged.output)
async def test_legacy_gallery_and_search_shapes_unchanged(self):
path = self.media("legacy.png", mtime=100)
self.media("other.png", mtime=99)
self._link_image(path, "dunes at dusk")
async def legacy():
bodies = []
for url, params in (
("/prompt_manager/images/output", {}),
("/prompt_manager/images/search", {"q": "dunes"}),
("/prompt_manager/images/all", {}),
):
response = await self.client.get(url, params=params)
self.assertEqual(response.status, 200, url)
bodies.append(await response.json())
return bodies
before = await legacy()
self.tag(path, ["dunes", "warm_light"])
self.tag(self.output_dir / "other.png", ["dunes"])
self.api.invalidate_gallery_cache()
self.assertEqual(await legacy(), before)
def test_link_facts_resolve_pairs_skip_ambiguous_and_chunk(self):
from database.models import normalize_image_path
db = self.api.db
plain = self._save_prompt("Plain dunes", category="Day", tags=["a", "b"])
positive = self._save_prompt("Pair positive", category="Night", tags=["p"])
pair = db.save_prompt_pair("Pair positive", "Noise", title="Pair")
owners = (self._save_prompt("Owner one"), self._save_prompt("Owner two"))
paths = {
name: str(self.output_dir / f"{name}.png")
for name in ("plain", "paired", "ambiguous")
}
db.link_image_to_prompt(plain, paths["plain"])
image = db.link_image_to_prompt(positive, paths["paired"])
db.link_image_to_pair(pair, image, {"source": "fixture"})
for owner in owners:
db.link_image_to_prompt(owner, paths["ambiguous"])
keys = {name: normalize_image_path(path) for name, path in paths.items()}
padding = [f"/nowhere/{index}.png" for index in range(600)]
facts = db.get_image_link_facts([*padding, *keys.values(), keys["plain"]])
self.assertEqual(
facts,
{
keys["plain"]: (True, None, plain, "Day", ("a", "b")),
keys["paired"]: (True, None, positive, "Night", ("p",)),
},
)
searched = db.get_image_link_facts(list(keys.values()), "dunes")
self.assertEqual(
{path: row[0] for path, row in searched.items()},
{keys["plain"]: True, keys["paired"]: False},
)
self.assertEqual(db.get_image_link_facts([]), {})
def test_new_link_fact_reads_use_indexes(self):
from database.operations import (
_LINK_FACTS_SQL,
_LINK_FILTER_HIT,
_POSITIVE_TAG_NAMES_SQL,
)
with self.api.db.model.get_connection() as conn:
conn.executemany(
"INSERT INTO prompts(id,text,hash) VALUES(?,?,?)",
[(i, f"prompt {i}", f"hash-{i}") for i in range(1, 3_001)],
)
conn.executemany(
"INSERT INTO generated_images(prompt_id,image_path,filename,"
"file_path) VALUES(?,?,?,?)",
[
(i, f"/x/{i}.png", f"{i}.png", f"/x/{i}.png")
for i in range(1, 3_001)
],
)
for hit, params in (
("1", {}),
(_LINK_FILTER_HIT, {"pattern": "%prompt 1%"}),
):
sql = _LINK_FACTS_SQL.format(hit=hit, marks=":path0,:path1")
bound = {**params, "path0": "/x/1.png", "path1": "/x/2.png"}
plan = [
row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, bound)
]
self.assertIn(
"SEARCH gi USING INDEX idx_image_file_path (file_path=?)", plan
)
self.assertFalse(
[s for s in plan if s.startswith(("SCAN gi", "SCAN other"))], plan
)
plan = [
row[3]
for row in conn.execute(
"EXPLAIN QUERY PLAN " + _POSITIVE_TAG_NAMES_SQL.format(marks="?,?"),
(1, 2),
)
]
self.assertFalse([step for step in plan if step.startswith("SCAN")], plan)
tests/test_autotag_engine.py
+153
View File
@@ -301,6 +301,159 @@ class TestGgufNormalisation(EngineTestCase):
)
SCORED_TAGS = [
{"name": "desert", "category": "0"},
{"name": "sandstone", "category": "0"},
{"name": "Sandstone", "category": "0"},
{"name": "hatsune_miku", "category": "4"},
{"name": "rating:safe", "category": "9"},
] + [{"name": f"tag_{i:03d}", "category": "0"} for i in range(120)]
SCORED_PROBABILITIES = [0.96, 0.98, 0.5, 0.9, 0.99] + [
0.9 - i * 0.004 for i in range(120)
]
class TestTagSuggestions(EngineTestCase):
"""``generate_tag_suggestions``: scored, ordered, bounded, distinct."""
def _wd14(self, tags=SCORED_TAGS, probabilities=SCORED_PROBABILITIES):
_touch_model_files(self.models_dir, "wd14-vit")
self.session = FakeOnnxSession(probabilities)
service = self._service(("wd14", (self.session, tags)))
service.load_model("wd14-vit")
return service
def _gguf(self, text):
_touch_model_files(self.models_dir, "gguf")
self.model = FakeGgufModel(text)
service = self._service(("gguf", self.model))
service.load_model("gguf")
return service
def test_wd14_suggestions_carry_real_probabilities_sorted(self):
service = self._wd14()
suggestions = service.generate_tag_suggestions(
str(self.image_path),
limit=10,
general_threshold=0.35,
character_threshold=0.85,
)
self.assertEqual(suggestions[:2], [("sandstone", 0.98), ("desert", 0.96)])
# The character tag passes its own threshold; rating tags never show.
self.assertIn(("hatsune_miku", 0.9), suggestions)
self.assertNotIn("rating:safe", [name for name, _ in suggestions])
scores = [score for _, score in suggestions]
self.assertEqual(scores, sorted(scores, reverse=True))
self.assertTrue(all(isinstance(score, float) for score in scores))
# "Sandstone" (0.5) repeats "sandstone" case-insensitively: dropped.
self.assertEqual(
[name.casefold() for name, _ in suggestions].count("sandstone"), 1
)
def test_suggestions_keep_the_requested_limit(self):
service = self._wd14()
for limit in (1, 10, 100):
with self.subTest(limit=limit):
suggestions = service.generate_tag_suggestions(
str(self.image_path),
limit=limit,
general_threshold=0.0,
character_threshold=0.0,
)
self.assertEqual(len(suggestions), limit)
for bad in (0, -1, 101, 2.5, "10", True, None):
with self.subTest(bad=bad), self.assertRaises(ValueError):
service.generate_tag_suggestions(str(self.image_path), limit=bad)
def test_thresholds_default_to_the_service_values(self):
service = self._wd14()
service.wd14_general_threshold = 0.95
service.wd14_character_threshold = 0.95
suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10)
self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)])
def test_llm_suggestions_have_null_confidence_in_output_order(self):
service = self._gguf("Zebra, apple, ZEBRA, mango tree, kiwi, plum")
suggestions = service.generate_tag_suggestions(
str(self.image_path), limit=3, prompt="tags please"
)
self.assertEqual(
suggestions, [("zebra", None), ("apple", None), ("mango_tree", None)]
)
def test_suggestions_skip_names_save_would_reject_before_the_limit(self):
from database.annotations import normalize_image_tags
too_long = "x" * 60
service = self._gguf(f"{too_long}, sky\ncloud, zebra, apple, mango, kiwi, plum")
suggestions = service.generate_tag_suggestions(
str(self.image_path), limit=3, prompt="tags please"
)
# The 60-character part and the part holding a newline are dropped
# before the count, so three valid chips still come back.
names = [name for name, _ in suggestions]
self.assertEqual(names, ["zebra", "apple", "mango"])
self.assertEqual(normalize_image_tags(names), names)
def test_generate_tags_return_shape_unchanged(self):
service = self._wd14()
self.assertEqual(
service.generate_tags(
str(self.image_path), general_threshold=0.95, character_threshold=0.95
),
["sandstone", "desert"],
)
gguf = self._gguf("Smile, smile, long hair")
self.assertEqual(
gguf.generate_tags(str(self.image_path)), ["smile", "long_hair"]
)
def test_explicit_prompt_does_not_mutate_custom_prompt(self):
service = self._gguf("tag")
service.custom_prompt = "saved legacy prompt"
service.generate_tag_suggestions(
str(self.image_path), limit=10, prompt="dialog instructions"
)
self.assertEqual(
self.model.calls[0]["messages"][1]["content"][0]["text"],
"dialog instructions",
)
self.assertEqual(service.custom_prompt, "saved legacy prompt")
def test_suggestions_need_a_loaded_model(self):
service = self._service(("gguf", FakeGgufModel("x")))
with self.assertRaises(RuntimeError):
service.generate_tag_suggestions(str(self.image_path), limit=10)
def test_image_file_is_closed_after_inference(self):
service = self._gguf("tag")
opened = []
real_open = autotag.Image.open
def tracking_open(*args, **kwargs):
image = real_open(*args, **kwargs)
opened.append(image)
return image
with patch.object(autotag.Image, "open", tracking_open):
service.generate_tag_suggestions(str(self.image_path), limit=10)
service.generate_tags(str(self.image_path))
self.assertEqual(len(opened), 2)
for image in opened:
self.assertIsNone(getattr(image, "fp", None))
class TestGenerateTagsErrors(EngineTestCase):
def test_no_model_loaded(self):
tests/test_config.py
+81
View File
@@ -19,6 +19,7 @@ sys.modules.setdefault("server", _mock_server)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py import config as config_module
from py.config import GalleryConfig, IntegrationConfig, PromptManagerConfig
@@ -404,5 +405,85 @@ class TestConfigPathOverride(unittest.TestCase):
self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 77)
class TestAutotagConfig(unittest.TestCase):
"""``autotag.defaults.suggestion_limit``: integer 1..100, default 10 (D68)."""
def setUp(self):
self.config_class = getattr(config_module, "AutotagConfig", None)
self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented")
self._orig = PromptManagerConfig.get_config()
def tearDown(self):
self.config_class.SUGGESTION_LIMIT = 10
PromptManagerConfig.update_config(self._orig)
def _load(self, value):
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "config.json")
with open(path, "w") as f:
# json writes NaN for float("nan"), which json.load reads back.
json.dump({"autotag": {"defaults": {"suggestion_limit": value}}}, f)
with self.assertLogs("prompt_manager.config", level="WARNING") as logs:
PromptManagerConfig.load_from_file(path)
return logs
def test_every_copy_of_the_tag_set_cap_is_one_value(self):
# A suggestion request is bounded by the image tag set cap; the
# copies that cannot import it (config loads before the database,
# the inspection envelope stays storage-free) are pinned to it here.
from database.annotations import MAX_IMAGE_TAGS
from py import autotag
from utils import inspection_response
self.assertEqual(autotag.MAX_SUGGESTION_LIMIT, MAX_IMAGE_TAGS)
self.assertEqual(config_module.MAX_SUGGESTION_LIMIT, MAX_IMAGE_TAGS)
self.assertEqual(inspection_response.MAX_IMAGE_TAGS, MAX_IMAGE_TAGS)
def test_default_is_ten_and_persisted_under_autotag_defaults(self):
self.assertEqual(config_module.DEFAULT_SUGGESTION_LIMIT, 10)
self.assertEqual(config_module.MAX_SUGGESTION_LIMIT, 100)
self.assertEqual(self.config_class.suggestion_limit(), 10)
self.assertEqual(
PromptManagerConfig.get_config()["autotag"],
{"defaults": {"suggestion_limit": 10}},
)
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "config.json")
PromptManagerConfig.update_config(
{"autotag": {"defaults": {"suggestion_limit": 25}}}
)
PromptManagerConfig.save_to_file(path)
self.config_class.SUGGESTION_LIMIT = 10
PromptManagerConfig.load_from_file(path)
self.assertEqual(self.config_class.suggestion_limit(), 25)
def test_valid_limits_are_kept(self):
for value in (1, 10, 100):
with self.subTest(value=value):
PromptManagerConfig.update_config(
{"autotag": {"defaults": {"suggestion_limit": value}}}
)
self.assertEqual(self.config_class.suggestion_limit(), value)
def test_suggestion_limit_setting_validated_on_load(self):
for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]):
with self.subTest(value=value):
self.config_class.SUGGESTION_LIMIT = 42
logs = self._load(value)
self.assertEqual(self.config_class.suggestion_limit(), 10)
self.assertIn("suggestion_limit", "\n".join(logs.output))
def test_missing_or_malformed_sections_keep_the_current_value(self):
self.config_class.SUGGESTION_LIMIT = 7
for config in ({}, {"autotag": {}}, {"autotag": {"defaults": {}}}):
PromptManagerConfig.update_config(config)
self.assertEqual(self.config_class.suggestion_limit(), 7)
with self.assertLogs("prompt_manager.config", level="WARNING"):
PromptManagerConfig.update_config({"autotag": "nope"})
with self.assertLogs("prompt_manager.config", level="WARNING"):
PromptManagerConfig.update_config({"autotag": {"defaults": [1]}})
self.assertEqual(self.config_class.suggestion_limit(), 7)
if __name__ == "__main__":
unittest.main()
tests/test_inspection_response.py
+65
View File
@@ -823,6 +823,71 @@ def _png_with_chunks(path, chunks):
return path
class TestImageTagsFacet(unittest.TestCase):
"""B3: the rooted envelope's ``image_tags`` facet is validated."""
def encoded(self, facet):
import inspect as signature_of
parameters = signature_of.signature(encode_inspection_response).parameters
if "image_tags" not in parameters:
self.fail("the image_tags facet is not implemented")
body = encode_inspection_response(
image=IMAGE,
root_scope="a" * 64,
file={
"observation": "current-file",
"source_label": "output",
"media_type": "image",
"extension": ".png",
"size": 100,
"modified_time": 100,
},
prompts=PROMPTS,
resolution={"state": "resolved", "sampler_id": "3", "reason": ""},
inspection=inspection({"prompt": graph_text()}),
library={"status": "none"},
image_tags=facet,
)
return json.loads(body)
def test_valid_facets_pass_through(self):
for facet in (
{"state": "current", "tags": ["dunes", "warm_light"]},
{"state": "current", "tags": []},
{"state": "none", "tags": []},
{"state": "unavailable", "tags": []},
{"state": "current", "tags": [f"tag {i}" for i in range(100)]},
):
with self.subTest(facet=facet):
self.assertEqual(self.encoded(facet)["image_tags"], facet)
def test_envelope_rejects_malformed_image_tags(self):
unavailable = {"state": "unavailable", "tags": []}
for facet in (
"current",
["dunes"],
{"state": "current"},
{"state": "stale", "tags": []},
{"state": "current", "tags": "dunes"},
{"state": "current", "tags": ["dunes", 3]},
{"state": "current", "tags": [""]},
{"state": "current", "tags": [" dunes"]},
{"state": "current", "tags": ["a,b"]},
{"state": "current", "tags": ["x" * 51]},
{"state": "current", "tags": [f"tag {i}" for i in range(101)]},
{"state": "none", "tags": ["dunes"]},
{"state": "unavailable", "tags": ["dunes"]},
{"state": "current", "tags": [], "locator": "/abs/path.png"},
):
with self.subTest(facet=facet):
self.assertEqual(self.encoded(facet)["image_tags"], unavailable)
def test_facet_is_absent_unless_given(self):
self.assertNotIn("image_tags", json.loads(encode(inspection({}))))
self.assertNotIn("image_tags", self.encoded(None))
class TestVideoProbeBoundsMatchPngBudgets(unittest.TestCase):
"""ffprobe escapes tag text; the stdout bounds must never gate a record
that the identical PNG reads within its decoded budget (I3)."""
tests/test_media_capture.py
+130 -1
View File
@@ -7,7 +7,9 @@ import struct
import tempfile
import threading
import time
import unicodedata
import unittest
import warnings
import zlib
from pathlib import Path
from unittest.mock import patch
@@ -239,6 +241,41 @@ class CaptureStoreTestCase(unittest.TestCase):
self.capture(store, "alias.png")
self.assert_empty(store)
def test_case_and_unicode_spellings_of_a_listed_file_are_denied(self):
# On a case- or normalization-insensitive filesystem (macOS APFS,
# Windows) these spellings open the listed file while realpath keeps
# the requested spelling, which would store its tags under a second
# locator. Every OS denies them; the listed spelling is the control.
nfc = unicodedata.normalize("NFC", "caf\u00e9.png")
nfd = unicodedata.normalize("NFD", nfc)
self.assertNotEqual(nfc, nfd)
self.write("a.png", png_bytes(1, 1))
self.write(nfc, png_bytes(1, 1))
self.write("Sub/b.png", png_bytes(1, 1))
store = self.make_store()
for relative in ("a.png", nfc, "Sub/b.png"):
with self.subTest(listed=relative):
cap = self.capture(store, relative)
self.assertEqual(
cap.locator,
os.path.normcase(
os.path.join(os.path.realpath(self.root), *relative.split("/"))
),
)
self.assertTrue(store.release(cap.token))
for alias in ("A.PNG", "A.png", nfd, "sub/b.png", "SUB/B.PNG", "Sub/B.png"):
with self.subTest(alias=alias):
with self.assertRaises(CaptureDenied):
self.capture(store, alias)
with self.assertRaises(CaptureDenied):
store.verify(self.root, alias)
# A missing file with no other spelling keeps its usual outcome.
with self.assertRaises(CaptureChanged):
self.capture(store, "Sub/missing.png")
self.assert_empty(store)
assert_closed(self, self.parent)
def test_locator_matches_listing_identity_for_symlinked_root(self):
real_root = self.parent / "real-output"
real_root.mkdir()
@@ -596,6 +633,98 @@ class CaptureStoreTestCase(unittest.TestCase):
cap.sha256, hashlib.sha256(b"\xff\xd8not really a jpeg").hexdigest()
)
def test_metadata_over_inspection_budgets_does_not_block_generate(self):
# B2 carry-over (A M1): inferable is decided from the image header
# alone, so prompt metadata over the inspection budgets (here a 3 MiB
# text chunk) leaves a decodable image inferable.
base = png_bytes(64, 32)
huge = _chunk(b"tEXt", b"prompt\x00" + b"x" * (3 * 1024 * 1024))
self.write("big-metadata.png", base[:-12] + huge + base[-12:])
store = self.make_store()
cap = self.capture(store, "big-metadata.png")
self.assertEqual((cap.width, cap.height), (64, 32))
self.assertTrue(cap.inferable)
self.assertIsNone(cap.inferable_reason)
assert_closed(self, self.parent)
def test_decompression_bomb_header_is_not_inferable(self):
self.write("bomb.png", png_bytes(20_000, 20_000))
store = self.make_store()
cap = self.capture(store, "bomb.png")
self.assertFalse(cap.inferable)
self.assertIn("4096", cap.inferable_reason)
assert_closed(self, self.parent)
def test_header_dimensions_reads_only_the_header(self):
path = self.write("dims.png", png_bytes(7, 5))
self.assertEqual(capture_module.header_dimensions(str(path)), (7, 5))
assert_closed(self, self.parent)
with self.assertRaises(capture_module.HeaderUnreadable):
capture_module.header_dimensions(
str(self.write("bad.png", b"\x89PNG\r\n\x1a\nnope"))
)
assert_closed(self, self.parent)
def test_header_dimensions_leaves_the_warning_filters_alone(self):
# The capture workers run in the ComfyUI process: touching the
# process-global warnings filters from them races (catch_warnings is
# not thread-safe) and can leave an "error" filter installed for every
# later Pillow open in the host.
from PIL import Image
path = str(self.write("dims.png", png_bytes(7, 5)))
before = list(warnings.filters)
seen_during = []
real_open = Image.open
def recording_open(*args, **kwargs):
seen_during.append(list(warnings.filters))
return real_open(*args, **kwargs)
with patch.object(Image, "open", recording_open):
self.assertEqual(capture_module.header_dimensions(path), (7, 5))
self.assertEqual(seen_during, [before])
self.assertEqual(list(warnings.filters), before)
barrier = threading.Barrier(8)
errors = []
def worker():
try:
barrier.wait()
for _ in range(50):
capture_module.header_dimensions(path)
except Exception as exc: # pragma: no cover - reported below
errors.append(exc)
threads = [threading.Thread(target=worker) for _ in range(8)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
self.assertEqual(errors, [])
self.assertEqual(list(warnings.filters), before)
assert_closed(self, self.parent)
def test_header_over_the_warning_size_is_refused_by_the_pixel_cap(self):
# 100 MP: past Pillow's warning size, under its error size. The
# header is read, then the inference pixel cap refuses it.
from PIL import Image
path = str(self.write("big.png", png_bytes(10_000, 10_000)))
with self.assertWarns(Image.DecompressionBombWarning):
self.assertEqual(capture_module.header_dimensions(path), (10_000, 10_000))
store = self.make_store()
with self.assertWarns(Image.DecompressionBombWarning):
cap = self.capture(store, "big.png")
self.assertFalse(cap.inferable)
self.assertIn("4096", cap.inferable_reason)
assert_closed(self, self.parent)
def test_video_capture_is_not_inferable(self):
data = b"\x00\x00\x00\x18ftypmp42" + os.urandom(32)
self.write("clip.mp4", data)
@@ -604,7 +733,7 @@ class CaptureStoreTestCase(unittest.TestCase):
def never(path): # pragma: no cover - guard only
raise AssertionError("video captures are never probed")
with patch.object(capture_module, "read_inspection_file", never):
with patch.object(capture_module, "header_dimensions", never):
cap = self.capture(store, "clip.mp4")
self.assertEqual((cap.media_type, cap.width, cap.height), ("video", None, None))
self.assertFalse(cap.inferable)
tests/test_model_gate.py
+217
View File
@@ -0,0 +1,217 @@
"""The shared, non-blocking model gate (#249 B2).
One owner at a time may load, swap, run, download or unload a tagging model.
A release requested while gated work still runs in a worker thread takes
effect only when that work really exits, so a client disconnect never lets a
second owner swap the model under a running inference.
"""
import asyncio
import os
import sys
import threading
import unittest
import unittest.mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try: # guarded so a missing module fails assertions, not test loading
from py import model_gate
except ImportError:
model_gate = None
WAIT = 30.0 # ceiling for cross-thread handoffs, never a short clock
class GateTestCase(unittest.TestCase):
def setUp(self):
self.assertIsNotNone(model_gate, "py.model_gate is not implemented")
self.gate = model_gate.ModelGate()
class TestGate(GateTestCase):
def test_gate_is_nonblocking_and_single_owner(self):
self.assertIsNone(self.gate.holder())
self.assertTrue(self.gate.try_acquire("legacy-batch"))
self.assertEqual(self.gate.holder(), "legacy-batch")
# A second owner is refused at once, never queued.
self.assertFalse(self.gate.try_acquire("image-annotations"))
self.assertFalse(self.gate.try_acquire("legacy-batch"))
self.assertEqual(self.gate.holder(), "legacy-batch")
# Only the holder can release.
self.gate.release("image-annotations")
self.assertEqual(self.gate.holder(), "legacy-batch")
self.gate.release("legacy-batch")
self.assertIsNone(self.gate.holder())
self.gate.release("legacy-batch") # idempotent
self.assertTrue(self.gate.try_acquire("image-annotations"))
self.assertEqual(self.gate.holder(), "image-annotations")
def test_release_waits_for_running_work(self):
self.assertTrue(self.gate.try_acquire("owner"))
self.gate.begin("owner")
self.gate.release("owner")
self.assertEqual(self.gate.holder(), "owner")
self.assertFalse(self.gate.try_acquire("other"))
# No new work is admitted once the owner let go.
with self.assertRaises(model_gate.GateNotHeld):
self.gate.begin("owner")
self.gate.end("owner")
self.assertIsNone(self.gate.holder())
def test_begin_requires_the_holder(self):
with self.assertRaises(model_gate.GateNotHeld):
self.gate.begin("nobody")
self.assertTrue(self.gate.try_acquire("owner"))
with self.assertRaises(model_gate.GateNotHeld):
self.gate.begin("other")
self.gate.end("other") # a stray end never disturbs the holder
self.assertEqual(self.gate.holder(), "owner")
def test_module_gate_is_one_process_wide_instance(self):
self.assertIsNone(model_gate.holder())
self.assertTrue(model_gate.try_acquire("a"))
try:
self.assertFalse(model_gate.try_acquire("b"))
self.assertEqual(model_gate.holder(), "a")
finally:
model_gate.release("a")
self.assertIsNone(model_gate.holder())
def test_new_owner_names_are_unique(self):
first = model_gate.new_owner("legacy-single")
second = model_gate.new_owner("legacy-single")
self.assertNotEqual(first, second)
self.assertTrue(first.startswith("legacy-single:"))
class TestGatedExecutor(GateTestCase):
def test_cancelled_waiter_keeps_the_gate_until_the_thread_exits(self):
started = threading.Event()
proceed = threading.Event()
finished = threading.Event()
def blocking_work():
started.set()
proceed.wait(WAIT)
finished.set()
return "done"
async def scenario():
self.assertTrue(self.gate.try_acquire("native"))
task = asyncio.ensure_future(self.gate.run("native", blocking_work))
try:
await asyncio.get_running_loop().run_in_executor(
None, started.wait, WAIT
)
task.cancel()
with self.assertRaises(asyncio.CancelledError):
await task
finally:
self.gate.release("native")
# The waiter is gone but the inference still runs: still held.
self.assertEqual(self.gate.holder(), "native")
self.assertFalse(self.gate.try_acquire("legacy-unload"))
proceed.set()
await asyncio.get_running_loop().run_in_executor(None, finished.wait, WAIT)
for _ in range(1000):
if self.gate.holder() is None:
break
await asyncio.sleep(0.01)
self.assertIsNone(self.gate.holder())
asyncio.run(scenario())
def test_run_returns_the_result_and_propagates_errors(self):
async def scenario():
self.assertTrue(self.gate.try_acquire("owner"))
try:
self.assertEqual(await self.gate.run("owner", lambda x: x * 2, 21), 42)
def boom():
raise ValueError("bad")
with self.assertRaises(ValueError):
await self.gate.run("owner", boom)
self.assertEqual(self.gate.holder(), "owner")
finally:
self.gate.release("owner")
self.assertIsNone(self.gate.holder())
with self.assertRaises(model_gate.GateNotHeld):
await self.gate.run("owner", lambda: None)
asyncio.run(scenario())
def test_work_that_never_starts_gives_its_admission_back(self):
async def scenario():
loop = asyncio.get_running_loop()
self.assertTrue(self.gate.try_acquire("owner"))
def refused(executor, func):
raise RuntimeError("executor shut down")
with unittest.mock.patch.object(loop, "run_in_executor", refused):
with self.assertRaises(RuntimeError):
await self.gate.run("owner", lambda: None)
def cancelled(executor, func):
future = loop.create_future()
future.cancel() # a loop shutdown drops the queued call
return future
with unittest.mock.patch.object(loop, "run_in_executor", cancelled):
with self.assertRaises(asyncio.CancelledError):
await self.gate.run("owner", lambda: None)
await asyncio.sleep(0) # let the done callback run
self.gate.release("owner")
self.assertIsNone(self.gate.holder())
asyncio.run(scenario())
def test_skipped_work_runs_its_cleanup_exactly_once(self):
async def scenario():
loop = asyncio.get_running_loop()
skipped = []
def cleanup():
skipped.append(True)
self.assertTrue(self.gate.try_acquire("owner"))
self.assertEqual(
await self.gate.run("owner", lambda: 7, on_skipped=cleanup), 7
)
self.assertEqual(skipped, [])
def refused(executor, func):
raise RuntimeError("executor shut down")
with unittest.mock.patch.object(loop, "run_in_executor", refused):
with self.assertRaises(RuntimeError):
await self.gate.run("owner", lambda: None, on_skipped=cleanup)
self.assertEqual(skipped, [True])
def cancelled(executor, func):
future = loop.create_future()
future.cancel()
return future
with unittest.mock.patch.object(loop, "run_in_executor", cancelled):
with self.assertRaises(asyncio.CancelledError):
await self.gate.run("owner", lambda: None, on_skipped=cleanup)
await asyncio.sleep(0) # let the done callback run
self.assertEqual(skipped, [True, True])
self.gate.release("owner")
with self.assertRaises(model_gate.GateNotHeld):
await self.gate.run("owner", lambda: None, on_skipped=cleanup)
self.assertEqual(skipped, [True, True, True])
self.assertIsNone(self.gate.holder())
asyncio.run(scenario())
if __name__ == "__main__":
unittest.main()
utils/inspection_response.py
+43 -1
View File
@@ -6,12 +6,17 @@ from decimal import Decimal, InvalidOperation
from typing import TYPE_CHECKING, Any, Literal, Optional
from .inspection_outcomes import MESSAGE_LENGTH, project_resolution
from .validators import validate_tags
if TYPE_CHECKING:
from .prompt_inspector import InspectionRead
MAX_RESPONSE_BYTES = 8 * 1024 * 1024
MAX_LIBRARY_BYTES = 32 * 1024
# The image tag set cap (database/annotations.py MAX_IMAGE_TAGS; pinned equal
# by tests/test_config.py).
MAX_IMAGE_TAGS = 100
_IMAGE_TAG_STATES = frozenset({"current", "none", "unavailable"})
_MAX_SAFE_INTEGER = 2**53 - 1
_RECORD_KEYS = frozenset(
{"prompt", "workflow", "parameters", "comfyassets_prompt_pairs"}
@@ -363,6 +368,35 @@ def _library_projection(library: dict) -> dict:
return result
def _image_tag_name(value: object) -> bool:
if not isinstance(value, str) or value != value.strip() or "," in value:
return False
try:
return validate_tags([value])
except ValueError:
return False
def project_image_tags(facet: object) -> dict:
"""The ``image_tags`` facet as the client accepts it, else unavailable.
``current`` carries the byte record's names (at most 100 valid vocabulary
names); ``none`` and ``unavailable`` carry no names. Anything else, extra
keys included, is reported as ``unavailable`` rather than passed on.
"""
unavailable = {"state": "unavailable", "tags": []}
if not isinstance(facet, dict) or set(facet) != {"state", "tags"}:
return unavailable
state, tags = facet["state"], facet["tags"]
if state not in _IMAGE_TAG_STATES or not isinstance(tags, list):
return unavailable
if state != "current" and tags:
return unavailable
if len(tags) > MAX_IMAGE_TAGS or not all(_image_tag_name(tag) for tag in tags):
return unavailable
return {"state": state, "tags": list(tags)}
def _bounded_message(value: Any) -> Optional[str]:
"""A role line as the client accepts it: 1 to 512 UTF-16 code units."""
if not isinstance(value, str) or not value:
@@ -394,6 +428,7 @@ def _encode_inspection(
resolution: dict,
inspection: "InspectionRead | None",
library: dict,
image_tags: dict | None = None,
) -> bytes:
"""Serialize exactly once per budget step, preserving identity and roles."""
metadata = _metadata(inspection)
@@ -407,6 +442,8 @@ def _encode_inspection(
"metadata": metadata,
"library": _library_projection(library),
}
if image_tags is not None:
response["image_tags"] = project_image_tags(image_tags)
body = _bytes(response)
# Omission order keeps prompt evidence longest: the workflow graph first,
# then the prompt graph, then all raw records atomically.
@@ -438,8 +475,12 @@ def encode_inspection_response(
resolution: dict,
inspection: "InspectionRead | None",
library: dict,
image_tags: dict | None = None,
) -> bytes:
"""Preserve the existing rooted wire and its bounded metadata projection."""
"""Preserve the existing rooted wire and its bounded metadata projection.
*image_tags*, when given, adds the validated ``image_tags`` facet (B3).
"""
return _encode_inspection(
identity={"root_scope": root_scope, "image": image},
file=file,
@@ -447,6 +488,7 @@ def encode_inspection_response(
resolution=resolution,
inspection=inspection,
library=library,
image_tags=image_tags,
)
utils/media_capture.py
+67 -14
View File
@@ -26,18 +26,15 @@ import stat as stat_module
import tempfile
import threading
import time
import warnings
import unicodedata
from dataclasses import dataclass
from pathlib import Path
from typing import Callable
from PIL import Image
try:
from .prompt_inspector import MetadataReadError, read_inspection_file
from .video_metadata import VIDEO_EXTENSIONS, is_video_path
except ImportError: # imported with the repo root on sys.path
from utils.prompt_inspector import MetadataReadError, read_inspection_file
from utils.video_metadata import VIDEO_EXTENSIONS, is_video_path
logger = logging.getLogger(__name__)
@@ -578,20 +575,17 @@ class CaptureStore:
def _probe(self, path: str, media_type: str):
"""Header-only dimensions of the capture; never decodes pixels.
Uses the inspection reader, so a PNG whose prompt metadata is over the
inspection budgets is reported as unreadable (not inferable) although
its pixels decode; B2 may switch to a dimensions-only header read.
Only the image header decides whether Generate is offered: prompt
metadata over the inspection budgets does not make a decodable image
unreadable (the metadata itself is read elsewhere, under its budgets).
"""
if media_type == "video":
return None, None, False, VIDEO_ONLY_REASON
try:
with warnings.catch_warnings():
warnings.simplefilter("ignore", Image.DecompressionBombWarning)
read = read_inspection_file(Path(path))
except (MetadataReadError, Image.DecompressionBombError):
return None, None, False, HEADER_REASON
width, height = read.width, read.height
if width is None or height is None:
width, height = header_dimensions(path)
except HeaderTooLarge:
return None, None, False, DIMENSION_REASON
except HeaderUnreadable:
return None, None, False, HEADER_REASON
if (
width > MAX_INFERENCE_SIDE
@@ -602,6 +596,36 @@ class CaptureStore:
return width, height, True, None
class HeaderUnreadable(Exception):
"""The file has no image header Pillow can read."""
class HeaderTooLarge(HeaderUnreadable):
"""The header declares a decompression bomb (no pixel is decoded)."""
def header_dimensions(path: str) -> tuple[int, int]:
"""``(width, height)`` from the image header alone; the file is closed.
``Image.open`` reads only the header. Pillow's decompression-bomb error
(twice its warning size) is ``HeaderTooLarge``; a size between the two
only warns, and the caller's inference pixel cap, far below Pillow's
warning size, refuses it. The process-global warnings filters are never
touched: this runs on capture workers inside the host process, and
``warnings.catch_warnings`` is not thread-safe.
"""
try:
with Image.open(path) as image:
width, height = image.size
except Image.DecompressionBombError as exc:
raise HeaderTooLarge(str(exc)) from None
except Exception as exc: # any decoder failure means "no usable header"
raise HeaderUnreadable(str(exc)) from None
if not width or not height:
raise HeaderUnreadable("The header declares no size")
return int(width), int(height)
def _relative_parts(relative) -> tuple[list[str], str]:
"""Split a safe relative POSIX media path; anything else is denied."""
if not isinstance(relative, str) or not relative:
@@ -629,4 +653,33 @@ def _resolve(root, parts: list[str]) -> tuple[str, str]:
locator = os.path.normcase(candidate)
if os.path.normcase(os.path.realpath(candidate)) != locator:
raise CaptureDenied("The file is not a listed rooted file")
_deny_spelling_aliases(real_root, parts)
return candidate, locator
def _spelling_key(name: str) -> str:
return unicodedata.normalize("NFC", name).casefold()
def _deny_spelling_aliases(real_root: str, parts: list[str]) -> None:
"""Deny a case or Unicode spelling of an entry the directory spells otherwise.
On a case- or normalization-insensitive filesystem (macOS APFS, Windows)
``A.PNG`` or an NFD name opens the listed ``a.png`` or NFC name, while
``realpath`` keeps the requested spelling (macOS), so the tags would be
stored under a second locator. Each part must be spelled as its directory
lists it. A part with no entry of any spelling is left to the open that
follows (a missing file keeps its usual outcome).
"""
directory = real_root
for part in parts:
try:
names = os.listdir(directory)
except OSError:
return
if part not in names:
key = _spelling_key(part)
if any(_spelling_key(name) == key for name in names):
raise CaptureDenied("The file is not a listed rooted file")
return
directory = os.path.join(directory, part)