diff --git a/TESTING.md b/TESTING.md index 6b05b1e..ce4bd6b 100644 --- a/TESTING.md +++ b/TESTING.md @@ -3826,7 +3826,9 @@ threshold off that grid is refused ("General-tag confidence must be a multiple o 0.05.", `validate_saved_threshold`; `draftProblem` holds back Preview tagging review and the batch Review with the same sentence) and falls back on load; an accepted one is stored exactly on the grid (#369). A per-request threshold keeps its own check: `/image-annotations/suggest` and -`/jobs/autotag` take 0–1; the legacy `/autotag/single` and `/autotag/start` only require a number. `GET /autotag/models` reports the +`/jobs/autotag` take 0–1; the legacy `/autotag/start` passes its options through +`validate_options` and starts the `autotag` job titled "Tag output images (legacy page)"; +the legacy `/autotag/single` only requires a number. `GET /autotag/models` reports the saved values (`wd14_general_threshold`, `wd14_character_threshold`, `suggestion_limit`, `default_prompt`, `model_type`, `existing`), so the Tag image and Auto-tag selected dialogs start from them; a preset (the bulk review, Preview tagging review) wins. @@ -3892,6 +3894,57 @@ images", "Tag output images (legacy page)", "Tagging one image (legacy page)", is pinned by `tests/test_autotag_schema.py` (forward-only: one empty `WITHOUT ROWID` table, optional in old backups, refused without the annotation group). +- **The legacy pages (#379, PR 2)** write image tags, never prompt tags. + `POST /autotag/start` (the admin and gallery batch) reads `model_type`, `prompt`, + `skip_tagged` (Skip or Merge), `keep_in_memory` and the thresholds (else the saved + ones), checks them with `validate_options` (WD14: both thresholds; GGUF/HF: + `instructions` = the non-blank `prompt`, else the saved template; the saved + "Suggestions per image") and starts the `autotag` job "Tag output images (legacy + page)" (gate owner `legacy-batch`) through the native start, so the safety backup and + the disk check apply. A second legacy batch answers JSON 409 `{success: false, error: + "Auto-tagging is already running"}` (also when another legacy start wins the race past + that check); `SseStream` hands a non-2xx JSON `error` to `onerror` as `serverMessage`, + which both pages show instead of "Connection lost during auto tagging". Any other refusal (a job in the way: " + already running", `model-unavailable`, `insufficient-space`, `model-busy`, a bad + option) is one `{type: "error", status: "error", message}` frame. The stream relays + the job: the phases as `{type: "progress", progress: 0 | 5 | 8, status: "Taking a + safety backup…" | "Loading model…" | "Listing output images…"}`, each image as + `{type: "progress", progress: 10 + progress × 85 / 100, status, processed, tagged, + skipped}` (at most every 0.5 s; `processed` = Tagged + Unchanged + Errors), and the end + as `complete` (`processed`, `tagged`, `skipped`, `errors`, `skipped_outside_roots: 0`, + `status` "Complete" or "No still images found", `model_status`, `model_loaded`), + `{type: "cancelled", message: "Cancelled. Tags already added were kept."}` or an + `error` frame with the job's status line. A job quiet for 5 s repeats its last frame, + and a closed stream cancels the job (the legacy pages have no other stop). + `POST /autotag/single` takes `existing` (`"skip"` or `"merge"`, default merge). Under + Skip, an image whose current bytes are auto-tagged answers `{success, already_tagged: + true, image_tags, prompt_id, filename}` after one stat, without the model gate. + Otherwise the image is captured and the model runs on the capture under the gate with + the saved cap; the answer adds `tags`, `sha256` (the captured bytes) and `image_tags`, + and `service.custom_prompt` is never set. Refusals: `409 model-unavailable`, `413 + too-large`, `404 file-not-found`, `409 file-changed`, `503 busy`, `415 not-a-still` + (the image is marked), `409 model-busy`. `POST /autotag/apply` takes exactly + `{image_id, sha256, tags}` or `{path, root_index?, sha256, tags}` (64 KiB), re-hashes + the file and unions the tags into that image (`add_image_tags`, marked): `{success, + added_tags, total_tags}`; an empty list only marks the image. A body with `prompt_id` + is `400` "This page is out of date. Reload it."; a changed file `409 file-changed` + ("The image changed since its tags were generated."); a full set `409 tag-limit`; a + restore `409 storage-changed`. The admin review asks each image with Skip first: an + auto-tagged one shows "Image Already Auto-tagged" / "A tagging model already tagged + this image. Generate new tags?" with its image tags (Skip, Skip All Tagged, Re-tag asks + again with Merge); the gallery review always merges. Apply & Next sends the image and + its `sha256` (linked or not). The batch copy: "Skip images already auto-tagged", + "Auto tagging complete! Tagged N images." ("1 image"), and under both progress bars + "Closing this window stops tagging. Tags already added are kept. Prompt Manager's + Activity page shows the time left.". Pinned by `tests/test_autotag_legacy.py`, + `tests/js/image-helpers.test.js` and `legacy-autotag.spec`. +- **Tag image dialog marks (DEC4).** A save after Generate suggestions ran on that + capture marks the image auto-tagged (also a replayed save, and a saved retry after + `tag-conflict`); a save without Generate does not. The API keeps the suggested capture + tokens; each new one drops the tokens whose capture is gone (`expires_in`, which never + resets a capture's idle clock), a release or a saved receipt drops its own, and + closing the captures clears them (`TestDialogMarks`). + No browser test downloads or runs a model: the injected `TestModel` answers, its `download_model` raises unless `/__test__/tagging/download {mode: hold | succeed | fail}` installs a fake that waits on a fixture event, and every tagging reset releases diff --git a/database/operations.py b/database/operations.py index 78ba02c..0f45097 100644 --- a/database/operations.py +++ b/database/operations.py @@ -1808,37 +1808,6 @@ class PromptDatabase( reader = DocumentReader(conn) return [self._image_row_to_dict(row, reader) for row in cursor] - def autotag_batch_candidates(self) -> List[Dict[str, Any]]: - """The images a Tag output images batch looks at (#328). - - ``{id, image_path, prompt_id, prompt_tags}`` for every image with a - stored path that is linked to a prompt, newest first (the - ``get_all_images`` order, ties newest row first). Never reads - ``workflow_data`` or the other documents: a batch needs none of them. - """ - sql = ( - "SELECT gi.id, gi.image_path, gi.prompt_id, " - "(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt " - "JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = gi.prompt_id) " - "AS prompt_tags " - "FROM generated_images gi " - "INNER JOIN prompts p ON gi.prompt_id = p.id " - "WHERE gi.image_path IS NOT NULL AND gi.image_path != '' " - "ORDER BY gi.generation_time DESC, gi.id DESC" - ) - with self.model.get_connection() as conn: - return [ - { - "id": row["id"], - "image_path": row["image_path"], - "prompt_id": row["prompt_id"], - "prompt_tags": ( - row["prompt_tags"].split("|||") if row["prompt_tags"] else [] - ), - } - for row in conn.execute(sql) - ] - def get_all_images( self, limit: int = 0, diff --git a/py/api/__init__.py b/py/api/__init__.py index d45f2c4..485578c 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -363,6 +363,9 @@ class PromptManagerAPI( # Private media captures of the annotation routes (py/api/annotations.py), # created on first use with their restore subscription. self._captures = None + # Capture tokens Generate suggestions ran on: a save of one marks the + # image auto-tagged (#379). + self._suggested_tokens = set() self._captures_unsubscribe = None # Set by close_annotation_captures (shutdown): no store is recreated. self._captures_closed = False diff --git a/py/api/annotations.py b/py/api/annotations.py index 8c799ff..47f905e 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -296,6 +296,7 @@ class AnnotationRoutesMixin: 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: @@ -515,6 +516,7 @@ class AnnotationRoutesMixin: 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"} ) @@ -564,6 +566,9 @@ class AnnotationRoutesMixin: 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, @@ -574,6 +579,7 @@ class AnnotationRoutesMixin: expected_epoch=expected_epoch, expected_tag_revision=revision, tags=tags, + mark_autotagged=suggested, ) except AnnotationCommandError as exc: current = exc.current or {} @@ -590,10 +596,23 @@ class AnnotationRoutesMixin: 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: @@ -675,6 +694,7 @@ class AnnotationRoutesMixin: finally: if not started: store.release_use(token) + self._remember_suggested(store, token) return web.json_response( { "success": True, diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 5d0b7fe..7314f0f 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -6,7 +6,6 @@ import functools import json import os import re -import threading import time from pathlib import Path from typing import Optional @@ -33,6 +32,8 @@ AUTOTAG_JOB_TITLE = "Tag output images" AUTOTAG_PREVIEW_LEASE = "Tag output images review" AUTOTAG_OWNER_KIND = "autotag-job" BATCH_PREVIEW_BODY_MAX = 1024 +APPLY_BODY_MAX = 64 * 1024 # POST /autotag/apply: one image and its tags +_HEX64 = re.compile(r"[0-9a-f]{64}") BATCH_JOB_BODY_MAX = 32 * 1024 REVIEW_REQUIRED = ( "Replacing tags needs a review of each image. Use Preview tagging review, " @@ -133,16 +134,6 @@ _MEDIA_EXTENSIONS = frozenset( _SCAN_MAX_FILES = 50_000 _SCAN_MAX_DEPTH = 12 -# One batch run at a time; a second autotag/start gets 409 while this is held. -_autotag_batch_lock = threading.Lock() - - -def _file_exists(path) -> bool: - try: - return Path(str(path)).exists() - except (OSError, ValueError): - return False - def path_is_within(candidate, root) -> bool: """True when ``candidate`` is strictly inside ``root`` on the real filesystem. @@ -178,27 +169,39 @@ async def _single(event): _monotonic = time.monotonic # the legacy frames' throttle clock (patched in tests) -# What one image came to in a batch (``on_item``): fixed strings, never an -# exception's text. -SKIP_DETAILS = { - "unlinked": "File not found", - "prompt already processed": "Prompt handled by another image", - "outside allowed directories": "Outside the output folders", - "file missing": "File not found", - "already tagged": "Prompt already tagged", +# The legacy pages' batch (#379): the ``autotag`` job, relayed as their frames. +LEGACY_BATCH_TITLE = "Tag output images (legacy page)" +LEGACY_FRAME_INTERVAL = 0.5 # progress frames at most this often +LEGACY_REPEAT_AFTER = 5.0 # a quiet job repeats its last frame (notices a closed page) +_LEGACY_PHASES = { + "backup": (0, "Taking a safety backup…"), + "model": (5, "Loading model…"), + "listing": (8, "Listing output images…"), } -def _added(count): - return f"Added {count} tag" if count == 1 else f"Added {count} tags" +def _legacy_busy(): + """The JSON 409 a legacy start gets while another legacy batch runs.""" + return web.json_response({"success": False, "error": _BUSY_ERROR}, status=409) -async def _refusal(title): - """The single error event of an auto-tag batch refused by another job.""" - message = f"{title} already running" +async def _error_frame(message): + """The single error frame of a legacy batch that could not start.""" yield {"type": "error", "status": "error", "message": message} +def _legacy_counts(job) -> dict: + """Legacy counters: ``processed`` = images the model ran on.""" + counters = job.counters + return { + "processed": counters.get("tagged", 0) + + counters.get("unchanged", 0) + + counters.get("errors", 0), + "tagged": counters.get("tagged", 0), + "skipped": counters.get("skipped", 0), + } + + def wd14_device(service) -> str: """Worker thread: ``"gpu"`` when WD14 runs (or would run) on CUDA, else ``"cpu"``; a service without the probe is ``"cpu"``.""" @@ -476,303 +479,181 @@ class AutotagRoutesMixin: params.update(body) return params - @staticmethod - def _real_tags(tags) -> list: - """Normalise a tag field to a list without bookkeeping markers.""" - if isinstance(tags, str): - tags = [t.strip() for t in tags.split(",") if t.strip()] - return [ - t - for t in (tags or []) - if t != "auto-scanned" - and not t.startswith("prepend:") - and not t.startswith("append:") - ] - - async def _apply_tags_to_prompt(self, prompt_id, tags) -> int: - """Merge ``tags`` into the prompt; how many new tags were added.""" - existing_prompt = await self._run_in_executor( - self.db.get_prompt_by_id, prompt_id - ) - if not existing_prompt: - return 0 - existing_tags = existing_prompt.get("tags", []) - if isinstance(existing_tags, str): - existing_tags = [t.strip() for t in existing_tags.split(",") if t.strip()] - new_tags = list(dict.fromkeys(t for t in tags if t not in existing_tags)) - if not new_tags: - return 0 - await self._run_in_executor( - self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags - ) - return len(new_tags) - - async def _autotag_batch_events(self, service, opts: dict): - """Yield SSE progress dicts while tagging every image linked to a prompt. - - Both batches run here: the legacy page's (its frames unchanged) and - the ``autotag`` job's. ``opts``: ``model_type``, ``keep_in_memory``, - ``existing`` (``skip`` leaves prompts with real tags alone, ``merge`` - adds to them), the two thresholds, and optionally ``gate_owner``, - ``prompt`` (the template, passed to each call; the legacy page sets - ``custom_prompt`` instead), ``cancelled`` (checked before each image: - a cancelled run returns without a ``complete`` frame) and - ``on_item(index, total, group, image_path, detail)`` (one outcome per - image, ``SKIP_DETAILS`` for the skips). - """ - model_type = opts["model_type"] - keep_in_memory = opts["keep_in_memory"] - cancelled = opts.get("cancelled") or (lambda: False) - on_item = opts.get("on_item") or (lambda *_args: None) - status = await self._run_in_executor(service.get_models_status) - if not status.get(model_type, {}).get("downloaded"): - yield {"type": "error", "message": f"Model {model_type} not downloaded"} - return - - if cancelled(): - return - yield {"type": "progress", "progress": 0, "status": "Loading model..."} - run = self._gated_runner(opts.get("gate_owner")) - try: - await run(service.load_model, model_type, True) - except Exception: - self.logger.exception("Failed to load model") - yield { - "type": "error", - "message": "Failed to load model. Check server logs for details.", - } - return - - if opts.get("custom_prompt"): - service.custom_prompt = opts["custom_prompt"] - tag_kwargs = { - "general_threshold": opts["general_threshold"], - "character_threshold": opts["character_threshold"], - } - if opts.get("prompt"): - tag_kwargs["prompt"] = opts["prompt"] - - yield { - "type": "progress", - "progress": 5, - "status": "Model loaded. Fetching all images from database...", - } - images = await self._run_in_executor(self.db.autotag_batch_candidates) - total = len(images) - if total == 0: - yield { - "type": "complete", - "processed": 0, - "tagged": 0, - "skipped": 0, - "errors": 0, - "skipped_outside_roots": 0, - "status": "No images with linked prompts found", - } - if not keep_in_memory: - service.unload_model() - return - - yield { - "type": "progress", - "progress": 10, - "status": f"Found {total} images. Processing...", - } - - counts = { - "processed": 0, - "tagged": 0, - "skipped": 0, - "errors": 0, - "skipped_outside_roots": 0, - } - tagged_prompt_ids = set() - last_update = _monotonic() - finished = False - - def progress_event(i, message): - return { - "type": "progress", - "progress": 10 + int((i + 1) / total * 85), - "status": message, - "processed": counts["processed"], - "tagged": counts["tagged"], - "skipped": counts["skipped"], - "skipped_outside_roots": counts["skipped_outside_roots"], - } - - try: - for i, image_data in enumerate(images): - if cancelled(): - return - image_path = image_data.get("image_path") - prompt_id = image_data.get("prompt_id") - skip_reason = None - allowed_path = None - if not image_path or not prompt_id: - skip_reason = "unlinked" - elif prompt_id in tagged_prompt_ids: - skip_reason = "prompt already processed" - else: - # Rows can point anywhere (restored DB, hand edits): only - # media files inside the served output roots are tagged. - allowed_path = self._resolve_allowed_image_path(str(image_path)) - if allowed_path is None: - if _file_exists(image_path): - skip_reason = "outside allowed directories" - counts["skipped_outside_roots"] += 1 - else: - skip_reason = "file missing" - elif opts["existing"] == "skip" and self._real_tags( - image_data.get("prompt_tags", []) - ): - tagged_prompt_ids.add(prompt_id) - skip_reason = "already tagged" - - if skip_reason is not None: - counts["skipped"] += 1 - on_item(i, total, "skipped", image_path, SKIP_DETAILS[skip_reason]) - now = _monotonic() - if (now - last_update) >= 0.5 or i == total - 1: - yield progress_event( - i, f"Skipping {i + 1}/{total} ({skip_reason})..." - ) - await asyncio.sleep(0.01) - last_update = now - continue - - try: - tags = await run(service.generate_tags, allowed_path, **tag_kwargs) - counts["processed"] += 1 - added = 0 - if tags: - added = await self._apply_tags_to_prompt(prompt_id, tags) - if added: - counts["tagged"] += 1 - on_item(i, total, "tagged", image_path, _added(added)) - else: - counts["skipped"] += 1 - on_item(i, total, "unchanged", image_path, "No new tags") - tagged_prompt_ids.add(prompt_id) - except Exception as img_err: - self.logger.error( - f"Error processing {os.path.basename(str(image_path))}: " - f"{img_err}" - ) - counts["errors"] += 1 - counts["processed"] += 1 - on_item(i, total, "errors", image_path, "Tagging failed") - - yield progress_event(i, f"Processing {i + 1}/{total}...") - await asyncio.sleep(0.01) - last_update = _monotonic() - - finished = True - if not keep_in_memory: - service.unload_model() - yield { - "type": "complete", - "progress": 100, - **counts, - "status": "Complete", - "model_status": ( - "Model kept in memory" if keep_in_memory else "Model unloaded" - ), - "model_loaded": keep_in_memory, - } - finally: - if not finished and not keep_in_memory: - service.unload_model() - async def start_autotag(self, request): - """Start batch auto-tagging with streaming progress (POST only).""" + """The legacy pages' batch (POST only): the ``autotag`` job as frames. + + Its parameters (``model_type``, ``prompt``, ``skip_tagged``, + ``keep_in_memory``, the thresholds or the saved ones) become the job's + options, checked by ``validate_options``. The job, titled "Tag output + images (legacy page)", tags each output image and never a prompt + (#379). A second legacy batch answers JSON 409; every other refusal is + one ``error`` frame (both pages read this route as a stream). Closing + the stream cancels the job. + """ + from ..autotag import MODELS + from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig + from .annotation_suggest import OptionError, validate_options + params = await self._read_start_params(request) if params is None: return self._bad_request("Request body must be a JSON object") - - opts = { - "model_type": params.get("model_type", "gguf"), - "custom_prompt": params.get("prompt", ""), - "existing": ( - "skip" - if str(params.get("skip_tagged", "true")).lower() == "true" - else "merge" - ), - "keep_in_memory": str(params.get("keep_in_memory", "true")).lower() - == "true", - "general_threshold": params.get("general_threshold"), - "character_threshold": params.get("character_threshold"), - } + model_type = params.get("model_type", "gguf") try: - opts["general_threshold"], opts["character_threshold"] = ( - _thresholds_or_saved( - opts["general_threshold"], opts["character_threshold"] - ) + general, character = _thresholds_or_saved( + params.get("general_threshold"), params.get("character_threshold") ) except (ValueError, TypeError): return self._bad_request( "general_threshold and character_threshold must be numeric" ) - - errors = [] # an error event ends the run; Activity lists it failed - - async def events(): - try: - from ..autotag import get_autotag_service - - async for event in self._autotag_batch_events( - get_autotag_service(), opts - ): - if event.get("type") == "error": - errors.append(event) - yield event - except Exception: - self.logger.exception("AutoTag error") - errors.append(_INTERNAL_ERROR) - yield { - "type": "error", - "message": _INTERNAL_ERROR, - } - - # The batch writes prompt tags: a records hold on the registry, so it - # never runs beside a scan, restore or the workspace tagging job. + existing = ( + "skip" + if str(params.get("skip_tagged", "true")).lower() == "true" + else "merge" + ) + keep_in_memory = str(params.get("keep_in_memory", "true")).lower() == "true" + if not isinstance(model_type, str) or model_type not in MODELS: + return await self._stream_sse( + request, _error_frame(f"Invalid model type: {model_type}") + ) + if model_type.startswith("wd14"): + options = {"general_threshold": general, "character_threshold": character} + else: + prompt = params.get("prompt") + template = prompt if isinstance(prompt, str) and prompt.strip() else None + options = {"instructions": template or AutotagConfig.TEMPLATE} + try: + _model, options = validate_options( + model_type, + options, + default_limit=AutotagConfig.suggestion_limit(), + max_limit=MAX_SUGGESTION_LIMIT, + ) + except OptionError as error: + return await self._stream_sse(request, _error_frame(str(error))) + running = self._jobs.active(AUTOTAG_BATCH_JOB_KIND) + if running is not None and running.title == LEGACY_BATCH_TITLE: + return _legacy_busy() + # A job in the way is named before the engine is asked anything. busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records") if busy is not None: - # A second legacy batch keeps its JSON 409; the workspace's Tag - # output images job is named in a stream frame. - if busy.kind == AUTOTAG_BATCH_JOB_KIND and busy.title != AUTOTAG_JOB_TITLE: - return web.json_response( - {"success": False, "error": _BUSY_ERROR}, status=409 - ) - # Both legacy pages read this route as a stream, where a 409 reads - # as a lost connection: refuse with one error frame instead (the - # admin page keys on ``type``, the gallery on ``status``). - return await self._stream_sse(request, _refusal(busy.title)) - async with self._jobs.hold( - AUTOTAG_BATCH_JOB_KIND, - OWNER_LABELS["legacy-batch"], - admission="records", - ) as job: - response = await self._start_autotag_batch(request, opts, events) - if response.status >= 400 or errors: - job.finish("failed") - return response - - async def _start_autotag_batch(self, request, opts, events): - if not _autotag_batch_lock.acquire(blocking=False): - return web.json_response( - {"success": False, "error": _BUSY_ERROR}, status=409 + return await self._stream_sse( + request, _error_frame(f"{busy.title} already running") ) try: - 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"]) + job, refusal = await self._start_autotag( + model_type=model_type, + existing=existing, + options=options, + keep_in_memory=keep_in_memory, + title=LEGACY_BATCH_TITLE, + owner_kind="legacy-batch", + ) + except Exception: + self.logger.exception("AutoTag error") + return await self._stream_sse(request, _error_frame(_INTERNAL_ERROR)) + if refusal is not None: + if refusal.message == f"{LEGACY_BATCH_TITLE} already running": + # Another legacy start won the race past the check above. + return _legacy_busy() + return await self._stream_sse(request, _error_frame(refusal.message)) + return await self._stream_sse( + request, self._legacy_batch_frames(job, keep_in_memory) + ) + + async def _legacy_batch_frames(self, job, keep_in_memory): + """The legacy frames of a running ``autotag`` job, then its end frame. + + Phases map to fixed frames; progress frames come at most every 0.5 s; + a job quiet for 5 s repeats the last frame, so a closed page is noticed + (aiohttp sees a disconnect only on a write). A stream closed before the + job ended cancels it. + """ + events = job.events() + pending = None + last = None + last_progress = None + try: + while True: + if pending is None: + pending = asyncio.ensure_future(events.__anext__()) + done, _ = await asyncio.wait({pending}, timeout=LEGACY_REPEAT_AFTER) + if not done: + if last is not None: + yield last + continue + try: + payload = pending.result() + except StopAsyncIteration: + pending = None + break + pending = None + frame = self._legacy_frame(job, payload) + if frame is None: + continue + if frame.get("phase") is None: + now = _monotonic() + if ( + last_progress is not None + and now - last_progress < LEGACY_FRAME_INTERVAL + ): + continue + last_progress = now + frame.pop("phase", None) + last = frame + yield frame + yield self._legacy_end_frame(job, keep_in_memory) finally: - _autotag_batch_lock.release() + if pending is not None: + pending.cancel() + with contextlib.suppress(BaseException): + await pending + with contextlib.suppress(BaseException): + await events.aclose() + if job.running: + self._jobs.cancel(job.id) + + @staticmethod + def _legacy_frame(job, payload): + """One job event as a legacy progress frame, or None (an end event).""" + kind = payload.get("type") + if kind == "phase": + progress, status = _LEGACY_PHASES.get(payload.get("phase"), (0, "")) + if not status: + return None + return { + "type": "progress", + "progress": progress, + "status": status, + "phase": True, + } + if kind != "progress" or not job.total: + return None + return { + "type": "progress", + "progress": 10 + job.progress * 85 // 100, + "status": job.status_line, + **_legacy_counts(job), + } + + @staticmethod + def _legacy_end_frame(job, keep_in_memory): + if job.state == "finished": + return { + "type": "complete", + "progress": 100, + **_legacy_counts(job), + "errors": job.counters.get("errors", 0), + "skipped_outside_roots": 0, + "status": "Complete" if job.total else "No still images found", + "model_status": ( + "Model kept in memory" if keep_in_memory else "Model unloaded" + ), + "model_loaded": keep_in_memory, + } + if job.state == "cancelled": + return {"type": "cancelled", "message": BATCH_CANCELLED} + return {"type": "error", "status": "error", "message": job.status_line} # ── Tag output images (#328) ──────────────────────────────────────── @@ -969,59 +850,69 @@ class AutotagRoutesMixin: status=403, ) - async def _select_autotag_target(self, data: dict): - """Pick the image to tag from an autotag/single body. + def _root_of(self, resolved: str): + """``(root_index, root, relative)`` of a resolved file inside a root. - Returns ``(image_path, prompt_id, None)`` on success or - ``(None, None, error_response)`` when the request must be rejected. + The first output root that holds it (as ``path_is_within`` checks); + ``relative`` is ``/``-separated below the root's real path, so the + capture store's alias check holds. None when no root holds it. + """ + for index, root in enumerate(self._get_all_output_dirs()): + if path_is_within(resolved, root): + real_root = os.path.realpath(str(root)) + relative = os.path.relpath(resolved, real_root).replace(os.sep, "/") + return index, root, relative + return None + + async def _select_autotag_target(self, data: dict): + """Pick the image a legacy single/apply body names. + + Returns ``(target, None)`` with ``target`` = ``{path, prompt_id, + root_index, root, relative}``, or ``(None, error_response)``. """ image_id = data.get("image_id") if image_id is not None: if isinstance(image_id, bool) or not isinstance(image_id, (int, str)): - return None, None, self._bad_request("image_id must be an integer") + return None, self._bad_request("image_id must be an integer") try: image_id = int(image_id) except (TypeError, ValueError): - return None, None, self._bad_request("image_id must be an integer") + return None, self._bad_request("image_id must be an integer") record = await self._run_in_executor(self.db.get_image_by_id, image_id) if not record or not record.get("image_path"): - return ( - None, - None, - web.json_response( - {"success": False, "error": f"Image {image_id} not found"}, - status=404, - ), + return None, web.json_response( + {"success": False, "error": f"Image {image_id} not found"}, + status=404, ) stored_path = record["image_path"] resolved = self._resolve_allowed_image_path(stored_path) if resolved is None: - return None, None, self._forbidden(stored_path) - return resolved, record.get("prompt_id"), None + return None, self._forbidden(stored_path) + return self._target_of(resolved, record.get("prompt_id"), stored_path) raw_path = data.get("path") or data.get("image_path") if not raw_path or not isinstance(raw_path, str): - return None, None, self._bad_request("image_id or path is required") + return None, self._bad_request("image_id or path is required") if "root_index" in data: root_index = data["root_index"] if isinstance(root_index, bool) or not isinstance(root_index, int): - return None, None, self._bad_request("root_index must be an integer") + return None, self._bad_request("root_index must be an integer") roots = self._get_all_output_dirs() if not 0 <= root_index < len(roots): - return None, None, self._forbidden(raw_path) + return None, self._forbidden(raw_path) root = roots[root_index] if Path(raw_path).is_absolute() or ".." in raw_path.replace( "\\", "/" ).split("/"): - return None, None, self._forbidden(raw_path) + return None, self._forbidden(raw_path) raw_path = str(Path(root) / raw_path) if not path_is_within(raw_path, root): - return None, None, self._forbidden(raw_path) + return None, self._forbidden(raw_path) resolved = self._resolve_allowed_image_path(raw_path) if resolved is None: - return None, None, self._forbidden(raw_path) + return None, self._forbidden(raw_path) prompt_id = None try: @@ -1034,125 +925,345 @@ class AutotagRoutesMixin: ) except Exception as e: self.logger.warning(f"Could not find linked prompt: {e}") - return resolved, prompt_id, None + return self._target_of(resolved, prompt_id, raw_path) + + def _target_of(self, resolved, prompt_id, shown): + located = self._root_of(resolved) + if located is None: + return None, self._forbidden(shown) + index, root, relative = located + return { + "path": resolved, + "prompt_id": prompt_id, + "root_index": index, + "root": root, + "relative": relative, + }, None @staticmethod def _bad_request(message: str): return web.json_response({"success": False, "error": message}, status=400) + @staticmethod + def _coded_error(status: int, code: str, message: str): + return web.json_response( + {"success": False, "code": code, "error": message}, status=status + ) + + def _capture_refusal(self, exc, target): + """The legacy answer for a file the capture store refused (#379).""" + from .image_tag_jobs import FILE_CHANGED, FILE_NOT_FOUND, TOO_LARGE + + name = type(exc).__name__ + if name == "CaptureTooLarge": + return self._coded_error(413, "too-large", TOO_LARGE) + if name == "CaptureCapacity": + return self._coded_error( + 503, "busy", "The file could not be read now. Try again." + ) + if name == "CaptureChanged" and os.path.lexists(target["path"]): + return self._coded_error(409, "file-changed", FILE_CHANGED) + if name in ("CaptureDenied", "CaptureChanged"): + return self._coded_error(404, "file-not-found", FILE_NOT_FOUND) + return self._coded_error(500, "read-failed", "The file could not be read.") + + def _single_options(self, data, model_type): + """The validated suggestion options of a legacy single body.""" + from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig + from .annotation_suggest import validate_options + + general, character = _thresholds_or_saved( + data.get("general_threshold"), data.get("character_threshold") + ) + if model_type.startswith("wd14"): + options = {"general_threshold": general, "character_threshold": character} + else: + prompt = data.get("prompt") + template = prompt if isinstance(prompt, str) and prompt.strip() else None + options = {"instructions": template or AutotagConfig.TEMPLATE} + use_gpu = data.get("use_gpu", True) + options["use_gpu"] = use_gpu if isinstance(use_gpu, bool) else True + return validate_options( + model_type, + options, + default_limit=AutotagConfig.suggestion_limit(), + max_limit=MAX_SUGGESTION_LIMIT, + )[1] + async def autotag_single(self, request): - """Generate tags for a single image selected by image_id or gallery path.""" + """Suggest tags for one image (the legacy review), never writing. + + ``existing: "skip"`` answers ``already_tagged`` without taking the + model gate when the image's current bytes are auto-tagged (#379). + Otherwise the image is captured and the model runs on the capture + under the gate, with at most the saved "Suggestions per image"; + the answer carries the captured ``sha256`` that ``/autotag/apply`` + writes against, and the image's current ``image_tags``. + """ + from ..autotag import MODELS + from .annotation_suggest import OptionError + try: data = await request.json() - if not isinstance(data, dict): - return self._bad_request("Request body must be a JSON object") + except (json.JSONDecodeError, UnicodeDecodeError): + return self._bad_request("Request body must be valid JSON") + if not isinstance(data, dict): + return self._bad_request("Request body must be a JSON object") + try: model_type = data.get("model_type", "gguf") - custom_prompt = data.get("prompt") - use_gpu = data.get("use_gpu", True) - general_threshold = data.get("general_threshold") - character_threshold = data.get("character_threshold") - try: - general_threshold, character_threshold = _thresholds_or_saved( - general_threshold, character_threshold + if not isinstance(model_type, str) or model_type not in MODELS: + return self._bad_request( + "model_type must be one of " + ", ".join(sorted(MODELS)) ) + existing = data.get("existing", "merge") + if existing not in ("skip", "merge"): + return self._bad_request('existing must be "skip" or "merge"') + try: + options = self._single_options(data, model_type) except (ValueError, TypeError): return self._bad_request( "general_threshold and character_threshold must be numeric" ) - - image_path, prompt_id, error = await self._select_autotag_target(data) + except OptionError as error: + return self._bad_request(str(error)) + target, error = await self._select_autotag_target(data) if error is not None: return error - - 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 custom_prompt: - service.custom_prompt = custom_prompt - - 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( - { - "success": True, - "tags": tags, - "prompt_id": prompt_id, - "filename": os.path.basename(image_path), - } - ) - - except json.JSONDecodeError: - return self._bad_request("Request body must be valid JSON") + base = { + "success": True, + "prompt_id": target["prompt_id"], + "filename": os.path.basename(target["path"]), + } + if existing == "skip": + marked = await self._run_in_executor(self._marked_tags, target) + if marked is not None: + return web.json_response( + { + **base, + "already_tagged": True, + "tags": [], + "image_tags": marked, + } + ) + return await self._suggest_single(target, model_type, options, base) except Exception as e: self.logger.error(f"AutoTag single error: {e}") return web.json_response( {"success": False, "error": self._public_error(e)}, status=500 ) - async def apply_autotag(self, request): - """Apply selected tags to a prompt.""" + def _locator_of(self, target) -> str: + return os.path.normcase( + os.path.join( + os.path.realpath(str(target["root"])), *target["relative"].split("/") + ) + ) + + def _marked_tags(self, target): + """Worker thread: the current image tags when the image is + auto-tagged, else None. One stat, no hash.""" + from .images import _current_fingerprint + + fingerprint = _current_fingerprint(target["path"]) + if fingerprint is None: + return None + locator = self._locator_of(target) + if not self.db.autotag_state({locator: fingerprint}).get(locator): + return None + return self.db.current_image_tags({locator: fingerprint}).get(locator, []) + + async def _suggest_single(self, target, model_type, options, base): + from ..autotag import MODELS, get_autotag_service + from .annotation_suggest import unavailable_message + from .annotations import _pinned_suggestions, _settled, _submit + try: - data = await request.json() - prompt_id = data.get("prompt_id") - tags = data.get("tags", []) + from ...utils.media_capture import CaptureError + except ImportError: + from utils.media_capture import CaptureError - if not prompt_id: - return web.json_response( - {"success": False, "error": "prompt_id is required"}, status=400 + service = get_autotag_service() + status = await self._run_in_executor(service.get_models_status) + if not (status.get(model_type) or {}).get("downloaded"): + message = unavailable_message(MODELS[model_type]["name"]) + return self._coded_error(409, "model-unavailable", message) + owner = model_gate.new_owner("legacy-single") + if not model_gate.try_acquire(owner): + return model_busy_response() + store = self._capture_store() + try: + epoch = await self._run_in_executor(self.db.get_execution_epoch) + scope = self._gallery_root_scope(target["root"], target["root_index"]) + try: + capture = await _settled( + _submit( + store.capture, + target["root"], + target["relative"], + epoch=epoch, + scope=scope, + root_index=target["root_index"], + ), + lambda made: made is not None and store.release(made.token), ) - - if not tags: - return web.json_response( - {"success": True, "message": "No tags to apply"} + except CaptureError as exc: + return self._capture_refusal(exc, target) + try: + if not capture.inferable: + await self._mark_untaggable(capture, epoch) + return self._coded_error( + 415, "not-a-still", "Not a supported still image" + ) + if store.acquire(capture.token) is None: + return self._coded_error( + 409, + "storage-changed", + "Prompt storage was restored. Reload the page.", + ) + pairs = await model_gate.run( + owner, + _pinned_suggestions, + store, + capture.token, + service, + capture, + model_type, + options, + on_skipped=functools.partial(store.release_use, capture.token), ) + finally: + store.release(capture.token) + finally: + model_gate.release(owner) + current = await self._run_in_executor( + self.db.read_image_tags, capture.locator, capture.sha256 + ) + return web.json_response( + { + **base, + "tags": [str(name) for name, _score in pairs], + "sha256": capture.sha256, + "image_tags": current["tags"], + } + ) - prompt = await self._run_in_executor(self.db.get_prompt_by_id, prompt_id) - if not prompt: - return web.json_response( - {"success": False, "error": f"Prompt {prompt_id} not found"}, - status=404, - ) - - existing_tags = prompt.get("tags", []) - if isinstance(existing_tags, str): - existing_tags = [ - t.strip() for t in existing_tags.split(",") if t.strip() - ] - - new_tags = [t for t in tags if t not in existing_tags] - all_tags = existing_tags + new_tags - + async def _mark_untaggable(self, facts, epoch): + try: await self._run_in_executor( - self.db.update_prompt_metadata, prompt_id, tags=all_tags + functools.partial( + self.db.add_image_tags, + locator=facts.locator, + sha256=facts.sha256, + size=facts.size, + mtime_ns=facts.mtime_ns, + expected_epoch=epoch, + tags=[], + mark_autotagged=True, + ) ) + except Exception: + self.logger.warning("An unsupported image could not be marked") - return web.json_response( - {"success": True, "added_tags": new_tags, "total_tags": len(all_tags)} + async def apply_autotag(self, request): + """Union reviewed tags into one image's tags and mark it (#379). + + The body names the image (``image_id``, or ``path`` with optional + ``root_index``), the ``sha256`` ``/autotag/single`` returned and the + ``tags``; it is written only onto those exact bytes. A body with + ``prompt_id`` (a page loaded before the update) is refused: no path + writes prompt tags. + """ + try: + from ...database.annotations import ImageTagLimitError, normalize_image_tags + from ...database.library_commands import LibraryCommandError + from ...utils.media_capture import CaptureError + except ImportError: + from database.annotations import ImageTagLimitError, normalize_image_tags + from database.library_commands import LibraryCommandError + from utils.media_capture import CaptureError + + data, error = await self._read_json_body(request, max_bytes=APPLY_BODY_MAX) + if error is not None: + return error + if "prompt_id" in data: + return self._bad_request("This page is out of date. Reload it.") + keys = set(data) + if keys not in ( + {"image_id", "sha256", "tags"}, + {"path", "sha256", "tags"}, + {"path", "root_index", "sha256", "tags"}, + ): + return self._bad_request( + "The body must be {image_id, sha256, tags} or " + "{path, root_index?, sha256, tags}" ) + sha256 = data["sha256"] + if not isinstance(sha256, str) or not _HEX64.fullmatch(sha256): + return self._bad_request( + "sha256 must be 64 lowercase hexadecimal characters" + ) + try: + tags = normalize_image_tags(data["tags"]) + except ImageTagLimitError as exc: + return self._coded_error(400, "tag-limit", str(exc)) + except ValueError as exc: + return self._bad_request(str(exc)) + try: + epoch = await self._run_in_executor(self.db.get_execution_epoch) + target, error = await self._select_autotag_target(data) + if error is not None: + return error + store = self._capture_store() + try: + facts = await self._run_in_executor( + store.verify, target["root"], target["relative"] + ) + except CaptureError as exc: + return self._capture_refusal(exc, target) + if facts.sha256 != sha256: + return self._coded_error( + 409, + "file-changed", + "The image changed since its tags were generated.", + ) + try: + receipt = await self._run_in_executor( + functools.partial( + self.db.add_image_tags, + locator=facts.locator, + sha256=facts.sha256, + size=facts.size, + mtime_ns=facts.mtime_ns, + expected_epoch=epoch, + tags=tags, + mark_autotagged=True, + ) + ) + except LibraryCommandError as exc: + if exc.code == "storage-changed": + return self._coded_error( + 409, + "storage-changed", + "Prompt storage was restored. Reload the page.", + ) + if exc.code == "tag-limit": + from .annotations import _MESSAGES + return self._coded_error(409, "tag-limit", _MESSAGES["tag-limit"]) + raise except Exception as e: self.logger.error(f"Apply autotag error: {e}") return web.json_response( {"success": False, "error": self._public_error(e)}, status=500 ) + added = receipt.get("added", 0) + return web.json_response( + { + "success": True, + "added_tags": receipt["tags"][-added:] if added else [], + "total_tags": len(receipt["tags"]), + } + ) async def unload_autotag_model(self, request): """Manually unload the AutoTag model from memory.""" diff --git a/tests/browser/legacy-autotag.spec.mjs b/tests/browser/legacy-autotag.spec.mjs index 1140a88..ff96f68 100644 --- a/tests/browser/legacy-autotag.spec.mjs +++ b/tests/browser/legacy-autotag.spec.mjs @@ -95,8 +95,14 @@ async function watchProgress(page) { await page.evaluate(() => { window.progressSeen = []; const node = document.getElementById('autoTagProgressPercent'); - new MutationObserver(() => window.progressSeen.push(node.textContent)) - .observe(node, { childList: true, characterData: true, subtree: true }); + // Each change from its own record: frames that land in one task share one + // callback, and the node's text by then is only the last of them. + new MutationObserver(records => { + for (const record of records) { + if (record.type === 'characterData') window.progressSeen.push(record.target.data); + for (const added of record.addedNodes) window.progressSeen.push(added.textContent); + } + }).observe(node, { childList: true, characterData: true, subtree: true }); }); } @@ -108,25 +114,58 @@ async function runGalleryBatch(page) { await page.locator('#startAutoTagBtn').click(); await done; await expect(page.locator('#autoTagProgressModal')).toBeHidden(); - await expect(page.getByText(/^Auto tagging complete! Applied tags to \d+ prompts\.$/)).toBeVisible(); + // #379: the batch counts images, never prompts. + await expect(page.getByText(/^Auto tagging complete! Tagged \d+ images?\.$/)).toBeVisible(); +} + +const CLOSE_NOTE = 'Closing this window stops tagging. Tags already added are kept. Prompt Manager\'s Activity page shows the ' + + 'time left.'; + +async function serverJobs(request) { + const response = await request.get('/prompt_manager/jobs'); + await expect(response).toBeOK(); + return (await response.json()).jobs; +} + +async function runNativeBatch(request) { + const started = await request.post('/prompt_manager/jobs/autotag', { data: { model_type: 'wd14-swinv2', existing: 'skip', + options: { general_threshold: 0.35, character_threshold: 0.85 } } }); + expect(started.status()).toBe(202); + const { job } = await started.json(); + await expect.poll(async () => (await serverJobs(request)).find(item => item.id === job.id)?.state).toBe('finished'); } test('gallery batch modal counts progress and closes on complete', async ({ page, request }) => { await page.goto('/prompt_manager/gallery'); await runGalleryBatch(page); const seen = await page.evaluate(() => window.progressSeen); - // The fixture's prompts are already tagged, so the batch skips them; its progress frames still move the bar. - expect(seen).toContain('95%'); + // The phase frames (model 5%, listing 8%) then the per-image frames (from 10%) move the bar. + expect(seen).toContain('8%'); + expect(seen.some(value => parseInt(value, 10) >= 10)).toBe(true); + await expect(page.locator('#autoTagCloseNote')).toHaveText(CLOSE_NOTE); + // The batch ran as the autotag job of the legacy page. + expect((await serverJobs(request)).filter(job => job.kind === 'autotag').map(job => [job.title, job.state])) + .toEqual([['Tag output images (legacy page)', 'finished']]); expect((await downloads(request)).length).toBe(0); }); +test('admin batch copy names images, not prompts (#379)', async ({ page }) => { + await page.goto('/prompt_manager/admin'); + await page.locator('#autoTagBtn').click(); + await expect(page.locator('input[name="autoTagMode"][value="skip"] + span')).toHaveText('Skip images already auto-tagged'); + await expect(page.getByText('"auto-scanned" is not counted as a real tag')).toHaveCount(0); + await expect(page.locator('#autoTagCloseNote')).toHaveText(CLOSE_NOTE); + await expect(page.locator('#retagConfirmModal h3')).toHaveText('Image Already Auto-tagged'); + await expect(page.locator('#retagConfirmModal')).toContainText('A tagging model already tagged this image. Generate new tags?'); +}); + test('admin review opens the first image', async ({ page }) => { await page.goto('/prompt_manager/admin'); await page.locator('#autoTagBtn').click(); await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check(); await page.locator('#startReviewBtn').click(); - // The first image's prompt is already tagged: confirm a new pass, then the review opens on that image. - await page.locator('#retagConfirmBtn').click(); + // #379: no image is auto-tagged yet (prompt tags no longer count), so the review opens on the first image. + await expect(page.locator('#retagConfirmModal')).toBeHidden(); await expect(page.locator('#autoTagReviewModal')).toBeVisible(); await expect(page.locator('#reviewImageName')).not.toHaveText(/^(image_name\.png|Unknown|)$/); await expect(page.getByText('Failed to start review mode')).toHaveCount(0); @@ -150,3 +189,70 @@ test('gallery unload works and is refused while the gate is held', async ({ page await unload.click(); await expect(page.getByText('WD14-SWINV2 model unloaded successfully')).toBeVisible(); }); + +/** The admin review lists every row, videos too: Skip any image that is not one until the confirmation shows. */ +async function untilAutoTagged(page, confirm) { + // Each image's skip check answers before its confirmation or its error shows. + const failed = page.locator('#reviewTagsVisible .text-pm-error'); + for (let tries = 0; tries < 5; tries++) { + await expect.poll(async () => (await confirm.isVisible()) || (await failed.count()) > 0).toBe(true); + if (await confirm.isVisible()) return; + const answered = page.waitForResponse(item => item.url().endsWith('/prompt_manager/autotag/single')); + await page.locator('#skipReviewBtn').click(); + await answered; + } + await expect(confirm).toBeVisible(); +} + +// #379: the admin review asks before re-tagging an image a model already tagged, and Apply & Next writes that image's +// tags (its bytes, never its prompt). +test('admin review: an auto-tagged image asks first; Re-tag then Apply & Next tags the image', async ({ page, request }) => { + await runNativeBatch(request); + await page.goto('/prompt_manager/admin'); + await page.locator('#autoTagBtn').click(); + await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check(); + await page.locator('#startReviewBtn').click(); + const confirm = page.locator('#retagConfirmModal'); + await untilAutoTagged(page, confirm); + await expect(confirm.locator('#retagExistingTags span').first()).toHaveText('sandstone'); + await page.locator('#retagConfirmBtn').click(); + await expect(page.locator('#autoTagReviewModal')).toBeVisible(); + await expect(page.locator('#reviewTagsVisible .tag-chip').first()).toContainText('sandstone'); + const applied = page.waitForRequest(item => item.url().endsWith('/prompt_manager/autotag/apply')); + const answer = page.waitForResponse(item => item.url().endsWith('/prompt_manager/autotag/apply')); + await page.locator('#applyReviewBtn').click(); + const body = (await applied).postDataJSON(); + expect(body.prompt_id).toBeUndefined(); + expect(body.sha256).toMatch(/^[0-9a-f]{64}$/); + expect((await answer).status()).toBe(200); +}); + +test('admin review: Skip All Tagged passes every auto-tagged image', async ({ page, request }) => { + await runNativeBatch(request); + await page.goto('/prompt_manager/admin'); + await page.locator('#autoTagBtn').click(); + await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check(); + await page.locator('#startReviewBtn').click(); + const confirm = page.locator('#retagConfirmModal'); + await untilAutoTagged(page, confirm); + await page.locator('#retagSkipAllBtn').click(); + await expect(page.locator('#retagConfirmModal')).toBeHidden(); + await expect(page.getByText('Review complete!')).toBeVisible(); +}); + +test('gallery review: Apply & Next tags the image by its path and bytes', async ({ page }) => { + await page.goto('/prompt_manager/gallery'); + await page.locator('#autoTagBtn').click(); + await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check(); + const suggested = page.waitForRequest(item => item.url().endsWith('/prompt_manager/autotag/single')); + await page.locator('#startReviewBtn').click(); + expect((await suggested).postDataJSON().existing).toBe('merge'); + await expect(page.locator('#reviewTagsContainer .tag-chip').first()).toContainText('sandstone'); + const applied = page.waitForRequest(item => item.url().endsWith('/prompt_manager/autotag/apply')); + const answer = page.waitForResponse(item => item.url().endsWith('/prompt_manager/autotag/apply')); + await page.locator('#applyReviewBtn').click(); + const body = (await applied).postDataJSON(); + expect(Object.keys(body).sort()).toEqual(['path', 'root_index', 'sha256', 'tags'].filter(key => key in body).sort()); + expect(body.prompt_id).toBeUndefined(); + expect((await answer).status()).toBe(200); +}); diff --git a/tests/browser/workflows.spec.mjs b/tests/browser/workflows.spec.mjs index 62d8645..e822d99 100644 --- a/tests/browser/workflows.spec.mjs +++ b/tests/browser/workflows.spec.mjs @@ -132,7 +132,8 @@ test('the gallery cache clear keeps preferences and other tools caches', async ( }))).toEqual({ theme: 'light', other: 'keep', prefixed: 'keep', notice: '3.2.0', caches: ['other-tool-gallery-thumbnails'] }); }); -test('actual gallery review generates and persists tags for the selected root', async ({ page, request }) => { +// #379: the reviewed tags are that image's own tags; no prompt changes. +test('actual gallery review generates and persists tags for the selected root\'s image', async ({ page, request }) => { await page.goto('/prompt_manager/gallery'); await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2); await page.locator('#autoTagBtn').click(); @@ -140,13 +141,15 @@ test('actual gallery review generates and persists tags for the selected root', await page.locator('#startReviewBtn').click(); const generation = await generated; expect(generation.request().postDataJSON().root_index).toBe(0); - await expect(page.locator('#reviewTagsContainer')).toContainText('generated'); + await expect(page.locator('#reviewTagsContainer')).toContainText('sandstone'); const result = page.waitForResponse(response => response.url().endsWith('/prompt_manager/autotag/apply')); await page.locator('#applyReviewBtn').click(); expect((await (await result).json()).success).toBe(true); const data = await (await request.get('/prompt_manager/recent?limit=20')).json(); - expect(data.results.find(prompt => prompt.text === 'Cat portrait').tags).toContain('generated'); - expect(data.results.find(prompt => prompt.text === 'Dog landscape').tags).not.toContain('generated'); + expect(data.results.find(prompt => prompt.text === 'Cat portrait').tags).toEqual(['cats']); + expect(data.results.find(prompt => prompt.text === 'Dog landscape').tags).toEqual(['dogs']); + const tagged = await (await request.get('/prompt_manager/images/output?mode=workspace&tags=sandstone')).json(); + expect(tagged.images.map(image => `${image.root_index}:${image.path}`)).toEqual(['0:shared/same.png']); }); test('actual Tags page reloads the selected filter and opens its image viewer', async ({ page }) => { diff --git a/tests/js/gallery-review.test.js b/tests/js/gallery-review.test.js index 32e27c1..52825f7 100644 --- a/tests/js/gallery-review.test.js +++ b/tests/js/gallery-review.test.js @@ -36,33 +36,48 @@ function fixture() { const reply = (index, data) => requests[index].resolve({ json: async () => data }); return { manager, requests, reply, notifications }; } -test('gallery review retains root identity and applies tags with the returned prompt ID', async () => { +const SHA = 'a'.repeat(64); +test('gallery review retains root identity and applies the tags to the image by its sha256 (#379)', async () => { const { manager, requests, reply } = fixture(); const load = manager.loadNextReviewImage(); assert.equal(requests[0].body.root_index, 1); - reply(0, { success: true, prompt_id: 42, tags: ['smile'] }); + assert.equal(requests[0].body.existing, 'merge'); + reply(0, { success: true, prompt_id: 42, tags: ['smile'], sha256: SHA, image_tags: [] }); await load; manager.loadNextReviewImage = () => {}; const apply = manager.applyReviewTags(); - assert.deepEqual(requests[1].body, { prompt_id: 42, tags: ['smile'] }); + assert.equal(requests[1].url, '/prompt_manager/autotag/apply'); + assert.deepEqual(requests[1].body, ImageHelpers.autotagApplyBody( + { id: 'digest', path: 'same.png', root_index: 1 }, SHA, ['smile'])); + assert.equal('prompt_id' in requests[1].body, false); reply(1, { success: true }); await apply; }); -test('unlinked images do not submit an invalid apply request or advance', async () => { +test('an unlinked image applies to the image like a linked one', async () => { const { manager, requests, reply } = fixture(); const load = manager.loadNextReviewImage(); - reply(0, { success: true, prompt_id: null, tags: ['smile'] }); + reply(0, { success: true, prompt_id: null, tags: ['smile'], sha256: SHA, image_tags: [] }); await load; + manager.loadNextReviewImage = () => {}; const apply = manager.applyReviewTags(); - if (requests[1]) reply(1, { success: false, error: "Unexpected apply" }); + assert.equal(requests[1].body.sha256, SHA); + assert.deepEqual(requests[1].body.tags, ['smile']); + reply(1, { success: true }); await apply; - assert.equal(requests.length, 1); - assert.equal(manager.autoTagState.reviewIndex, 0); +}); +test('an answer without a captured sha256 never sends an apply', async () => { + const { manager, requests, reply } = fixture(); + const load = manager.loadNextReviewImage(); + reply(0, { success: true, prompt_id: 42, tags: ['smile'] }); + await load; + manager.loadNextReviewImage = () => {}; + await manager.applyReviewTags(); + assert.equal(requests.filter(r => r.url === '/prompt_manager/autotag/apply').length, 0); }); test('a failed apply keeps the current image available for retry', async () => { const { manager, requests, reply } = fixture(); const load = manager.loadNextReviewImage(); - reply(0, { success: true, prompt_id: 42, tags: ['smile'] }); + reply(0, { success: true, prompt_id: 42, tags: ['smile'], sha256: SHA, image_tags: [] }); await load; const apply = manager.applyReviewTags(); reply(1, { success: false, error: 'Try again' }); diff --git a/tests/js/image-helpers.test.js b/tests/js/image-helpers.test.js index 0439330..a353109 100644 --- a/tests/js/image-helpers.test.js +++ b/tests/js/image-helpers.test.js @@ -5,6 +5,8 @@ const { formatImageCaption, workflowDocumentHtml, autotagSingleBody, + autotagApplyBody, + autotagCompleteMessage, fallbackImageSource, } = require("../../web/js/image-helpers.js"); @@ -121,3 +123,27 @@ test('metadata lookup preserves encoded filename characters and works behind a p assert.equal(metadataLookupUrl('/prompt_manager/images/serve/a.png?v=123'), '/prompt_manager/images/prompt/a.png'); assert.equal(metadataLookupUrl('/elsewhere/a.png'), null); }); + +test("autotagSingleBody asks to skip auto-tagged images only when told (#379)", () => { + assert.equal(autotagSingleBody({ id: 1 }, WD14).existing, "merge"); + assert.equal(autotagSingleBody({ id: 1 }, { ...WD14, existing: "skip" }).existing, "skip"); + assert.equal(autotagSingleBody({ id: 1 }, { ...WD14, existing: "odd" }).existing, "merge"); +}); + +test("autotagApplyBody names the image, its bytes and the tags, never a prompt (#379)", () => { + const sha = "a".repeat(64); + assert.deepEqual(autotagApplyBody({ id: 4, prompt_id: 7 }, sha, ["bird"]), { image_id: 4, sha256: sha, tags: ["bird"] }); + assert.deepEqual(autotagApplyBody({ id: "digest", path: "sub/a.png", root_index: 1 }, sha, []), + { path: "sub/a.png", root_index: 1, sha256: sha, tags: [] }); + assert.deepEqual(autotagApplyBody({ id: 2 }, undefined, null), { image_id: 2, sha256: "", tags: [] }); + const tags = ["x"]; + autotagApplyBody({ id: 1 }, sha, tags).tags.push("y"); + assert.deepEqual(tags, ["x"]); +}); + +test("autotagCompleteMessage counts images (#379)", () => { + assert.equal(autotagCompleteMessage(2), "Auto tagging complete! Tagged 2 images."); + assert.equal(autotagCompleteMessage(1), "Auto tagging complete! Tagged 1 image."); + assert.equal(autotagCompleteMessage(0), "Auto tagging complete! Tagged 0 images."); + assert.equal(autotagCompleteMessage(undefined), "Auto tagging complete! Tagged 0 images."); +}); diff --git a/tests/js/sse-stream.test.js b/tests/js/sse-stream.test.js index 63e2f99..13cd052 100644 --- a/tests/js/sse-stream.test.js +++ b/tests/js/sse-stream.test.js @@ -97,6 +97,30 @@ test("connect reports a non-2xx response through onerror and never onmessage", a assert.equal(source.readyState, source.CLOSED); }); +test("a non-2xx JSON answer's error sentence reaches onerror as serverMessage", async () => { + let error = null; + const fetchImpl = async () => new Response( + JSON.stringify({ success: false, error: "Auto-tagging is already running" }), + { status: 409, headers: { "Content-Type": "application/json" } }, + ); + const source = connect("/x", { fetchImpl }); + source.onerror = (err) => { error = err; }; + await until(() => error !== null); + assert.match(error.message, /409/); + assert.equal(error.status, 409); + assert.equal(error.serverMessage, "Auto-tagging is already running"); +}); + +test("a non-2xx answer without a JSON error has no serverMessage", async () => { + let error = null; + const fetchImpl = async () => new Response("nope", { status: 500 }); + const source = connect("/x", { fetchImpl }); + source.onerror = (err) => { error = err; }; + await until(() => error !== null); + assert.equal(error.status, 500); + assert.equal(error.serverMessage, undefined); +}); + test("connect reports a network failure through onerror", async () => { let error = null; const source = connect("/x", { fetchImpl: async () => { throw new TypeError("Failed to fetch"); } }); diff --git a/tests/test_api_annotation_suggest.py b/tests/test_api_annotation_suggest.py index 29727ee..de097d6 100644 --- a/tests/test_api_annotation_suggest.py +++ b/tests/test_api_annotation_suggest.py @@ -363,23 +363,20 @@ class TestModelGate(SuggestTestCase): 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 + # The legacy batch is the autotag job (#379): it suggests per image. + self.tagger.block = (started, proceed) 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) + calls = len(self.tagger.calls) 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(len(self.tagger.calls), calls) self.assertEqual(self.tagger.loaded, "wd14-vit") # not swapped finally: + self.tagger.block = None proceed.set() response = await batch await response.read() diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index da42059..e2d2dff 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -5,6 +5,7 @@ no network request is made. """ import asyncio +import hashlib import json import os import shutil @@ -34,6 +35,7 @@ from aiohttp.test_utils import AioHTTPTestCase from database.operations import PromptDatabase from py.api import PromptManagerAPI +from py.api.autotag_routes import MODEL_BUSY_ERROR from py.api.jobs import JobBusy from utils.hashing import generate_prompt_hash @@ -70,6 +72,7 @@ class FakeAutotagService: self.status_calls = 0 self.status_threads = [] self.status_paths = {} + self.suggest_calls = 0 def get_models_status(self): self.status_calls += 1 @@ -115,6 +118,25 @@ class FakeAutotagService: raise self.generate_error return list(self.tags) + def generate_tag_suggestions( + self, + image_path, + *, + limit, + prompt=None, + general_threshold=None, + character_threshold=None, + ): + """The suggestion entry point (#379): the fixed tags, at most ``limit``.""" + self.suggest_calls += 1 + tags = self.generate_tags( + image_path, + prompt=prompt, + general_threshold=general_threshold, + character_threshold=character_threshold, + ) + return [(name, None) for name in tags[:limit]] + class AutotagAPITestCase(AioHTTPTestCase): """Stands up the PromptManager routes with a temp DB, temp output dir and a @@ -175,8 +197,17 @@ class AutotagAPITestCase(AioHTTPTestCase): ) def _make_image(self, directory, name="img.png"): + """A small, valid image (the legacy review captures and probes it).""" + from PIL import Image + path = directory / name - path.write_bytes(b"\x89PNG\r\n\x1a\n") + path.parent.mkdir(parents=True, exist_ok=True) + shade = sum(name.encode()) % 256 + fmt = {".jpg": "JPEG", ".jpeg": "JPEG", ".webp": "WEBP"}.get( + path.suffix.lower(), "PNG" + ) + with Image.new("RGB", (4, 4), (shade, 64, 128)) as image: + image.save(path, fmt) return path def _link_image(self, prompt_id, path): @@ -216,10 +247,9 @@ class TestAutotagSingleContainment(AutotagAPITestCase): self.assertEqual(data["tags"], ["1girl", "smile"]) self.assertEqual(data["prompt_id"], prompt_id) self.assertEqual(len(self.service.generate_calls), 1) - self.assertEqual( - os.path.realpath(self.service.generate_calls[0][0]), - os.path.realpath(str(image)), - ) + # #379: the model reads a private capture of the image, never the file. + self.assertNotIn(str(self.output_dir), self.service.generate_calls[0][0]) + self.assertEqual(data["sha256"], hashlib.sha256(image.read_bytes()).hexdigest()) async def test_path_inside_output_dir_is_tagged(self): image = self._make_image(self.output_dir) @@ -449,238 +479,6 @@ class TestAutotagSideEffectRoutesArePostOnly(AutotagAPITestCase): self.assertTrue((await resp.json())["success"]) -class TestAutotagBatch(AutotagAPITestCase): - """Batch tagging over the DB image list with SSE progress reporting.""" - - def _seed(self): - """Three prompts: A (one image), B (two images), C (missing file).""" - a = self._save_prompt("prompt a") - b = self._save_prompt("prompt b") - c = self._save_prompt("prompt c") - self._link_image(a, self._make_image(self.output_dir, "a1.png")) - self._link_image(b, self._make_image(self.output_dir, "b1.png")) - self._link_image(b, self._make_image(self.output_dir, "b2.png")) - self._link_image(c, self.output_dir / "missing.png") - return a, b, c - - async def _start(self, query="model_type=gguf", body=None): - if body is None: - resp = await self.client.request( - "POST", f"/prompt_manager/autotag/start?{query}" - ) - else: - resp = await self._post_json(f"/prompt_manager/autotag/start?{query}", body) - return resp, self._sse_events(await resp.text()) - - async def test_batch_tags_each_prompt_once_and_reports_counts(self): - a, b, c = self._seed() - - resp, events = await self._start() - - self.assertEqual(resp.status, 200) - done = events[-1] - self.assertEqual(done["type"], "complete") - self.assertEqual(done["processed"], 2) - self.assertEqual(done["tagged"], 2) - self.assertEqual(done["skipped"], 2) - self.assertEqual(done["errors"], 0) - self.assertEqual(len(self.service.generate_calls), 2) - self.assertEqual(self.api.db.get_prompt_by_id(a)["tags"], ["1girl", "smile"]) - self.assertEqual(self.api.db.get_prompt_by_id(b)["tags"], ["1girl", "smile"]) - self.assertEqual(self.api.db.get_prompt_by_id(c)["tags"], []) - progress = [e for e in events if e["type"] == "progress"] - self.assertTrue(all(0 <= e["progress"] <= 100 for e in progress)) - self.assertTrue(any("processed" in e for e in progress)) - - async def test_rows_outside_output_dirs_are_skipped_and_counted(self): - inside = self._save_prompt("inside") - outside = self._save_prompt("outside") - self._link_image(inside, self._make_image(self.output_dir, "in.png")) - secret = self._make_image(self.secret_dir, "password.png") - self._link_image(outside, secret) - - _, events = await self._start() - - done = events[-1] - self.assertEqual(done["type"], "complete") - self.assertEqual(done["processed"], 1) - self.assertEqual(done["tagged"], 1) - self.assertEqual(done["skipped"], 1) - self.assertEqual(done["skipped_outside_roots"], 1) - self.assertEqual(len(self.service.generate_calls), 1) - self.assertEqual( - os.path.realpath(self.service.generate_calls[0][0]), - os.path.realpath(str(self.output_dir / "in.png")), - ) - self.assertEqual(self.api.db.get_prompt_by_id(outside)["tags"], []) - progress = [e for e in events if e["type"] == "progress" and "skipped" in e] - self.assertTrue(progress) - self.assertTrue(all("skipped_outside_roots" in e for e in progress)) - body = json.dumps(events) - self.assertNotIn(str(self.secret_dir), body) - - async def test_second_start_while_running_is_409(self): - self._seed() - started = threading.Event() - release = threading.Event() - - def slow_generate(image_path, **kwargs): - started.set() - release.wait(5) - return ["slow"] - - self.service.generate_tags = slow_generate - first = asyncio.ensure_future(self._start()) - try: - await asyncio.get_event_loop().run_in_executor(None, started.wait, 5) - self.assertTrue(started.is_set()) - - resp = await self.client.request( - "POST", "/prompt_manager/autotag/start?model_type=gguf" - ) - - self.assertEqual(resp.status, 409) - self.assertEqual( - await resp.json(), - {"success": False, "error": "Auto-tagging is already running"}, - ) - finally: - release.set() - resp1, events = await first - self.assertEqual(resp1.status, 200) - self.assertEqual(events[-1]["type"], "complete") - - async def test_lock_is_released_after_a_run(self): - self._seed() - - resp1, _ = await self._start() - resp2, events = await self._start() - - self.assertEqual(resp1.status, 200) - self.assertEqual(resp2.status, 200) - self.assertEqual(events[-1]["type"], "complete") - - async def test_lock_is_released_when_the_run_fails(self): - self._seed() - self.service.load_error = RuntimeError("no model") - - _, events = await self._start() - self.assertEqual(events[-1]["type"], "error") - - self.service.load_error = None - resp, events = await self._start() - self.assertEqual(resp.status, 200) - self.assertEqual(events[-1]["type"], "complete") - - async def test_already_tagged_prompts_are_skipped_unless_asked(self): - p = self._save_prompt("tagged", tags=["cat"]) - self._link_image(p, self._make_image(self.output_dir, "t.png")) - - _, events = await self._start("model_type=gguf&skip_tagged=true") - self.assertEqual(events[-1]["skipped"], 1) - self.assertEqual(self.service.generate_calls, []) - - _, events = await self._start("model_type=gguf&skip_tagged=false") - self.assertEqual(events[-1]["tagged"], 1) - self.assertEqual( - self.api.db.get_prompt_by_id(p)["tags"], ["cat", "1girl", "smile"] - ) - - async def test_prompt_already_holding_the_tags_counts_as_skipped(self): - p = self._save_prompt("same", tags=["1girl", "smile"]) - self._link_image(p, self._make_image(self.output_dir, "s.png")) - - _, events = await self._start("model_type=gguf&skip_tagged=false") - - self.assertEqual(events[-1]["skipped"], 1) - self.assertEqual(events[-1]["tagged"], 0) - - async def test_generation_errors_are_counted_and_do_not_abort(self): - self._seed() - self.service.generate_error = RuntimeError("boom") - - _, events = await self._start() - - self.assertEqual(events[-1]["type"], "complete") - # a1, b1 and b2 each fail: a failed prompt is retried on its next image - self.assertEqual(events[-1]["errors"], 3) - self.assertEqual(events[-1]["tagged"], 0) - - async def test_json_body_parameters_are_accepted(self): - self._seed() - - _, events = await self._start( - "", {"model_type": "wd14-vit", "general_threshold": 0.5, "prompt": "p"} - ) - - self.assertEqual(events[-1]["type"], "complete") - self.assertEqual(self.service.load_calls, [("wd14-vit", True)]) - self.assertEqual(self.service.generate_calls[0][2], 0.5) - self.assertEqual(self.service.custom_prompt, "p") - - async def test_keep_in_memory_false_unloads_after_batch(self): - self._seed() - - _, events = await self._start("model_type=gguf&keep_in_memory=false") - - self.assertEqual(events[-1]["model_loaded"], False) - self.assertEqual(self.service.unload_calls, 1) - - async def test_model_not_downloaded_is_an_error_event(self): - self.service.downloaded = False - - _, events = await self._start() - - self.assertEqual(events[-1]["type"], "error") - self.assertEqual(self.service.load_calls, []) - - async def test_load_failure_is_an_error_event_without_details(self): - self.service.load_error = RuntimeError("/srv/models/secret.gguf is corrupt") - - _, events = await self._start() - - self.assertEqual(events[-1]["type"], "error") - self.assertNotIn("/srv/models", events[-1]["message"]) - - async def test_non_numeric_threshold_is_400(self): - resp = await self.client.request( - "POST", "/prompt_manager/autotag/start?general_threshold=abc" - ) - self.assertEqual(resp.status, 400) - - async def test_invalid_json_body_is_400(self): - resp = await self.client.request( - "POST", - "/prompt_manager/autotag/start", - data="{not json", - headers={"Content-Type": "application/json"}, - ) - self.assertEqual(resp.status, 400) - - async def test_client_disconnect_stops_the_batch_and_unloads(self): - self._seed() - real_write = web.StreamResponse.write - state = {"writes": 0} - - async def flaky_write(response, data): - state["writes"] += 1 - if b"Processing " in data: - raise ConnectionResetError("Cannot write to closing transport") - return await real_write(response, data) - - with patch.object(web.StreamResponse, "write", flaky_write): - try: - resp = await self.client.request( - "POST", "/prompt_manager/autotag/start?keep_in_memory=false" - ) - await resp.read() - except Exception: - pass - - self.assertEqual(len(self.service.generate_calls), 1) - self.assertEqual(self.service.unload_calls, 1) - - class TestAutotagModelsAndDownloadErrors(AutotagAPITestCase): async def test_models_error_is_500(self): @@ -834,7 +632,15 @@ class TestAutotagSingleEdgeCases(AutotagAPITestCase): self.assertEqual(resp.status, 200) self.assertEqual(self.service.load_calls, [("wd14-vit", True)]) - self.assertEqual(self.service.custom_prompt, "terse") + # #379: the request's prompt is for this call only. + self.assertEqual(self.service.custom_prompt, "") + resp = await self._post_json( + "/prompt_manager/autotag/single", + {"path": str(image), "model_type": "gguf", "prompt": "terse"}, + ) + self.assertEqual(resp.status, 200) + self.assertEqual(self.service.generate_calls[-1][1], "terse") + self.assertEqual(self.service.custom_prompt, "") async def test_image_id_without_stored_path_is_404(self): prompt_id = self._save_prompt() @@ -850,64 +656,6 @@ class TestAutotagSingleEdgeCases(AutotagAPITestCase): self.assertEqual(resp.status, 404) -class TestApplyAutotag(AutotagAPITestCase): - - async def test_missing_prompt_id_is_400(self): - resp = await self._post_json("/prompt_manager/autotag/apply", {"tags": ["a"]}) - self.assertEqual(resp.status, 400) - - async def test_no_tags_is_a_noop(self): - resp = await self._post_json("/prompt_manager/autotag/apply", {"prompt_id": 1}) - data = await resp.json() - self.assertTrue(data["success"]) - self.assertEqual(data["message"], "No tags to apply") - - async def test_unknown_prompt_is_404(self): - resp = await self._post_json( - "/prompt_manager/autotag/apply", {"prompt_id": 999, "tags": ["a"]} - ) - self.assertEqual(resp.status, 404) - - async def test_tags_are_merged_without_duplicates(self): - prompt_id = self._save_prompt(tags=["cat"]) - - resp = await self._post_json( - "/prompt_manager/autotag/apply", - {"prompt_id": prompt_id, "tags": ["cat", "dog"]}, - ) - - data = await resp.json() - self.assertEqual(data["added_tags"], ["dog"]) - self.assertEqual(data["total_tags"], 2) - self.assertEqual( - self.api.db.get_prompt_by_id(prompt_id)["tags"], ["cat", "dog"] - ) - - async def test_string_tags_from_db_are_split(self): - prompt_id = self._save_prompt() - with patch.object( - self.api.db, - "get_prompt_by_id", - return_value={"id": prompt_id, "tags": "cat, dog"}, - ): - resp = await self._post_json( - "/prompt_manager/autotag/apply", - {"prompt_id": prompt_id, "tags": ["dog", "owl"]}, - ) - data = await resp.json() - self.assertEqual(data["added_tags"], ["owl"]) - self.assertEqual(data["total_tags"], 3) - - async def test_db_error_is_500(self): - with patch.object( - self.api.db, "get_prompt_by_id", side_effect=RuntimeError("db") - ): - resp = await self._post_json( - "/prompt_manager/autotag/apply", {"prompt_id": 1, "tags": ["a"]} - ) - self.assertEqual(resp.status, 500) - - class TestAutotagErrorBodiesHideServerPaths(AutotagAPITestCase): """Every autotag handler reports exceptions through the shared public helper.""" @@ -932,9 +680,12 @@ class TestAutotagErrorBodiesHideServerPaths(AutotagAPITestCase): await self._assert_hidden(resp) async def test_apply_error(self): - with patch.object(self.api.db, "get_prompt_by_id", side_effect=self._locked()): + image = self._make_image(self.output_dir) + digest = hashlib.sha256(image.read_bytes()).hexdigest() + with patch.object(self.api.db, "add_image_tags", side_effect=self._locked()): resp = await self._post_json( - "/prompt_manager/autotag/apply", {"prompt_id": 1, "tags": ["a"]} + "/prompt_manager/autotag/apply", + {"path": str(image), "sha256": digest, "tags": ["a"]}, ) await self._assert_hidden(resp) @@ -1008,10 +759,13 @@ class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase): "POST", "/prompt_manager/autotag/download/gguf" ) ) - await self._busy( - await self.client.request( - "POST", "/prompt_manager/autotag/start?model_type=gguf" - ) + # The legacy batch is read as a stream: one error frame (#379). + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start?model_type=gguf" + ) + self.assertEqual( + self._sse_events(await resp.text()), + [{"type": "error", "status": "error", "message": MODEL_BUSY_ERROR}], ) await self._busy( await self.client.request("POST", "/prompt_manager/autotag/unload") @@ -1026,7 +780,7 @@ class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase): 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. + # A refused start left no job behind: the next start runs. resp = await self.client.request( "POST", "/prompt_manager/autotag/start?model_type=gguf" ) @@ -1238,7 +992,7 @@ if __name__ == "__main__": class TestGalleryReviewContract(AutotagAPITestCase): - async def test_root_relative_gallery_targets_generate_and_apply_to_their_prompt( + async def test_root_relative_gallery_targets_generate_and_apply_to_their_image( self, ): second = self.tmp_root / "second-output" @@ -1255,16 +1009,22 @@ class TestGalleryReviewContract(AutotagAPITestCase): self.assertEqual(response.status, 200) data = await response.json() self.assertEqual(data["prompt_id"], pid) - self.assertEqual( - self.service.generate_calls[-1][0], str((root / "same.png").resolve()) - ) response = await self._post_json( "/prompt_manager/autotag/apply", - {"prompt_id": data["prompt_id"], "tags": data["tags"]}, + { + "path": "same.png", + "root_index": index, + "sha256": data["sha256"], + "tags": data["tags"], + }, ) self.assertEqual(response.status, 200) + # The image's own tags; the prompt is never written (#379). + self.assertEqual(self.api.db.get_prompt_by_id(pid)["tags"], []) + locator = os.path.normcase(os.path.realpath(str(root / "same.png"))) self.assertEqual( - set(self.api.db.get_prompt_by_id(pid)["tags"]), set(data["tags"]) + self.api.db.read_image_tags(locator, data["sha256"])["tags"], + data["tags"], ) async def test_gallery_unlinked_image_returns_no_prompt_id(self): @@ -2058,108 +1818,6 @@ class TestTagOutputImagesJob(BatchJobTestCase): self.assertTrue(inside) -# The legacy batch's per-image frames over BatchJobTestCase's library, newest -# first, with every throttle check passing (pinned from main's generator). -GOLDEN_LEGACY_ROWS = [ - { - "type": "progress", - "progress": 24, - "status": "Processing 1/6...", - "processed": 1, - "tagged": 1, - "skipped": 0, - "skipped_outside_roots": 0, - }, - { - "type": "progress", - "progress": 38, - "status": "Skipping 2/6 (outside allowed directories)...", - "processed": 1, - "tagged": 1, - "skipped": 1, - "skipped_outside_roots": 1, - }, - { - "type": "progress", - "progress": 52, - "status": "Skipping 3/6 (file missing)...", - "processed": 1, - "tagged": 1, - "skipped": 2, - "skipped_outside_roots": 1, - }, - { - "type": "progress", - "progress": 66, - "status": "Processing 4/6...", - "processed": 2, - "tagged": 2, - "skipped": 2, - "skipped_outside_roots": 1, - }, - { - "type": "progress", - "progress": 80, - "status": "Skipping 5/6 (prompt already processed)...", - "processed": 2, - "tagged": 2, - "skipped": 3, - "skipped_outside_roots": 1, - }, - { - "type": "progress", - "progress": 95, - "status": "Skipping 6/6 (already tagged)...", - "processed": 2, - "tagged": 2, - "skipped": 4, - "skipped_outside_roots": 1, - }, -] - - -class TestLegacyBatchFrames(BatchJobTestCase): - - async def test_legacy_batch_events_unchanged(self): - """Golden frames of the legacy batch over the seeded library.""" - self._seed() - ticks = iter(range(1000)) # every throttle check sees a second pass - with patch("py.api.autotag_routes._monotonic", side_effect=lambda: next(ticks)): - resp = await self.client.post( - "/prompt_manager/autotag/start?model_type=wd14-vit&skip_tagged=true" - ) - frames = self._sse_events(await resp.text()) # the stream runs here - self.assertEqual( - frames, - [ - {"type": "progress", "progress": 0, "status": "Loading model..."}, - { - "type": "progress", - "progress": 5, - "status": "Model loaded. Fetching all images from database...", - }, - { - "type": "progress", - "progress": 10, - "status": "Found 6 images. Processing...", - }, - *GOLDEN_LEGACY_ROWS, - { - "type": "complete", - "progress": 100, - "processed": 2, - "tagged": 2, - "skipped": 4, - "errors": 0, - "skipped_outside_roots": 1, - "status": "Complete", - "model_status": "Model kept in memory", - "model_loaded": True, - }, - ], - ) - - class TestTagOutputImagesEdges(BatchJobTestCase): async def test_cancel_during_backup_loads_no_model(self): diff --git a/tests/test_autotag_legacy.py b/tests/test_autotag_legacy.py new file mode 100644 index 0000000..126ef67 --- /dev/null +++ b/tests/test_autotag_legacy.py @@ -0,0 +1,617 @@ +"""The legacy pages and the Tag image dialog tag images, never prompts (#379, PR 2). + +- ``POST /autotag/start`` (legacy admin and gallery batch) starts the + ``autotag`` job titled "Tag output images (legacy page)" and relays it as + the legacy frames; closing the stream cancels it. +- ``POST /autotag/single`` suggests on a capture with the saved cap and + answers the captured ``sha256``; ``existing: "skip"`` answers + ``already_tagged`` for an auto-tagged image without the model gate. +- ``POST /autotag/apply`` unions into that image's tags and marks it; a + ``prompt_id`` body is refused. +- A Tag image dialog save after Generate suggestions marks the image. + +No real model runs: the per-image fake from test_autotag_batch answers. +""" + +import asyncio +import hashlib +import json +import sqlite3 +import threading +from unittest.mock import patch + +from aiohttp import web + +try: + from test_autotag_batch import BatchTestCase, _absolute, _strings +except ImportError: + from tests.test_autotag_batch import BatchTestCase, _absolute, _strings + +from database.execution_history_schema import rotate_history_epoch +from py import model_gate +from py.api import autotag_routes, image_tag_jobs +from py.config import AutotagConfig + +START = "/prompt_manager/autotag/start" +SINGLE = "/prompt_manager/autotag/single" +APPLY = "/prompt_manager/autotag/apply" +SUGGEST = "/prompt_manager/image-annotations/suggest" +OBSERVE = "/prompt_manager/image-annotations/observe" +TAGS = "/prompt_manager/image-annotations/tags" +WAIT = 30.0 +LEGACY_TITLE = "Tag output images (legacy page)" + + +def frames_of(text): + return [ + json.loads(line[len("data: ") :]) + for line in text.splitlines() + if line.startswith("data: ") + ] + + +class LegacyTestCase(BatchTestCase): + async def legacy(self, query="model_type=wd14-vit&skip_tagged=true"): + response = await self.client.post(f"{START}?{query}") + return response, frames_of(await response.text()) + + def digest(self, path): + return hashlib.sha256((self.output_dir / path).read_bytes()).hexdigest() + + async def single(self, body): + response = await self.client.post(SINGLE, json=body) + return response.status, await response.json() + + async def apply(self, body): + response = await self.client.post(APPLY, json=body) + return response.status, await response.json() + + def snapshot_prompts(self): + conn = sqlite3.connect(self._temp_db.name) + try: + return [ + conn.execute(f"SELECT * FROM {t} ORDER BY 1,2").fetchall() + for t in ("prompts", "prompt_tags") + ] + finally: + conn.close() + + +class TestLegacyBatch(LegacyTestCase): + async def test_the_legacy_batch_tags_each_image_and_relays_frames(self): + prompt = self.prompt("one prompt", tags=["beach"]) + a = self.image("a.png", (1, 2, 3), sees=["bird"]) + b = self.image("b.png", (4, 5, 6), sees=["waves"]) + self.link(prompt, a) + self.link(prompt, b) + before = self.snapshot_prompts() + ticks = iter(range(1000)) + with patch.object( + autotag_routes, "_monotonic", side_effect=lambda: next(ticks) + ): + response, frames = await self.legacy() + self.assertEqual(response.status, 200) + self.assertEqual(self.stored("a.png")["tags"], ["bird"]) + self.assertEqual(self.stored("b.png")["tags"], ["waves"]) + self.assertEqual(self.snapshot_prompts(), before) + phases = [f["status"] for f in frames if f.get("progress") in (0, 5, 8)] + self.assertEqual( + phases, + sorted( + phases, + key=[ + "Taking a safety backup…", + "Loading model…", + "Listing output images…", + ].index, + ), + ) + self.assertIn("Listing output images…", phases) + rows = [f for f in frames if f["type"] == "progress" and "processed" in f] + self.assertTrue(rows) + for row in rows: + self.assertEqual( + set(row), + {"type", "progress", "status", "processed", "tagged", "skipped"}, + ) + self.assertTrue(10 <= row["progress"] <= 95, row) + self.assertEqual( + frames[-1], + { + "type": "complete", + "progress": 100, + "processed": 2, + "tagged": 2, + "skipped": 0, + "errors": 0, + "skipped_outside_roots": 0, + "status": "Complete", + "model_status": "Model kept in memory", + "model_loaded": True, + }, + ) + job = self.api._jobs.list()[0] + self.assertEqual( + (job.title, job.kind, job.state), (LEGACY_TITLE, "autotag", "finished") + ) + self.assertEqual([s for s in _strings(frames) if _absolute(s)], []) + + # Skip: the marked images are skipped before the model runs. + self.tagger.calls.clear() + _, again = await self.legacy() + self.assertEqual(again[-1]["skipped"], 2) + self.assertEqual(again[-1]["processed"], 0) + self.assertEqual(self.tagger.calls, []) + + async def test_no_still_images_and_keep_in_memory_false(self): + _, frames = await self.legacy("model_type=wd14-vit&keep_in_memory=false") + self.assertEqual(frames[-1]["status"], "No still images found") + self.assertEqual(frames[-1]["model_loaded"], False) + self.assertEqual(frames[-1]["model_status"], "Model unloaded") + self.assertEqual(self.tagger.unload_calls, 1) + + async def test_language_model_options_come_from_the_page_or_the_template(self): + self.image("a.png", (1, 1, 1)) + await self.legacy("model_type=gguf&prompt=Tags%20only") + self.assertEqual(self.tagger.calls[-1]["prompt"], "Tags only") + self.assertEqual(self.tagger.calls[-1]["limit"], 10) + self.image("b.png", (2, 2, 2)) + await self.relist() + await self.legacy("model_type=gguf&prompt=%20") + self.assertEqual(self.tagger.calls[-1]["prompt"], AutotagConfig.TEMPLATE) + self.assertEqual(self.tagger.custom_prompt, "legacy prompt") + + async def test_each_refusal_is_one_error_frame(self): + self.tagger.downloaded["wd14-vit"] = False + _, frames = await self.legacy() + self.assertEqual(len(frames), 1) + self.assertEqual(frames[0]["type"], "error") + self.assertEqual(frames[0]["status"], "error") + self.assertIn("WD14 ViT is not downloaded", frames[0]["message"]) + self.tagger.downloaded["wd14-vit"] = True + + _, frames = await self.legacy("model_type=wd14-vit&general_threshold=2") + self.assertEqual( + frames, + [ + { + "type": "error", + "status": "error", + "message": "general_threshold must be between 0 and 1", + } + ], + ) + + owner = model_gate.new_owner("image-annotations") + self.assertTrue(model_gate.try_acquire(owner)) + try: + _, frames = await self.legacy() + finally: + model_gate.release(owner) + self.assertEqual(frames[0]["message"], autotag_routes.MODEL_BUSY_ERROR) + self.assertEqual(self.api._jobs.list(), []) + self.assertIsNone(model_gate.holder()) + + response = await self.client.post(f"{START}?general_threshold=abc") + self.assertEqual(response.status, 400) + + async def test_busy_jobs_refuse_the_legacy_batch(self): + self.image("a.png", (1, 1, 1)) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + status, data = await self.start_batch() + self.assertEqual(status, 202, data) + job = self.api._jobs.get(data["job"]["id"]) + try: + await self.wait_for(entered) + _, frames = await self.legacy() + self.assertEqual( + frames, + [ + { + "type": "error", + "status": "error", + "message": "Tag output images already running", + } + ], + ) + finally: + self.tagger.block = None + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + + # A second legacy batch keeps its JSON 409. + self.image("b.png", (2, 2, 2)) + await self.relist() + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + first = asyncio.ensure_future( + self.legacy("model_type=wd14-vit&skip_tagged=false") + ) + try: + await self.wait_for(entered) + second = await self.client.post(f"{START}?model_type=wd14-vit") + self.assertEqual(second.status, 409) + self.assertEqual( + await second.json(), + {"success": False, "error": "Auto-tagging is already running"}, + ) + finally: + self.tagger.block = None + proceed.set() + _, frames = await first + self.assertEqual(frames[-1]["type"], "complete") + + async def test_a_legacy_start_that_loses_the_race_gets_the_json_409(self): + # Two starts can both pass the title check before either registers its + # job; the loser's busy refusal names the legacy batch and answers the + # same JSON 409 as a second legacy batch. + refusal = image_tag_jobs.Refusal( + 409, + "busy", + f"{autotag_routes.LEGACY_BATCH_TITLE} already running", + web.json_response({}, status=409), + ) + + async def refused(**_kwargs): + return None, refusal + + with patch.object(self.api, "_start_autotag", side_effect=refused): + response = await self.client.post(f"{START}?model_type=wd14-vit") + self.assertEqual(response.status, 409) + self.assertEqual( + await response.json(), + {"success": False, "error": "Auto-tagging is already running"}, + ) + + async def test_closing_the_stream_cancels_the_job(self): + for n in range(3): + self.image(f"{n}.png", (n, n, n)) + real_write = web.StreamResponse.write + + async def flaky_write(response, data): + if b'"processed"' in data: + raise ConnectionResetError("Cannot write to closing transport") + return await real_write(response, data) + + with patch.object(web.StreamResponse, "write", flaky_write): + try: + response = await self.client.post(f"{START}?model_type=wd14-vit") + await response.read() + except Exception: # nosec B110 - the client sees the dropped stream + pass + job = self.api._jobs.list()[0] + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(job.state, "cancelled") + self.assertLess(job.counters["tagged"], 3) + + async def test_a_quiet_job_repeats_its_last_frame(self): + self.image("a.png", (1, 1, 1)) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + with patch.object(autotag_routes, "LEGACY_REPEAT_AFTER", 0.05): + run = asyncio.ensure_future(self.legacy()) + await self.wait_for(entered) + await asyncio.sleep(0.3) + self.tagger.block = None + proceed.set() + _, frames = await run + statuses = [f.get("status") for f in frames] + repeated = [s for s in statuses if s and statuses.count(s) > 1] + self.assertTrue(repeated, statuses) + self.assertEqual(frames[-1]["type"], "complete") + + async def test_a_cancel_from_activity_ends_with_a_cancelled_frame(self): + for n in range(2): + self.image(f"{n}.png", (n, n, n)) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + run = asyncio.ensure_future(self.legacy()) + try: + await self.wait_for(entered) + job = self.api._jobs.active("autotag") + await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel") + finally: + self.tagger.block = None + proceed.set() + _, frames = await run + self.assertEqual( + frames[-1], + { + "type": "cancelled", + "message": "Cancelled. Tags already added were kept.", + }, + ) + + +class TestLegacySingleAndApply(LegacyTestCase): + async def test_single_suggests_on_a_capture_and_apply_tags_only_that_image(self): + prompt = self.prompt("shared prompt", tags=["keep"]) + a = self.image("a.png", (1, 2, 3), sees=["bird", "sky"]) + b = self.image("b.png", (4, 5, 6), sees=["cat"]) + image_a = self.link(prompt, a) + self.link(prompt, b) + self.saved_set("b.png", b.read_bytes(), ["hand"]) + before = self.snapshot_prompts() + status, data = await self.single( + {"image_id": image_a, "model_type": "wd14-vit"} + ) + self.assertEqual(status, 200, data) + self.assertEqual(data["tags"], ["bird", "sky"]) + self.assertEqual(data["sha256"], self.digest("a.png")) + self.assertEqual(data["image_tags"], []) + self.assertEqual(data["prompt_id"], prompt) + self.assertEqual(self.tagger.custom_prompt, "legacy prompt") + # Inference read the private capture, not the file. + self.assertNotIn(str(self.output_dir), self.tagger.read_paths[-1]) + status, applied = await self.apply( + {"image_id": image_a, "sha256": data["sha256"], "tags": ["bird"]} + ) + self.assertEqual( + (status, applied), + (200, {"success": True, "added_tags": ["bird"], "total_tags": 1}), + ) + self.assertEqual(self.stored("a.png")["tags"], ["bird"]) + self.assertEqual(self.stored("b.png")["tags"], ["hand"]) + self.assertEqual(self.snapshot_prompts(), before) + self.assertEqual(self.marked(), {self.locator("a.png")}) + + async def test_skip_answers_already_tagged_without_the_gate(self): + a = self.image("a.png", (1, 2, 3), sees=["bird"]) + await self.run_batch() + self.tagger.calls.clear() + owner = model_gate.new_owner("image-tag-job") + self.assertTrue(model_gate.try_acquire(owner)) + try: + status, data = await self.single( + { + "path": "a.png", + "root_index": 0, + "existing": "skip", + "model_type": "wd14-vit", + } + ) + finally: + model_gate.release(owner) + self.assertEqual(status, 200, data) + self.assertEqual(data["already_tagged"], True) + self.assertEqual(data["image_tags"], ["bird"]) + self.assertEqual(self.tagger.calls, []) + # Hand tags alone do not count; merge always suggests. + b = self.image("b.png", (9, 9, 9)) + self.saved_set("b.png", b.read_bytes(), ["hand"]) + status, data = await self.single( + {"path": str(b), "existing": "skip", "model_type": "wd14-vit"} + ) + self.assertNotIn("already_tagged", data) + self.assertEqual(data["image_tags"], ["hand"]) + status, data = await self.single( + {"path": str(a), "existing": "merge", "model_type": "wd14-vit"} + ) + self.assertNotIn("already_tagged", data) + + async def test_single_uses_the_saved_cap_and_the_refusal_table(self): + self.image("a.png", (1, 2, 3)) + AutotagConfig.SUGGESTION_LIMIT = 3 + status, data = await self.single( + {"path": "a.png", "root_index": 0, "model_type": "wd14-vit"} + ) + self.assertEqual((status, len(data["tags"])), (200, 3)) + (self.output_dir / "broken.png").write_bytes(b"not an image") + await self.relist() + status, data = await self.single( + {"path": "broken.png", "root_index": 0, "model_type": "wd14-vit"} + ) + self.assertEqual((status, data["code"]), (415, "not-a-still")) + self.assertIn(self.locator("broken.png"), self.marked()) + self.tagger.downloaded["gguf"] = False + status, data = await self.single( + {"path": "a.png", "root_index": 0, "model_type": "gguf"} + ) + self.assertEqual((status, data["code"]), (409, "model-unavailable")) + self.api._capture_store().max_file = 16 + status, data = await self.single( + {"path": "a.png", "root_index": 0, "model_type": "wd14-vit"} + ) + self.assertEqual((status, data["code"]), (413, "too-large")) + self.api._capture_store().max_file = 2**20 + owner = model_gate.new_owner("image-annotations") + self.assertTrue(model_gate.try_acquire(owner)) + try: + status, data = await self.single( + {"path": "a.png", "root_index": 0, "model_type": "wd14-vit"} + ) + finally: + model_gate.release(owner) + self.assertEqual((status, data["code"]), (409, "model-busy")) + + async def test_apply_body_and_answer_table(self): + a = self.image("a.png", (1, 2, 3)) + digest = self.digest("a.png") + status, data = await self.apply({"prompt_id": 1, "tags": ["a"]}) + self.assertEqual( + (status, data["error"]), (400, "This page is out of date. Reload it.") + ) + status, _ = await self.apply( + {"path": "a.png", "root_index": 0, "sha256": "x", "tags": ["a"]} + ) + self.assertEqual(status, 400) + status, _ = await self.apply( + {"path": "a.png", "root_index": 0, "sha256": digest, "tags": ["bad,name"]} + ) + self.assertEqual(status, 400) + status, _ = await self.apply({"path": "a.png", "sha256": digest, "tags": []}) + self.assertEqual(status, 403) # a bare relative path is not inside a root + # An empty reviewed list only marks the image. + status, data = await self.apply({"path": str(a), "sha256": digest, "tags": []}) + self.assertEqual((status, data["added_tags"]), (200, [])) + self.assertEqual(self.marked(), {self.locator("a.png")}) + # The file changed since its tags were generated. + status, data = await self.apply( + {"path": str(a), "sha256": "0" * 64, "tags": ["x"]} + ) + self.assertEqual((status, data["code"]), (409, "file-changed")) + # Too many tags. + full = self.image("full.png", (7, 7, 7)) + self.saved_set("full.png", full.read_bytes(), [f"t{i}" for i in range(100)]) + status, data = await self.apply( + {"path": str(full), "sha256": self.digest("full.png"), "tags": ["one_more"]} + ) + self.assertEqual((status, data["code"]), (409, "tag-limit")) + # Gone. + status, data = await self.apply( + {"path": "gone.png", "root_index": 0, "sha256": digest, "tags": ["x"]} + ) + self.assertIn(status, (403, 404)) + + async def test_apply_after_a_restore_writes_nothing(self): + a = self.image("a.png", (1, 2, 3)) + digest = self.digest("a.png") + real = self.api.db.get_execution_epoch + + def stale_epoch(): + epoch = real() + conn = sqlite3.connect(self._temp_db.name) + try: + rotate_history_epoch(conn) + finally: + conn.close() + return epoch + + with patch.object(self.api.db, "get_execution_epoch", side_effect=stale_epoch): + status, data = await self.apply( + {"path": str(a), "sha256": digest, "tags": ["x"]} + ) + self.assertEqual((status, data["code"]), (409, "storage-changed")) + self.assertEqual(self.stored("a.png")["tags"], []) + + async def test_image_id_and_path_bodies_reach_the_same_bytes(self): + prompt = self.prompt("p") + a = self.image("sub/a.png", (1, 2, 3), sees=["x"]) + image_id = self.link(prompt, a) + _, by_id = await self.single({"image_id": image_id, "model_type": "wd14-vit"}) + _, by_path = await self.single( + {"path": "sub/a.png", "root_index": 0, "model_type": "wd14-vit"} + ) + self.assertEqual(by_id["sha256"], by_path["sha256"]) + status, _ = await self.apply( + {"image_id": image_id, "sha256": by_id["sha256"], "tags": ["x"]} + ) + self.assertEqual(status, 200) + status, data = await self.apply( + { + "path": "sub/a.png", + "root_index": 0, + "sha256": by_path["sha256"], + "tags": ["y"], + } + ) + self.assertEqual((status, data["total_tags"]), (200, 2)) + + +class TestDialogMarks(LegacyTestCase): + """DEC4: a Tag image dialog save after Generate suggestions marks the image.""" + + def target(self, path): + return { + "root_index": 0, + "root_scope": self.api._gallery_root_scope(self.output_dir, 0), + "path": path, + } + + async def observe(self, path): + response = await self.client.post(OBSERVE, json={"target": self.target(path)}) + data = await response.json() + self.assertEqual(response.status, 200, data) + return data + + async def suggest(self, observed, path): + response = await self.client.post( + SUGGEST, + json={ + "target": self.target(path), + "capture_token": observed["capture"]["token"], + "expected_epoch": observed["library_epoch"], + "model_type": "wd14-vit", + "options": {"general_threshold": 0.35, "character_threshold": 0.85}, + }, + ) + self.assertEqual(response.status, 200, await response.text()) + + async def save(self, observed, path, tags, revision=None): + response = await self.client.put( + TAGS, + json={ + "target": self.target(path), + "capture_token": observed["capture"]["token"], + "sha256": observed["capture"]["sha256"], + "expected_epoch": observed["library_epoch"], + "expected_tag_revision": ( + observed["tag_revision"] if revision is None else revision + ), + "tags": tags, + }, + ) + return response.status, await response.json() + + async def test_save_after_generate_marks_and_a_skip_batch_leaves_it(self): + self.image("d.png", (1, 2, 3), sees=["bird", "sky"]) + observed = await self.observe("d.png") + await self.suggest(observed, "d.png") + status, _ = await self.save(observed, "d.png", ["sky"]) # bird removed + self.assertEqual(status, 200) + self.assertEqual(self.marked(), {self.locator("d.png")}) + job = await self.run_batch(existing="skip") + self.assertEqual(self.rows(job, "skipped"), [("d.png", "Already auto-tagged")]) + self.assertEqual(self.stored("d.png")["tags"], ["sky"]) + + async def test_save_without_generate_does_not_mark(self): + self.image("e.png", (1, 2, 3)) + observed = await self.observe("e.png") + status, _ = await self.save(observed, "e.png", ["hand"]) + self.assertEqual(status, 200) + self.assertEqual(self.marked(), set()) + + async def test_a_conflict_then_a_saved_retry_marks(self): + f = self.image("f.png", (1, 2, 3)) + observed = await self.observe("f.png") + await self.suggest(observed, "f.png") + self.saved_set("f.png", f.read_bytes(), ["other"]) # another writer + status, data = await self.save(observed, "f.png", ["mine"]) + self.assertEqual((status, data["code"]), (409, "tag-conflict")) + self.assertEqual(self.marked(), set()) + status, _ = await self.save( + observed, "f.png", ["mine"], revision=data["tag_revision"] + ) + self.assertEqual(status, 200) + self.assertEqual(self.marked(), {self.locator("f.png")}) + + async def test_another_capture_does_not_inherit_the_mark(self): + self.image("g.png", (1, 2, 3)) + self.image("h.png", (4, 5, 6)) + first = await self.observe("g.png") + await self.suggest(first, "g.png") + second = await self.observe("h.png") + status, _ = await self.save(second, "h.png", ["x"]) + self.assertEqual(status, 200) + self.assertEqual(self.marked(), set()) + self.assertIn(first["capture"]["token"], self.api._suggested_tokens) + + async def test_gone_tokens_are_dropped_on_the_next_insert_without_touching_clocks( + self, + ): + self.image("i.png", (1, 2, 3)) + self.image("j.png", (4, 5, 6)) + first = await self.observe("i.png") + await self.suggest(first, "i.png") + store = self.api._capture_store() + store.release(first["capture"]["token"]) + second = await self.observe("j.png") + seen = [] + real_get = store.get + with patch.object( + store, "get", side_effect=lambda t: seen.append(t) or real_get(t) + ): + await self.suggest(second, "j.png") + self.assertEqual(seen, []) + self.assertEqual(self.api._suggested_tokens, {second["capture"]["token"]}) diff --git a/tests/test_database.py b/tests/test_database.py index 09fe6e9..92935ce 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1398,50 +1398,6 @@ class TestModelHousekeeping(DatabaseTestCase): self.assertIsNot(self.db.model.get_connection(), failing) -class TestAutotagBatchCandidates(DatabaseTestCase): - """#328: the batch's candidate read never loads workflow documents.""" - - def test_autotag_candidates_skip_workflow_data(self): - from utils.hashing import generate_prompt_hash - - first = self.db.save_prompt( - text="a", tags=["one", "two"], prompt_hash=generate_prompt_hash("a") - ) - second = self.db.save_prompt(text="b", prompt_hash=generate_prompt_hash("b")) - old = self.db.link_image_to_prompt( - first, "/out/a.png", {"workflow": {"big": "x" * 1000}} - ) - new = self.db.link_image_to_prompt(second, "/out/b.png") - statements = [] - conn = self.db.model.get_connection() # this thread's own connection - conn.set_trace_callback(statements.append) - try: - rows = self.db.autotag_batch_candidates() - finally: - conn.set_trace_callback(None) - self.assertEqual( - rows[0], - { - "id": new, - "image_path": "/out/b.png", - "prompt_id": second, - "prompt_tags": [], - }, - ) - self.assertEqual( - {**rows[1], "prompt_tags": sorted(rows[1]["prompt_tags"])}, - { - "id": old, - "image_path": "/out/a.png", - "prompt_id": first, - "prompt_tags": ["one", "two"], - }, - ) - self.assertEqual(len(rows), 2) - self.assertTrue(statements) - self.assertFalse(any("workflow_data" in sql for sql in statements)) - - if __name__ == "__main__": unittest.main() diff --git a/tests/test_job_admission.py b/tests/test_job_admission.py index 44d47ff..4e8e558 100644 --- a/tests/test_job_admission.py +++ b/tests/test_job_admission.py @@ -535,7 +535,9 @@ class TestHeldJobStates(_LegacyHoldCase): self.assertIn("event: error", text) self.assertEqual(await self._listed(), [("thumbnails", "failed")]) - async def test_legacy_autotag_that_ends_in_error_is_failed(self): + async def test_legacy_autotag_refused_at_start_leaves_no_job(self): + # #379: the legacy batch is the autotag job; a model that is not + # downloaded refuses the start (one error frame) and registers nothing. service = MagicMock() service.get_models_status.return_value = {} # model not downloaded with patch("py.autotag.get_autotag_service", return_value=service): @@ -545,7 +547,7 @@ class TestHeldJobStates(_LegacyHoldCase): text = await resp.text() self.assertEqual(resp.status, 200) self.assertIn('"type": "error"', text) - self.assertEqual(await self._listed(), [("autotag", "failed")]) + self.assertEqual(await self._listed(), []) if __name__ == "__main__": diff --git a/web/admin.html b/web/admin.html index dee1d8a..f80b356 100644 --- a/web/admin.html +++ b/web/admin.html @@ -1187,14 +1187,13 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr <div class="space-y-1.5"> <label class="flex items-center p-2.5 bg-pm-input rounded-pm-md cursor-pointer hover:bg-pm-hover border border-pm"> <input type="radio" name="autoTagMode" value="skip" checked class="w-3.5 h-3.5 accent-[var(--pm-accent)]"> - <span class="ml-2.5 text-pm text-xs">Skip images with existing tags</span> + <span class="ml-2.5 text-pm text-xs">Skip images already auto-tagged</span> </label> <label class="flex items-center p-2.5 bg-pm-input rounded-pm-md cursor-pointer hover:bg-pm-hover border border-pm"> <input type="radio" name="autoTagMode" value="retag" class="w-3.5 h-3.5 accent-[var(--pm-accent)]"> <span class="ml-2.5 text-pm text-xs">Re-tag all images</span> </label> </div> - <p class="text-pm-muted text-[11px] mt-1">Note: "auto-scanned" is not counted as a real tag</p> </div> <!-- Action Buttons --> @@ -1293,6 +1292,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr <div class="w-full bg-pm-input rounded-full h-1.5"> <div id="autoTagProgressBar" class="bg-pm-accent h-1.5 rounded-full transition-all duration-300" style="width: 0%"></div> </div> + <p id="autoTagCloseNote" class="text-pm-muted text-[11px] mt-1.5">Closing this window stops tagging. Tags already added are kept. Prompt Manager's Activity page shows the time left.</p> </div> <div class="grid grid-cols-3 gap-3 mb-4"> @@ -1332,7 +1332,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr <!-- Re-tag Confirmation Modal --> <div id="retagConfirmModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50"> <div class="bg-pm-surface rounded-pm-lg p-4 max-w-lg w-full mx-4 border border-pm shadow-pm"> - <h3 class="text-sm font-semibold text-pm mb-3">Image Already Tagged</h3> + <h3 class="text-sm font-semibold text-pm mb-3">Image Already Auto-tagged</h3> <div class="mb-3"> <img id="retagPreviewImage" src="" alt="Preview" class="w-full h-40 object-contain bg-pm-primary rounded-pm-md"> @@ -1345,7 +1345,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr </div> </div> - <p class="text-pm-secondary text-xs mb-4">This image already has tags. Would you like to generate new tags?</p> + <p class="text-pm-secondary text-xs mb-4">A tagging model already tagged this image. Generate new tags?</p> <div class="flex gap-2 justify-end"> <button id="retagSkipBtn" diff --git a/web/gallery.html b/web/gallery.html index 00b890a..5613555 100644 --- a/web/gallery.html +++ b/web/gallery.html @@ -682,6 +682,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr <div class="w-full bg-pm-input rounded-full h-2"> <div id="autoTagProgressBar" class="bg-pm-accent h-2 rounded-full transition-all duration-300" style="width: 0%"></div> </div> + <p id="autoTagCloseNote" class="text-pm-muted text-xs mt-1.5">Closing this window stops tagging. Tags already added are kept. Prompt Manager's Activity page shows the time left.</p> </div> <!-- Stats --> diff --git a/web/js/admin.js b/web/js/admin.js index 200e462..e48c098 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -3801,7 +3801,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; } else if (data.type === 'complete') { this.autoTagState.eventSource.close(); this.hideModal("autoTagProgressModal"); - this.showNotification(`Auto tagging complete! Applied tags to ${data.tagged || 0} prompts.`, 'success'); + this.showNotification(ImageHelpers.autotagCompleteMessage(data.tagged), 'success'); this.refreshList(); // Refresh the prompt list } else if (data.type === 'error') { this.autoTagState.eventSource.close(); @@ -3819,7 +3819,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; console.error('EventSource readyState:', this.autoTagState.eventSource.readyState); this.autoTagState.eventSource.close(); this.hideModal("autoTagProgressModal"); - this.showNotification('Connection lost during auto tagging', 'error'); + this.showNotification((error && error.serverMessage) || 'Connection lost during auto tagging', 'error'); }; } catch (error) { console.error('Auto tag error:', error); @@ -3894,10 +3894,14 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; } } - // Check if tags array has "real" tags (excluding auto-scanned and prepend:* metadata) - getRealTags(tags) { - if (!tags || !Array.isArray(tags)) return []; - return tags.filter(tag => tag !== 'auto-scanned' && !tag.startsWith('prepend:') && !tag.startsWith('append:')); + /** POST /autotag/single for `image` (#379: `existing` "skip" answers already_tagged for an auto-tagged image). */ + async requestReviewTags(image, existing) { + const response = await ApiClient.fetchUrl('/prompt_manager/autotag/single', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(ImageHelpers.autotagSingleBody(image, { ...this.autoTagState, existing })) + }); + return response.json(); } // Show the re-tag confirmation modal @@ -3949,26 +3953,48 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; // always provides a servable url; never rebuild it from the path. const imageUrl = this.getImageUrl(image); - // Check if image already has real tags (excluding auto-scanned) - const realTags = this.getRealTags(image.prompt_tags); + this.autoTagState.currentImage = image; + this.autoTagState.currentSha256 = null; + this.autoTagState.tagsExpanded = false; // Reset accordion for new image - if (realTags.length > 0) { - // Skip if "Skip All Tagged" was selected - if (this.autoTagState.skipAllTagged) { - this.autoTagState.reviewIndex++; - return this.loadNextReviewImage(); + let data; + try { + // Each image is first asked with "skip": an image a tagging model already tagged answers + // already_tagged (its image tags, never its prompt's) without running the model (#379). + data = await this.requestReviewTags(image, 'skip'); + if (data.success && data.already_tagged) { + if (this.autoTagState.skipAllTagged) { + this.autoTagState.reviewIndex++; + return this.loadNextReviewImage(); + } + const choice = await this.showRetagConfirmation(image, imageUrl, data.image_tags || []); + if (choice === 'skip' || choice === 'skipAll') { + this.autoTagState.reviewIndex++; + return this.loadNextReviewImage(); + } + // choice === 'retag': generate new tags for it + this.showReviewImage(image, imageUrl); + data = await this.requestReviewTags(image, 'merge'); + } else { + this.showReviewImage(image, imageUrl); } - - // Show confirmation modal - const choice = await this.showRetagConfirmation(image, imageUrl, realTags); - - if (choice === 'skip' || choice === 'skipAll') { - this.autoTagState.reviewIndex++; - return this.loadNextReviewImage(); + if (data.success) { + this.autoTagState.currentTags = data.tags; + this.autoTagState.currentSha256 = data.sha256 || null; + this.renderReviewTags(); + } else { + document.getElementById('reviewTagsVisible').innerHTML = + `<div class="text-pm-error">Error: ${escapeHtml(data.error)}</div>`; } - // choice === 'retag' - continue with tagging + } catch (error) { + console.error('Error generating tags:', error); + this.showReviewImage(image, imageUrl); + document.getElementById('reviewTagsVisible').innerHTML = + '<div class="text-pm-error">Failed to generate tags</div>'; } + } + showReviewImage(image, imageUrl) { document.getElementById('reviewImage').src = imageUrl; // The record's file name, else the last segment of its path (#328: `filename` was never declared). const filename = image.filename || String(image.image_path || image.path || '').split(/[\\/]/).pop(); @@ -3977,37 +4003,6 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; document.getElementById('reviewTagsHidden').innerHTML = ''; document.getElementById('reviewTagsAccordion').classList.add('hidden'); document.getElementById('reviewTagsToggle').classList.add('hidden'); - - // Store prompt_id from database for applying tags later - this.autoTagState.currentPromptId = image.prompt_id; - this.autoTagState.tagsExpanded = false; // Reset accordion for new image - - try { - // Review rows come from /images/all (database ids), so this resolves to image_id - const requestBody = ImageHelpers.autotagSingleBody(image, this.autoTagState); - - const response = await ApiClient.fetchUrl('/prompt_manager/autotag/single', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(requestBody) - }); - - const data = await response.json(); - - if (data.success) { - this.autoTagState.currentTags = data.tags; - // Use prompt_id from database, fallback to API response - this.autoTagState.currentPromptId = image.prompt_id || data.prompt_id; - this.renderReviewTags(); - } else { - document.getElementById('reviewTagsVisible').innerHTML = - `<div class="text-pm-error">Error: ${escapeHtml(data.error)}</div>`; - } - } catch (error) { - console.error('Error generating tags:', error); - document.getElementById('reviewTagsVisible').innerHTML = - '<div class="text-pm-error">Failed to generate tags</div>'; - } } /** /prompt_manager/images/all is capped per call; collect every page. */ @@ -4126,15 +4121,9 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; } async applyReviewTags() { - if (this.autoTagState.currentTags.length === 0) { - this.skipReviewImage(); - return; - } - - const promptId = this.autoTagState.currentPromptId; - - if (!promptId) { - this.showNotification('No linked prompt found for this image', 'warning'); + // #379: the reviewed tags go to this image (linked or not), never its prompt; an empty + // list still marks the image as reviewed by a tagging model. + if (!this.autoTagState.currentSha256) { this.skipReviewImage(); return; } @@ -4143,10 +4132,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; const response = await ApiClient.fetchUrl('/prompt_manager/autotag/apply', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - prompt_id: promptId, - tags: this.autoTagState.currentTags - }) + body: JSON.stringify(ImageHelpers.autotagApplyBody( + this.autoTagState.currentImage, this.autoTagState.currentSha256, this.autoTagState.currentTags)) }); const data = await response.json(); diff --git a/web/js/gallery.js b/web/js/gallery.js index 51838c6..0c4e90d 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -3560,19 +3560,24 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; } else if (data.type === 'complete') { this.autoTagState.eventSource.close(); this.hideModal("autoTagProgressModal"); - this.showNotification(`Auto tagging complete! Applied tags to ${data.tagged || 0} prompts.`, 'success'); + this.showNotification(ImageHelpers.autotagCompleteMessage(data.tagged), 'success'); this.loadImages(); // Refresh the gallery } else if (data.type === 'error') { this.autoTagState.eventSource.close(); this.hideModal("autoTagProgressModal"); this.showNotification(`Auto tag error: ${data.message}`, 'error'); + } else if (data.type === 'cancelled') { + // Cancelled in Prompt Manager's Activity (#379): as its own cancel. + this.autoTagState.eventSource.close(); + this.hideModal("autoTagProgressModal"); + this.showNotification('Auto tagging cancelled', 'info'); } }; - this.autoTagState.eventSource.onerror = () => { + this.autoTagState.eventSource.onerror = (error) => { this.autoTagState.eventSource.close(); this.hideModal("autoTagProgressModal"); - this.showNotification('Connection lost during auto tagging', 'error'); + this.showNotification((error && error.serverMessage) || 'Connection lost during auto tagging', 'error'); }; } catch (error) { console.error('Auto tag error:', error); @@ -3629,6 +3634,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; this.reviewRequest = request; this.autoTagState.currentTags = []; this.autoTagState.reviewPromptId = null; + this.autoTagState.reviewImage = null; + this.autoTagState.reviewSha256 = null; if (this.autoTagState.reviewIndex >= this.autoTagState.reviewImages.length) { this.hideModal("autoTagReviewModal"); this.showNotification('Review complete!', 'success'); @@ -3643,8 +3650,9 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; document.getElementById('reviewTagsContainer').innerHTML = '<div class="text-pm-secondary">Generating tags...</div>'; try { - // Gallery entries come from /images/output: id is a path digest, so address by path - const requestBody = ImageHelpers.autotagSingleBody(image, this.autoTagState); + // Gallery entries come from /images/output: id is a path digest, so address by path. + // The gallery review always suggests (it never skipped tagged images, #379). + const requestBody = ImageHelpers.autotagSingleBody(image, { ...this.autoTagState, existing: 'merge' }); const response = await ApiClient.fetchUrl('/prompt_manager/autotag/single', { method: 'POST', @@ -3657,6 +3665,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; if (this.reviewRequest !== request) return; if (data.success) { this.autoTagState.reviewPromptId = data.prompt_id ?? null; + this.autoTagState.reviewImage = image; + this.autoTagState.reviewSha256 = data.sha256 || null; this.autoTagState.currentTags = data.tags; this.renderReviewTags(); } else { @@ -3697,16 +3707,12 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; } async applyReviewTags() { - if (this.autoTagState.currentTags.length === 0) { + // #379: the reviewed tags go to this image (linked or not), never its prompt; an empty + // list still marks the image as reviewed by a tagging model. + if (!this.autoTagState.reviewSha256) { this.skipReviewImage(); return; } - - const promptId = this.autoTagState.reviewPromptId; - if (!promptId) { - this.showNotification('This image has no linked prompt. Scan the output folder first.', 'warning'); - return; - } const request = this.reviewRequest; if (this.reviewApplying === request) return; this.reviewApplying = request; @@ -3715,10 +3721,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; const response = await ApiClient.fetchUrl('/prompt_manager/autotag/apply', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - prompt_id: promptId, - tags: this.autoTagState.currentTags - }) + body: JSON.stringify(ImageHelpers.autotagApplyBody( + this.autoTagState.reviewImage, this.autoTagState.reviewSha256, this.autoTagState.currentTags)) }); const data = await response.json(); diff --git a/web/js/image-helpers.js b/web/js/image-helpers.js index fc019d7..7ab5047 100644 --- a/web/js/image-helpers.js +++ b/web/js/image-helpers.js @@ -70,10 +70,28 @@ * path digest as id, which would 404, so those are addressed by path. * * @param {{id?: unknown, path?: unknown}} image - * @param {{modelType: string, generalThreshold?: number, characterThreshold?: number, prompt?: string}} settings + * `existing` (#379): "skip" answers `already_tagged` for an image a model + * already tagged; anything else is "merge" (always suggest). + * + * @param {{id?: unknown, path?: unknown}} image + * @param {{modelType: string, generalThreshold?: number, characterThreshold?: number, prompt?: string, existing?: string}} settings * @returns {object} JSON-serialisable request body */ function autotagSingleBody(image, settings) { + const target = autotagTarget(image); + const existing = settings.existing === "skip" ? "skip" : "merge"; + const modelType = String(settings.modelType || ""); + const params = modelType.startsWith("wd14") + ? { + general_threshold: settings.generalThreshold, + character_threshold: settings.characterThreshold, + } + : { prompt: settings.prompt }; + return { ...target, model_type: modelType, existing, ...params }; + } + + /** The image a legacy single/apply body names: its row id, or its path (and root). */ + function autotagTarget(image) { const entry = image || {}; const target = typeof entry.path === "string" && entry.path && !Number.isInteger(entry.id) @@ -82,14 +100,22 @@ if ("path" in target && Number.isInteger(entry.root_index)) { target.root_index = entry.root_index; } - const modelType = String(settings.modelType || ""); - const params = modelType.startsWith("wd14") - ? { - general_threshold: settings.generalThreshold, - character_threshold: settings.characterThreshold, - } - : { prompt: settings.prompt }; - return { ...target, model_type: modelType, ...params }; + return target; + } + + /** + * Body for POST /prompt_manager/autotag/apply (#379): the image, the + * `sha256` /autotag/single answered and the reviewed tags. The tags go to + * that image only, never to its prompt; an empty list marks it reviewed. + */ + function autotagApplyBody(image, sha256, tags) { + return { ...autotagTarget(image), sha256: String(sha256 || ""), tags: Array.isArray(tags) ? [...tags] : [] }; + } + + /** The legacy batch's completion toast (#379): images, never prompts. */ + function autotagCompleteMessage(tagged) { + const count = Number.isSafeInteger(tagged) && tagged > 0 ? tagged : 0; + return `Auto tagging complete! Tagged ${count} ${count === 1 ? "image" : "images"}.`; } /** @@ -122,7 +148,15 @@ } } - const api = { formatImageCaption, workflowDocumentHtml, autotagSingleBody, fallbackImageSource, metadataLookupUrl }; + const api = { + formatImageCaption, + workflowDocumentHtml, + autotagSingleBody, + autotagApplyBody, + autotagCompleteMessage, + fallbackImageSource, + metadataLookupUrl, + }; if (typeof module !== "undefined" && module.exports) module.exports = api; else root.ImageHelpers = api; })(typeof window !== "undefined" ? window : globalThis); diff --git a/web/js/sse-stream.js b/web/js/sse-stream.js index 9ee98ba..ee3372d 100644 --- a/web/js/sse-stream.js +++ b/web/js/sse-stream.js @@ -59,6 +59,20 @@ * @returns {{onmessage: Function|null, onerror: Function|null, onclose: Function|null, * readyState: number, close: Function, CONNECTING: number, OPEN: number, CLOSED: number}} */ + // A non-2xx answer as an Error carrying its status and, when the body is + // the API's JSON {error}, that sentence as serverMessage for the page. + async function httpError(response) { + const error = new Error(`HTTP ${response.status}`); + error.status = response.status; + try { + const data = await response.json(); + if (data && typeof data.error === "string" && data.error) error.serverMessage = data.error; + } catch (_) { + // Not JSON: the status alone. + } + return error; + } + function connect(url, { method = "GET", body, fetchImpl } = {}) { const doFetch = fetchImpl || ((...args) => globalThis.fetch(...args)); const controller = new AbortController(); @@ -103,7 +117,7 @@ if (body !== undefined) init.body = body; const requestUrl = root.ApiClient ? root.ApiClient.resolveUrl(url) : url; const response = await doFetch(requestUrl, init); - if (!response.ok) throw new Error(`HTTP ${response.status}`); + if (!response.ok) throw await httpError(response); if (source.readyState === CLOSED) return; source.readyState = OPEN;