diff --git a/database/operations.py b/database/operations.py index 0f45097..78ba02c 100644 --- a/database/operations.py +++ b/database/operations.py @@ -1808,6 +1808,37 @@ 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/autotag_routes.py b/py/api/autotag_routes.py index 5ba1ec0..8749339 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -7,12 +7,20 @@ import json import os import re import threading +import time from pathlib import Path from typing import Optional from urllib.parse import quote from .job_routes import busy_response -from .jobs import AUTOTAG_BATCH_JOB_KIND, MODEL_DOWNLOAD_JOB_KIND, JobBusy +from .jobs import ( + AUTOTAG_BATCH_JOB_KIND, + AUTOTAG_GROUP_NAMES, + MODEL_DOWNLOAD_JOB_KIND, + JobBusy, +) +from .admin import _plural, _scan_location, _scan_roots +from .maintenance import disk_plan, short_space_message from .prompts import thumbnail_url_for from . import thumbnail_settings from .. import model_gate @@ -20,6 +28,18 @@ from .. import model_gate from aiohttp import web _INTERNAL_ERROR = "An internal error occurred. Check server logs for details." + +# Tag output images (#328): the native prompt-tag batch. +AUTOTAG_JOB_TITLE = "Tag output images" +AUTOTAG_PREVIEW_LEASE = "Tag output images review" +AUTOTAG_OWNER_KIND = "autotag-job" +BATCH_PREVIEW_BODY_MAX = 1024 +BATCH_JOB_BODY_MAX = 32 * 1024 +REVIEW_REQUIRED = ( + "Replacing tags needs a review of each image. Use Preview tagging review, " + "or choose Skip or Merge for a batch run." +) +BATCH_CANCELLED = "Cancelled. Tags already added were kept." _BUSY_ERROR = "Auto-tagging is already running" MODEL_BUSY_ERROR = ( "A tagging model is busy with another operation. Try again when it finishes." @@ -157,6 +177,23 @@ async def _single(event): yield 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", +} + + +def _added(count): + return f"Added {count} tag" if count == 1 else f"Added {count} tags" + + async def _refusal(title): """The single error event of an auto-tag batch refused by another job.""" message = f"{title} already running" @@ -175,6 +212,10 @@ class AutotagRoutesMixin: async def download_autotag_model_route(request): return await self.download_autotag_model(request) + @routes.post("/prompt_manager/autotag/batch/preview") + async def preview_autotag_batch_route(request): + return await self.preview_autotag_batch(request) + @routes.post("/prompt_manager/autotag/start") async def start_autotag_route(request): return await self.start_autotag(request) @@ -437,31 +478,42 @@ class AutotagRoutesMixin: and not t.startswith("append:") ] - async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool: - """Merge ``tags`` into the prompt; True when something new was added.""" + 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 False + 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 = [t for t in tags if t not in existing_tags] + new_tags = list(dict.fromkeys(t for t in tags if t not in existing_tags)) if not new_tags: - return False + return 0 await self._run_in_executor( self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags ) - return True + 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.""" - import time as _time + """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, ``BATCH_DETAILS``). + """ model_type = opts["model_type"] keep_in_memory = opts["keep_in_memory"] - status = service.get_models_status() + 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 @@ -478,15 +530,21 @@ class AutotagRoutesMixin: } return - if opts["custom_prompt"]: + 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.get_all_images) + images = await self._run_in_executor(self.db.autotag_batch_candidates) total = len(images) if total == 0: yield { @@ -498,7 +556,8 @@ class AutotagRoutesMixin: "skipped_outside_roots": 0, "status": "No images with linked prompts found", } - service.unload_model() + if not keep_in_memory: + service.unload_model() return yield { @@ -515,7 +574,7 @@ class AutotagRoutesMixin: "skipped_outside_roots": 0, } tagged_prompt_ids = set() - last_update = _time.monotonic() + last_update = _monotonic() finished = False def progress_event(i, message): @@ -531,6 +590,8 @@ class AutotagRoutesMixin: 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 @@ -549,7 +610,7 @@ class AutotagRoutesMixin: counts["skipped_outside_roots"] += 1 else: skip_reason = "file missing" - elif opts["skip_tagged"] and self._real_tags( + elif opts["existing"] == "skip" and self._real_tags( image_data.get("prompt_tags", []) ): tagged_prompt_ids.add(prompt_id) @@ -557,7 +618,8 @@ class AutotagRoutesMixin: if skip_reason is not None: counts["skipped"] += 1 - now = _time.monotonic() + 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})..." @@ -567,17 +629,17 @@ class AutotagRoutesMixin: continue try: - tags = await run( - service.generate_tags, - allowed_path, - general_threshold=opts["general_threshold"], - character_threshold=opts["character_threshold"], - ) + tags = await run(service.generate_tags, allowed_path, **tag_kwargs) counts["processed"] += 1 - if tags and await self._apply_tags_to_prompt(prompt_id, tags): + 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( @@ -586,10 +648,11 @@ class AutotagRoutesMixin: ) 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 = _time.monotonic() + last_update = _monotonic() finished = True if not keep_in_memory: @@ -617,7 +680,11 @@ class AutotagRoutesMixin: opts = { "model_type": params.get("model_type", "gguf"), "custom_prompt": params.get("prompt", ""), - "skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true", + "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"), @@ -692,6 +759,269 @@ class AutotagRoutesMixin: finally: _autotag_batch_lock.release() + # ── Tag output images (#328) ──────────────────────────────────────── + + async def _batch_body(self, request, keys, max_bytes): + """``(data, None)`` for a JSON object with exactly ``keys``.""" + data, error = await self._read_json_body(request, max_bytes=max_bytes) + if error is not None: + return None, error + if set(data) != set(keys): + return None, self._bad_request( + "The body must be exactly {" + ", ".join(keys) + "}" + ) + return data, None + + @staticmethod + def _batch_choice(data): + """The model's engine entry and the policy, or a 400 message.""" + from ..autotag import MODELS + from ..config import EXISTING_POLICIES + + model_type, existing = data.get("model_type"), data.get("existing") + if not isinstance(model_type, str) or model_type not in MODELS: + return None, "model_type must be one of " + ", ".join(sorted(MODELS)) + if not isinstance(existing, str) or existing not in EXISTING_POLICIES: + return None, "existing must be one of " + ", ".join(EXISTING_POLICIES) + return MODELS[model_type], None + + async def preview_autotag_batch(self, request): + """Review Tag output images: counts, model, backup and disk; no write. + + Opens no image file. Refused (409 ``busy``) while an exclusive job runs. + """ + try: + async with self._jobs.lease(AUTOTAG_PREVIEW_LEASE): + return await self._preview_autotag_batch(request) + except JobBusy as busy: + return busy_response(busy) + + async def _preview_autotag_batch(self, request): + from ..autotag import get_autotag_service + + data, error = await self._batch_body( + request, ("model_type", "existing"), BATCH_PREVIEW_BODY_MAX + ) + if error is not None: + return error + config, message = self._batch_choice(data) + if message: + return self._bad_request(message) + model_type, existing = data["model_type"], data["existing"] + policy = self._safety_backup_policy() + try: + candidates = await self._run_in_executor(self.db.autotag_batch_candidates) + status = await self._run_in_executor( + get_autotag_service().get_models_status + ) + disk, _short, _size = await self._run_in_executor( + functools.partial( + disk_plan, + self.db.model.db_path, + backup=policy != "off", + vacuum=False, + ) + ) + except Exception as e: + self.logger.exception("Tag output images could not be reviewed") + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) + prompts = {row["prompt_id"]: row["prompt_tags"] for row in candidates} + tagged = sum(1 for tags in prompts.values() if self._real_tags(tags)) + to_tag = len(prompts) - tagged if existing == "skip" else len(prompts) + return web.json_response( + { + "success": True, + "images": len(candidates), + "prompts": len(prompts), + "already_tagged": tagged, + "to_tag": to_tag, + "model": { + "type": model_type, + "label": config["label"], + "size_label": size_label(config["size"]), + "downloaded": bool( + (status.get(model_type) or {}).get("downloaded") + ), + }, + "existing": existing, + "review_required": existing == "replace_after_review", + "backup": {"policy": policy}, + "disk": disk, + } + ) + + async def start_autotag_job(self, request): + """``POST /jobs/autotag``: start Tag output images (202), or refuse. + + In order: 400 body or options, 400 ``review-required``; 409 + ``model-unavailable``; 507 ``insufficient-space``; 409 ``busy`` + (another records job, the legacy batch's hold included); 409 + ``model-busy``. No await follows the busy check. + """ + from ..autotag import get_autotag_service + from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig + from .annotation_suggest import OptionError, validate_options + + data, error = await self._batch_body( + request, ("model_type", "existing", "options"), BATCH_JOB_BODY_MAX + ) + if error is not None: + return error + config, message = self._batch_choice(data) + if message: + return self._bad_request(message) + model_type, existing, options = ( + data["model_type"], + data["existing"], + data["options"], + ) + if isinstance(options, dict) and {"suggestion_limit", "use_gpu"} & set(options): + return self._bad_request("options take no suggestion_limit or use_gpu") + try: + _model, options = validate_options( + model_type, + options, + default_limit=AutotagConfig.suggestion_limit(), + max_limit=MAX_SUGGESTION_LIMIT, + ) + except OptionError as error: + return self._bad_request(str(error)) + if existing == "replace_after_review": + return web.json_response( + {"success": False, "code": "review-required", "error": REVIEW_REQUIRED}, + status=400, + ) + service = get_autotag_service() + backup_on = self._safety_backup_policy() == "on" + try: + status = await self._run_in_executor(service.get_models_status) + candidates = await self._run_in_executor(self.db.autotag_batch_candidates) + _disk, short, _size = await self._run_in_executor( + functools.partial( + disk_plan, self.db.model.db_path, backup=backup_on, vacuum=False + ) + ) + except Exception as e: + self.logger.exception("Tag output images could not start") + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) + if not (status.get(model_type) or {}).get("downloaded"): + return web.json_response( + { + "success": False, + "code": "model-unavailable", + "error": f"{config['label']} is not downloaded. " + "Download it in System › Auto-tagging.", + }, + status=409, + ) + if backup_on and short is not None: + return web.json_response( + { + "success": False, + "code": "insufficient-space", + "error": short_space_message(short), + }, + status=507, + ) + busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records") + if busy is not None: + return busy_response(JobBusy(busy)) + owner = model_gate.new_owner(AUTOTAG_OWNER_KIND) + if not model_gate.try_acquire(owner): + return model_busy_response() + opts = { + "model_type": model_type, + "existing": existing, + "keep_in_memory": True, + "gate_owner": owner, + "general_threshold": options.get("general_threshold"), + "character_threshold": options.get("character_threshold"), + "prompt": options.get("instructions"), + } + try: + job = self._jobs.start( + AUTOTAG_BATCH_JOB_KIND, + AUTOTAG_JOB_TITLE, + admission="records", + cancellable=True, + scope={ + "model_type": model_type, + "existing": existing, + "images": len(candidates), + "label": config["label"], + }, + groups=AUTOTAG_GROUP_NAMES, + runner=functools.partial( + self._run_autotag_job, service, opts, backup_on + ), + ) + except BaseException as refused: + model_gate.release(owner) + if isinstance(refused, JobBusy): + return busy_response(refused) + raise + job.task.add_done_callback(lambda _task: model_gate.release(owner)) + return web.json_response({"success": True, "job": job.summary()}, status=202) + + async def _run_autotag_job(self, service, opts, backup_on, job): + ok, _backup = await self._safety_backup_step(job, self.db, backup_on) + if not ok: + return + roots = _scan_roots(self._get_all_output_dirs()) + job.log_line( + f"Tagging started · {job.scope['label']} · " + + ("skip prompts already tagged" if opts["existing"] == "skip" else "merge") + ) + job.update(status_line="Loading the tagging model") + + def on_item(index, total, group, image_path, detail): + path, root, root_index = _scan_location(image_path or "", roots) + job.update( + processed=index + 1, + total=total, + progress=100 * (index + 1) // total, + status_line=f"Tagging {index + 1} of {total}", + ) + job.record( + group, + { + "path": path, + "root": root, + "root_index": root_index, + "detail": detail, + }, + f"{path}: {detail}", + ) + job.publish({"type": "progress", "progress": job.progress}) + + final = None + try: + async for event in self._autotag_batch_events( + service, + {**opts, "cancelled": lambda: job.cancelled, "on_item": on_item}, + ): + if event.get("type") in ("complete", "error"): + final = event + if final is None and job.cancelled: + job.update(status_line=BATCH_CANCELLED) + job.log_line(BATCH_CANCELLED) + job.finish("cancelled") + elif final is not None and final["type"] == "error": + job.update(status_line=final["message"]) + job.log_line(f"Tagging failed · {final['message']}") + job.finish("failed") + else: + done = f"Tagged {_plural(job.counters.get('tagged', 0), 'prompt')}" + job.update(status_line=done, progress=100) + job.log_line(done) + job.finish("finished") + finally: + await self._prune_backup_pool(job, self.db) + def _gated_runner(self, owner): """Executor calls under the model gate *owner* holds (plain when None).""" if owner is None: diff --git a/py/api/job_routes.py b/py/api/job_routes.py index 59cadb6..82dcbfc 100644 --- a/py/api/job_routes.py +++ b/py/api/job_routes.py @@ -134,6 +134,10 @@ class JobRoutesMixin: async def start_model_download_job_route(request): return await self.start_model_download_job(request) + @routes.post("/prompt_manager/jobs/autotag") + async def start_autotag_job_route(request): + return await self.start_autotag_job(request) + @routes.post("/prompt_manager/jobs/thumbnails") async def start_thumbnail_job_route(request): return await self.start_thumbnail_job(request) diff --git a/py/api/jobs.py b/py/api/jobs.py index a385f3b..39b74b3 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -87,6 +87,8 @@ SCAN_GROUPS = { # Record-changing: it writes image tag sets and new vocabulary names. IMAGE_TAG_JOB_KIND = "image-tags" IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors") +# Tag output images (#328) counts its prompts the same way. +AUTOTAG_GROUP_NAMES = IMAGE_TAG_GROUP_NAMES # Thumbnail generation (#249 System S5): previews written, already fresh, failed. THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors") diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 2ff8b0c5..01fac1d 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -1806,3 +1806,495 @@ class TestSavedTaggingDefaults(AutotagAPITestCase): self.assertEqual( [call[2:] for call in self.service.generate_calls], [(0.45, 0.75)] ) + + +class BatchJobTestCase(AutotagAPITestCase): + """Fixtures for Tag output images: four prompts, five linked images. + + ``tagged`` already has a real tag; ``fresh`` has two images (the second + is "handled by another image"); ``missing`` points at no file; + ``outside`` is a file outside the output folders; ``marked`` has only + bookkeeping markers (not a real tag). + """ + + def setUp(self): + super().setUp() + from py.config import PromptManagerConfig + + backup = PromptManagerConfig.SAFETY_BACKUP + self.addCleanup(setattr, PromptManagerConfig, "SAFETY_BACKUP", backup) + PromptManagerConfig.SAFETY_BACKUP = False + + def _seed(self): + ids = { + "tagged": self._save_prompt("tagged prompt", tags=["existing"]), + "fresh": self._save_prompt("fresh prompt"), + "missing": self._save_prompt("missing prompt"), + "outside": self._save_prompt("outside prompt"), + "marked": self._save_prompt("marked prompt", tags=["auto-scanned"]), + } + self._link_image(ids["tagged"], self._make_image(self.output_dir, "t.png")) + self._link_image(ids["fresh"], self._make_image(self.output_dir, "f1.png")) + self._link_image(ids["fresh"], self._make_image(self.output_dir, "f2.png")) + self._link_image(ids["missing"], self.output_dir / "gone.png") + self._link_image(ids["outside"], self._make_image(self.secret_dir, "o.png")) + self._link_image(ids["marked"], self._make_image(self.output_dir, "m.png")) + return ids + + def _tags(self, prompt_id): + return self.api.db.get_prompt_by_id(prompt_id)["tags"] + + async def _preview(self, model_type="wd14-vit", existing="skip"): + return await self._post_json( + "/prompt_manager/autotag/batch/preview", + {"model_type": model_type, "existing": existing}, + ) + + async def _start_job(self, existing="skip", model_type="wd14-vit", options=None): + if options is None: + options = ( + {"general_threshold": 0.3, "character_threshold": 0.8} + if model_type.startswith("wd14") + else {"instructions": "Tags please."} + ) + return await self._post_json( + "/prompt_manager/jobs/autotag", + {"model_type": model_type, "existing": existing, "options": options}, + ) + + async def _run_job(self, **kwargs): + resp = await self._start_job(**kwargs) + self.assertEqual(resp.status, 202, await resp.text()) + job = self.api._jobs.get((await resp.json())["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), 30) + return job + + +class TestTagOutputImagesPreview(BatchJobTestCase): + + async def test_preview_writes_nothing(self): + ids = self._seed() + # The database and its WAL are byte-identical afterwards (-shm holds the + # WAL index, which any reader updates), and no prompt changed. + files = [self._temp_db.name + suffix for suffix in ("", "-wal")] + + def snapshot(): + return ( + [Path(f).read_bytes() if os.path.exists(f) else None for f in files], + {pid: self._tags(pid) for pid in ids.values()}, + ) + + before = snapshot() + resp = await self._preview() + self.assertEqual(resp.status, 200, await resp.text()) + self.assertEqual(snapshot(), before) + self.assertEqual(self.service.generate_calls, []) + self.assertEqual(self.service.load_calls, []) + + async def test_preview_counts_skip_and_merge(self): + self._seed() + skip = await (await self._preview(existing="skip")).json() + self.assertEqual( + {k: skip[k] for k in ("images", "prompts", "already_tagged", "to_tag")}, + {"images": 6, "prompts": 5, "already_tagged": 1, "to_tag": 4}, + ) + self.assertEqual( + skip["model"], + { + "type": "wd14-vit", + "label": "WD14 ViT", + "size_label": "~400 MB", + "downloaded": True, + }, + ) + self.assertEqual(skip["existing"], "skip") + self.assertFalse(skip["review_required"]) + self.assertEqual(skip["backup"], {"policy": "off"}) + self.assertEqual( + set(skip["disk"]), {"free_bytes", "required_bytes", "sufficient"} + ) + merge = await (await self._preview(existing="merge")).json() + self.assertEqual(merge["to_tag"], 5) + leaked = [s for s in _strings(skip) if _absolute(s)] + self.assertEqual(leaked, []) + + async def test_preview_flags_replace_as_review_required(self): + self._seed() + data = await (await self._preview(existing="replace_after_review")).json() + self.assertTrue(data["review_required"]) + self.assertEqual(data["prompts"], 5) + + async def test_preview_refuses_bad_bodies(self): + for body in ( + {"model_type": "nope", "existing": "skip"}, + {"model_type": "wd14-vit", "existing": "replace"}, + {"model_type": "wd14-vit"}, + {"model_type": "wd14-vit", "existing": "skip", "x": 1}, + ): + resp = await self._post_json("/prompt_manager/autotag/batch/preview", body) + self.assertEqual(resp.status, 400, body) + + +class TestTagOutputImagesJob(BatchJobTestCase): + + async def test_job_refuses_replace_after_review(self): + self._seed() + resp = await self._start_job(existing="replace_after_review") + data = await resp.json() + self.assertEqual(resp.status, 400) + self.assertEqual(data["code"], "review-required") + self.assertEqual( + data["error"], + "Replacing tags needs a review of each image. Use Preview tagging " + "review, or choose Skip or Merge for a batch run.", + ) + self.assertEqual(self.api._jobs.list(), []) + + async def test_job_refuses_bad_options(self): + for options in ( + {"general_threshold": 0.3}, + {"general_threshold": 0.3, "character_threshold": 0.8, "use_gpu": True}, + { + "general_threshold": 0.3, + "character_threshold": 0.8, + "suggestion_limit": 5, + }, + {"general_threshold": 2, "character_threshold": 0.8}, + ): + resp = await self._start_job(options=options) + self.assertEqual(resp.status, 400, options) + self.assertEqual(self.api._jobs.list(), []) + self.assertIsNone(model_gate.holder()) + + async def test_unavailable_model_is_409(self): + self.service.downloaded = False + resp = await self._start_job() + data = await resp.json() + self.assertEqual(resp.status, 409) + self.assertEqual(data["code"], "model-unavailable") + self.assertEqual( + data["error"], + "WD14 ViT is not downloaded. Download it in System › Auto-tagging.", + ) + + async def test_job_writes_what_legacy_writes(self): + ids = self._seed() + await self._run_job(existing="skip") + native = {name: sorted(self._tags(pid)) for name, pid in ids.items()} + # The same library through the legacy batch, in a fresh database. + for name, pid in ids.items(): + self.api.db.update_prompt_metadata( + pid, + tags=( + ["existing"] + if name == "tagged" + else (["auto-scanned"] if name == "marked" else []) + ), + ) + resp = await self.client.post( + "/prompt_manager/autotag/start?model_type=wd14-vit&skip_tagged=true" + ) + self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete") + legacy = {name: sorted(self._tags(pid)) for name, pid in ids.items()} + self.assertEqual(native, legacy) + self.assertEqual(native["tagged"], ["existing"]) + self.assertEqual(native["fresh"], ["1girl", "smile"]) + + async def test_job_outcomes_follow_the_mapping(self): + ids = self._seed() + self.api.db.update_prompt_metadata( + ids["marked"], tags=["auto-scanned", "1girl", "smile"] + ) + + def details(job): + return { + group: sorted(row["detail"] for row in rows) + for group, rows in job.outcomes.items() + } + + merged = await self._run_job(existing="merge") + self.assertEqual(merged.state, "finished") + self.assertEqual( + details(merged), + { + "tagged": ["Added 2 tags", "Added 2 tags"], + "unchanged": ["No new tags"], + "skipped": [ + "File not found", + "Outside the output folders", + "Prompt handled by another image", + ], + "errors": [], + }, + ) + self.assertEqual(merged.status_line, "Tagged 2 prompts") + + skipped = await self._run_job(existing="skip") + self.assertEqual(details(skipped)["tagged"], []) + self.assertEqual(details(skipped)["skipped"].count("Prompt already tagged"), 3) + self.assertEqual(skipped.status_line, "Tagged 0 prompts") + + async def test_tagging_errors_are_fixed_strings(self): + self._seed() + self.service.generate_error = RuntimeError("/srv/secret/model exploded") + job = await self._run_job(existing="merge") + self.assertTrue(job.outcomes["errors"]) + self.assertEqual( + {r["detail"] for r in job.outcomes["errors"]}, {"Tagging failed"} + ) + self.assertNotIn("/srv/secret", json.dumps(job.report())) + + async def test_job_progress_is_determinate(self): + self._seed() + seen = [] + original = self.service.generate_tags + + def generate(*args, **kwargs): + job = self.api._jobs.active("autotag") + seen.append((job.processed, job.total)) + return original(*args, **kwargs) + + self.service.generate_tags = generate + job = await self._run_job(existing="merge") + self.assertEqual((job.processed, job.total, job.progress), (6, 6, 100)) + self.assertTrue(all(total == 6 for _p, total in seen[1:]), seen) + self.assertTrue(job.cancellable) + + async def test_job_cancel_stops_before_the_next_image_and_keeps_writes(self): + ids = self._seed() + started = threading.Event() + proceed = threading.Event() + calls = [] + + def slow(image_path, **kwargs): + calls.append(image_path) + started.set() + proceed.wait(30) + return ["slow"] + + self.service.generate_tags = slow + resp = await self._start_job(existing="merge") + job = self.api._jobs.get((await resp.json())["job"]["id"]) + try: + await asyncio.get_running_loop().run_in_executor(None, started.wait, 30) + cancel = await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel") + self.assertTrue((await cancel.json())["cancellation_requested"]) + finally: + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), 30) + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.status_line, "Cancelled. Tags already added were kept.") + self.assertEqual(len(calls), 1) + tagged = [n for n, pid in ids.items() if "slow" in self._tags(pid)] + self.assertEqual(len(tagged), 1) + self.assertIsNone(model_gate.holder()) + self.assertEqual(self.service.unload_calls, 0) + + async def test_job_refused_while_legacy_batch_holds(self): + async with self.api._jobs.hold( + "autotag", "Tag output images (legacy page)", admission="records" + ): + resp = await self._start_job() + data = await resp.json() + self.assertEqual(resp.status, 409) + self.assertEqual(data["code"], "busy") + self.assertEqual(data["running"]["title"], "Tag output images (legacy page)") + self.assertIsNone(model_gate.holder()) + + async def test_legacy_refused_while_job_runs(self): + self._seed() + started, proceed = threading.Event(), threading.Event() + + def slow(image_path, **kwargs): + started.set() + proceed.wait(30) + return ["slow"] + + self.service.generate_tags = slow + resp = await self._start_job(existing="merge") + job = self.api._jobs.get((await resp.json())["job"]["id"]) + try: + await asyncio.get_running_loop().run_in_executor(None, started.wait, 30) + self.assertTrue(model_gate.holder().startswith("autotag-job:")) + legacy = await self.client.post("/prompt_manager/autotag/start") + self.assertEqual(legacy.status, 409) + unload = await self.client.post("/prompt_manager/autotag/unload") + self.assertEqual((await unload.json())["operation"], "Tag output images") + finally: + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), 30) + + async def test_job_releases_gate_when_start_is_refused(self): + class Running: + id = "other" + title = "Database restore" + + with patch.object(self.api._jobs, "start", side_effect=JobBusy(Running())): + resp = await self._start_job() + self.assertEqual(resp.status, 409) + self.assertIsNone(model_gate.holder()) + + async def test_model_busy_when_gate_held(self): + owner = model_gate.new_owner("image-annotations") + self.assertTrue(model_gate.try_acquire(owner)) + try: + resp = await self._start_job() + data = await resp.json() + finally: + model_gate.release(owner) + self.assertEqual(resp.status, 409) + self.assertEqual(data["code"], "model-busy") + self.assertEqual(self.api._jobs.list(), []) + + async def test_backup_failure_aborts_before_write(self): + from py.config import PromptManagerConfig + from database.safety_backup import SafetyBackupFailed + + ids = self._seed() + PromptManagerConfig.SAFETY_BACKUP = True + with patch( + "py.api.maintenance.take_safety_backup", + side_effect=SafetyBackupFailed("disk full"), + ): + job = await self._run_job(existing="merge") + self.assertEqual(job.state, "failed") + self.assertEqual(job.scope.get("failure"), "backup-failed") + self.assertEqual(self.service.generate_calls, []) + self.assertEqual(self._tags(ids["fresh"]), []) + + async def test_short_disk_is_507_insufficient_space(self): + from py.config import PromptManagerConfig + + PromptManagerConfig.SAFETY_BACKUP = True + short = ( + {"free_bytes": 1, "required_bytes": 10, "sufficient": False}, + ("database folder", "the safety backup", 10, 1), + 10, + ) + with patch("py.api.autotag_routes.disk_plan", return_value=short): + resp = await self._start_job() + data = await resp.json() + self.assertEqual(resp.status, 507) + self.assertEqual(data["code"], "insufficient-space") + self.assertIsNone(model_gate.holder()) + self.assertEqual(self.api._jobs.list(), []) + + async def test_job_rows_carry_no_absolute_paths(self): + self._seed() + job = await self._run_job(existing="merge") + report = job.report() + self.assertEqual([s for s in _strings(report) if _absolute(s)], []) + rows = [row for rows in job.outcomes.values() for row in rows] + self.assertTrue(rows) + inside = [row for row in rows if row["root_index"] == 0] + self.assertTrue(inside) + + async def test_native_batch_leaves_custom_prompt_unchanged(self): + self._seed() + self.service.custom_prompt = "legacy page prompt" + await self._run_job(existing="merge", model_type="gguf") + self.assertEqual(self.service.custom_prompt, "legacy page prompt") + prompts = {call[1] for call in self.service.generate_calls} + self.assertEqual(prompts, {"Tags please."}) + + +# 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, + }, + ], + ) diff --git a/tests/test_database.py b/tests/test_database.py index 92935ce..09fe6e9 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1398,6 +1398,50 @@ 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()