Vito
b5a56bee0b
* feat(images): store images' own categories in the library database (#332) One new table, image_categories, and its category index, created empty by a forward-only migration that reads and writes only its own objects. Startup checks the structure; restore also checks the values. The store sets, clears and reads own categories by file key (the deepest output source holding the file and the path below it), and copies a removed duplicate's category to a kept copy that has none. * feat(images): duplicate cleanup keeps images' own categories (#332) The kept copy's own category wins; a kept copy without one receives the removed copy's, keyed by the deepest output source holding each file. The removed copy's row is never deleted, so a give-back or a failed delete loses nothing. The job log reports how many were kept. * feat(images): set, remove and list images' own categories (#332) PUT /prompt_manager/images/category sets the own category of up to 500 listed files, DELETE removes it (listed or not), and GET /prompt_manager/images/categories lists the values whose file is listed. Targets are the Images listing identity of an output source, checked by media_target (shared with Collections). The routes never touch prompts, recipes or links; POST /prompt_manager/bulk/category is unchanged. * fix(images): key symlinked output folders and keep key work off the event loop (#332) Review fixes: a source matches by its configured spelling too, so a detected output folder that is a symlink keys its files (routes and duplicate cleanup). File keys are computed in the executor. A path no source holds is a plain 403, a lone surrogate a 400. The install repairs a table that lacks its category index instead of refusing to open. * feat(images): set or remove a category on all N matching images (#332) The category routes also accept the Images listing filter of "Select all N matching" (#392) instead of explicit items, resolved now and capped at 500 images like an explicit selection. * test(vendor): count the #332 host case in the execution matrix image-attribution.spec.mjs gains the image category route case, so the execution config lists 92 tests (3 attribution cases per project). * test(legacy): record the image category routes in the legacy contract snapshot (#332) #332 adds PUT/DELETE /prompt_manager/images/category and GET /prompt_manager/images/categories; the legacy pages' routes are unchanged.
774 lines
30 KiB
Python
774 lines
30 KiB
Python
"""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
|
|
import threading
|
|
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
|
|
|
|
# The target check is shared with the collection routes (#331); both names
|
|
# stay importable from here.
|
|
from .media_target import MAX_TARGET_PATH # noqa: F401
|
|
from .media_target import media_target as _target
|
|
from .images import _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 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 CaptureError, CaptureStore
|
|
|
|
# Request bodies of the annotation routes are small (a target plus a tag set).
|
|
MAX_BODY_BYTES = 64 * 1024
|
|
CAPTURE_URL = "/prompt_manager/image-annotations/capture/"
|
|
_PREVIEW_CHUNK = 64 * 1024
|
|
|
|
_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}")
|
|
# The store is created on first use by the event loop, the scan's executor or
|
|
# the monitor's worker (Needs review), so its creation is serialized.
|
|
_CAPTURE_STORE_LOCK = threading.Lock()
|
|
|
|
_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 _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``).
|
|
Thread-safe: Needs review recording calls it off the event loop.
|
|
Raises ``CaptureError`` after ``close_annotation_captures``.
|
|
"""
|
|
with _CAPTURE_STORE_LOCK:
|
|
if getattr(self, "_captures_closed", False):
|
|
# A recording still in flight at shutdown must not create a
|
|
# new temp directory after the cleanup removed the old one.
|
|
raise CaptureError("The capture store is closed")
|
|
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, *, final: bool = False) -> None:
|
|
"""Remove every capture file and the owned directory. Idempotent.
|
|
|
|
Taken under the store lock. ``final`` (the shutdown cleanup) also makes
|
|
``_capture_store`` refuse to create another store, so a recording
|
|
still in flight cannot leak a new temp directory after it.
|
|
"""
|
|
with _CAPTURE_STORE_LOCK:
|
|
if final:
|
|
self._captures_closed = True
|
|
unsubscribe, self._captures_unsubscribe = self._captures_unsubscribe, None
|
|
store, self._captures = self._captures, None
|
|
self._suggested_tokens = set()
|
|
if unsubscribe is not None:
|
|
unsubscribe()
|
|
if store is not None:
|
|
store.close()
|
|
|
|
async def cleanup_annotation_captures(self, app: web.Application) -> None:
|
|
self.close_annotation_captures(final=True)
|
|
|
|
# -- 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
|
|
try:
|
|
store = self._capture_store()
|
|
except CaptureError as exc: # the server is shutting down
|
|
return _capture_error(exc)
|
|
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)
|
|
self._suggested_tokens.discard(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")
|
|
# A save after Generate suggestions ran on this capture marks the
|
|
# image auto-tagged (#379), so a Skip batch leaves the review alone.
|
|
suggested = token in self._suggested_tokens
|
|
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,
|
|
mark_autotagged=suggested,
|
|
)
|
|
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.")
|
|
# Dropped only after a saved or replayed receipt: a retry after a
|
|
# tag-conflict still marks.
|
|
self._suggested_tokens.discard(token)
|
|
return web.json_response(
|
|
{"success": True, **receipt}, headers={"Cache-Control": "no-store"}
|
|
)
|
|
|
|
def _remember_suggested(self, store, token: str) -> None:
|
|
"""Note that Generate ran on *token*; forget tokens whose capture is
|
|
gone. ``expires_in`` does not reset a capture's idle clock."""
|
|
self._suggested_tokens = {
|
|
kept
|
|
for kept in self._suggested_tokens
|
|
if store.expires_in(kept) is not None
|
|
}
|
|
self._suggested_tokens.add(token)
|
|
|
|
# -- 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)
|
|
self._remember_suggested(store, 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
|