feat(autotag): tag output images as a reviewed Activity job
This commit is contained in:
Vito Sansevero
@@ -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,
|
||||
|
||||
+355
-25
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user