Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ac10be9dfd | ||
|
|
efbf635165 | ||
|
|
8b2b091098 | ||
|
|
ecc3189649 | ||
|
|
121e586fcb | ||
|
|
e0f909a6bb | ||
|
|
3c4e49708f | ||
|
|
2e60ce797e | ||
|
|
cb7e0f5c48 | ||
|
|
55f0e5c56b | ||
|
|
04f2c921cf |
@@ -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,
|
||||
|
||||
@@ -10,12 +10,15 @@ filesystem or the database.
|
||||
|
||||
import math
|
||||
|
||||
WD14_MODELS = frozenset({"wd14-swinv2", "wd14-vit"})
|
||||
LLM_MODELS = frozenset({"gguf", "hf"})
|
||||
MODEL_TYPES = WD14_MODELS | LLM_MODELS
|
||||
|
||||
MAX_INSTRUCTION_CHARS = 4000
|
||||
MAX_INSTRUCTION_BYTES = 16 * 1024
|
||||
from ..tagging_defaults import ( # noqa: F401 - re-exported for the route
|
||||
LLM_MODEL_TYPES as LLM_MODELS,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES as WD14_MODELS,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
|
||||
SUGGEST_FIELDS = frozenset(
|
||||
{"target", "capture_token", "expected_epoch", "model_type", "options"}
|
||||
@@ -32,27 +35,17 @@ class OptionError(ValueError):
|
||||
|
||||
|
||||
def _threshold(name: str, value) -> float:
|
||||
if type(value) not in (int, float) or not math.isfinite(value):
|
||||
raise OptionError(f"{name} must be a finite number")
|
||||
if not 0 <= value <= 1:
|
||||
raise OptionError(f"{name} must be between 0 and 1")
|
||||
return float(value)
|
||||
try:
|
||||
return validate_threshold(name, value)
|
||||
except ValueError as error:
|
||||
raise OptionError(str(error)) from None
|
||||
|
||||
|
||||
def _instructions(value) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise OptionError("instructions must be non-empty text")
|
||||
if len(value) > MAX_INSTRUCTION_CHARS:
|
||||
raise OptionError(
|
||||
f"instructions must be at most {MAX_INSTRUCTION_CHARS} characters"
|
||||
)
|
||||
try:
|
||||
encoded = value.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
raise OptionError("instructions must be valid Unicode text") from None
|
||||
if len(encoded) > MAX_INSTRUCTION_BYTES:
|
||||
raise OptionError("instructions are too long")
|
||||
return value
|
||||
return validate_template(value)
|
||||
except ValueError as error:
|
||||
raise OptionError(str(error)) from None
|
||||
|
||||
|
||||
def validate_options(model_type, options, *, default_limit: int, max_limit: int):
|
||||
|
||||
+394
-36
@@ -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."
|
||||
@@ -141,11 +161,39 @@ def path_is_within(candidate, root) -> bool:
|
||||
return child != parent and child.is_relative_to(parent)
|
||||
|
||||
|
||||
def _thresholds_or_saved(general, character):
|
||||
"""The request's WD14 thresholds as floats; one it omits is the saved
|
||||
default (``autotag.defaults``, OQ4). Raises ValueError/TypeError."""
|
||||
from ..config import AutotagConfig
|
||||
|
||||
return (
|
||||
AutotagConfig.GENERAL_THRESHOLD if general is None else float(general),
|
||||
AutotagConfig.CHARACTER_THRESHOLD if character is None else float(character),
|
||||
)
|
||||
|
||||
|
||||
async def _single(event):
|
||||
"""A stream of one 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"
|
||||
@@ -164,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)
|
||||
@@ -200,12 +252,15 @@ class AutotagRoutesMixin:
|
||||
{
|
||||
"success": True,
|
||||
"models": public_models(status, loaded_type),
|
||||
"default_prompt": service.default_prompt,
|
||||
# The saved tagging defaults (#328, System › Auto-tagging).
|
||||
"default_prompt": AutotagConfig.TEMPLATE,
|
||||
"model_loaded": loaded,
|
||||
"loaded_model_type": loaded_type,
|
||||
"wd14_general_threshold": service.wd14_general_threshold,
|
||||
"wd14_character_threshold": service.wd14_character_threshold,
|
||||
"wd14_general_threshold": AutotagConfig.GENERAL_THRESHOLD,
|
||||
"wd14_character_threshold": AutotagConfig.CHARACTER_THRESHOLD,
|
||||
"suggestion_limit": AutotagConfig.suggestion_limit(),
|
||||
"model_type": AutotagConfig.MODEL_TYPE,
|
||||
"existing": AutotagConfig.EXISTING,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -423,35 +478,48 @@ 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, ``SKIP_DETAILS`` for the skips).
|
||||
"""
|
||||
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
|
||||
|
||||
if cancelled():
|
||||
return
|
||||
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
|
||||
run = self._gated_runner(opts.get("gate_owner"))
|
||||
try:
|
||||
@@ -464,15 +532,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 {
|
||||
@@ -484,7 +558,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 {
|
||||
@@ -501,7 +576,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):
|
||||
@@ -517,6 +592,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
|
||||
@@ -535,7 +612,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)
|
||||
@@ -543,7 +620,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})..."
|
||||
@@ -553,17 +631,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(
|
||||
@@ -572,10 +650,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:
|
||||
@@ -603,16 +682,22 @@ 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"),
|
||||
"character_threshold": params.get("character_threshold"),
|
||||
}
|
||||
try:
|
||||
for key in ("general_threshold", "character_threshold"):
|
||||
if opts[key] is not None:
|
||||
opts[key] = float(opts[key])
|
||||
opts["general_threshold"], opts["character_threshold"] = (
|
||||
_thresholds_or_saved(
|
||||
opts["general_threshold"], opts["character_threshold"]
|
||||
)
|
||||
)
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
@@ -642,7 +727,9 @@ class AutotagRoutesMixin:
|
||||
# never runs beside a scan, restore or the workspace tagging job.
|
||||
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
|
||||
if busy is not None:
|
||||
if busy.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload
|
||||
# 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
|
||||
)
|
||||
@@ -676,6 +763,278 @@ 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
|
||||
try:
|
||||
await self._tag_output_images(service, opts, job)
|
||||
except Exception:
|
||||
self.logger.exception("Tag output images failed")
|
||||
job.update(status_line=_INTERNAL_ERROR)
|
||||
job.log_line(f"Tagging failed · {_INTERNAL_ERROR}")
|
||||
job.finish("failed")
|
||||
finally:
|
||||
if backup_on:
|
||||
await self._prune_backup_pool(job, self.db)
|
||||
|
||||
async def _tag_output_images(self, service, opts, job):
|
||||
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
|
||||
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 not None and final["type"] == "error":
|
||||
job.update(status_line=final["message"])
|
||||
job.log_line(f"Tagging failed · {final['message']}")
|
||||
job.finish("failed")
|
||||
elif job.cancelled: # a cancel wins over a run that then completed
|
||||
job.update(status_line=BATCH_CANCELLED)
|
||||
job.log_line(BATCH_CANCELLED)
|
||||
job.finish("cancelled")
|
||||
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")
|
||||
|
||||
def _gated_runner(self, owner):
|
||||
"""Executor calls under the model gate *owner* holds (plain when None)."""
|
||||
if owner is None:
|
||||
@@ -792,10 +1151,9 @@ class AutotagRoutesMixin:
|
||||
general_threshold = data.get("general_threshold")
|
||||
character_threshold = data.get("character_threshold")
|
||||
try:
|
||||
if general_threshold is not None:
|
||||
general_threshold = float(general_threshold)
|
||||
if character_threshold is not None:
|
||||
character_threshold = float(character_threshold)
|
||||
general_threshold, character_threshold = _thresholds_or_saved(
|
||||
general_threshold, character_threshold
|
||||
)
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -47,11 +47,23 @@ _SETTINGS_FILE_KEYS = {
|
||||
"THUMBNAIL_FORMAT": ("thumbnails", "format"),
|
||||
"THUMBNAIL_QUALITY": ("thumbnails", "quality"),
|
||||
"THUMBNAIL_EDGE": ("thumbnails", "edge"),
|
||||
# The one tagging configuration (#328), published to AutotagConfig.
|
||||
"AUTOTAG_MODEL_TYPE": ("autotag", "defaults", "model_type"),
|
||||
"AUTOTAG_EXISTING": ("autotag", "defaults", "existing"),
|
||||
"AUTOTAG_GENERAL_THRESHOLD": ("autotag", "defaults", "general_threshold"),
|
||||
"AUTOTAG_CHARACTER_THRESHOLD": ("autotag", "defaults", "character_threshold"),
|
||||
"AUTOTAG_TEMPLATE": ("autotag", "defaults", "template"),
|
||||
"AUTOTAG_SUGGESTION_LIMIT": ("autotag", "defaults", "suggestion_limit"),
|
||||
}
|
||||
# The thumbnail settings a running Convert outdated job fixes (#317).
|
||||
_THUMBNAIL_SETTINGS = ("THUMBNAIL_FORMAT", "THUMBNAIL_QUALITY", "THUMBNAIL_EDGE")
|
||||
# Attributes that live on GalleryConfig; the others on PromptManagerConfig.
|
||||
_GALLERY_SETTINGS = frozenset({"MONITORING_ENABLED"})
|
||||
# Attributes that live on AutotagConfig, named without the AUTOTAG_ prefix.
|
||||
_AUTOTAG_SETTINGS = frozenset(
|
||||
name for name in _SETTINGS_FILE_KEYS if name.startswith("AUTOTAG_")
|
||||
)
|
||||
_AUTOTAG_PREFIX = "AUTOTAG_"
|
||||
SOURCE_CHANGE_JOB_KIND = "source_change"
|
||||
SOURCE_CHANGE_JOB_TITLE = "Source change"
|
||||
|
||||
@@ -151,6 +163,63 @@ def _set_nested(config, sections, key, value):
|
||||
config[sections[0]] = merged
|
||||
|
||||
|
||||
def _validate_autotag_settings(data):
|
||||
"""The ``autotag_*`` settings as ``{AUTOTAG_ATTR: value}``.
|
||||
|
||||
The first bad value raises ValueError with a message that names the
|
||||
setting by its label in System › Auto-tagging.
|
||||
"""
|
||||
from ..config import (
|
||||
EXISTING_POLICIES,
|
||||
MAX_SUGGESTION_LIMIT,
|
||||
MODEL_TYPES,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
|
||||
updates = {}
|
||||
if "autotag_model_type" in data:
|
||||
value = data["autotag_model_type"]
|
||||
if not isinstance(value, str) or value not in MODEL_TYPES:
|
||||
raise ValueError(
|
||||
"Tagging model must be one of " + ", ".join(sorted(MODEL_TYPES)) + "."
|
||||
)
|
||||
updates["AUTOTAG_MODEL_TYPE"] = value
|
||||
if "autotag_existing" in data:
|
||||
value = data["autotag_existing"]
|
||||
if not isinstance(value, str) or value not in EXISTING_POLICIES:
|
||||
raise ValueError(
|
||||
"Existing tags must be skip, merge or replace_after_review."
|
||||
)
|
||||
updates["AUTOTAG_EXISTING"] = value
|
||||
for key, label in (
|
||||
("autotag_general_threshold", "General-tag confidence"),
|
||||
("autotag_character_threshold", "Character-tag confidence"),
|
||||
):
|
||||
if key in data:
|
||||
try:
|
||||
value = validate_threshold(label, data[key])
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{error}.") from None
|
||||
updates[key.upper()] = value
|
||||
if "autotag_template" in data:
|
||||
try:
|
||||
updates["AUTOTAG_TEMPLATE"] = validate_template(
|
||||
data["autotag_template"], "LLM tag template"
|
||||
)
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{error}.") from None
|
||||
if "autotag_suggestion_limit" in data:
|
||||
value = data["autotag_suggestion_limit"]
|
||||
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
|
||||
raise ValueError(
|
||||
"Suggestions per image must be a whole number from 1 to "
|
||||
f"{MAX_SUGGESTION_LIMIT}."
|
||||
)
|
||||
updates["AUTOTAG_SUGGESTION_LIMIT"] = value
|
||||
return updates
|
||||
|
||||
|
||||
async def _read_json(request):
|
||||
"""Parse a JSON object body through the package's capped reader."""
|
||||
from . import _read_json_body
|
||||
@@ -177,7 +246,7 @@ class SettingsRoutesMixin:
|
||||
async def get_settings(self, request):
|
||||
"""Get current settings."""
|
||||
try:
|
||||
from ..config import PromptManagerConfig, GalleryConfig
|
||||
from ..config import AutotagConfig, PromptManagerConfig, GalleryConfig
|
||||
|
||||
root_paths = [
|
||||
self._public_path(d) for d in GalleryConfig.MONITORING_DIRECTORIES
|
||||
@@ -203,6 +272,14 @@ class SettingsRoutesMixin:
|
||||
"thumbnail_format": PromptManagerConfig.THUMBNAIL_FORMAT,
|
||||
"thumbnail_quality": PromptManagerConfig.THUMBNAIL_QUALITY,
|
||||
"thumbnail_edge": PromptManagerConfig.THUMBNAIL_EDGE,
|
||||
"autotag_model_type": AutotagConfig.MODEL_TYPE,
|
||||
"autotag_existing": AutotagConfig.EXISTING,
|
||||
"autotag_general_threshold": AutotagConfig.GENERAL_THRESHOLD,
|
||||
"autotag_character_threshold": (
|
||||
AutotagConfig.CHARACTER_THRESHOLD
|
||||
),
|
||||
"autotag_template": AutotagConfig.TEMPLATE,
|
||||
"autotag_suggestion_limit": AutotagConfig.SUGGESTION_LIMIT,
|
||||
"cpu_count": PromptManagerConfig.max_worker_threads(),
|
||||
"gallery_root_paths": root_paths,
|
||||
"gallery_root_path": root_paths[0] if root_paths else "",
|
||||
@@ -425,7 +502,12 @@ class SettingsRoutesMixin:
|
||||
"""
|
||||
from ..config import GalleryConfig, PromptManagerConfig
|
||||
|
||||
from ..config import AutotagConfig
|
||||
|
||||
for name, value in updates.items():
|
||||
if name in _AUTOTAG_SETTINGS:
|
||||
setattr(AutotagConfig, name[len(_AUTOTAG_PREFIX) :], value)
|
||||
continue
|
||||
owner = GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig
|
||||
setattr(owner, name, value)
|
||||
restart_fields = []
|
||||
@@ -564,6 +646,7 @@ class SettingsRoutesMixin:
|
||||
updates["THUMBNAIL_EDGE"] = thumbnail_settings.validate_edge(
|
||||
data["thumbnail_edge"]
|
||||
)
|
||||
updates.update(_validate_autotag_settings(data))
|
||||
return updates
|
||||
|
||||
def _parse_gallery_roots(self, data):
|
||||
|
||||
+24
-27
@@ -25,6 +25,23 @@ except ImportError:
|
||||
from database.annotations import MAX_IMAGE_TAGS
|
||||
from database.vocabulary import validate_vocabulary_tag_name
|
||||
|
||||
# The tag template and thresholds live with the saved tagging defaults
|
||||
# (#328), so config.py never imports PIL.
|
||||
try:
|
||||
from .tagging_defaults import DEFAULT_TAG_TEMPLATE
|
||||
except ImportError:
|
||||
from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
|
||||
|
||||
|
||||
def _autotag_config():
|
||||
"""The saved defaults; imported late, as config needs ComfyUI's server."""
|
||||
try:
|
||||
from .config import AutotagConfig
|
||||
except ImportError:
|
||||
from py.config import AutotagConfig
|
||||
return AutotagConfig
|
||||
|
||||
|
||||
# Try to import logging from utils, fallback to standard logging
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
@@ -90,17 +107,7 @@ MODELS = {
|
||||
}
|
||||
|
||||
# Default prompts
|
||||
DEFAULT_PROMPT = (
|
||||
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
||||
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
||||
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
||||
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
||||
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
||||
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
||||
"Prioritize visual accuracy.\n"
|
||||
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
||||
"descriptive."
|
||||
)
|
||||
DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE
|
||||
|
||||
# Upper bound of one suggestion request: the image tag set cap.
|
||||
MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS
|
||||
@@ -184,10 +191,6 @@ class AutoTagService:
|
||||
self._current_model_type: Optional[str] = None
|
||||
self._custom_prompt: str = DEFAULT_PROMPT
|
||||
|
||||
# WD14 threshold defaults
|
||||
self._wd14_general_threshold: float = 0.35
|
||||
self._wd14_character_threshold: float = 0.85
|
||||
|
||||
@property
|
||||
def models_config(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get the models configuration dictionary."""
|
||||
@@ -210,19 +213,13 @@ class AutoTagService:
|
||||
|
||||
@property
|
||||
def wd14_general_threshold(self) -> float:
|
||||
return self._wd14_general_threshold
|
||||
|
||||
@wd14_general_threshold.setter
|
||||
def wd14_general_threshold(self, value: float):
|
||||
self._wd14_general_threshold = max(0.0, min(1.0, value))
|
||||
"""The saved General-tag confidence (``autotag.defaults``)."""
|
||||
return _autotag_config().GENERAL_THRESHOLD
|
||||
|
||||
@property
|
||||
def wd14_character_threshold(self) -> float:
|
||||
return self._wd14_character_threshold
|
||||
|
||||
@wd14_character_threshold.setter
|
||||
def wd14_character_threshold(self, value: float):
|
||||
self._wd14_character_threshold = max(0.0, min(1.0, value))
|
||||
"""The saved Character-tag confidence (``autotag.defaults``)."""
|
||||
return _autotag_config().CHARACTER_THRESHOLD
|
||||
|
||||
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get availability status for all model types.
|
||||
@@ -865,12 +862,12 @@ class AutoTagService:
|
||||
gen_thresh = (
|
||||
general_threshold
|
||||
if general_threshold is not None
|
||||
else self._wd14_general_threshold
|
||||
else _autotag_config().GENERAL_THRESHOLD
|
||||
)
|
||||
char_thresh = (
|
||||
character_threshold
|
||||
if character_threshold is not None
|
||||
else self._wd14_character_threshold
|
||||
else _autotag_config().CHARACTER_THRESHOLD
|
||||
)
|
||||
return self._generate_wd14(
|
||||
session, tags_list, image, gen_thresh, char_thresh
|
||||
|
||||
+96
-16
@@ -563,18 +563,90 @@ class IntegrationConfig:
|
||||
DEFAULT_SUGGESTION_LIMIT = 10
|
||||
MAX_SUGGESTION_LIMIT = 100 # pinned to MAX_IMAGE_TAGS by tests/test_config.py
|
||||
|
||||
# The tagging engines, options, template and validators (server-free, #328).
|
||||
try:
|
||||
from .tagging_defaults import ( # noqa: F401 - re-exported
|
||||
DEFAULT_TAG_TEMPLATE,
|
||||
EXISTING_POLICIES,
|
||||
LLM_MODEL_TYPES,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
except ImportError:
|
||||
from py.tagging_defaults import ( # noqa: F401 - re-exported
|
||||
DEFAULT_TAG_TEMPLATE,
|
||||
EXISTING_POLICIES,
|
||||
LLM_MODEL_TYPES,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
|
||||
|
||||
def _valid_model_type(value: Any) -> str:
|
||||
if not isinstance(value, str) or value not in MODEL_TYPES:
|
||||
raise ValueError("must be one of " + ", ".join(sorted(MODEL_TYPES)))
|
||||
return value
|
||||
|
||||
|
||||
def _valid_existing(value: Any) -> str:
|
||||
if not isinstance(value, str) or value not in EXISTING_POLICIES:
|
||||
raise ValueError("must be one of " + ", ".join(EXISTING_POLICIES))
|
||||
return value
|
||||
|
||||
|
||||
def _valid_suggestion_limit(value: Any) -> int:
|
||||
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
|
||||
raise ValueError(f"must be an integer from 1 to {MAX_SUGGESTION_LIMIT}")
|
||||
return value
|
||||
|
||||
|
||||
class AutotagConfig:
|
||||
"""Persisted tagging defaults under ``autotag.defaults`` in config.json.
|
||||
|
||||
``suggestion_limit`` is how many suggestions Generate returns (the top N
|
||||
by score), an integer in ``1..MAX_SUGGESTION_LIMIT``. A hand-edited value
|
||||
outside that range, or of another type, falls back to the default with a
|
||||
warning rather than failing the load.
|
||||
One tagging configuration (#328): the model, what a batch does with
|
||||
prompts already tagged, the two WD14 thresholds, the language-model
|
||||
template and how many suggestions Generate returns. Each key is checked
|
||||
alone on load; a hand-edited value that fails its check falls back to
|
||||
its default with one warning rather than failing the load.
|
||||
"""
|
||||
|
||||
MODEL_TYPE = "wd14-swinv2"
|
||||
EXISTING = "skip"
|
||||
GENERAL_THRESHOLD = 0.35
|
||||
CHARACTER_THRESHOLD = 0.85
|
||||
TEMPLATE = DEFAULT_TAG_TEMPLATE
|
||||
SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
|
||||
|
||||
# config.json key -> (attribute, validator, default)
|
||||
_KEYS = {
|
||||
"model_type": ("MODEL_TYPE", _valid_model_type, "wd14-swinv2"),
|
||||
"existing": ("EXISTING", _valid_existing, "skip"),
|
||||
"general_threshold": (
|
||||
"GENERAL_THRESHOLD",
|
||||
lambda v: validate_threshold("value", v),
|
||||
0.35,
|
||||
),
|
||||
"character_threshold": (
|
||||
"CHARACTER_THRESHOLD",
|
||||
lambda v: validate_threshold("value", v),
|
||||
0.85,
|
||||
),
|
||||
"template": ("TEMPLATE", validate_template, DEFAULT_TAG_TEMPLATE),
|
||||
"suggestion_limit": (
|
||||
"SUGGESTION_LIMIT",
|
||||
_valid_suggestion_limit,
|
||||
DEFAULT_SUGGESTION_LIMIT,
|
||||
),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def suggestion_limit(cls) -> int:
|
||||
return cls.SUGGESTION_LIMIT
|
||||
@@ -587,7 +659,14 @@ class AutotagConfig:
|
||||
|
||||
@classmethod
|
||||
def get_config(cls) -> Dict[str, Any]:
|
||||
return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}}
|
||||
"""The saved defaults; the template only when it is not the built-in
|
||||
one, so a later built-in template still reaches who never set one."""
|
||||
defaults = {
|
||||
key: getattr(cls, attr) for key, (attr, _v, _d) in cls._KEYS.items()
|
||||
}
|
||||
if defaults["template"] == DEFAULT_TAG_TEMPLATE:
|
||||
del defaults["template"]
|
||||
return {"defaults": defaults}
|
||||
|
||||
@classmethod
|
||||
def update_config(cls, new_config: Any) -> None:
|
||||
@@ -598,17 +677,18 @@ class AutotagConfig:
|
||||
if not isinstance(defaults, dict):
|
||||
config_logger.warning("Ignoring autotag.defaults that is not an object")
|
||||
return
|
||||
if "suggestion_limit" not in defaults:
|
||||
return
|
||||
value = defaults["suggestion_limit"]
|
||||
if cls.valid_suggestion_limit(value):
|
||||
cls.SUGGESTION_LIMIT = value
|
||||
return
|
||||
config_logger.warning(
|
||||
"autotag.defaults.suggestion_limit must be an integer from 1 to "
|
||||
f"{MAX_SUGGESTION_LIMIT}; using {DEFAULT_SUGGESTION_LIMIT}"
|
||||
)
|
||||
cls.SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
|
||||
for key, (attr, validate, default) in cls._KEYS.items():
|
||||
if key not in defaults:
|
||||
continue
|
||||
try:
|
||||
setattr(cls, attr, validate(defaults[key]))
|
||||
except ValueError as error:
|
||||
config_logger.warning(
|
||||
f"autotag.defaults.{key} {error}; using {default!r}"
|
||||
if key != "template"
|
||||
else f"autotag.defaults.template {error}; using the built-in one"
|
||||
)
|
||||
setattr(cls, attr, default)
|
||||
|
||||
|
||||
class PromptManagerConfig:
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
"""The tagging configuration's constants and validators (#328).
|
||||
|
||||
Server-free on purpose: ``py.api`` must import without ComfyUI's ``server``
|
||||
module, which ``py/config.py`` needs; config re-exports all of this.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
# The tagging engines (``py/autotag.py`` MODELS) and the options they take.
|
||||
WD14_MODEL_TYPES = frozenset({"wd14-swinv2", "wd14-vit"})
|
||||
LLM_MODEL_TYPES = frozenset({"gguf", "hf"})
|
||||
MODEL_TYPES = WD14_MODEL_TYPES | LLM_MODEL_TYPES
|
||||
MAX_INSTRUCTION_CHARS = 4000
|
||||
MAX_INSTRUCTION_BYTES = 16 * 1024
|
||||
# What a batch does with a prompt that already has tags (#328).
|
||||
EXISTING_POLICIES = ("skip", "merge", "replace_after_review")
|
||||
|
||||
# The language-model tag template every page starts from until one is saved.
|
||||
DEFAULT_TAG_TEMPLATE = (
|
||||
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
||||
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
||||
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
||||
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
||||
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
||||
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
||||
"Prioritize visual accuracy.\n"
|
||||
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
||||
"descriptive."
|
||||
)
|
||||
|
||||
|
||||
def validate_threshold(name: str, value: Any) -> float:
|
||||
"""A WD14 confidence threshold in [0, 1], or ValueError naming ``name``."""
|
||||
if type(value) not in (int, float) or not math.isfinite(value):
|
||||
raise ValueError(f"{name} must be a finite number")
|
||||
if not 0 <= value <= 1:
|
||||
raise ValueError(f"{name} must be between 0 and 1")
|
||||
return float(value)
|
||||
|
||||
|
||||
def validate_template(value: Any, name: str = "instructions") -> str:
|
||||
"""Non-empty language-model instructions within the character and byte
|
||||
caps, or ValueError naming ``name``."""
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"{name} must be non-empty text")
|
||||
if len(value) > MAX_INSTRUCTION_CHARS:
|
||||
raise ValueError(f"{name} must be at most {MAX_INSTRUCTION_CHARS} characters")
|
||||
try:
|
||||
encoded = value.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
raise ValueError(f"{name} must be valid Unicode text") from None
|
||||
if len(encoded) > MAX_INSTRUCTION_BYTES:
|
||||
raise ValueError(f"{name} {'are' if name == 'instructions' else 'is'} too long")
|
||||
return value
|
||||
@@ -87,3 +87,66 @@ test('a failed download closes the modal and says so', async ({ page, request })
|
||||
await expect(status.getByRole('button', { name: 'Download', exact: true })).toBeVisible();
|
||||
expect((await downloads(request)).map(job => job.state)).toEqual(['failed']);
|
||||
});
|
||||
|
||||
// (wd14-vit is not downloaded in these tests: beforeEach marks it so the download tests can offer it.)
|
||||
// The gallery's batch reader keyed on `data.status` values the server never sends, so its modal never moved or
|
||||
// closed (#328 Problem 5). It now reads `type` frames as the admin page does.
|
||||
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 });
|
||||
});
|
||||
}
|
||||
|
||||
async function runGalleryBatch(page) {
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await watchProgress(page);
|
||||
const done = page.waitForResponse(response => response.url().includes('/prompt_manager/autotag/start'));
|
||||
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();
|
||||
}
|
||||
|
||||
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%');
|
||||
expect((await downloads(request)).length).toBe(0);
|
||||
});
|
||||
|
||||
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();
|
||||
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);
|
||||
});
|
||||
|
||||
test('gallery unload works and is refused while the gate is held', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await runGalleryBatch(page); // keeps the model in memory
|
||||
await page.locator('#autoTagBtn').click();
|
||||
const unload = page.locator('#unloadModelBtn');
|
||||
await expect(unload).toBeVisible();
|
||||
|
||||
// A held download owns the model gate: Unload is refused with the server's error.
|
||||
const held = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'gguf' } });
|
||||
expect(held.status()).toBe(202);
|
||||
await unload.click();
|
||||
await expect(page.getByText(/^Failed to unload model: A tagging model is busy/)).toBeVisible();
|
||||
|
||||
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
|
||||
await expect.poll(async () => (await downloads(request)).map(job => job.state)).toEqual(['finished']);
|
||||
await unload.click();
|
||||
await expect(page.getByText('WD14-SWINV2 model unloaded successfully')).toBeVisible();
|
||||
});
|
||||
|
||||
@@ -203,6 +203,8 @@ def register_tagging_commands(
|
||||
model.unavailable = model.unavailable - {model_type}
|
||||
return download["result"]
|
||||
|
||||
_set_defaults(model, DEFAULTS) # the server starts on the standard ones
|
||||
|
||||
@routes.post("/__test__/tagging/reset")
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
data = await request.json() if request.can_read_body else {}
|
||||
@@ -331,12 +333,19 @@ DEFAULTS = (0.35, 0.85, 10)
|
||||
|
||||
|
||||
def _set_defaults(model, values) -> None:
|
||||
"""Set the saved tagging defaults (``AutotagConfig``, #328); a reset also
|
||||
puts back the saved model, batch policy and template (the fake model's
|
||||
short instructions, which the concept comparisons show)."""
|
||||
from py.config import AutotagConfig
|
||||
|
||||
general, character, limit = values
|
||||
model.wd14_general_threshold = general
|
||||
model.wd14_character_threshold = character
|
||||
AutotagConfig.GENERAL_THRESHOLD = general
|
||||
AutotagConfig.CHARACTER_THRESHOLD = character
|
||||
AutotagConfig.SUGGESTION_LIMIT = limit
|
||||
if values == DEFAULTS:
|
||||
AutotagConfig.MODEL_TYPE = "wd14-swinv2"
|
||||
AutotagConfig.EXISTING = "skip"
|
||||
AutotagConfig.TEMPLATE = model.default_prompt
|
||||
|
||||
|
||||
def _valid_defaults(data: object) -> bool:
|
||||
|
||||
@@ -225,3 +225,26 @@ test('a failed duplicate cleanup says the server reason (S7 review M3)', async (
|
||||
assert.equal(finishToast(failed), `Duplicate cleanup failed · ${reason}`);
|
||||
assert.equal(finishToast({ ...failed, status_line: '' }), 'Duplicate cleanup failed · 0 removed · 1 error');
|
||||
});
|
||||
|
||||
// #328: Tag output images (`autotag`) reuses the image tagging view: Tagged / Unchanged / Skipped / Errors and Cancel.
|
||||
// The legacy page's hold of the same kind counts nothing and is not cancellable, so it keeps the generic held view.
|
||||
test('autotag gets Tagged / Unchanged / Skipped / Errors and Cancel; its legacy hold keeps the generic view', async () => {
|
||||
const { kindOf, viewOf, canCancel } = await source;
|
||||
const { startToast, finishToast } = await store;
|
||||
const counters = { tagged: 3, unchanged: 1, skipped: 2, errors: 0 };
|
||||
const batch = job({ kind: 'autotag', title: 'Tag output images', counters, scope: { images: 6 } });
|
||||
const view = viewOf(batch);
|
||||
assert.deepEqual(view.tiles.map(([group, label]) => [group, label]),
|
||||
[['tagged', 'Tagged'], ['unchanged', 'Unchanged'], ['skipped', 'Skipped'], ['errors', 'Errors']]);
|
||||
assert.equal(kindOf('autotag').pillVerb, 'Tagging');
|
||||
assert.equal(kindOf('autotag').cancelLabel, 'Cancel tagging');
|
||||
assert.equal(canCancel(batch), true);
|
||||
assert.equal(startToast(batch), 'Tag output images started.');
|
||||
assert.match(finishToast({ ...batch, state: 'finished', finished_at: 2 }), /^Tag output images finished/);
|
||||
|
||||
const held = job({ kind: 'autotag', title: 'Tag output images (legacy page)', counters: {}, cancellable: false,
|
||||
progress: 0, processed: 0, total: 0 });
|
||||
const legacy = viewOf(held);
|
||||
assert.deepEqual([legacy.tiles, legacy.outcomes, [...legacy.openable]], [[], [], []]);
|
||||
assert.equal(canCancel(held), false);
|
||||
});
|
||||
|
||||
@@ -197,6 +197,52 @@ test('instructions default to the server prompt', async () => {
|
||||
assert.equal(c.state.options.suggestion_limit, 7, 'the suggestion count survives a model change');
|
||||
});
|
||||
|
||||
// #328: one tagging configuration. The dialogs start on the model saved in System › Auto-tagging, unless a preset
|
||||
// (a review session's choice) or the user already picked one; an unknown saved model keeps the first option.
|
||||
test('dialog starts on the saved model unless a preset or the user chose one', async () => {
|
||||
const saved = { [`GET ${MODELS}`]: () => response(200, models({ model_type: 'gguf' })) };
|
||||
const { c } = await opened(saved);
|
||||
assert.equal(c.state.model, 'gguf');
|
||||
assert.equal(c.state.options.instructions, DEFAULT_PROMPT);
|
||||
|
||||
const { c: preset } = await opened(saved, { preset: { model: 'wd14-vit', options: {} } });
|
||||
assert.equal(preset.state.model, 'wd14-vit');
|
||||
|
||||
const { createImageTaggingController, createTaggingOptionsController } = await load();
|
||||
const h = harness(defaultRoutes(saved));
|
||||
const picked = createImageTaggingController({ api: h.api, target: TARGET, source: 'images' });
|
||||
picked.setModel('hf');
|
||||
await picked.open();
|
||||
await settle();
|
||||
assert.equal(picked.state.model, 'hf', 'the user picked a model before the read landed');
|
||||
|
||||
// A pick made while the models read is still in flight wins over the saved model that lands after it.
|
||||
const gate = deferred();
|
||||
const late = harness(defaultRoutes({ [`GET ${MODELS}`]: () => gate.promise }));
|
||||
const racing = createImageTaggingController({ api: late.api, target: TARGET, source: 'images' });
|
||||
const opening = racing.open();
|
||||
await settle();
|
||||
assert.equal(racing.state.phase, 'observing');
|
||||
racing.setModel('wd14-vit');
|
||||
gate.resolve(response(200, models({ model_type: 'gguf' })));
|
||||
await opening;
|
||||
await settle();
|
||||
assert.equal(racing.state.model, 'wd14-vit');
|
||||
racing.dispose();
|
||||
|
||||
const { c: unknown } = await opened({ [`GET ${MODELS}`]: () => response(200, models({ model_type: 'wd14-convnext' })) });
|
||||
assert.equal(unknown.state.model, 'wd14-swinv2');
|
||||
|
||||
const options = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api });
|
||||
await options.open();
|
||||
assert.equal(options.state.model, 'gguf');
|
||||
const chosen = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api });
|
||||
chosen.setModel('wd14-vit');
|
||||
await chosen.open();
|
||||
assert.equal(chosen.state.model, 'wd14-vit');
|
||||
for (const controller of [c, preset, picked, unknown, options, chosen]) controller.dispose();
|
||||
});
|
||||
|
||||
test('an unreadable or malformed models read falls back to the documented defaults', async () => {
|
||||
for (const route of [() => fail(500, null), () => response(200, models({ default_prompt: 4, suggestion_limit: 0, wd14_general_threshold: 2, wd14_character_threshold: 'x' }))]) {
|
||||
const { c } = await opened({ [`GET ${MODELS}`]: route });
|
||||
|
||||
@@ -1145,8 +1145,10 @@ class TestEnvelopes(AnnotationAPITestCase):
|
||||
sorted(await self.json_of(resp)),
|
||||
[
|
||||
"default_prompt",
|
||||
"existing", # #328 adds the saved batch policy and model
|
||||
"loaded_model_type",
|
||||
"model_loaded",
|
||||
"model_type",
|
||||
"models",
|
||||
"success",
|
||||
"suggestion_limit", # B2 adds it (D68); nothing else changes
|
||||
|
||||
@@ -1708,3 +1708,646 @@ class TestLegacyDownloadAttaches(DownloadJobTestCase):
|
||||
self.assertEqual(native.status, 500)
|
||||
self.assertFalse((await native.json())["success"])
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
|
||||
class TestSavedTaggingDefaults(AutotagAPITestCase):
|
||||
"""#328: /models reports the saved defaults; legacy requests without
|
||||
thresholds use the saved ones (OQ4)."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
from py.config import AutotagConfig
|
||||
|
||||
self.config = AutotagConfig
|
||||
for attr in (
|
||||
"MODEL_TYPE",
|
||||
"EXISTING",
|
||||
"GENERAL_THRESHOLD",
|
||||
"CHARACTER_THRESHOLD",
|
||||
"TEMPLATE",
|
||||
"SUGGESTION_LIMIT",
|
||||
):
|
||||
self.addCleanup(setattr, AutotagConfig, attr, getattr(AutotagConfig, attr))
|
||||
|
||||
async def test_models_reports_saved_defaults(self):
|
||||
from py.config import DEFAULT_TAG_TEMPLATE
|
||||
|
||||
data = await (await self.client.get("/prompt_manager/autotag/models")).json()
|
||||
self.assertEqual(
|
||||
{k: data[k] for k in ("model_type", "existing", "default_prompt")},
|
||||
{
|
||||
"model_type": "wd14-swinv2",
|
||||
"existing": "skip",
|
||||
"default_prompt": DEFAULT_TAG_TEMPLATE,
|
||||
},
|
||||
)
|
||||
self.config.MODEL_TYPE = "gguf"
|
||||
self.config.EXISTING = "replace_after_review"
|
||||
self.config.GENERAL_THRESHOLD = 0.4
|
||||
self.config.CHARACTER_THRESHOLD = 0.7
|
||||
self.config.TEMPLATE = "Saved template"
|
||||
self.config.SUGGESTION_LIMIT = 15
|
||||
data = await (await self.client.get("/prompt_manager/autotag/models")).json()
|
||||
self.assertEqual(
|
||||
{
|
||||
k: data[k]
|
||||
for k in (
|
||||
"model_type",
|
||||
"existing",
|
||||
"wd14_general_threshold",
|
||||
"wd14_character_threshold",
|
||||
"default_prompt",
|
||||
"suggestion_limit",
|
||||
)
|
||||
},
|
||||
{
|
||||
"model_type": "gguf",
|
||||
"existing": "replace_after_review",
|
||||
"wd14_general_threshold": 0.4,
|
||||
"wd14_character_threshold": 0.7,
|
||||
"default_prompt": "Saved template",
|
||||
"suggestion_limit": 15,
|
||||
},
|
||||
)
|
||||
|
||||
async def test_single_without_thresholds_uses_saved_ones(self):
|
||||
image = self._make_image(self.output_dir)
|
||||
self.config.GENERAL_THRESHOLD = 0.4
|
||||
self.config.CHARACTER_THRESHOLD = 0.7
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single",
|
||||
{"path": str(image), "model_type": "wd14-vit"},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single",
|
||||
{
|
||||
"path": str(image),
|
||||
"model_type": "wd14-vit",
|
||||
"general_threshold": 0.2,
|
||||
"character_threshold": 0.3,
|
||||
},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(
|
||||
[call[2:] for call in self.service.generate_calls],
|
||||
[(0.4, 0.7), (0.2, 0.3)],
|
||||
)
|
||||
|
||||
async def test_legacy_batch_without_thresholds_uses_saved_ones(self):
|
||||
prompt = self._save_prompt("p")
|
||||
self._link_image(prompt, self._make_image(self.output_dir, "a.png"))
|
||||
self.config.GENERAL_THRESHOLD = 0.45
|
||||
self.config.CHARACTER_THRESHOLD = 0.75
|
||||
resp = await self.client.post(
|
||||
"/prompt_manager/autotag/start?model_type=wd14-vit"
|
||||
)
|
||||
self.assertEqual(self._sse_events(await resp.text())[-1]["type"], "complete")
|
||||
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()
|
||||
with patch("PIL.Image.open") as image_open:
|
||||
resp = await self._preview()
|
||||
self.assertEqual(resp.status, 200, await resp.text())
|
||||
image_open.assert_not_called() # opens no image file
|
||||
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")
|
||||
# Both legacy pages read a stream: the refusal is one error frame
|
||||
# that names the running job.
|
||||
self.assertEqual(legacy.status, 200)
|
||||
self.assertEqual(
|
||||
self._sse_events(await legacy.text()),
|
||||
[
|
||||
{
|
||||
"type": "error",
|
||||
"status": "error",
|
||||
"message": "Tag output images already running",
|
||||
}
|
||||
],
|
||||
)
|
||||
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,
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestTagOutputImagesEdges(BatchJobTestCase):
|
||||
|
||||
async def test_cancel_during_backup_loads_no_model(self):
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
self._seed()
|
||||
PromptManagerConfig.SAFETY_BACKUP = True
|
||||
started, proceed = threading.Event(), threading.Event()
|
||||
|
||||
def slow_backup(model, label="backup"):
|
||||
started.set()
|
||||
proceed.wait(30)
|
||||
return "prompts.db.backup_x"
|
||||
|
||||
with patch("py.api.maintenance.take_safety_backup", side_effect=slow_backup):
|
||||
resp = await self._start_job(existing="merge")
|
||||
job = self.api._jobs.get((await resp.json())["job"]["id"])
|
||||
await asyncio.get_running_loop().run_in_executor(None, started.wait, 30)
|
||||
await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel")
|
||||
proceed.set()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), 30)
|
||||
self.assertEqual(job.state, "cancelled")
|
||||
self.assertEqual(self.service.load_calls, [])
|
||||
self.assertEqual(self.service.generate_calls, [])
|
||||
|
||||
async def test_unexpected_error_fails_the_job_with_a_reason(self):
|
||||
self._seed()
|
||||
with patch.object(
|
||||
self.api.db, "autotag_batch_candidates", side_effect=[[], RuntimeError("x")]
|
||||
):
|
||||
job = await self._run_job(existing="merge")
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(
|
||||
job.status_line,
|
||||
"An internal error occurred. Check server logs for details.",
|
||||
)
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
@@ -19,8 +19,15 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from py import autotag
|
||||
from py.autotag import MODELS, AutoTagService
|
||||
# py.autotag reads its saved defaults from py.config, which imports ComfyUI's
|
||||
# ``server`` at load: stand in for it when this module runs alone.
|
||||
_mock_server = MagicMock()
|
||||
_mock_server.PromptServer.instance.routes = MagicMock()
|
||||
sys.modules.setdefault("server", _mock_server)
|
||||
|
||||
from py import autotag # noqa: E402
|
||||
from py.autotag import MODELS, AutoTagService # noqa: E402
|
||||
from py.config import AutotagConfig # noqa: E402
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
@@ -254,12 +261,19 @@ class TestWd14Thresholding(EngineTestCase):
|
||||
self.assertEqual(list(feeds.keys()), ["input"])
|
||||
self.assertEqual(tuple(feeds["input"].shape), (1, 8, 8, 3))
|
||||
|
||||
def test_service_threshold_setters_clamp(self):
|
||||
service = self._loaded([0.0] * 6)
|
||||
service.wd14_general_threshold = 1.7
|
||||
service.wd14_character_threshold = -2
|
||||
self.assertEqual(service.wd14_general_threshold, 1.0)
|
||||
self.assertEqual(service.wd14_character_threshold, 0.0)
|
||||
def test_service_thresholds_are_the_saved_defaults(self):
|
||||
# #328: one configuration; the service has no thresholds of its own.
|
||||
service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5])
|
||||
with (
|
||||
patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.25),
|
||||
patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.75),
|
||||
):
|
||||
self.assertEqual(
|
||||
(service.wd14_general_threshold, service.wd14_character_threshold),
|
||||
(0.25, 0.75),
|
||||
)
|
||||
tags = service.generate_tags(str(self.image_path))
|
||||
self.assertEqual(tags, ["1girl", "hatsune_miku", "long_hair", "smile"])
|
||||
|
||||
|
||||
class TestGgufNormalisation(EngineTestCase):
|
||||
@@ -369,10 +383,13 @@ class TestTagSuggestions(EngineTestCase):
|
||||
|
||||
def test_thresholds_default_to_the_service_values(self):
|
||||
service = self._wd14()
|
||||
service.wd14_general_threshold = 0.95
|
||||
service.wd14_character_threshold = 0.95
|
||||
|
||||
suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10)
|
||||
with (
|
||||
patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.95),
|
||||
patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.95),
|
||||
):
|
||||
suggestions = service.generate_tag_suggestions(
|
||||
str(self.image_path), limit=10
|
||||
)
|
||||
|
||||
self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)])
|
||||
|
||||
|
||||
+103
-2
@@ -413,9 +413,21 @@ class TestAutotagConfig(unittest.TestCase):
|
||||
self.config_class = getattr(config_module, "AutotagConfig", None)
|
||||
self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented")
|
||||
self._orig = PromptManagerConfig.get_config()
|
||||
self._defaults = {
|
||||
attr: getattr(self.config_class, attr, None)
|
||||
for attr in (
|
||||
"MODEL_TYPE",
|
||||
"EXISTING",
|
||||
"GENERAL_THRESHOLD",
|
||||
"CHARACTER_THRESHOLD",
|
||||
"TEMPLATE",
|
||||
"SUGGESTION_LIMIT",
|
||||
)
|
||||
}
|
||||
|
||||
def tearDown(self):
|
||||
self.config_class.SUGGESTION_LIMIT = 10
|
||||
for attr, value in self._defaults.items():
|
||||
setattr(self.config_class, attr, value)
|
||||
PromptManagerConfig.update_config(self._orig)
|
||||
|
||||
def _load(self, value):
|
||||
@@ -446,8 +458,23 @@ class TestAutotagConfig(unittest.TestCase):
|
||||
self.assertEqual(self.config_class.suggestion_limit(), 10)
|
||||
self.assertEqual(
|
||||
PromptManagerConfig.get_config()["autotag"],
|
||||
{"defaults": {"suggestion_limit": 10}},
|
||||
{
|
||||
"defaults": {
|
||||
"model_type": "wd14-swinv2",
|
||||
"existing": "skip",
|
||||
"general_threshold": 0.35,
|
||||
"character_threshold": 0.85,
|
||||
"suggestion_limit": 10,
|
||||
}
|
||||
},
|
||||
)
|
||||
# The built-in template is not written; a saved one is.
|
||||
self.config_class.TEMPLATE = "Saved."
|
||||
self.assertEqual(
|
||||
PromptManagerConfig.get_config()["autotag"]["defaults"]["template"],
|
||||
"Saved.",
|
||||
)
|
||||
self.config_class.TEMPLATE = config_module.DEFAULT_TAG_TEMPLATE
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = os.path.join(tmpdir, "config.json")
|
||||
PromptManagerConfig.update_config(
|
||||
@@ -466,6 +493,80 @@ class TestAutotagConfig(unittest.TestCase):
|
||||
)
|
||||
self.assertEqual(self.config_class.suggestion_limit(), value)
|
||||
|
||||
def test_autotag_defaults_fall_back_per_key_on_load(self):
|
||||
"""#328: each saved tagging default is checked alone on load; a bad
|
||||
hand-edited value falls back to its default with one warning and
|
||||
leaves the other keys as saved."""
|
||||
good = {
|
||||
"model_type": "gguf",
|
||||
"existing": "merge",
|
||||
"general_threshold": 0.4,
|
||||
"character_threshold": 0.7,
|
||||
"template": "Tags only.",
|
||||
"suggestion_limit": 15,
|
||||
}
|
||||
fallback = {
|
||||
"model_type": "wd14-swinv2",
|
||||
"existing": "skip",
|
||||
"general_threshold": 0.35,
|
||||
"character_threshold": 0.85,
|
||||
"template": config_module.DEFAULT_TAG_TEMPLATE,
|
||||
"suggestion_limit": 10,
|
||||
}
|
||||
bad_values = {
|
||||
"model_type": ["wd14-convnext", 3, None],
|
||||
"existing": ["replace", "", True],
|
||||
"general_threshold": [1.5, -0.1, float("nan"), "0.4", True],
|
||||
"character_threshold": [2, float("inf"), None],
|
||||
# 4000 characters are at most 16,000 UTF-8 bytes, so the 16 KiB
|
||||
# byte cap can only be met past the character cap.
|
||||
"template": ["", " ", "x" * 4001, "\U0001f600" * 4097, "\ud800", 7],
|
||||
}
|
||||
for key, values in bad_values.items():
|
||||
for value in values:
|
||||
with self.subTest(key=key, value=value):
|
||||
saved = dict(good, **{key: value})
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = os.path.join(tmpdir, "config.json")
|
||||
with open(path, "w") as f:
|
||||
json.dump({"autotag": {"defaults": saved}}, f)
|
||||
with self.assertLogs(
|
||||
"prompt_manager.config", level="WARNING"
|
||||
) as logs:
|
||||
PromptManagerConfig.load_from_file(path)
|
||||
self.assertEqual(len(logs.records), 1, logs.output)
|
||||
self.assertIn(f"autotag.defaults.{key}", logs.output[0])
|
||||
loaded = {
|
||||
k: getattr(self.config_class, attr)
|
||||
for k, (attr, _v, _d) in self.config_class._KEYS.items()
|
||||
}
|
||||
self.assertEqual(loaded, dict(good, **{key: fallback[key]}))
|
||||
|
||||
def test_tagging_validators_moved_with_their_messages(self):
|
||||
validate_threshold = config_module.validate_threshold
|
||||
validate_template = config_module.validate_template
|
||||
self.assertEqual(validate_threshold("general_threshold", 1), 1.0)
|
||||
for value, message in (
|
||||
(float("nan"), "general_threshold must be a finite number"),
|
||||
("0.4", "general_threshold must be a finite number"),
|
||||
(1.01, "general_threshold must be between 0 and 1"),
|
||||
):
|
||||
with self.assertRaises(ValueError) as caught:
|
||||
validate_threshold("general_threshold", value)
|
||||
self.assertEqual(str(caught.exception), message)
|
||||
self.assertEqual(validate_template("Tags."), "Tags.")
|
||||
for value, message in (
|
||||
(" ", "instructions must be non-empty text"),
|
||||
("x" * 4001, "instructions must be at most 4000 characters"),
|
||||
("\ud800", "instructions must be valid Unicode text"),
|
||||
):
|
||||
with self.assertRaises(ValueError) as caught:
|
||||
validate_template(value)
|
||||
self.assertEqual(str(caught.exception), message)
|
||||
self.assertEqual(
|
||||
config_module.MODEL_TYPES, {"gguf", "hf", "wd14-swinv2", "wd14-vit"}
|
||||
)
|
||||
|
||||
def test_suggestion_limit_setting_validated_on_load(self):
|
||||
for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]):
|
||||
with self.subTest(value=value):
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -976,3 +976,129 @@ class TestSourceFormError(unittest.TestCase):
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
class TestAutotagDefaultsSettings(SourcesTestCase):
|
||||
"""#328: the tagging defaults save through POST /settings like the others."""
|
||||
|
||||
ATTRS = (
|
||||
"MODEL_TYPE",
|
||||
"EXISTING",
|
||||
"GENERAL_THRESHOLD",
|
||||
"CHARACTER_THRESHOLD",
|
||||
"TEMPLATE",
|
||||
"SUGGESTION_LIMIT",
|
||||
)
|
||||
SAVED = {
|
||||
"autotag_model_type": "gguf",
|
||||
"autotag_existing": "merge",
|
||||
"autotag_general_threshold": 0.4,
|
||||
"autotag_character_threshold": 0.8,
|
||||
"autotag_template": "Tags only.",
|
||||
"autotag_suggestion_limit": 15,
|
||||
}
|
||||
|
||||
async def asyncSetUp(self):
|
||||
await super().asyncSetUp()
|
||||
from py.config import AutotagConfig
|
||||
|
||||
self.autotag = AutotagConfig
|
||||
for attr in self.ATTRS:
|
||||
self.addCleanup(setattr, AutotagConfig, attr, getattr(AutotagConfig, attr))
|
||||
|
||||
def _runtime(self):
|
||||
return {attr: getattr(self.autotag, attr) for attr in self.ATTRS}
|
||||
|
||||
async def test_autotag_defaults_save_round_trip(self):
|
||||
self._write_config({"custom": {"keep": True}})
|
||||
revision = (await self._get())["revision"]
|
||||
status, data = await self._post(
|
||||
"/prompt_manager/settings", {"expected_revision": revision, **self.SAVED}
|
||||
)
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(data["restart_required_fields"], [])
|
||||
saved = self._read_config()
|
||||
self.assertEqual(
|
||||
saved["autotag"]["defaults"],
|
||||
{key[len("autotag_") :]: value for key, value in self.SAVED.items()},
|
||||
)
|
||||
self.assertEqual(saved["custom"], {"keep": True})
|
||||
self.assertEqual(
|
||||
self._runtime(),
|
||||
{
|
||||
"MODEL_TYPE": "gguf",
|
||||
"EXISTING": "merge",
|
||||
"GENERAL_THRESHOLD": 0.4,
|
||||
"CHARACTER_THRESHOLD": 0.8,
|
||||
"TEMPLATE": "Tags only.",
|
||||
"SUGGESTION_LIMIT": 15,
|
||||
},
|
||||
)
|
||||
settings = (await self._get())["settings"]
|
||||
self.assertEqual({key: settings[key] for key in self.SAVED}, self.SAVED)
|
||||
|
||||
async def test_stale_revision_changes_no_default(self):
|
||||
revision = (await self._get())["revision"]
|
||||
self._write_config({"custom": {"edited": "elsewhere"}})
|
||||
before = (self._config_bytes(), self._runtime())
|
||||
status, data = await self._post(
|
||||
"/prompt_manager/settings", {"expected_revision": revision, **self.SAVED}
|
||||
)
|
||||
self.assertEqual(status, 409, data)
|
||||
self.assertEqual(data["code"], "settings-changed")
|
||||
self.assertEqual((self._config_bytes(), self._runtime()), before)
|
||||
|
||||
async def test_bad_default_message_names_the_setting(self):
|
||||
cases = (
|
||||
(
|
||||
"autotag_general_threshold",
|
||||
1.5,
|
||||
"General-tag confidence must be between 0 and 1.",
|
||||
),
|
||||
(
|
||||
"autotag_character_threshold",
|
||||
"0.8",
|
||||
"Character-tag confidence must be a finite number.",
|
||||
),
|
||||
(
|
||||
"autotag_character_threshold",
|
||||
True,
|
||||
"Character-tag confidence must be a finite number.",
|
||||
),
|
||||
(
|
||||
"autotag_model_type",
|
||||
"wd14-convnext",
|
||||
"Tagging model must be one of gguf, hf, wd14-swinv2, wd14-vit.",
|
||||
),
|
||||
(
|
||||
"autotag_existing",
|
||||
"replace",
|
||||
"Existing tags must be skip, merge or replace_after_review.",
|
||||
),
|
||||
("autotag_template", " ", "LLM tag template must be non-empty text."),
|
||||
(
|
||||
"autotag_template",
|
||||
"x" * 4001,
|
||||
"LLM tag template must be at most 4000 characters.",
|
||||
),
|
||||
(
|
||||
"autotag_suggestion_limit",
|
||||
0,
|
||||
"Suggestions per image must be a whole number from 1 to 100.",
|
||||
),
|
||||
(
|
||||
"autotag_suggestion_limit",
|
||||
2.5,
|
||||
"Suggestions per image must be a whole number from 1 to 100.",
|
||||
),
|
||||
)
|
||||
self._write_config({"custom": {"keep": True}})
|
||||
before = (self._config_bytes(), self._runtime())
|
||||
for key, value, message in cases:
|
||||
with self.subTest(key=key, value=value):
|
||||
status, data = await self._post(
|
||||
"/prompt_manager/settings", {**self.SAVED, key: value}
|
||||
)
|
||||
self.assertEqual(status, 400, data)
|
||||
self.assertEqual(data["error"], message)
|
||||
self.assertEqual((self._config_bytes(), self._runtime()), before)
|
||||
|
||||
@@ -3970,6 +3970,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
}
|
||||
|
||||
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();
|
||||
document.getElementById('reviewImageName').textContent = filename || 'Unknown';
|
||||
document.getElementById('reviewTagsVisible').innerHTML = '<div class="text-pm-secondary">Generating tags...</div>';
|
||||
document.getElementById('reviewTagsHidden').innerHTML = '';
|
||||
|
||||
+13
-17
@@ -3547,29 +3547,25 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
this.autoTagState.eventSource.onmessage = (event) => {
|
||||
const data = JSON.parse(event.data);
|
||||
|
||||
if (data.status === 'loading_model') {
|
||||
document.getElementById('autoTagCurrentFile').textContent = 'Loading model...';
|
||||
} else if (data.status === 'processing') {
|
||||
const percent = Math.round((data.processed / data.total) * 100);
|
||||
document.getElementById('autoTagCurrentFile').textContent = data.current_file || 'Processing...';
|
||||
document.getElementById('autoTagProgressPercent').textContent = `${percent}%`;
|
||||
document.getElementById('autoTagProgressBar').style.width = `${percent}%`;
|
||||
document.getElementById('autoTagProcessed').textContent = data.processed;
|
||||
document.getElementById('autoTagApplied').textContent = data.tags_applied;
|
||||
document.getElementById('autoTagSkipped').textContent = data.skipped;
|
||||
} else if (data.status === 'complete') {
|
||||
// The server sends `type` frames, as the admin page reads them (#328).
|
||||
if (data.type === 'progress') {
|
||||
document.getElementById('autoTagCurrentFile').textContent = data.status || 'Processing...';
|
||||
document.getElementById('autoTagProgressPercent').textContent = `${data.progress || 0}%`;
|
||||
document.getElementById('autoTagProgressBar').style.width = `${data.progress || 0}%`;
|
||||
if (data.processed !== undefined) {
|
||||
document.getElementById('autoTagProcessed').textContent = data.processed;
|
||||
document.getElementById('autoTagApplied').textContent = data.tagged || 0;
|
||||
document.getElementById('autoTagSkipped').textContent = data.skipped || 0;
|
||||
}
|
||||
} else if (data.type === 'complete') {
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification(`Auto tagging complete! Applied tags to ${data.tags_applied} prompts.`, 'success');
|
||||
this.showNotification(`Auto tagging complete! Applied tags to ${data.tagged || 0} prompts.`, 'success');
|
||||
this.loadImages(); // Refresh the gallery
|
||||
} else if (data.status === 'error') {
|
||||
} else if (data.type === 'error') {
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification(`Auto tag error: ${data.message}`, 'error');
|
||||
} else if (data.status === 'cancelled') {
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification('Auto tagging cancelled', 'info');
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -104,6 +104,24 @@ function maintenanceFinish(job) {
|
||||
// D85: an image tagging job says which of its two modes started.
|
||||
const IMAGE_TAG_START_TOASTS = Object.freeze({ add: 'Tagging started.', auto: 'Auto-tagging started.' });
|
||||
|
||||
/** D89: the image tagging job's view, shared with Tag output images (#328). */
|
||||
const IMAGE_TAGS = Object.freeze({
|
||||
...GENERIC,
|
||||
pillVerb: 'Tagging',
|
||||
startToast: job => IMAGE_TAG_START_TOASTS[job.scope?.mode] ?? titleStarted(job),
|
||||
finishToast: namedFinish(job => job.title, taggedOf, tagLine),
|
||||
line: tagLine,
|
||||
kept: taggedOf,
|
||||
scope: imagesOf,
|
||||
tiles: tiles([['tagged', 'Tagged', 'added'], ['unchanged', 'Unchanged', ''], ['skipped', 'Skipped', 'warn'],
|
||||
['errors', 'Errors', 'bad']]),
|
||||
outcomes: tiles([['tagged', 'Tagged'], ['errors', 'Errors'], ['skipped', 'Skipped'], ['unchanged', 'Unchanged']]),
|
||||
openable: groups('tagged', 'unchanged'),
|
||||
log: 'Tagging log',
|
||||
cancelLabel: 'Cancel tagging',
|
||||
cancel: 'Cancel tagging',
|
||||
});
|
||||
|
||||
/** Every kind with a descriptor of its own. */
|
||||
export const KINDS = Object.freeze({
|
||||
__proto__: null,
|
||||
@@ -125,22 +143,10 @@ export const KINDS = Object.freeze({
|
||||
cancel: 'Cancel scan',
|
||||
}),
|
||||
// D89: Tagged (green), Unchanged, Skipped (warning) and Errors (red); Tagged and Unchanged rows open the viewer.
|
||||
'image-tags': Object.freeze({
|
||||
...GENERIC,
|
||||
pillVerb: 'Tagging',
|
||||
startToast: job => IMAGE_TAG_START_TOASTS[job.scope?.mode] ?? titleStarted(job),
|
||||
finishToast: namedFinish(job => job.title, taggedOf, tagLine),
|
||||
line: tagLine,
|
||||
kept: taggedOf,
|
||||
scope: imagesOf,
|
||||
tiles: tiles([['tagged', 'Tagged', 'added'], ['unchanged', 'Unchanged', ''], ['skipped', 'Skipped', 'warn'],
|
||||
['errors', 'Errors', 'bad']]),
|
||||
outcomes: tiles([['tagged', 'Tagged'], ['errors', 'Errors'], ['skipped', 'Skipped'], ['unchanged', 'Unchanged']]),
|
||||
openable: groups('tagged', 'unchanged'),
|
||||
log: 'Tagging log',
|
||||
cancelLabel: 'Cancel tagging',
|
||||
cancel: 'Cancel tagging',
|
||||
}),
|
||||
'image-tags': IMAGE_TAGS,
|
||||
// #328 (extends D111): Tag output images, the prompt-tag batch, counts its images the same way. Its legacy-page
|
||||
// hold counts nothing and is not cancellable, so `viewOf` gives it the generic held view.
|
||||
autotag: IMAGE_TAGS,
|
||||
// D111: Generated (green), Skipped (already fresh) and Errors (red); a row with a preview opens the viewer.
|
||||
thumbnails: Object.freeze({
|
||||
...GENERIC,
|
||||
|
||||
@@ -179,14 +179,15 @@ function validateReceipt(value) {
|
||||
}
|
||||
|
||||
/**
|
||||
* Thresholds, instructions and suggestion count from the server, else the documented defaults, and the models'
|
||||
* download status (`models`, null when unknown).
|
||||
* Thresholds, instructions and suggestion count from the server, else the documented defaults, the models'
|
||||
* download status (`models`, null when unknown) and the saved tagging model (`model`, null when unknown; #328).
|
||||
*/
|
||||
async function readDefaults(api, signal) {
|
||||
try {
|
||||
const value = await api.get(MODELS_URL, {signal, cache: 'no-store'});
|
||||
if (!object(value)) return {defaults: DEFAULTS, models: null};
|
||||
if (!object(value)) return {defaults: DEFAULTS, models: null, model: null};
|
||||
return {
|
||||
model: typeof value.model_type === 'string' && modelOf(value.model_type) ? value.model_type : null,
|
||||
defaults: {
|
||||
general_threshold: unit(value.wd14_general_threshold) ? value.wd14_general_threshold : DEFAULTS.general_threshold,
|
||||
character_threshold: unit(value.wd14_character_threshold) ? value.wd14_character_threshold : DEFAULTS.character_threshold,
|
||||
@@ -197,7 +198,7 @@ async function readDefaults(api, signal) {
|
||||
};
|
||||
} catch {
|
||||
// The defaults keep the dialog usable; Generate reports a real failure itself.
|
||||
return {defaults: DEFAULTS, models: null};
|
||||
return {defaults: DEFAULTS, models: null, model: null};
|
||||
}
|
||||
}
|
||||
|
||||
@@ -263,6 +264,7 @@ export function createImageTaggingController({api, target, source = 'images', on
|
||||
requireValue(SOURCES.has(source), 'Unknown tagging source');
|
||||
requireValue(typeof onLateSave === 'function', 'onLateSave must be a function');
|
||||
const chosen = validatePreset(preset);
|
||||
let modelPicked = false; // the user chose a model: the saved one no longer applies
|
||||
|
||||
let state = freeze(derive(initialState(source, chosen)));
|
||||
const listeners = new Set();
|
||||
@@ -336,8 +338,11 @@ export function createImageTaggingController({api, target, source = 'images', on
|
||||
}
|
||||
capture = {token: value.capture.token, sha256: value.capture.sha256};
|
||||
epoch = value.library_epoch;
|
||||
({defaults} = await defaultsRead);
|
||||
const read = await defaultsRead;
|
||||
({defaults} = read);
|
||||
if (closed || serial !== openSerial) return;
|
||||
// The saved model (System › Auto-tagging) unless a preset or the user already chose one (#328).
|
||||
const savedModel = !chosen && !modelPicked && read.model ? {model: read.model} : {};
|
||||
observedOnce = true;
|
||||
const {width, height, media_type: mediaType, preview_url: previewUrl, generate} = value.capture;
|
||||
// A kept draft stays as typed unless another writer changed the set meanwhile: then the set leads it (I2).
|
||||
@@ -350,7 +355,7 @@ export function createImageTaggingController({api, target, source = 'images', on
|
||||
savedTags: value.tags.slice(), tagRevision: value.tag_revision,
|
||||
manual: draft === null ? value.tags.join(', ') : changed ? mergeCurrent(value.tags, draft) : draft,
|
||||
draftRevision: state.draftRevision + 1,
|
||||
options: {...defaults, ...chosen?.options}, optionsRevision: state.optionsRevision + 1,
|
||||
...savedModel, options: {...defaults, ...chosen?.options}, optionsRevision: state.optionsRevision + 1,
|
||||
generateSupported: generate.supported, generateReason: generate.reason,
|
||||
statusText: changed ? `${COPY.changedElsewhere} ${savedNow(value.tags)}${reason ? ` ${reason}` : ''}` : reason,
|
||||
});
|
||||
@@ -362,7 +367,9 @@ export function createImageTaggingController({api, target, source = 'images', on
|
||||
|
||||
function setModel(id) {
|
||||
requireValue(modelOf(id), 'Unknown tagging model');
|
||||
if (closed || id === state.model) return;
|
||||
if (closed) return;
|
||||
modelPicked = true;
|
||||
if (id === state.model) return;
|
||||
stopGenerating();
|
||||
const phase = state.phase === 'generating' || state.phase === 'generate-error' ? 'ready' : state.phase;
|
||||
set({
|
||||
@@ -535,6 +542,7 @@ export function createTaggingOptionsController({api} = {}) {
|
||||
let defaults = DEFAULTS;
|
||||
let models = null;
|
||||
let abort = null;
|
||||
let modelPicked = false;
|
||||
const deriveOptions = value => ({...value, canGenerate: false, unavailable: modelUnavailable(models, value.model),
|
||||
canStart: value.phase === 'ready' && !modelUnavailable(models, value.model)});
|
||||
let state = freeze(deriveOptions({phase: 'loading', model: MODEL_OPTIONS[0].id, options: {...DEFAULTS}, optionsRevision: 0}));
|
||||
@@ -550,11 +558,15 @@ export function createTaggingOptionsController({api} = {}) {
|
||||
const read = await readDefaults(api, abort.signal);
|
||||
if (closed) return;
|
||||
({defaults, models} = read);
|
||||
set({phase: 'ready', options: {...defaults}, optionsRevision: state.optionsRevision + 1});
|
||||
// The saved model (System › Auto-tagging) unless the user already chose one (#328).
|
||||
const savedModel = !modelPicked && read.model ? {model: read.model} : {};
|
||||
set({phase: 'ready', ...savedModel, options: {...defaults}, optionsRevision: state.optionsRevision + 1});
|
||||
}
|
||||
function setModel(id) {
|
||||
requireValue(modelOf(id), 'Unknown tagging model');
|
||||
if (closed || id === state.model) return;
|
||||
if (closed) return;
|
||||
modelPicked = true;
|
||||
if (id === state.model) return;
|
||||
// As in the Tag image dialog: the model's fields start from the defaults again; the count is kept (D91).
|
||||
set({model: id, optionsRevision: state.optionsRevision + 1,
|
||||
options: {...state.options, general_threshold: defaults.general_threshold,
|
||||
|
||||
Reference in New Issue
Block a user