diff --git a/py/api/__init__.py b/py/api/__init__.py index b45b3d58..fbebf5c3 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -7,6 +7,7 @@ Split into domain-specific mixins for maintainability: - LoggingRoutesMixin: log management endpoints - AutotagRoutesMixin: auto-tagging model management and tagging - JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py) +- ImageTagJobsMixin: bulk Add tags / Auto-tag as image-tags jobs """ import asyncio @@ -32,6 +33,7 @@ from .history import HistoryRoutesMixin from .images import ImageRoutesMixin from .admin import AdminRoutesMixin from .job_routes import JobRoutesMixin +from .image_tag_jobs import ImageTagJobsMixin from .jobs import JobRegistry from .logging_routes import LoggingRoutesMixin from .autotag_routes import AutotagRoutesMixin @@ -260,6 +262,7 @@ class PromptManagerAPI( ImageRoutesMixin, AdminRoutesMixin, JobRoutesMixin, + ImageTagJobsMixin, LoggingRoutesMixin, AutotagRoutesMixin, LoraIntegrationMixin, diff --git a/py/api/image_tag_jobs.py b/py/api/image_tag_jobs.py new file mode 100644 index 00000000..24951165 --- /dev/null +++ b/py/api/image_tag_jobs.py @@ -0,0 +1,446 @@ +"""Bulk image tagging as an Activity job: Add tags and Auto-tag (#249 D1). + +``POST /prompt_manager/jobs/image-tags`` starts one ``image-tags`` job over +1..500 listed images, named by the listing identity ``{root_index, +root_scope, path}``. The job is record-changing (it writes image tag sets and +new vocabulary names), so it never runs beside the metadata scan. + +- **Add tags** (``mode: "add"``) hashes each current file and unions a fixed + tag set into the record of exactly those bytes. +- **Auto-tag** (``mode: "auto"``) copies each still image into a private + capture, asks one legacy engine for its top N tags on the capture, re-hashes + the current file and unions the suggestions only when the bytes are still + the captured ones. It holds the shared model gate for the whole job, loads + the model once and never downloads one. + +Every step reuses the Tag dialog's pieces: the capture store, the model gate, +the suggestion worker and the storage call ``add_image_tags``. Files are +handled one at a time and the job stops at a file boundary when cancelled, +keeping what it already wrote. Rows and log lines carry root-relative paths +and public root labels only. +""" + +import functools +import os +from dataclasses import dataclass + +from aiohttp import web + +from .. import model_gate +from .admin import _plural +from .annotation_suggest import ( + OptionError, + unavailable_message, + validate_options, +) +from .annotations import ( + _MESSAGES, + _error, + _pinned_suggestions, + _settled, + _submit, + _target, +) +from .autotag_routes import model_busy_response +from .images import _public_path +from .job_routes import busy_response +from .jobs import IMAGE_TAG_GROUP_NAMES, IMAGE_TAG_JOB_KIND, JobBusy +from .prompts import bad_request + +try: + from ...database.annotations import ImageTagLimitError, normalize_image_tags + from ...database.library_commands import LibraryCommandError + from ...utils.media_capture import ( + CaptureCapacity, + CaptureChanged, + CaptureDenied, + CaptureError, + CaptureTooLarge, + ) + from ...utils.video_metadata import is_video_path +except ImportError: + from database.annotations import ImageTagLimitError, normalize_image_tags + from database.library_commands import LibraryCommandError + from utils.media_capture import ( + CaptureCapacity, + CaptureChanged, + CaptureDenied, + CaptureError, + CaptureTooLarge, + ) + from utils.video_metadata import is_video_path + +MAX_JOB_TARGETS = 500 +# 500 targets of up to ~1.1 KiB each plus the options: the annotation routes' +# 64 KiB bound is for one target, too small for a full selection. +MAX_JOB_BODY_BYTES = 1024 * 1024 + +ADD_TITLE = "Add tags to images" +AUTO_TITLE = "Auto-tag images" +GATE_OWNER_KIND = "image-tag-job" + +_FIELDS = { + "add": frozenset({"mode", "targets", "tags"}), + "auto": frozenset({"mode", "targets", "model_type", "options"}), +} +_VERBS = {"add": "Add tags", "auto": "Auto-tag"} +_LOG_PREFIXES = {"tagged": "+", "unchanged": "=", "skipped": "-", "errors": "!"} + +NO_NEW_TAGS = "No new tags" +VIDEO_SKIPPED = "Video: manual tags only" +TOO_LARGE = "Too large to verify" +FILE_CHANGED = "File changed" +FILE_NOT_FOUND = "File not found" +NOT_A_STILL = "Not a supported still image" +INFERENCE_FAILED = "Auto-tagging failed. Check server logs for details." +SAVE_FAILED = "The tags could not be saved. Check server logs for details." +READ_FAILED = "The file could not be read." +SOURCE_CHANGED = "The source changed. Refresh the image collection." +STORAGE_RESTORED = "Prompt storage was restored; start the job again." + + +class _JobStopped(Exception): + """No later file can be written (prompt storage was restored).""" + + +@dataclass(frozen=True) +class _Spec: + """A validated start request.""" + + mode: str + targets: tuple + tags: tuple = () + model_type: str | None = None + options: dict | None = None + + @property + def title(self) -> str: + return ADD_TITLE if self.mode == "add" else AUTO_TITLE + + def scope(self) -> dict: + scope = {"mode": self.mode, "images": len(self.targets)} + if self.mode == "auto": + scope["model_type"] = self.model_type + return scope + + +def _parse_targets(raw): + """Unique validated targets in request order, or ``(None, error)``.""" + if not isinstance(raw, list) or not 1 <= len(raw) <= MAX_JOB_TARGETS: + return None, bad_request(f"targets must list 1 to {MAX_JOB_TARGETS} images") + seen, targets = set(), [] + for item in raw: + target, error = _target(item) + if error is not None: + return None, error + if target not in seen: + seen.add(target) + targets.append(target) + return tuple(targets), None + + +def _parse_add(targets, data): + 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)) + if not tags: + return None, bad_request("tags must name at least one tag") + return _Spec("add", targets, tags=tuple(tags)), None + + +def _parse_auto(targets, data): + from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig + + 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 None, bad_request(str(exc)) + return _Spec("auto", targets, model_type=model_type, options=options), None + + +def parse_image_tag_request(data: dict): + """A :class:`_Spec` for a start body, or ``(None, error response)``. + + ``{mode:"add", targets, tags}`` or ``{mode:"auto", targets, model_type, + options}`` exactly; options follow the Tag dialog's Generate rules. + """ + fields = _FIELDS.get(data.get("mode")) + if fields is None: + return None, bad_request('mode must be "add" or "auto"') + if set(data) != fields: + return None, bad_request( + "Request must contain exactly " + ", ".join(sorted(fields)) + ) + targets, error = _parse_targets(data["targets"]) + if error is not None: + return None, error + if data["mode"] == "add": + return _parse_add(targets, data) + return _parse_auto(targets, data) + + +def _capture_outcome(exc: CaptureError, root, relative: str): + """The (group, detail) of a file the capture store refused to read.""" + if isinstance(exc, CaptureTooLarge): + return "skipped", TOO_LARGE + if isinstance(exc, CaptureDenied): + return "skipped", FILE_NOT_FOUND + if isinstance(exc, CaptureChanged): + path = os.path.join(root, *relative.split("/")) + return "skipped", FILE_CHANGED if os.path.lexists(path) else FILE_NOT_FOUND + if isinstance(exc, CaptureCapacity): + return "errors", _MESSAGES["review-capacity"] + return "errors", READ_FAILED + + +def _load_model(service, model_type: str, use_gpu: bool) -> None: + """Worker thread: load (or swap to) *model_type*; never downloads.""" + if service.is_model_loaded() and service.get_loaded_model_type() == model_type: + return + if service.load_model(model_type, use_gpu) is False: + raise RuntimeError(f"{model_type} could not be loaded") + + +class ImageTagJobsMixin: + """The ``image-tags`` job kind: start route and runner.""" + + async def start_image_tag_job(self, request: web.Request) -> web.Response: + """Start bulk tagging (202), or 409 ``busy`` / ``model-busy``. + + Every target's root scope is checked before the job starts + (``source-changed``). Auto-tag takes the model gate here, so a start + while anything else holds it is refused at once. + """ + data, error = await self._read_json_body(request, max_bytes=MAX_JOB_BODY_BYTES) + if error is not None: + return error + spec, error = parse_image_tag_request(data) + if error is not None: + return error + for index, scope in dict.fromkeys((t[0], t[1]) for t in spec.targets): + _root, error = self._scoped_root(index, scope) + if error is not None: + return error + epoch = await self._run_in_executor(self.db.get_execution_epoch) + # A running job in the way answers first: it says which job to wait for. + blocker = self._jobs.blocker(IMAGE_TAG_JOB_KIND, record_changing=True) + if blocker is not None: + return busy_response(JobBusy(blocker)) + owner = None + if spec.mode == "auto": + owner = model_gate.new_owner(GATE_OWNER_KIND) + if not model_gate.try_acquire(owner): + return model_busy_response() + try: + job = self._jobs.start( + IMAGE_TAG_JOB_KIND, + spec.title, + record_changing=True, + scope=spec.scope(), + runner=functools.partial(self._run_image_tag_job, spec, epoch, owner), + groups=IMAGE_TAG_GROUP_NAMES, + ) + except JobBusy as busy: + if owner is not None: + model_gate.release(owner) + return busy_response(busy) + if owner is not None: + # Also when the task is cancelled before its runner ever starts. + job.task.add_done_callback(lambda _task: model_gate.release(owner)) + return web.json_response({"success": True, "job": job.summary()}, status=202) + + # -- runner ---------------------------------------------------------------- + + async def _run_image_tag_job(self, spec: _Spec, epoch: str, owner, job) -> None: + verb = _VERBS[spec.mode] + total = len(spec.targets) + job.update(total=total, status_line=f"Processing file 0 of {total:,}") + started = f"{verb} started · {_plural(total, 'image')}" + if spec.mode == "auto": + started += f" · {spec.model_type}" + job.log_line(started) + try: + service = None + if spec.mode == "auto": + service = await self._prepare_tagging_model(job, spec, owner) + if service is None: + job.finish("failed") + return + if not await self._tag_each_file(job, spec, epoch, owner, service): + job.log_line(f"{verb} stopped · {STORAGE_RESTORED}") + job.finish("failed") + return + state = "cancelled" if job.cancelled else "finished" + job.log_line(f"{verb} {state} · {_plural(job.processed, 'file')} processed") + job.publish({"type": state, "progress": job.progress}) + job.finish(state) + finally: + if owner is not None: + model_gate.release(owner) + + async def _tag_each_file(self, job, spec, epoch, owner, service) -> bool: + """Tag the targets in order; False when the job had to stop.""" + total = len(spec.targets) + for number, target in enumerate(spec.targets, 1): + if job.cancelled: + break + job.update(status_line=f"Processing file {number:,} of {total:,}") + job.publish({"type": "progress", "progress": job.progress}) + stopped = False + try: + group, detail = await self._tag_one_file( + spec, epoch, owner, service, target + ) + except _JobStopped: + stopped, group, detail = True, "errors", STORAGE_RESTORED + except Exception: + self.logger.exception("Image tagging failed for one file") + group, detail = "errors", SAVE_FAILED + self._record_tag_outcome(job, target, group, detail) + job.update(processed=number, progress=number * 100 // total) + if stopped: + return False + return True + + async def _prepare_tagging_model(self, job, spec: _Spec, owner): + """The loaded service, or None (logged) when the job cannot run.""" + from ..autotag import MODELS, get_autotag_service + + service = get_autotag_service() + status = await self._run_in_executor(service.get_models_status) + if not status.get(spec.model_type, {}).get("downloaded"): + name = MODELS.get(spec.model_type, {}).get("name", spec.model_type) + message = unavailable_message(name) + job.update(status_line=message) + job.log_line(f"! {message}") + return None + try: + await model_gate.run( + owner, + _load_model, + service, + spec.model_type, + spec.options["use_gpu"], + ) + except Exception: + self.logger.exception("The auto-tag job could not load its model") + job.update(status_line=INFERENCE_FAILED) + job.log_line(f"! {INFERENCE_FAILED}") + return None + return service + + async def _tag_one_file(self, spec, epoch, owner, service, target): + """The (group, detail) of one target; raises _JobStopped on a restore.""" + index, scope, relative = target + root, error = self._scoped_root(index, scope) + if error is not None: + return "errors", SOURCE_CHANGED + store = self._capture_store() + if spec.mode == "add": + try: + facts = await self._run_in_executor(store.verify, root, relative) + except CaptureError as exc: + return _capture_outcome(exc, root, relative) + return await self._union_image_tags(facts, list(spec.tags), epoch) + if is_video_path(relative): + return "skipped", VIDEO_SKIPPED + try: + capture = await _settled( + _submit( + store.capture, + root, + relative, + epoch=epoch, + scope=scope, + root_index=index, + ), + lambda made: made is not None and store.release(made.token), + ) + except CaptureError as exc: + return _capture_outcome(exc, root, relative) + try: + return await self._auto_tag_capture( + store, capture, spec, epoch, owner, service, root, relative + ) + finally: + store.release(capture.token) + + async def _auto_tag_capture( + self, store, capture, spec, epoch, owner, service, root, relative + ): + """Suggest on the capture, then union only onto the same current bytes.""" + if not capture.inferable: + return "skipped", NOT_A_STILL + if store.acquire(capture.token) is None: + raise _JobStopped # every capture was invalidated (restore) + try: + pairs = await model_gate.run( + owner, + _pinned_suggestions, + store, + capture.token, + service, + capture, + spec.model_type, + spec.options, + on_skipped=functools.partial(store.release_use, capture.token), + ) + except Exception: + self.logger.exception("Auto-tagging failed for one file") + return "errors", INFERENCE_FAILED + try: + facts = await self._run_in_executor(store.verify, root, relative) + except CaptureError as exc: + return _capture_outcome(exc, root, relative) + if facts.locator != capture.locator or facts.sha256 != capture.sha256: + return "skipped", FILE_CHANGED + names = [str(name) for name, _score in pairs] + if not names: + return "unchanged", NO_NEW_TAGS + return await self._union_image_tags(facts, names, epoch) + + async def _union_image_tags(self, facts, tags: list, epoch: str): + """``add_image_tags`` on the verified bytes, as a (group, detail).""" + try: + receipt = await self._run_in_executor( + self.db.add_image_tags, + locator=facts.locator, + sha256=facts.sha256, + size=facts.size, + mtime_ns=facts.mtime_ns, + expected_epoch=epoch, + tags=tags, + ) + except LibraryCommandError as exc: + if exc.code == "storage-changed": + raise _JobStopped from None + if exc.code == "tag-limit": + return "errors", _MESSAGES["tag-limit"] + self.logger.warning(f"Image tags not saved ({exc.code})") + return "errors", SAVE_FAILED + except ValueError: + self.logger.exception("Image tags were refused") + return "errors", SAVE_FAILED + added = receipt.get("added", 0) + if added: + return "tagged", f"Added {_plural(added, 'tag')}" + return "unchanged", NO_NEW_TAGS + + def _record_tag_outcome(self, job, target, group: str, detail: str) -> None: + index, _scope, relative = target + roots = list(self._get_all_output_dirs()) + label = _public_path(roots[index]) if index < len(roots) else "" + where = f"{label}/{relative}" if label else relative + job.record( + group, + {"path": relative, "root": label, "root_index": index, "detail": detail}, + f"{_LOG_PREFIXES[group]} {where} · {detail}", + ) diff --git a/py/api/job_routes.py b/py/api/job_routes.py index 71b35cb0..f8556601 100644 --- a/py/api/job_routes.py +++ b/py/api/job_routes.py @@ -1,8 +1,9 @@ """Session job routes (Activity): list, inspect, follow and cancel jobs. -The registry itself is ``py/api/jobs.py``; the scan job's runner is in -``py/api/admin.py``. Every response leaves out absolute server paths: job -rows already carry root-relative paths and public root labels. +The registry itself is ``py/api/jobs.py``. The scan job's runner is in +``py/api/admin.py``, the image tagging job's in ``py/api/image_tag_jobs.py``. +Every response leaves out absolute server paths: job rows already carry +root-relative paths and public root labels. """ import contextlib @@ -42,6 +43,20 @@ async def _scan_scope(request): return scope if scope in SCAN_SCOPES else None +def busy_response(busy: JobBusy): + """``409 busy``, naming the running job that holds the slot.""" + running = busy.running + return web.json_response( + { + "success": False, + "code": "busy", + "error": str(busy), + "running": {"id": running.id, "title": running.title}, + }, + status=409, + ) + + def _summary_event(job): return f"data: {json.dumps(job.summary())}\n\n".encode("utf-8") @@ -58,6 +73,10 @@ class JobRoutesMixin: async def start_scan_job_route(request): return await self.start_scan_job(request) + @routes.post("/prompt_manager/jobs/image-tags") + async def start_image_tag_job_route(request): + return await self.start_image_tag_job(request) + @routes.get("/prompt_manager/jobs/{job_id}") async def job_detail_route(request): return await self.job_detail(request) @@ -121,16 +140,7 @@ class JobRoutesMixin: try: job = self._start_scan_job(files) except JobBusy as busy: - running = busy.running - return web.json_response( - { - "success": False, - "code": "busy", - "error": str(busy), - "running": {"id": running.id, "title": running.title}, - }, - status=409, - ) + return busy_response(busy) return web.json_response({"success": True, "job": job.summary()}, status=202) async def job_events(self, request): diff --git a/py/api/jobs.py b/py/api/jobs.py index a5a4b0e5..27e4ecb5 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -9,8 +9,9 @@ as a snapshot and then every event published after it, until the job ends. Jobs live only in server memory and are gone after a restart. The registry keeps the running jobs and this session's last ``MAX_FINISHED_JOBS`` finished -ones. Record-changing jobs (the metadata scan, and later maintenance, cleanup -and restore) never run together, and one job of a kind runs at a time. +ones. Record-changing jobs (the metadata scan, bulk image tagging, and later +maintenance, cleanup and restore) never run together, and one job of a kind +runs at a time. """ import asyncio @@ -54,6 +55,11 @@ SCAN_GROUPS = { "failed": "errors", } +# Bulk image tagging (#249 D1): Add tags and Auto-tag over selected images. +# Record-changing: it writes image tag sets and new vocabulary names. +IMAGE_TAG_JOB_KIND = "image-tags" +IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors") + def _relative_path(path) -> str: """``path`` as a root-relative POSIX path; an absolute one keeps its name.""" @@ -292,13 +298,9 @@ class JobRegistry: running or when ``record_changing`` and another record-changing job runs. Must be called from the event loop. """ - same_kind = self.active(kind) - if same_kind is not None: - raise JobBusy(same_kind) - if record_changing: - for job in self._jobs.values(): - if job.running and job.record_changing: - raise JobBusy(job) + blocker = self.blocker(kind, record_changing=record_changing) + if blocker is not None: + raise JobBusy(blocker) self._prune() job = Job( kind, title, record_changing=record_changing, scope=scope, groups=groups @@ -307,6 +309,17 @@ class JobRegistry: job.task = asyncio.ensure_future(self._run(job, runner)) return job + def blocker(self, kind, *, record_changing): + """The running job a new ``kind`` job would be refused for, or None.""" + same_kind = self.active(kind) + if same_kind is not None: + return same_kind + if record_changing: + for job in self._jobs.values(): + if job.running and job.record_changing: + return job + return None + async def _run(self, job, runner): try: await runner(job) diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 14f7c464..89718109 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -462,6 +462,15 @@ class TestUiRoutes(CoreRoutesTestCase): ): self.assertIn(entry, table) + async def test_image_tag_job_route_is_registered_beside_the_scan_job(self): + table = { + (route.method, route.resource.canonical) + for route in self.app.router.routes() + if route.resource is not None + } + self.assertIn(("POST", "/prompt_manager/jobs/scan"), table) + self.assertIn(("POST", "/prompt_manager/jobs/image-tags"), table) + class TestStaticRoutes(CoreRoutesTestCase): diff --git a/tests/test_image_tag_jobs.py b/tests/test_image_tag_jobs.py new file mode 100644 index 00000000..089cfab2 --- /dev/null +++ b/tests/test_image_tag_jobs.py @@ -0,0 +1,648 @@ +"""Bulk image tagging as an Activity job: Add tags and Auto-tag (#249 D1). + +``POST /prompt_manager/jobs/image-tags`` starts one record-changing +``image-tags`` job over 1..500 listed images. Add tags unions a fixed set +into each file's byte record; Auto-tag runs one legacy engine on a private +capture of each still image (under the shared model gate, never +downloading) and unions its top N suggestions. No real model is loaded: the +engine is B2's fake service. +""" + +import asyncio +import hashlib +import json +import sqlite3 +import threading +import unittest +from unittest.mock import patch + +try: + from test_api_annotation_suggest import SCORED, WD14, FakeTagger, SuggestTestCase + from test_api_images import make_png +except ImportError: + from tests.test_api_annotation_suggest import ( + SCORED, + WD14, + FakeTagger, + SuggestTestCase, + ) + from tests.test_api_images import make_png + +from database.execution_history_schema import rotate_history_epoch +from py import model_gate +from py.api import jobs as jobs_module +from py.config import AutotagConfig + +START = "/prompt_manager/jobs/image-tags" +WAIT = 30.0 +KIND = getattr(jobs_module, "IMAGE_TAG_JOB_KIND", "image-tags") +GROUPS = ("tagged", "unchanged", "skipped", "errors") + + +class HookedTagger(FakeTagger): + """The B2 fake plus a hook that runs inside each inference.""" + + def __init__(self): + super().__init__() + self.on_generate = None + + def generate_tag_suggestions(self, image_path, **kwargs): + result = super().generate_tag_suggestions(image_path, **kwargs) + if self.on_generate is not None: + self.on_generate(image_path) + return result + + +class ImageTagJobTestCase(SuggestTestCase): + async def get_application(self): + app = await super().get_application() + self.tagger = HookedTagger() + patcher = patch("py.autotag.get_autotag_service", return_value=self.tagger) + patcher.start() + self.addCleanup(patcher.stop) + return app + + async def tearDownAsync(self): + # Every job of the test has ended before the app goes away. + for job in self.api._jobs.list(): + if job.task is not None and not job.task.done(): + job.task.cancel() + await asyncio.gather(job.task, return_exceptions=True) + await super().tearDownAsync() + + # -- helpers --------------------------------------------------------- + + def png(self, path, size=(8, 6), color=(200, 120, 40)): + make_png(self.output_dir / path, size=size, color=color) + return self.output_dir / path + + def digest(self, path): + return hashlib.sha256((self.output_dir / path).read_bytes()).hexdigest() + + def stored(self, path): + return self.api.db.read_image_tags(self.locator(path), self.digest(path)) + + def add_body(self, paths, tags, **changes): + return { + "mode": "add", + "targets": [self.target(p) for p in paths], + "tags": list(tags), + **changes, + } + + def auto_body(self, paths, model_type="wd14-vit", options=None, **changes): + return { + "mode": "auto", + "targets": [self.target(p) for p in paths], + "model_type": model_type, + "options": dict(WD14) if options is None else options, + **changes, + } + + async def start(self, body): + response = await self.client.post(START, json=body) + return response.status, await self.json_of(response) + + async def start_ok(self, body): + status, data = await self.start(body) + self.assertEqual(status, 202, data) + return self.api._jobs.get(data["job"]["id"]) + + async def run_job(self, body): + job = await self.start_ok(body) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + return job + + def rows(self, job, group): + return [(row["path"], row["detail"]) for row in job.outcomes[group]] + + async def wait_for(self, event): + ok = await asyncio.get_running_loop().run_in_executor(None, event.wait, WAIT) + self.assertTrue(ok, "the fake engine was never entered") + + +class TestAddTags(ImageTagJobTestCase): + async def test_add_job_unions_tags_on_each_target_and_records_outcomes(self): + for name in ("a.png", "b.png", "c.png"): + self.png(name, color=(len(name), ord(name[0]), 9)) + self.saved_set("a.png", (self.output_dir / "a.png").read_bytes(), ["keep_me"]) + self.saved_set( + "c.png", (self.output_dir / "c.png").read_bytes(), ["Sandstone", "desert"] + ) + + job = await self.run_job( + self.add_body( + ["a.png", "b.png", "c.png", "gone.png"], ["sandstone", "desert"] + ) + ) + + self.assertEqual(job.kind, KIND) + self.assertEqual(job.title, "Add tags to images") + self.assertEqual(job.state, "finished") + self.assertTrue(job.record_changing) + self.assertEqual(job.scope, {"mode": "add", "images": 4}) + self.assertEqual( + job.counters, {"tagged": 2, "unchanged": 1, "skipped": 1, "errors": 0} + ) + self.assertEqual( + self.stored("a.png")["tags"], ["keep_me", "sandstone", "desert"] + ) + self.assertEqual(self.stored("b.png")["tags"], ["sandstone", "desert"]) + # The stored spelling wins; nothing new means no write. + self.assertEqual(self.stored("c.png")["tags"], ["Sandstone", "desert"]) + self.assertEqual(self.stored("c.png")["tag_revision"], 1) + self.assertEqual( + self.rows(job, "tagged"), + [("a.png", "Added 2 tags"), ("b.png", "Added 2 tags")], + ) + self.assertEqual(self.rows(job, "unchanged"), [("c.png", "No new tags")]) + self.assertEqual(self.rows(job, "skipped"), [("gone.png", "File not found")]) + self.assertEqual(job.outcomes["tagged"][0]["root_index"], 0) + self.assertEqual((job.processed, job.total, job.progress), (4, 4, 100)) + self.assertTrue(any("+ " in line and "a.png" in line for line in job.log)) + self.assertTrue(job.log[-1].endswith("Add tags finished · 4 files processed")) + # Add tags never needs the model. + self.assertEqual(self.tagger.load_calls, []) + self.assertIsNone(model_gate.holder()) + + async def test_add_tags_vocabulary_names_become_global_tags(self): + self.png("a.png") + + await self.run_job(self.add_body(["a.png"], ["brand_new_tag"])) + + names = set(self.api.db.get_all_tags()) + self.assertIn("brand_new_tag", names) + + async def test_add_over_the_tag_cap_is_an_error_row_with_zero_writes(self): + self.png("a.png") + full = [f"t{i:03d}" for i in range(100)] + self.saved_set("a.png", (self.output_dir / "a.png").read_bytes(), full) + + job = await self.run_job(self.add_body(["a.png"], ["one_more"])) + + self.assertEqual(job.counters["errors"], 1) + self.assertEqual( + self.rows(job, "errors"), + [("a.png", "An image can carry at most 100 tags.")], + ) + self.assertEqual(self.stored("a.png")["tags"], full) + + async def test_video_skipped_in_auto_and_tagged_in_add(self): + (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 8) + + added = await self.run_job(self.add_body(["clip.mp4"], ["motion"])) + self.assertEqual(self.rows(added, "tagged"), [("clip.mp4", "Added 1 tag")]) + self.assertEqual(self.stored("clip.mp4")["tags"], ["motion"]) + + auto = await self.run_job(self.auto_body(["clip.mp4"])) + self.assertEqual( + self.rows(auto, "skipped"), [("clip.mp4", "Video: manual tags only")] + ) + self.assertEqual(self.tagger.calls, []) + self.assertEqual(self.stored("clip.mp4")["tags"], ["motion"]) + + async def test_not_a_still_image_and_too_large_are_skipped(self): + (self.output_dir / "broken.png").write_bytes(b"not an image") + self.png("big.png") + self.api._capture_store().max_file = 16 + + auto = await self.run_job(self.auto_body(["big.png"])) + self.assertEqual( + self.rows(auto, "skipped"), [("big.png", "Too large to verify")] + ) + + self.api._capture_store().max_file = 2**20 + auto = await self.run_job(self.auto_body(["broken.png"])) + self.assertEqual( + self.rows(auto, "skipped"), [("broken.png", "Not a supported still image")] + ) + + self.api._capture_store().max_file = 16 + added = await self.run_job(self.add_body(["big.png"], ["x"])) + self.assertEqual( + self.rows(added, "skipped"), [("big.png", "Too large to verify")] + ) + self.assertEqual(self.tagger.calls, []) + + +class TestAutoTag(ImageTagJobTestCase): + async def test_auto_job_applies_the_requested_count_per_file_without_removing_existing( # noqa: E501 + self, + ): + self.png("a.png") + self.png("b.png", color=(1, 2, 3)) + self.saved_set("a.png", (self.output_dir / "a.png").read_bytes(), ["keep_me"]) + options = {**WD14, "suggestion_limit": 2} + + job = await self.run_job(self.auto_body(["a.png", "b.png"], options=options)) + + self.assertEqual(job.title, "Auto-tag images") + self.assertEqual( + job.scope, {"mode": "auto", "images": 2, "model_type": "wd14-vit"} + ) + self.assertEqual(job.state, "finished") + self.assertEqual( + self.api.db.read_image_tags(self.locator("a.png"), self.digest("a.png"))[ + "tags" + ], + ["keep_me", "sandstone", "desert"], + ) + self.assertEqual(self.stored("b.png")["tags"], ["sandstone", "desert"]) + self.assertEqual( + job.counters, {"tagged": 2, "unchanged": 0, "skipped": 0, "errors": 0} + ) + # Inference read the private captures only, never the current files. + for path in self.tagger.read_paths: + self.assertNotIn(str(self.output_dir), path) + self.assertEqual([call["limit"] for call in self.tagger.calls], [2, 2]) + + # limit 3: at most three names are added to a file. + self.png("c.png", color=(9, 9, 9)) + options = {**WD14, "suggestion_limit": 3} + job = await self.run_job(self.auto_body(["c.png"], options=options)) + self.assertEqual(self.stored("c.png")["tags"], [name for name, _ in SCORED[:3]]) + self.assertEqual(self.rows(job, "tagged"), [("c.png", "Added 3 tags")]) + + # Omitted: the configured default. + AutotagConfig.SUGGESTION_LIMIT = 4 + self.png("d.png", color=(7, 7, 7)) + await self.run_job(self.auto_body(["d.png"])) + self.assertEqual(self.tagger.calls[-1]["limit"], 4) + self.assertEqual(len(self.stored("d.png")["tags"]), 4) + + # Suggestions the file already carries add nothing. + job = await self.run_job(self.auto_body(["d.png"])) + self.assertEqual(self.rows(job, "unchanged"), [("d.png", "No new tags")]) + + async def test_auto_job_holds_gate_and_never_downloads(self): + self.png("a.png") + self.png("b.png", color=(1, 2, 3)) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + + job = await self.start_ok(self.auto_body(["a.png", "b.png"])) + await self.wait_for(entered) + holder = model_gate.holder() + self.assertIsNotNone(holder) + self.assertTrue(holder.startswith("image-tag-job:"), holder) + # A legacy route cannot swap the model under the running job. + response = await self.client.post("/prompt_manager/autotag/unload") + self.assertEqual(response.status, 409) + self.assertEqual((await response.json())["code"], "model-busy") + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + + self.assertEqual(job.counters["tagged"], 2) + self.assertEqual(self.tagger.load_calls, [("wd14-vit", True)]) + self.tagger.download_model.assert_not_called() + self.assertIsNone(model_gate.holder()) + + async def test_model_unavailable_fails_the_job_before_the_first_file(self): + self.png("a.png") + self.tagger.downloaded["wd14-vit"] = False + + job = await self.run_job(self.auto_body(["a.png"])) + + self.assertEqual(job.state, "failed") + self.assertEqual(job.processed, 0) + self.assertEqual(job.counters, dict.fromkeys(GROUPS, 0)) + self.assertTrue( + any( + line.endswith( + "is not downloaded. Download it from the legacy Auto-tag page." + ) + for line in job.log + ), + job.log, + ) + self.assertEqual(self.tagger.load_calls, []) + self.tagger.download_model.assert_not_called() + self.assertIsNone(model_gate.holder()) + self.assertEqual(self.api._capture_store().stats()["captures"], 0) + + async def test_model_load_failure_fails_the_job_and_frees_the_gate(self): + self.png("a.png") + self.tagger.load_error = RuntimeError("/secret/models/wd14 exploded") + + job = await self.run_job(self.auto_body(["a.png"])) + + self.assertEqual(job.state, "failed") + self.assertEqual(job.processed, 0) + self.assertNotIn("/secret", json.dumps(job.detail())) + self.assertIsNone(model_gate.holder()) + + async def test_inference_failure_is_an_error_row_and_the_job_goes_on(self): + self.png("a.png") + self.png("b.png", color=(1, 2, 3)) + failing = {"first": True} + + def explode(_path): + if failing.pop("first", False): + raise RuntimeError("/secret/path boom") + + self.tagger.on_generate = explode + + job = await self.run_job(self.auto_body(["a.png", "b.png"])) + + self.assertEqual(job.counters["errors"], 1) + self.assertEqual(job.counters["tagged"], 1) + self.assertEqual( + self.rows(job, "errors"), + [("a.png", "Auto-tagging failed. Check server logs for details.")], + ) + self.assertNotIn("/secret", json.dumps(job.report())) + self.assertIsNone(model_gate.holder()) + + async def test_changed_file_is_skipped_not_tagged(self): + self.png("a.png") + before = self.digest("a.png") + + def rewrite(_capture_path): + make_png(self.output_dir / "a.png", size=(8, 6), color=(1, 1, 1)) + + self.tagger.on_generate = rewrite + + job = await self.run_job(self.auto_body(["a.png"])) + + self.assertEqual(self.rows(job, "skipped"), [("a.png", "File changed")]) + self.assertEqual(job.counters["tagged"], 0) + locator = self.locator("a.png") + self.assertEqual(self.api.db.read_image_tags(locator, before)["tags"], []) + self.assertEqual(self.stored("a.png")["tags"], []) + + async def test_captures_released_and_handles_closed(self): + try: + from open_handles import assert_closed + except ImportError: + from tests.open_handles import assert_closed + for index in range(3): + self.png(f"f{index}.png", color=(index, 0, 0)) + + await self.run_job(self.auto_body([f"f{index}.png" for index in range(3)])) + + store = self.api._capture_store() + self.assertEqual(store.stats()["captures"], 0) + assert_closed(self, store.directory) + assert_closed(self, str(self.output_dir)) + + +class TestStartAndExclusion(ImageTagJobTestCase): + async def test_start_returns_model_busy_or_job_busy(self): + self.png("a.png") + other = model_gate.new_owner("legacy-batch") + self.assertTrue(model_gate.try_acquire(other)) + try: + status, body = await self.start(self.auto_body(["a.png"])) + self.assertEqual((status, body["code"]), (409, "model-busy")) + self.assertEqual(self.api._jobs.list(), []) + # Add tags needs no model: it starts while the gate is held. + job = await self.run_job(self.add_body(["a.png"], ["x"])) + self.assertEqual(job.state, "finished") + finally: + model_gate.release(other) + + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + job = await self.start_ok(self.auto_body(["a.png"])) + await self.wait_for(entered) + try: + for body in (self.add_body(["a.png"], ["y"]), self.auto_body(["a.png"])): + resp = await self.client.post(START, json=body) + body = await resp.json() + self.assertEqual((resp.status, body["code"]), (409, "busy")) + self.assertEqual( + body["running"], {"id": job.id, "title": "Auto-tag images"} + ) + finally: + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + # A refused start never leaves the gate taken. + self.assertIsNone(model_gate.holder()) + + async def test_record_changing_excludes_running_scan(self): + self.png("a.png") + gate = asyncio.Event() + + async def scan_runner(job): + await gate.wait() + + scan = self.api._jobs.start( + "scan", + "Scan for metadata", + record_changing=True, + scope={"files": "new"}, + runner=scan_runner, + ) + try: + status, body = await self.start(self.add_body(["a.png"], ["x"])) + self.assertEqual((status, body["code"]), (409, "busy")) + self.assertEqual(body["running"]["title"], "Scan for metadata") + status, body = await self.start(self.auto_body(["a.png"])) + self.assertEqual((status, body["code"]), (409, "busy")) + self.assertIsNone(model_gate.holder()) + finally: + gate.set() + await scan.task + + # The other way round: a running tagging job keeps the scan out. + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + job = await self.start_ok(self.auto_body(["a.png"])) + await self.wait_for(entered) + try: + resp = await self.client.post("/prompt_manager/jobs/scan", json={}) + self.assertEqual(resp.status, 409) + self.assertEqual((await resp.json())["running"]["id"], job.id) + finally: + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + + async def test_a_start_that_loses_the_slot_race_frees_the_gate(self): + self.png("a.png") + gate = asyncio.Event() + + async def scan_runner(job): + await gate.wait() + + scan = self.api._jobs.start( + "scan", + "Scan for metadata", + record_changing=True, + scope={}, + runner=scan_runner, + ) + try: + # The pre-check saw a free slot; registry.start still refuses. + with patch.object(self.api._jobs, "blocker", side_effect=[None, scan]): + status, body = await self.start(self.auto_body(["a.png"])) + self.assertEqual((status, body["code"]), (409, "busy")) + self.assertIsNone(model_gate.holder()) + finally: + gate.set() + await scan.task + + async def test_cancel_stops_at_file_boundary_and_keeps_written_files(self): + for index in range(3): + self.png(f"f{index}.png", color=(index, 5, 5)) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + + job = await self.start_ok(self.auto_body([f"f{i}.png" for i in range(3)])) + await self.wait_for(entered) + response = await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel") + self.assertTrue((await response.json())["cancellation_requested"]) + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.counters["tagged"], 1) + self.assertEqual(job.processed, 1) + self.assertEqual(len(self.stored("f0.png")["tags"]), 10) + self.assertEqual(self.stored("f1.png")["tags"], []) + self.assertEqual(len(self.tagger.calls), 1) + self.assertTrue(job.log[-1].endswith("Auto-tag cancelled · 1 file processed")) + self.assertIsNone(model_gate.holder()) + self.assertEqual(self.api._capture_store().stats()["captures"], 0) + + async def test_stale_root_scope_refuses_the_start(self): + self.png("a.png") + body = self.add_body(["a.png"], ["x"]) + body["targets"][0]["root_scope"] = "0" * 64 + + status, data = await self.start(body) + + self.assertEqual((status, data["code"]), (409, "source-changed")) + self.assertEqual(self.api._jobs.list(), []) + + async def test_restore_during_the_job_stops_it_without_writes(self): + self.png("a.png") + self.png("b.png", color=(1, 2, 3)) + + def rotate(_path): + conn = sqlite3.connect(self._temp_db.name) + try: + rotate_history_epoch(conn) + finally: + conn.close() + + self.tagger.on_generate = rotate + + job = await self.run_job(self.auto_body(["a.png", "b.png"])) + + self.assertEqual(job.state, "failed") + self.assertEqual(job.counters["tagged"], 0) + self.assertEqual( + self.rows(job, "errors"), + [("a.png", "Prompt storage was restored; start the job again.")], + ) + self.assertEqual(self.stored("a.png")["tags"], []) + self.assertEqual(len(self.tagger.calls), 1) + self.assertIsNone(model_gate.holder()) + + async def test_invalid_bodies_rejected(self): + self.png("a.png") + target = self.target("a.png") + bad_bodies = [ + {"mode": "add", "targets": [], "tags": ["x"]}, + {"mode": "add", "targets": [target] * 501, "tags": ["x"]}, + {"mode": "add", "targets": "a.png", "tags": ["x"]}, + {"mode": "add", "targets": [{**target, "path": "../a.png"}], "tags": ["x"]}, + {"mode": "add", "targets": [{**target, "root_index": True}], "tags": ["x"]}, + {"mode": "add", "targets": [{**target, "extra": 1}], "tags": ["x"]}, + {"mode": "add", "targets": [{**target, "path": "a.txt"}], "tags": ["x"]}, + {"mode": "add", "targets": [target], "tags": []}, + {"mode": "add", "targets": [target], "tags": "x"}, + {"mode": "add", "targets": [target], "tags": ["bad\nname"]}, + {"mode": "add", "targets": [target]}, + {"mode": "add", "targets": [target], "tags": ["x"], "extra": 1}, + {"mode": "add", "targets": [target], "tags": ["x"], "model_type": "hf"}, + {"mode": "remove", "targets": [target], "tags": ["x"]}, + {"targets": [target], "tags": ["x"]}, + {"mode": "auto", "targets": [target], "model_type": "wd14-vit"}, + {"mode": "auto", "targets": [target], "model_type": "x", "options": {}}, + { + "mode": "auto", + "targets": [target], + "model_type": "wd14-vit", + "options": {}, + }, + { + "mode": "auto", + "targets": [target], + "model_type": "wd14-vit", + "options": {**WD14, "general_threshold": 2}, + }, + { + "mode": "auto", + "targets": [target], + "model_type": "wd14-vit", + "options": {**WD14, "suggestion_limit": 0}, + }, + { + "mode": "auto", + "targets": [target], + "model_type": "gguf", + "options": dict(WD14), + }, + { + "mode": "auto", + "targets": [target], + "model_type": "wd14-vit", + "options": dict(WD14), + "tags": ["x"], + }, + ] + for body in bad_bodies: + with self.subTest(body=json.dumps(body)[:120]): + status, data = await self.start(body) + self.assertIn(status, (400, 403), data) + self.assertFalse(data["success"]) + too_many = { + "mode": "add", + "targets": [target], + "tags": [f"t{i}" for i in range(101)], + } + status, data = await self.start(too_many) + self.assertEqual((status, data["code"]), (400, "tag-limit")) + response = await self.client.post(START, data="not json") + self.assertEqual(response.status, 400) + self.assertEqual(self.api._jobs.list(), []) + self.assertIsNone(model_gate.holder()) + + async def test_duplicate_targets_are_tagged_once(self): + self.png("a.png") + + job = await self.run_job(self.add_body(["a.png", "a.png"], ["x"])) + + self.assertEqual(job.scope["images"], 1) + self.assertEqual(job.counters["tagged"], 1) + + async def test_rows_and_report_carry_no_absolute_paths(self): + self.png("sub/a.png") + self.png("sub/b.png", color=(3, 3, 3)) + job = await self.run_job( + self.auto_body(["sub/a.png", "sub/b.png", "sub/missing.png"]) + ) + + for url in ( + START.rsplit("/", 1)[0], + f"/prompt_manager/jobs/{job.id}", + f"/prompt_manager/jobs/{job.id}/report", + ): + response = await self.client.get(url) + self.assertEqual(response.status, 200) + body = json.loads(await response.text()) + self.assert_no_paths(body) + detail = job.detail() + self.assertEqual( + [row["path"] for row in detail["outcomes"]["tagged"]], + ["sub/a.png", "sub/b.png"], + ) + self.assertTrue( + all(row["root"] == "output" for row in detail["outcomes"]["tagged"]) + ) + self.assertTrue(any("output/sub/a.png" in line for line in detail["log"])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 7f87f85e..f561be08 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -12,6 +12,7 @@ import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from py.api import jobs as jobs_module # noqa: E402 from py.api.jobs import ( # noqa: E402 DEFAULT_SCAN_SCOPE, LOG_CAP_LINE, @@ -415,6 +416,14 @@ class TestJobOutcomesAndLog(unittest.TestCase): SCAN_GROUP_NAMES, ("added", "known", "review", "no_metadata", "errors") ) + def test_image_tag_job_kind_and_groups(self): + # #249 D1: bulk Add tags / Auto-tag, counter-tile order. + self.assertEqual(getattr(jobs_module, "IMAGE_TAG_JOB_KIND", None), "image-tags") + self.assertEqual( + getattr(jobs_module, "IMAGE_TAG_GROUP_NAMES", None), + ("tagged", "unchanged", "skipped", "errors"), + ) + if __name__ == "__main__": unittest.main()