From 88fa13689ab25dfb7a53d955bd3d57a2cbc494ba Mon Sep 17 00:00:00 2001 From: Vito Date: Mon, 26 Oct 2026 23:49:15 -0700 Subject: [PATCH 1/6] feat(system): model status, a named model gate and downloads as Activity jobs (#328, PR 1) (#365) * feat(autotag): model status without paths, and a busy answer that names the operation * feat(autotag): model downloads run as Activity jobs * fix(autotag): keep each download route's error shape and harden the fake download --- py/api/autotag_routes.py | 293 +++++++++++--- py/api/job_routes.py | 4 + py/api/jobs.py | 2 + py/autotag.py | 7 +- tests/browser/legacy-autotag.spec.mjs | 89 +++++ tests/browser/server.py | 5 +- tests/browser/workspace-activity.spec.mjs | 30 ++ tests/browser/workspace_tagging_fixture.py | 52 +++ tests/test_api_autotag.py | 426 ++++++++++++++++++++- tests/test_job_admission.py | 2 +- web/admin.html | 3 +- web/gallery.html | 3 +- web/js/admin.js | 18 +- web/js/gallery.js | 18 +- 14 files changed, 876 insertions(+), 76 deletions(-) create mode 100644 tests/browser/legacy-autotag.spec.mjs diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 6114bb6d..e5521dd5 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -1,14 +1,18 @@ """AutoTag API routes for PromptManager.""" import asyncio +import contextlib +import functools import json import os +import re import threading from pathlib import Path from typing import Optional from urllib.parse import quote -from .jobs import AUTOTAG_BATCH_JOB_KIND +from .job_routes import busy_response +from .jobs import AUTOTAG_BATCH_JOB_KIND, MODEL_DOWNLOAD_JOB_KIND, JobBusy from .prompts import thumbnail_url_for from . import thumbnail_settings from .. import model_gate @@ -22,14 +26,86 @@ MODEL_BUSY_ERROR = ( ) +# Who holds the model gate, as a 409 model-busy names it (``operation``). An +# owner is ``":"`` (``model_gate.new_owner``); a download names +# its own job, whose title says which model. +MODEL_DOWNLOAD_OWNER_KIND = "model-download" +OWNER_LABELS = { + "autotag-job": "Tag output images", + "legacy-batch": "Tag output images (legacy page)", + "legacy-single": "Tagging one image (legacy page)", + "legacy-unload": "Unloading the tagging model", + "image-tag-job": "Auto-tag selected images", + "image-annotations": "Generate suggestions", +} +ANOTHER_OPERATION = "Another tagging operation" +_DOWNLOAD_FAILED = "The download failed. Check server logs for details." +# POST /jobs/model-download carries only {model_type}. +MODEL_DOWNLOAD_BODY_MAX = 1024 +# The running download's gate owner and job title. Only one download runs at a +# time and each needs the gate, so a new start replaces the entry. +_download_titles = {} + + +def _holder_label(holder: Optional[str]) -> str: + if not holder: + return ANOTHER_OPERATION + kind = holder.split(":", 1)[0] + if kind == MODEL_DOWNLOAD_OWNER_KIND: + return _download_titles.get(holder, ANOTHER_OPERATION) + return OWNER_LABELS.get(kind, ANOTHER_OPERATION) + + def model_busy_response() -> web.Response: - """``409 model-busy``: another operation holds the shared model gate.""" + """``409 model-busy``: another operation holds the shared model gate. + + ``operation`` names the holder (``OWNER_LABELS``), or "Another tagging + operation" when the gate was freed meanwhile or the kind is unknown. + """ return web.json_response( - {"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR}, + { + "success": False, + "code": "model-busy", + "error": MODEL_BUSY_ERROR, + "operation": _holder_label(model_gate.holder()), + }, status=409, ) +_SIZE = re.compile(r"^(~?[0-9.]+)\s*([KMGT]?B)$") + + +def size_label(size: str) -> str: + """The engine table's size with a space before the unit (``~400 MB``).""" + match = _SIZE.match(str(size or "").strip()) + return f"{match.group(1)} {match.group(2)}" if match else str(size or "") + + +def public_models(status: dict, loaded_type: Optional[str]) -> dict: + """Each engine model's public entry: names and sizes from ``MODELS``, the + flags from ``status``; never a path (``model_path`` is dropped).""" + from ..autotag import MODELS + + models = {} + for model_type, config in MODELS.items(): + flags = status.get(model_type) or {} + entry = { + "name": config["name"], + "label": config["label"], + "description": config["description"], + "size": config["size"], + "size_label": size_label(config["size"]), + "downloaded": bool(flags.get("downloaded")), + "loaded": loaded_type == model_type, + } + for key in ("model_exists", "mmproj_exists"): + if key in flags: + entry[key] = bool(flags[key]) + models[model_type] = entry + return models + + # Only these may be handed to the tagging engine or listed by the scanner. _MEDIA_EXTENSIONS = frozenset( {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"} @@ -65,6 +141,11 @@ def path_is_within(candidate, root) -> bool: return child != parent and child.is_relative_to(parent) +async def _single(event): + """A stream of one event.""" + yield event + + async def _refusal(title): """The single error event of an auto-tag batch refused by another job.""" message = f"{title} already running" @@ -110,15 +191,18 @@ class AutotagRoutesMixin: from ..config import AutotagConfig service = get_autotag_service() - models_status = service.get_models_status() + # Off the loop: the HuggingFace check walks the whole HF cache. + status = await self._run_in_executor(service.get_models_status) + loaded = service.is_model_loaded() + loaded_type = service.get_loaded_model_type() if loaded else None return web.json_response( { "success": True, - "models": models_status, + "models": public_models(status, loaded_type), "default_prompt": service.default_prompt, - "model_loaded": service.is_model_loaded(), - "loaded_model_type": service.get_loaded_model_type(), + "model_loaded": loaded, + "loaded_model_type": loaded_type, "wd14_general_threshold": service.wd14_general_threshold, "wd14_character_threshold": service.wd14_character_threshold, "suggestion_limit": AutotagConfig.suggestion_limit(), @@ -159,55 +243,156 @@ class AutotagRoutesMixin: return response async def download_autotag_model(self, request): - """Download an AutoTag model with streaming progress (POST only).""" + """Download a tagging model, streaming the legacy events (POST only). + + The download is the model_download Activity job: this request + attaches to a running download of the same model or starts one, and + streams the legacy frames until the job ends. A client disconnect + ends only the stream; the download goes on. A running download of + another model answers 409 model-busy (it holds the gate). + """ + from ..autotag import MODELS, get_autotag_service + model_type = request.match_info.get("model_type") - owner = model_gate.new_owner("legacy-download") - - async def events(): - try: - from ..autotag import get_autotag_service - - service = get_autotag_service() - if model_type not in service.models_config: - yield { - "type": "error", - "message": f"Invalid model type: {model_type}", - } - return - - yield { - "type": "progress", - "progress": 0, - "status": "Starting download...", - } - - def progress_callback(status: str, progress: float): - pass - - success = await model_gate.run( - owner, service.download_model, model_type, progress_callback - ) - if success: - yield { - "type": "complete", - "progress": 100, - "status": "Download complete", - } - else: - yield {"type": "error", "message": "Download failed"} - except Exception: - self.logger.exception("Download model error") - yield { - "type": "error", - "message": _INTERNAL_ERROR, - } - - if not model_gate.try_acquire(owner): - return model_busy_response() + if model_type not in MODELS: + invalid = {"type": "error", "message": f"Invalid model type: {model_type}"} + return await self._stream_sse(request, _single(invalid)) try: - return await self._stream_sse(request, events()) - finally: + job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND) + if job is not None and job.scope.get("model_type") != model_type: + return model_busy_response() + if job is None: + job, refusal = self._start_model_download( + model_type, get_autotag_service() + ) + if refusal is not None: + return refusal + except Exception: + self.logger.exception("Download model error") + failed = {"type": "error", "message": _INTERNAL_ERROR} + return await self._stream_sse(request, _single(failed)) + return await self._stream_sse(request, self._download_frames(job)) + + @staticmethod + async def _download_frames(job): + """The legacy download frames for job: start, then its end.""" + yield {"type": "progress", "progress": 0, "status": "Starting download..."} + async with contextlib.aclosing(job.events()) as events: + async for _payload in events: + pass + if job.state == "finished": + yield {"type": "complete", "progress": 100, "status": "Download complete"} + else: + yield {"type": "error", "message": "Download failed"} + + # ── model download job (#328) ─────────────────────────────────────── + + async def start_model_download_job(self, request): + """``POST /jobs/model-download {model_type}``: 202 with the job. + + 400 for a bad body, 409 ``busy`` while a download (or an exclusive + job) runs, 409 ``model-busy`` while another operation holds the gate. + """ + data, error = await self._read_json_body( + request, max_bytes=MODEL_DOWNLOAD_BODY_MAX + ) + if error is not None: + return error + if set(data) != {"model_type"}: + return self._bad_request("The body must be exactly {model_type}") + from ..autotag import MODELS, get_autotag_service + + model_type = data["model_type"] + if not isinstance(model_type, str) or model_type not in MODELS: + return self._bad_request( + "model_type must be one of " + ", ".join(sorted(MODELS)) + ) + try: + job, refusal = self._start_model_download(model_type, get_autotag_service()) + except Exception as e: + self.logger.exception("Model download start error") + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) + if refusal is not None: + return refusal + return web.json_response({"success": True, "job": job.summary()}, status=202) + + def _start_model_download(self, model_type, service): + """Register the download job: ``(job, None)`` or ``(None, refusal)``. + + Synchronous (no await between the checks and the start). The gate is + taken before the job registers and released when the runner ends, + also when the task is cancelled before the runner ever starts. + """ + from ..autotag import MODELS + + busy = self._jobs.blocker(MODEL_DOWNLOAD_JOB_KIND, admission="read") + if busy is not None: + return None, busy_response(JobBusy(busy)) + owner = model_gate.new_owner(MODEL_DOWNLOAD_OWNER_KIND) + if not model_gate.try_acquire(owner): + return None, model_busy_response() + config = MODELS[model_type] + label = config["label"] + title = f"Tagging model download · {label}" + _download_titles.clear() + _download_titles[owner] = title + try: + job = self._jobs.start( + MODEL_DOWNLOAD_JOB_KIND, + title, + admission="read", + cancellable=False, + scope={ + "model_type": model_type, + "label": label, + "size_label": size_label(config["size"]), + }, + runner=functools.partial( + self._run_model_download, service, model_type, owner + ), + ) + except BaseException as refused: + # JobBusy, or anything else the registry raised: give the gate back. model_gate.release(owner) + _download_titles.pop(owner, None) + if isinstance(refused, JobBusy): + return None, busy_response(refused) + raise + + def ended(_task): + model_gate.release(owner) + if model_gate.holder() != owner: # not while a thread still runs + _download_titles.pop(owner, None) + + job.task.add_done_callback(ended) + return job, None + + async def _run_model_download(self, service, model_type, owner, job): + label = job.scope["label"] + job.update(status_line=f"Downloading {label} ({job.scope['size_label']})") + job.log_line(f"Download started: {label}") + job.publish({"type": "status", "status": job.status_line}) + try: + ok = await model_gate.run(owner, service.download_model, model_type, None) + except Exception: + self.logger.exception(f"Tagging model download failed: {model_type}") + ok = False + finally: + # Before the job ends, so a reader that sees the end finds the + # gate free (a running download thread defers the release). + model_gate.release(owner) + if ok: + job.update(status_line=f"Downloaded {label}") + job.log_line(f"Download finished: {label}") + job.publish({"type": "status", "status": job.status_line}) + job.finish("finished") + return + job.update(status_line=_DOWNLOAD_FAILED) + job.log_line(f"Download failed: {label}") + job.publish({"type": "status", "status": job.status_line}) + job.finish("failed") async def _read_start_params(self, request) -> Optional[dict]: """Merge query-string and optional JSON-body parameters for autotag/start. @@ -467,7 +652,7 @@ class AutotagRoutesMixin: return await self._stream_sse(request, _refusal(busy.title)) async with self._jobs.hold( AUTOTAG_BATCH_JOB_KIND, - "Auto-tag images (legacy page)", + OWNER_LABELS["legacy-batch"], admission="records", ) as job: response = await self._start_autotag_batch(request, opts, events) diff --git a/py/api/job_routes.py b/py/api/job_routes.py index 47056a11..59cadb6a 100644 --- a/py/api/job_routes.py +++ b/py/api/job_routes.py @@ -130,6 +130,10 @@ class JobRoutesMixin: async def start_image_tag_job_route(request): return await self.start_image_tag_job(request) + @routes.post("/prompt_manager/jobs/model-download") + async def start_model_download_job_route(request): + return await self.start_model_download_job(request) + @routes.post("/prompt_manager/jobs/thumbnails") async def start_thumbnail_job_route(request): return await self.start_thumbnail_job(request) diff --git a/py/api/jobs.py b/py/api/jobs.py index 9e5c463b..a385f3b3 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -109,6 +109,8 @@ DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup" THUMBNAILS_JOB_KIND = "thumbnails" LORA_IMPORT_JOB_KIND = "lora_import" AUTOTAG_BATCH_JOB_KIND = "autotag" +# A tagging model download (#328): read admission, not cancellable. +MODEL_DOWNLOAD_JOB_KIND = "model_download" # Duplicate image files (#249 System S6, py/api/duplicate_files.py): the # find is a cancellable ``read`` job; the reviewed removal is ``records``. DUPLICATES_JOB_KIND = "duplicates" diff --git a/py/autotag.py b/py/autotag.py index 68c72a5b..4d6d4da8 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -45,10 +45,12 @@ except ImportError: return logger -# Model configurations - same as standalone_tagger.py +# Model configurations - same as standalone_tagger.py. ``label`` is the short +# name System, Activity and a 409 model-busy show (#328). MODELS = { "gguf": { "name": "GGUF (Recommended)", + "label": "GGUF JoyCaption", "description": "Quantized model, ~4.5GB, faster loading", "size": "~4.5GB", "repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF", @@ -59,6 +61,7 @@ MODELS = { }, "hf": { "name": "HuggingFace", + "label": "HuggingFace JoyCaption", "description": "Full precision model, ~16GB, higher quality", "size": "~16GB", "repo": "fancyfeast/llama-joycaption-beta-one-hf-llava", @@ -66,6 +69,7 @@ MODELS = { }, "wd14-swinv2": { "name": "WD14 SwinV2", + "label": "WD14 SwinV2", "description": "Fast classifier tagger, ~400MB, no prompt needed", "size": "~400MB", "repo": "SmilingWolf/wd-swinv2-tagger-v3", @@ -75,6 +79,7 @@ MODELS = { }, "wd14-vit": { "name": "WD14 ViT", + "label": "WD14 ViT", "description": "Fast classifier tagger, ~400MB, ViT architecture", "size": "~400MB", "repo": "SmilingWolf/wd-vit-tagger-v3", diff --git a/tests/browser/legacy-autotag.spec.mjs b/tests/browser/legacy-autotag.spec.mjs new file mode 100644 index 00000000..a928ac04 --- /dev/null +++ b/tests/browser/legacy-autotag.spec.mjs @@ -0,0 +1,89 @@ +// #328: the legacy admin and gallery auto-tag tools keep working beside System › Auto-tagging. The model download is +// the `model_download` Activity job; the pages' stream attaches to it. No model is ever downloaded: the fixture's +// fake download waits until the test releases it (`/__test__/tagging/download`). +import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs'; +import { test, expect } from './fixtures.mjs'; + +const TITLE = 'Tagging model download · WD14 ViT'; + +async function post(request, url, data) { + await expect(await request.post(url, data === undefined ? {} : { data })).toBeOK(); +} + +async function downloads(request) { + const response = await request.get('/prompt_manager/jobs'); + await expect(response).toBeOK(); + return (await response.json()).jobs.filter(job => job.kind === 'model_download'); +} + +test.beforeEach(async ({ page, request, browserName }) => { + await startCoverage(page, browserName); + // The 3.2.0 update notice would cover the admin page's toolbar. + await page.addInitScript(() => { localStorage.setItem('pm_update_notice_dismissed', '3.2.0'); }); + await post(request, '/__test__/jobs/restart'); + await post(request, '/__test__/tagging/reset'); + await post(request, '/__test__/tagging/models', { unavailable: ['wd14-vit'] }); + await post(request, '/__test__/tagging/download', { mode: 'hold' }); +}); + +test.afterEach(async ({ page, request, browserName }) => { + await saveCoverage(page, browserName); + // Releases a download a failed test left held, and puts the raising default back. + await post(request, '/__test__/tagging/reset'); + await post(request, '/__test__/jobs/restart'); +}); + +async function startDownload(page, request) { + await page.locator('#autoTagBtn').click(); + const status = page.locator('#wd14VitModelStatus'); + await status.getByRole('button', { name: 'Download', exact: true }).click(); + const modal = page.locator('#autoTagDownloadModal'); + await expect(modal).toBeVisible(); + await expect(page.locator('#downloadStatus')).toHaveText('Downloading…'); + await expect(page.locator('#downloadProgressBar')).toHaveClass(/animate-pulse/); + await expect(modal).toContainText('The download continues in Activity.'); + await expect(page.locator('#cancelDownloadBtn')).toHaveText('Close'); + const running = (await downloads(request)).filter(job => job.state === 'running'); + expect(running.map(job => [job.title, job.cancellable])).toEqual([[TITLE, false]]); + return { modal, status }; +} + +test('admin download modal closes on complete and Activity lists the job without Cancel', async ({ page, request }) => { + await page.goto('/prompt_manager/admin'); + const { modal, status } = await startDownload(page, request); + + await page.goto('/prompt_manager/workspace#/activity'); + const item = page.locator('#workspace-main section.surface > .activity-item') + .filter({ has: page.getByRole('link', { name: new RegExp(`${TITLE}$`) }) }); + await expect(item.locator('.chip')).toHaveText('Running'); + await expect(item.locator('progress')).not.toHaveAttribute('value'); + await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0); + + // Back on the admin page a second download of the same model attaches to the running job. + await page.goto('/prompt_manager/admin'); + await startDownload(page, request); + await post(request, '/__test__/tagging/download', { mode: 'succeed' }); + await expect(modal).toBeHidden(); + await expect(status).toContainText('Downloaded'); + expect((await downloads(request)).map(job => [job.title, job.state])).toEqual([[TITLE, 'finished']]); +}); + +test('gallery download modal closes on complete', async ({ page, request }) => { + await page.goto('/prompt_manager/gallery'); + const { modal, status } = await startDownload(page, request); + await post(request, '/__test__/tagging/download', { mode: 'succeed' }); + await expect(modal).toBeHidden(); + await expect(page.locator('#autoTagModal')).toBeVisible(); + await expect(status).toContainText('Downloaded'); + expect((await downloads(request)).map(job => job.state)).toEqual(['finished']); +}); + +test('a failed download closes the modal and says so', async ({ page, request }) => { + await page.goto('/prompt_manager/gallery'); + const { modal, status } = await startDownload(page, request); + await post(request, '/__test__/tagging/download', { mode: 'fail' }); + await expect(modal).toBeHidden(); + await expect(page.getByText('Download failed: Download failed')).toBeVisible(); + await expect(status.getByRole('button', { name: 'Download', exact: true })).toBeVisible(); + expect((await downloads(request)).map(job => job.state)).toEqual(['failed']); +}); diff --git a/tests/browser/server.py b/tests/browser/server.py index abccbfda..707634de 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -72,6 +72,8 @@ class TestModel: # Scored names the concept comparisons answer instead of the fixed list # (``/__test__/tagging/suggestions``); None keeps the fixed list. suggestions: list | None = None + # The engine's model table (``py/autotag.py`` MODELS), set at startup. + models_config: dict = {} def get_models_status(self) -> dict: return { @@ -79,7 +81,7 @@ class TestModel: "downloaded": name not in self.unavailable, "loaded": self.loaded == name, } - for name in ("gguf", "hf", "wd14-vit", "wd14-convnext", "wd14-swinv2") + for name in self.models_config } def is_model_loaded(self) -> bool: @@ -207,6 +209,7 @@ def create_app(root: Path) -> web.Application: from py import autotag from utils.hashing import generate_prompt_hash + TestModel.models_config = autotag.MODELS model = TestModel() autotag.get_autotag_service = lambda: model api = PromptManagerAPI() diff --git a/tests/browser/workspace-activity.spec.mjs b/tests/browser/workspace-activity.spec.mjs index 5f16414f..035b1620 100644 --- a/tests/browser/workspace-activity.spec.mjs +++ b/tests/browser/workspace-activity.spec.mjs @@ -1088,6 +1088,36 @@ test('a held legacy job and the health check show no tiles, no Results and no ma await expect(toast(page)).toHaveText('Health check finished.'); }); +// #328: a tagging model download is a job of its own kind with no byte progress and no Cancel; Activity shows it +// through the generic held view (its title, status line and an indeterminate bar), then its end toast. +test('a held model download shows its title, an indeterminate bar and no Cancel; its end toast follows', async ({ page, request }) => { + await page.setViewportSize({ width: 1600, height: 1100 }); + const title = 'Tagging model download · WD14 ViT'; + await expect(await request.post('/__test__/tagging/reset')).toBeOK(); + await expect(await request.post('/__test__/tagging/download', { data: { mode: 'hold' } })).toBeOK(); + try { + const response = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'wd14-vit' } }); + expect(response.status()).toBe(202); + const job = (await response.json()).job; + expect([job.kind, job.title, job.cancellable]).toEqual(['model_download', title, false]); + await page.goto(`${workspace}#/activity`); + const item = items(page).filter({ has: page.getByRole('link', { name: new RegExp(`${title}$`) }) }); + await expect(item.locator('progress')).not.toHaveAttribute('value'); + await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0); + await item.getByRole('link').click(); + await expect(heading(page).locator('h1')).toHaveText(title); + await expect(progressSection(page).locator('.panel-heading > span.small')).toHaveText('Downloading WD14 ViT (~400 MB)'); + await expect(progressSection(page).locator('progress')).not.toHaveAttribute('value'); + await expect(progressSection(page).locator('dl.stat-row')).toHaveCount(0); + await expect(heading(page).getByRole('button', { name: /^Cancel/ })).toHaveCount(0); + await expect(await request.post('/__test__/tagging/download', { data: { mode: 'succeed' } })).toBeOK(); + await expect(toast(page)).toHaveText(`${title} finished.`); + await expect(heading(page).locator('p.muted')).toHaveText(/^Finished \d\d:\d\d$/); + } finally { + await request.post('/__test__/tagging/reset'); + } +}); + // S5 review I3: a thumbnail job with nothing to generate ends on the server's final line ("No previews to generate · N // already up to date"), never on a stale "Processing file 0 of 0" under the Completed chip. test('a finished thumbnail job with nothing to do says so', async ({ page, request }) => { diff --git a/tests/browser/workspace_tagging_fixture.py b/tests/browser/workspace_tagging_fixture.py index 6fa32198..387fa829 100644 --- a/tests/browser/workspace_tagging_fixture.py +++ b/tests/browser/workspace_tagging_fixture.py @@ -12,6 +12,7 @@ import json import os from pathlib import Path import shutil +import threading import time from typing import TYPE_CHECKING @@ -181,6 +182,26 @@ def register_tagging_commands( routes: web.RouteTableDef, api: PromptManagerAPI, root: Path, model ) -> None: """Register fixed seed, snapshot and fake-model commands; no paths or SQL input.""" + # The fake download (``/__test__/tagging/download``): waits on ``event``, + # then answers ``result``. Every reset sets the event and puts back the + # model's raising default, so no test can leave a worker hanging. + # ``generation`` grows on every reset: a download a reset released (or one + # that timed out) fails without touching the next test's model state. + download = {"event": threading.Event(), "result": True, "generation": 0} + + def restore_download() -> None: + download["generation"] += 1 + download["event"].set() + model.__dict__.pop("download_model", None) + + def fake_download(model_type: str, progress_callback=None) -> bool: + event, generation = download["event"], download["generation"] + released = event.wait(DOWNLOAD_HOLD_LIMIT_S) + if not released or download["generation"] != generation: + return False + if download["result"]: + model.unavailable = model.unavailable - {model_type} + return download["result"] @routes.post("/__test__/tagging/reset") async def reset(request: web.Request) -> web.Response: @@ -195,6 +216,7 @@ def register_tagging_commands( model.unavailable = set() model.suggestions = None model.loaded = None + restore_download() _set_defaults(model, DEFAULTS) # A test that ends with the dialog open leaves its capture behind; drop # every capture so the next test starts below the store's limits. @@ -223,6 +245,33 @@ def register_tagging_commands( model.unavailable = set(data["unavailable"]) return web.json_response({"success": True}) + @routes.post("/__test__/tagging/download") + async def download_mode(request: web.Request) -> web.Response: + """Switch the model's download to a fake (never a real download). + + ``{"mode": "hold"}``: the next download waits until a later + ``succeed`` or ``fail`` (or a reset) releases it. ``succeed`` and + ``fail`` release a held download with that result and answer every + later one at once. A successful fake marks the model downloaded. + """ + data = await request.json() if request.can_read_body else None + if ( + request.query + or not isinstance(data, dict) + or set(data) != {"mode"} + or data["mode"] not in ("hold", "succeed", "fail") + ): + raise web.HTTPBadRequest() + if data["mode"] == "hold": + download["event"].set() + download["event"] = threading.Event() + download["result"] = True + else: + download["result"] = data["mode"] == "succeed" + download["event"].set() + model.download_model = fake_download + return web.json_response({"success": True}) + @routes.post("/__test__/tagging/suggestions") async def suggestions(request: web.Request) -> web.Response: """Answer Generate with fixed scored names (the concept's own sample). @@ -273,6 +322,9 @@ def register_tagging_commands( return web.json_response({"captures": count}) +# A held fake download gives up after this long, whatever the test does. +DOWNLOAD_HOLD_LIMIT_S = 60 + # The standard tagging defaults: WD14 general and character thresholds, and # the suggestion count (``AutotagConfig``). DEFAULTS = (0.35, 0.85, 10) diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index d2a9b153..3d907864 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -34,6 +34,7 @@ from aiohttp.test_utils import AioHTTPTestCase from database.operations import PromptDatabase from py.api import PromptManagerAPI +from py.api.jobs import JobBusy from utils.hashing import generate_prompt_hash try: # guarded so a missing module fails assertions, not test loading @@ -42,6 +43,10 @@ except ImportError: model_gate = None +# The engine's four model types (py/autotag.py MODELS), in its own order. +MODEL_TYPES = ("gguf", "hf", "wd14-swinv2", "wd14-vit") + + class FakeAutotagService: """Stand-in for AutoTagService that records calls and returns fixed tags.""" @@ -56,15 +61,24 @@ class FakeAutotagService: self.default_prompt = "default prompt" self.wd14_general_threshold = 0.35 self.wd14_character_threshold = 0.85 - self.models_config = {"gguf": {}, "wd14-vit": {}} + self.models_config = {name: {} for name in MODEL_TYPES} self._loaded_type = None self.generate_error = None self.download_result = True self.load_error = None + # Call counters and the thread each status read ran on. + self.status_calls = 0 + self.status_threads = [] + self.status_paths = {} def get_models_status(self): + self.status_calls += 1 + self.status_threads.append(threading.get_ident()) return { - name: {"downloaded": self.downloaded, "model_path": None} + name: { + "downloaded": self.downloaded, + "model_path": self.status_paths.get(name), + } for name in self.models_config } @@ -1055,7 +1069,9 @@ class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase): self.assertIsNone(model_gate.holder()) self.assertEqual(len(seen), 3) - self.assertTrue(all(owner and owner.startswith("legacy-") for owner in seen)) + # The download runs as the model_download job, which owns the gate. + kinds = [owner.split(":", 1)[0] for owner in seen] + self.assertEqual(kinds, ["legacy-single", "legacy-single", "model-download"]) async def test_validation_errors_win_over_the_gate(self): self.assertTrue(model_gate.try_acquire("test-holder")) @@ -1288,3 +1304,407 @@ class TestGalleryReviewContract(AutotagAPITestCase): "/prompt_manager/autotag/single", {"path": "link.png", "root_index": 0} ) self.assertEqual(response.status, 403) + + +def _strings(value): + """Every string inside a JSON value (keys included).""" + if isinstance(value, str): + yield value + elif isinstance(value, dict): + for key, item in value.items(): + yield key + yield from _strings(item) + elif isinstance(value, list): + for item in value: + yield from _strings(item) + + +def _absolute(text): + import ntpath + import posixpath + import re + + return ( + posixpath.isabs(text) + or ntpath.isabs(text) + or re.match(r"^[A-Za-z]:[\\/]", text) is not None + ) + + +class TestModelStatus(AutotagAPITestCase): + """GET /autotag/models: labels, sizes and the loaded flag, never a path.""" + + async def _models(self): + resp = await self.client.get("/prompt_manager/autotag/models") + self.assertEqual(resp.status, 200) + return await resp.json() + + async def test_models_status_has_no_absolute_paths(self): + self.service.status_paths = { + "wd14-vit": "/srv/models/WD14-ViT/model.onnx", + "gguf": "C:\\models\\GGUF\\joycaption.gguf", + "hf": "D:/hf-cache/models--fancyfeast", + } + data = await self._models() + for model in data["models"].values(): + self.assertNotIn("model_path", model) + leaked = [text for text in _strings(data) if _absolute(text)] + self.assertEqual(leaked, []) + + async def test_models_status_adds_label_size_label_and_loaded(self): + self.service._loaded_type = "wd14-vit" + data = await self._models() + models = data["models"] + self.assertEqual(list(models), list(MODEL_TYPES)) + self.assertEqual( + {name: (m["label"], m["size_label"]) for name, m in models.items()}, + { + "gguf": ("GGUF JoyCaption", "~4.5 GB"), + "hf": ("HuggingFace JoyCaption", "~16 GB"), + "wd14-swinv2": ("WD14 SwinV2", "~400 MB"), + "wd14-vit": ("WD14 ViT", "~400 MB"), + }, + ) + self.assertEqual( + {name: m["loaded"] for name, m in models.items()}, + {"gguf": False, "hf": False, "wd14-swinv2": False, "wd14-vit": True}, + ) + # Kept fields legacy pages and the workspace read. + for model in models.values(): + self.assertTrue({"name", "description", "size", "downloaded"} <= set(model)) + self.assertTrue(data["model_loaded"]) + self.assertEqual(data["loaded_model_type"], "wd14-vit") + + async def test_opening_models_never_loads_or_downloads(self): + await self._models() + self.assertEqual(self.service.status_calls, 1) + self.assertEqual(self.service.load_calls, []) + self.assertEqual(self.service.download_calls, []) + self.assertEqual(self.service.unload_calls, 0) + self.assertEqual(self.service.generate_calls, []) + + async def test_models_status_reads_off_the_event_loop(self): + await self._models() + self.assertEqual(len(self.service.status_threads), 1) + self.assertNotEqual(self.service.status_threads[0], threading.get_ident()) + + +class TestModelBusyNamesTheOperation(AutotagAPITestCase): + """A 409 model-busy names who holds the gate (``operation``).""" + + LABELS = { + "autotag-job": "Tag output images", + "legacy-batch": "Tag output images (legacy page)", + "legacy-single": "Tagging one image (legacy page)", + "legacy-unload": "Unloading the tagging model", + "image-tag-job": "Auto-tag selected images", + "image-annotations": "Generate suggestions", + } + + async def test_model_busy_names_the_operation(self): + self.service._loaded_type = "wd14-vit" + for kind, label in self.LABELS.items(): + owner = model_gate.new_owner(kind) + self.assertTrue(model_gate.try_acquire(owner)) + try: + resp = await self.client.post("/prompt_manager/autotag/unload") + data = await resp.json() + finally: + model_gate.release(owner) + self.assertEqual(resp.status, 409, kind) + self.assertEqual(data["code"], "model-busy", kind) + self.assertEqual(data["operation"], label, kind) + self.assertIn("busy", data["error"].lower()) + self.assertEqual(self.service.unload_calls, 0) + + async def test_unknown_or_freed_holder_is_another_operation(self): + from py.api.autotag_routes import model_busy_response + + self.assertIsNone(model_gate.holder()) + freed = json.loads(model_busy_response().body) + self.assertEqual(freed["operation"], "Another tagging operation") + self.assertTrue(model_gate.try_acquire("test-holder")) + try: + unknown = json.loads(model_busy_response().body) + finally: + model_gate.release("test-holder") + self.assertEqual(unknown["operation"], "Another tagging operation") + + async def test_unload_during_generate_is_409_model_busy(self): + self.service._loaded_type = "wd14-vit" + owner = model_gate.new_owner("image-annotations") + self.assertTrue(model_gate.try_acquire(owner)) + try: + resp = await self.client.post("/prompt_manager/autotag/unload") + data = await resp.json() + finally: + model_gate.release(owner) + self.assertEqual(resp.status, 409) + self.assertEqual(data["operation"], "Generate suggestions") + self.assertEqual(self.service.unload_calls, 0) + self.assertEqual(self.service.get_loaded_model_type(), "wd14-vit") + + +class DownloadJobTestCase(AutotagAPITestCase): + """Helpers for the ``model_download`` job: a download that waits for a + signal, so a test sees the job while it runs.""" + + TITLE = "Tagging model download · WD14 ViT" + + def setUp(self): + super().setUp() + self.started = threading.Event() + self.proceed = threading.Event() + + def _hold_download(self, result=True, error=None): + def download(model_type, progress_callback=None): + self.service.download_calls.append(model_type) + self.started.set() + self.proceed.wait(30) + if error is not None: + raise error + return result + + self.service.download_model = download + + async def _started(self): + loop = asyncio.get_running_loop() + self.assertTrue(await loop.run_in_executor(None, self.started.wait, 30)) + + async def _start(self, model_type="wd14-vit"): + return await self._post_json( + "/prompt_manager/jobs/model-download", {"model_type": model_type} + ) + + async def _until(self, condition, limit=10): + """Yield to the loop until ``condition()`` holds (fails after ``limit`` s).""" + deadline = asyncio.get_running_loop().time() + limit + while not condition(): + self.assertLess(asyncio.get_running_loop().time(), deadline) + await asyncio.sleep(0) + + async def _ended(self, job_id): + job = self.api._jobs.get(job_id) + self.proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), 30) + return job + + +class TestModelDownloadJob(DownloadJobTestCase): + + async def test_download_job_is_indeterminate_and_not_cancellable(self): + self._hold_download() + resp = await self._start() + self.assertEqual(resp.status, 202) + summary = (await resp.json())["job"] + try: + await self._started() + self.assertEqual(summary["kind"], "model_download") + self.assertEqual(summary["title"], self.TITLE) + self.assertFalse(summary["cancellable"]) + self.assertEqual(summary["admission"], "read") + self.assertEqual( + summary["scope"], + { + "model_type": "wd14-vit", + "label": "WD14 ViT", + "size_label": "~400 MB", + }, + ) + detail = ( + await ( + await self.client.get(f"/prompt_manager/jobs/{summary['id']}") + ).json() + )["job"] + self.assertEqual((detail["progress"], detail["total"]), (0, 0)) + self.assertEqual(detail["status_line"], "Downloading WD14 ViT (~400 MB)") + cancel = await self.client.post( + f"/prompt_manager/jobs/{summary['id']}/cancel" + ) + self.assertFalse((await cancel.json())["cancellation_requested"]) + finally: + job = await self._ended(summary["id"]) + self.assertEqual(job.state, "finished") + self.assertEqual(job.status_line, "Downloaded WD14 ViT") + self.assertEqual(self.service.download_calls, ["wd14-vit"]) + self.assertEqual(len(job.log), 2) + + async def test_start_body_is_checked_before_anything_runs(self): + for body in ( + {"model_type": "nope"}, + {"model_type": "wd14-vit", "extra": 1}, + {}, + ["wd14-vit"], + {"model_type": 3}, + ): + resp = await self._post_json("/prompt_manager/jobs/model-download", body) + self.assertEqual(resp.status, 400, body) + resp = await self._post_json( + "/prompt_manager/jobs/model-download", + {"model_type": "wd14-vit", "pad": "x" * 2048}, + ) + self.assertIn(resp.status, (400, 413)) + self.assertEqual(self.service.download_calls, []) + self.assertIsNone(model_gate.holder()) + self.assertEqual(self.api._jobs.list(), []) + + async def test_second_download_is_409_busy(self): + self._hold_download() + first = (await (await self._start()).json())["job"] + try: + await self._started() + resp = await self._start("gguf") + data = await resp.json() + self.assertEqual(resp.status, 409) + self.assertEqual(data["code"], "busy") + self.assertEqual(data["running"]["title"], self.TITLE) + finally: + await self._ended(first["id"]) + self.assertEqual(self.service.download_calls, ["wd14-vit"]) + + async def test_download_holds_gate_until_runner_exits(self): + self._hold_download() + self.service._loaded_type = "wd14-swinv2" + first = (await (await self._start()).json())["job"] + try: + await self._started() + self.assertTrue(model_gate.holder().startswith("model-download:")) + resp = await self.client.post("/prompt_manager/autotag/unload") + data = await resp.json() + self.assertEqual(resp.status, 409) + self.assertEqual(data["operation"], self.TITLE) + finally: + await self._ended(first["id"]) + self.assertIsNone(model_gate.holder()) + self.assertEqual(self.service.unload_calls, 0) + + async def test_download_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() + self.assertEqual(resp.status, 409) + self.assertEqual((await resp.json())["code"], "busy") + self.assertIsNone(model_gate.holder()) + self.assertEqual(self.service.download_calls, []) + + async def test_gate_held_elsewhere_is_409_model_busy(self): + owner = model_gate.new_owner("image-annotations") + self.assertTrue(model_gate.try_acquire(owner)) + try: + resp = await self._start() + data = await resp.json() + finally: + model_gate.release(owner) + self.assertEqual(resp.status, 409) + self.assertEqual(data["operation"], "Generate suggestions") + self.assertEqual(self.api._jobs.list(), []) + + async def test_failed_download_fails_the_job_without_paths(self): + for outcome in ({"result": False}, {"error": OSError("/srv/models is full")}): + self.started.clear() + self.proceed.clear() + self._hold_download(**outcome) + summary = (await (await self._start()).json())["job"] + await self._started() + job = await self._ended(summary["id"]) + self.assertEqual(job.state, "failed", outcome) + self.assertEqual( + job.status_line, "The download failed. Check server logs for details." + ) + report = await self.client.get( + f"/prompt_manager/jobs/{summary['id']}/report" + ) + text = await report.text() + self.assertNotIn("/srv/models", text) + leaked = [s for s in _strings(json.loads(text)) if _absolute(s)] + self.assertEqual(leaked, []) + self.assertIsNone(model_gate.holder()) + + +class TestLegacyDownloadAttaches(DownloadJobTestCase): + """POST /autotag/download/{type}: the legacy stream rides on the job.""" + + GOLDEN_OK = [ + {"type": "progress", "progress": 0, "status": "Starting download..."}, + {"type": "complete", "progress": 100, "status": "Download complete"}, + ] + GOLDEN_FAILED = [ + {"type": "progress", "progress": 0, "status": "Starting download..."}, + {"type": "error", "message": "Download failed"}, + ] + + async def test_legacy_download_events_unchanged(self): + resp = await self.client.post("/prompt_manager/autotag/download/wd14-vit") + self.assertEqual(resp.status, 200) + self.assertEqual(self._sse_events(await resp.text()), self.GOLDEN_OK) + # The gate is free by the time the stream says the download ended. + self.assertIsNone(model_gate.holder()) + jobs = self.api._jobs.list() + self.assertEqual( + [(j.kind, j.title, j.state) for j in jobs], + [("model_download", self.TITLE, "finished")], + ) + + self.service.download_result = False + resp = await self.client.post("/prompt_manager/autotag/download/wd14-vit") + self.assertEqual(self._sse_events(await resp.text()), self.GOLDEN_FAILED) + + async def test_legacy_download_attaches_to_running_job(self): + self._hold_download() + summary = (await (await self._start()).json())["job"] + await self._started() + legacy = asyncio.ensure_future( + self.client.post("/prompt_manager/autotag/download/wd14-vit") + ) + try: + job = self.api._jobs.get(summary["id"]) + await self._until(lambda: job.subscriber_count == 1) + finally: + await self._ended(summary["id"]) + resp = await legacy + self.assertEqual(self._sse_events(await resp.text()), self.GOLDEN_OK) + self.assertEqual(self.service.download_calls, ["wd14-vit"]) + self.assertEqual(len(self.api._jobs.list()), 1) + + async def test_legacy_download_of_other_type_is_409_model_busy(self): + self._hold_download() + summary = (await (await self._start()).json())["job"] + try: + await self._started() + resp = await self.client.post("/prompt_manager/autotag/download/gguf") + data = await resp.json() + finally: + await self._ended(summary["id"]) + self.assertEqual(resp.status, 409) + self.assertEqual(data["code"], "model-busy") + self.assertEqual(data["operation"], self.TITLE) + self.assertEqual(self.service.download_calls, ["wd14-vit"]) + + async def test_legacy_disconnect_keeps_the_job_running(self): + self._hold_download() + resp = await self.client.post("/prompt_manager/autotag/download/wd14-vit") + first = await resp.content.readline() + self.assertTrue(first.startswith(b"data: ")) + await self._started() + resp.close() # the page closed the modal + (job,) = self.api._jobs.list() + # The closed stream lets go of the job; the download goes on. + await self._until(lambda: job.subscriber_count == 0) + self.assertTrue(job.running) + self.assertTrue(model_gate.holder().startswith("model-download:")) + job = await self._ended(job.id) + self.assertEqual(job.state, "finished") + self.assertIsNone(model_gate.holder()) + + async def test_unexpected_start_error_keeps_each_route_shape(self): + with patch("py.autotag.get_autotag_service", side_effect=RuntimeError("x")): + legacy = await self.client.post("/prompt_manager/autotag/download/wd14-vit") + events = self._sse_events(await legacy.text()) + native = await self._start() + self.assertEqual(events[-1]["type"], "error") + self.assertEqual(native.status, 500) + self.assertFalse((await native.json())["success"]) + self.assertIsNone(model_gate.holder()) diff --git a/tests/test_job_admission.py b/tests/test_job_admission.py index 0bfb11c4..44d47ff4 100644 --- a/tests/test_job_admission.py +++ b/tests/test_job_admission.py @@ -460,7 +460,7 @@ class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase): async def test_second_legacy_autotag_batch_keeps_its_json_409(self): from py.api.autotag_routes import _BUSY_ERROR - self._run_job("autotag", "Auto-tag images (legacy page)") + self._run_job("autotag", "Tag output images (legacy page)") resp = await self.client.request( "POST", "/prompt_manager/autotag/start", json={} ) diff --git a/web/admin.html b/web/admin.html index bbd32c6a..dee1d8a9 100644 --- a/web/admin.html +++ b/web/admin.html @@ -1428,12 +1428,13 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
+

The download continues in Activity.

diff --git a/web/gallery.html b/web/gallery.html index ce1c8547..00b890a5 100644 --- a/web/gallery.html +++ b/web/gallery.html @@ -730,13 +730,14 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
+

The download continues in Activity.

diff --git a/web/js/admin.js b/web/js/admin.js index b02fb906..be782501 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -3689,6 +3689,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; document.getElementById('downloadStatus').textContent = 'Preparing...'; document.getElementById('downloadProgressPercent').textContent = '0%'; document.getElementById('downloadProgressBar').style.width = '0%'; + document.getElementById('downloadProgressBar').classList.remove('animate-pulse'); this.hideModal("autoTagModal"); this.showModal("autoTagDownloadModal"); @@ -3699,18 +3700,21 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; this.autoTagState.downloadEventSource.onmessage = (event) => { const data = JSON.parse(event.data); - if (data.status === 'downloading') { - const percent = Math.round(data.progress * 100); - document.getElementById('downloadStatus').textContent = data.message || 'Downloading...'; - document.getElementById('downloadProgressPercent').textContent = `${percent}%`; - document.getElementById('downloadProgressBar').style.width = `${percent}%`; - } else if (data.status === 'complete') { + // The server sends `type` (progress / complete / error); the download is an + // Activity job with no byte progress, so the bar is indeterminate (#328). + if (data.type === 'progress') { + document.getElementById('downloadStatus').textContent = 'Downloading\u2026'; + document.getElementById('downloadProgressPercent').textContent = ''; + const bar = document.getElementById('downloadProgressBar'); + bar.style.width = '100%'; + bar.classList.add('animate-pulse'); + } else if (data.type === 'complete') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal"); this.checkAutoTagModels(); this.showNotification('Model downloaded successfully!', 'success'); - } else if (data.status === 'error') { + } else if (data.type === 'error') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal"); diff --git a/web/js/gallery.js b/web/js/gallery.js index 97834b04..daf45ecc 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -3421,6 +3421,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; document.getElementById('downloadStatus').textContent = 'Preparing...'; document.getElementById('downloadProgressPercent').textContent = '0%'; document.getElementById('downloadProgressBar').style.width = '0%'; + document.getElementById('downloadProgressBar').classList.remove('animate-pulse'); this.hideModal("autoTagModal"); this.showModal("autoTagDownloadModal"); @@ -3431,18 +3432,21 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; this.autoTagState.downloadEventSource.onmessage = (event) => { const data = JSON.parse(event.data); - if (data.status === 'downloading') { - const percent = Math.round(data.progress * 100); - document.getElementById('downloadStatus').textContent = data.message || 'Downloading...'; - document.getElementById('downloadProgressPercent').textContent = `${percent}%`; - document.getElementById('downloadProgressBar').style.width = `${percent}%`; - } else if (data.status === 'complete') { + // The server sends `type` (progress / complete / error); the download is an + // Activity job with no byte progress, so the bar is indeterminate (#328). + if (data.type === 'progress') { + document.getElementById('downloadStatus').textContent = 'Downloading\u2026'; + document.getElementById('downloadProgressPercent').textContent = ''; + const bar = document.getElementById('downloadProgressBar'); + bar.style.width = '100%'; + bar.classList.add('animate-pulse'); + } else if (data.type === 'complete') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal"); this.checkAutoTagModels(); this.showNotification('Model downloaded successfully!', 'success'); - } else if (data.status === 'error') { + } else if (data.type === 'error') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal"); From 4032f4b6ab21a219a6b0616b89c03e10d416014a Mon Sep 17 00:00:00 2001 From: Vito Date: Tue, 27 Oct 2026 00:20:13 -0700 Subject: [PATCH 2/6] feat(system): one tagging configuration and Tag output images as a reviewed job (#328, PR 2) (#366) * feat(autotag): model status without paths, and a busy answer that names the operation * feat(autotag): model downloads run as Activity jobs * fix(autotag): keep each download route's error shape and harden the fake download * feat(autotag): one server-side tagging configuration * fix(autotag): the legacy batch and review work again * feat(autotag): tag output images as a reviewed Activity job * test(autotag): the legacy batch, review and unload page workflows * fix(autotag): server-free tagging defaults, cancel before load, and named legacy refusals * fix(autotag): name only the workspace batch in a legacy refusal frame --- database/operations.py | 31 + py/api/annotation_suggest.py | 39 +- py/api/autotag_routes.py | 430 +++++++++++- py/api/job_routes.py | 4 + py/api/jobs.py | 2 + py/api/settings_routes.py | 85 ++- py/autotag.py | 51 +- py/config.py | 112 ++- py/tagging_defaults.py | 55 ++ tests/browser/legacy-autotag.spec.mjs | 63 ++ tests/browser/workspace_tagging_fixture.py | 13 +- tests/js/workspace-activity-kinds.test.js | 23 + .../js/workspace-image-tagging-state.test.js | 46 ++ tests/test_api_annotations.py | 2 + tests/test_api_autotag.py | 643 ++++++++++++++++++ tests/test_autotag_engine.py | 41 +- tests/test_config.py | 105 ++- tests/test_database.py | 44 ++ tests/test_settings_sources.py | 126 ++++ web/js/admin.js | 2 + web/js/gallery.js | 30 +- web/js/workspace/activity-kinds.mjs | 38 +- web/js/workspace/image-tagging-state.mjs | 30 +- 23 files changed, 1854 insertions(+), 161 deletions(-) create mode 100644 py/tagging_defaults.py diff --git a/database/operations.py b/database/operations.py index 0f45097f..78ba02cd 100644 --- a/database/operations.py +++ b/database/operations.py @@ -1808,6 +1808,37 @@ class PromptDatabase( reader = DocumentReader(conn) return [self._image_row_to_dict(row, reader) for row in cursor] + def autotag_batch_candidates(self) -> List[Dict[str, Any]]: + """The images a Tag output images batch looks at (#328). + + ``{id, image_path, prompt_id, prompt_tags}`` for every image with a + stored path that is linked to a prompt, newest first (the + ``get_all_images`` order, ties newest row first). Never reads + ``workflow_data`` or the other documents: a batch needs none of them. + """ + sql = ( + "SELECT gi.id, gi.image_path, gi.prompt_id, " + "(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt " + "JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = gi.prompt_id) " + "AS prompt_tags " + "FROM generated_images gi " + "INNER JOIN prompts p ON gi.prompt_id = p.id " + "WHERE gi.image_path IS NOT NULL AND gi.image_path != '' " + "ORDER BY gi.generation_time DESC, gi.id DESC" + ) + with self.model.get_connection() as conn: + return [ + { + "id": row["id"], + "image_path": row["image_path"], + "prompt_id": row["prompt_id"], + "prompt_tags": ( + row["prompt_tags"].split("|||") if row["prompt_tags"] else [] + ), + } + for row in conn.execute(sql) + ] + def get_all_images( self, limit: int = 0, diff --git a/py/api/annotation_suggest.py b/py/api/annotation_suggest.py index 563bdc64..49e67a42 100644 --- a/py/api/annotation_suggest.py +++ b/py/api/annotation_suggest.py @@ -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): diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index e5521dd5..b69c5210 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -7,12 +7,20 @@ import json import os import re import threading +import time from pathlib import Path from typing import Optional from urllib.parse import quote from .job_routes import busy_response -from .jobs import AUTOTAG_BATCH_JOB_KIND, MODEL_DOWNLOAD_JOB_KIND, JobBusy +from .jobs import ( + AUTOTAG_BATCH_JOB_KIND, + AUTOTAG_GROUP_NAMES, + MODEL_DOWNLOAD_JOB_KIND, + JobBusy, +) +from .admin import _plural, _scan_location, _scan_roots +from .maintenance import disk_plan, short_space_message from .prompts import thumbnail_url_for from . import thumbnail_settings from .. import model_gate @@ -20,6 +28,18 @@ from .. import model_gate from aiohttp import web _INTERNAL_ERROR = "An internal error occurred. Check server logs for details." + +# Tag output images (#328): the native prompt-tag batch. +AUTOTAG_JOB_TITLE = "Tag output images" +AUTOTAG_PREVIEW_LEASE = "Tag output images review" +AUTOTAG_OWNER_KIND = "autotag-job" +BATCH_PREVIEW_BODY_MAX = 1024 +BATCH_JOB_BODY_MAX = 32 * 1024 +REVIEW_REQUIRED = ( + "Replacing tags needs a review of each image. Use Preview tagging review, " + "or choose Skip or Merge for a batch run." +) +BATCH_CANCELLED = "Cancelled. Tags already added were kept." _BUSY_ERROR = "Auto-tagging is already running" MODEL_BUSY_ERROR = ( "A tagging model is busy with another operation. Try again when it finishes." @@ -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" diff --git a/py/api/job_routes.py b/py/api/job_routes.py index 59cadb6a..82dcbfcb 100644 --- a/py/api/job_routes.py +++ b/py/api/job_routes.py @@ -134,6 +134,10 @@ class JobRoutesMixin: async def start_model_download_job_route(request): return await self.start_model_download_job(request) + @routes.post("/prompt_manager/jobs/autotag") + async def start_autotag_job_route(request): + return await self.start_autotag_job(request) + @routes.post("/prompt_manager/jobs/thumbnails") async def start_thumbnail_job_route(request): return await self.start_thumbnail_job(request) diff --git a/py/api/jobs.py b/py/api/jobs.py index a385f3b3..39b74b36 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -87,6 +87,8 @@ SCAN_GROUPS = { # Record-changing: it writes image tag sets and new vocabulary names. IMAGE_TAG_JOB_KIND = "image-tags" IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors") +# Tag output images (#328) counts its prompts the same way. +AUTOTAG_GROUP_NAMES = IMAGE_TAG_GROUP_NAMES # Thumbnail generation (#249 System S5): previews written, already fresh, failed. THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors") diff --git a/py/api/settings_routes.py b/py/api/settings_routes.py index 466ba41d..c38274c4 100644 --- a/py/api/settings_routes.py +++ b/py/api/settings_routes.py @@ -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): diff --git a/py/autotag.py b/py/autotag.py index 4d6d4da8..9a6d55b1 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -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 diff --git a/py/config.py b/py/config.py index 46b43641..2a14023a 100644 --- a/py/config.py +++ b/py/config.py @@ -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: diff --git a/py/tagging_defaults.py b/py/tagging_defaults.py new file mode 100644 index 00000000..dedaf9b6 --- /dev/null +++ b/py/tagging_defaults.py @@ -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 diff --git a/tests/browser/legacy-autotag.spec.mjs b/tests/browser/legacy-autotag.spec.mjs index a928ac04..1140a887 100644 --- a/tests/browser/legacy-autotag.spec.mjs +++ b/tests/browser/legacy-autotag.spec.mjs @@ -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(); +}); diff --git a/tests/browser/workspace_tagging_fixture.py b/tests/browser/workspace_tagging_fixture.py index 387fa829..843e2684 100644 --- a/tests/browser/workspace_tagging_fixture.py +++ b/tests/browser/workspace_tagging_fixture.py @@ -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: diff --git a/tests/js/workspace-activity-kinds.test.js b/tests/js/workspace-activity-kinds.test.js index 84831a61..68d3afc4 100644 --- a/tests/js/workspace-activity-kinds.test.js +++ b/tests/js/workspace-activity-kinds.test.js @@ -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); +}); diff --git a/tests/js/workspace-image-tagging-state.test.js b/tests/js/workspace-image-tagging-state.test.js index 604716f8..5607c0fe 100644 --- a/tests/js/workspace-image-tagging-state.test.js +++ b/tests/js/workspace-image-tagging-state.test.js @@ -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 }); diff --git a/tests/test_api_annotations.py b/tests/test_api_annotations.py index 6476ea7b..769afa89 100644 --- a/tests/test_api_annotations.py +++ b/tests/test_api_annotations.py @@ -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 diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 3d907864..444594c7 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -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()) diff --git a/tests/test_autotag_engine.py b/tests/test_autotag_engine.py index c76217ab..53e0bcac 100644 --- a/tests/test_autotag_engine.py +++ b/tests/test_autotag_engine.py @@ -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)]) diff --git a/tests/test_config.py b/tests/test_config.py index cf1da2fc..7a9e8ba4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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): diff --git a/tests/test_database.py b/tests/test_database.py index 92935ce1..09fe6e97 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1398,6 +1398,50 @@ class TestModelHousekeeping(DatabaseTestCase): self.assertIsNot(self.db.model.get_connection(), failing) +class TestAutotagBatchCandidates(DatabaseTestCase): + """#328: the batch's candidate read never loads workflow documents.""" + + def test_autotag_candidates_skip_workflow_data(self): + from utils.hashing import generate_prompt_hash + + first = self.db.save_prompt( + text="a", tags=["one", "two"], prompt_hash=generate_prompt_hash("a") + ) + second = self.db.save_prompt(text="b", prompt_hash=generate_prompt_hash("b")) + old = self.db.link_image_to_prompt( + first, "/out/a.png", {"workflow": {"big": "x" * 1000}} + ) + new = self.db.link_image_to_prompt(second, "/out/b.png") + statements = [] + conn = self.db.model.get_connection() # this thread's own connection + conn.set_trace_callback(statements.append) + try: + rows = self.db.autotag_batch_candidates() + finally: + conn.set_trace_callback(None) + self.assertEqual( + rows[0], + { + "id": new, + "image_path": "/out/b.png", + "prompt_id": second, + "prompt_tags": [], + }, + ) + self.assertEqual( + {**rows[1], "prompt_tags": sorted(rows[1]["prompt_tags"])}, + { + "id": old, + "image_path": "/out/a.png", + "prompt_id": first, + "prompt_tags": ["one", "two"], + }, + ) + self.assertEqual(len(rows), 2) + self.assertTrue(statements) + self.assertFalse(any("workflow_data" in sql for sql in statements)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_settings_sources.py b/tests/test_settings_sources.py index b964101d..e85549fc 100644 --- a/tests/test_settings_sources.py +++ b/tests/test_settings_sources.py @@ -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) diff --git a/web/js/admin.js b/web/js/admin.js index be782501..200e462a 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -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 = '
Generating tags...
'; document.getElementById('reviewTagsHidden').innerHTML = ''; diff --git a/web/js/gallery.js b/web/js/gallery.js index daf45ecc..51838c69 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -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'); } }; diff --git a/web/js/workspace/activity-kinds.mjs b/web/js/workspace/activity-kinds.mjs index 4bed30c7..8e7d0793 100644 --- a/web/js/workspace/activity-kinds.mjs +++ b/web/js/workspace/activity-kinds.mjs @@ -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, diff --git a/web/js/workspace/image-tagging-state.mjs b/web/js/workspace/image-tagging-state.mjs index e85b9597..b835f1fb 100644 --- a/web/js/workspace/image-tagging-state.mjs +++ b/web/js/workspace/image-tagging-state.mjs @@ -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, From ae755d2ce0436c13cd9b318efb05ee3005e97e72 Mon Sep 17 00:00:00 2001 From: Vito Date: Tue, 27 Oct 2026 00:59:39 -0700 Subject: [PATCH 3/6] feat(system): Auto-tagging screen, native in System (#328, PR 3) (#368) * feat(system): configure auto-tagging natively * test(system): compare Auto-tagging with the concept against the real backend --- TESTING.md | 125 +++- py/api/annotation_suggest.py | 2 +- tests/browser/system-tagging-concept.spec.mjs | 385 ++++++++++++ tests/browser/tagging-concept.spec.mjs | 2 +- .../browser/workspace-system-tagging.spec.mjs | 356 +++++++++++ tests/browser/workspace-system.spec.mjs | 6 +- tests/browser/workspace-tagging-bulk.spec.mjs | 2 +- tests/browser/workspace-tagging.spec.mjs | 2 +- tests/browser/workspace.spec.mjs | 12 +- tests/js/api-client.test.js | 10 + .../js/workspace-image-tagging-state.test.js | 2 +- .../js/workspace-images-tagging-bulk.test.js | 4 +- tests/js/workspace-system-state.test.js | 16 + tests/js/workspace-system-tagging.test.js | 583 ++++++++++++++++++ tests/test_api_annotation_suggest.py | 2 +- tests/test_image_tag_jobs.py | 2 +- tools/vendor/frontend-coverage-baseline.json | 6 + web/js/api-client.js | 6 + web/js/workspace/image-tagging-state.mjs | 2 +- web/js/workspace/system-integrations.mjs | 24 +- web/js/workspace/system-state.mjs | 11 +- web/js/workspace/system-tagging.mjs | 503 +++++++++++++++ web/js/workspace/system.mjs | 32 +- web/js/workspace/vocabulary.mjs | 3 +- web/js/workspace/workspace-dialogs.mjs | 29 + web/js/workspace/workspace.css | 7 + 26 files changed, 2080 insertions(+), 54 deletions(-) create mode 100644 tests/browser/system-tagging-concept.spec.mjs create mode 100644 tests/browser/workspace-system-tagging.spec.mjs create mode 100644 tests/js/workspace-system-tagging.test.js create mode 100644 web/js/workspace/system-tagging.mjs diff --git a/TESTING.md b/TESTING.md index 648ef497..edf48220 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1504,8 +1504,8 @@ targets of the longest legal path (1024 code points, even escaped as `\uXXXX`) f Its fields show the server's defaults (`GET /prompt_manager/autotag/models`, the D91 Suggestions count included) once that read lands; Start waits for it. A model the status lists as not downloaded keeps Start `aria-disabled` and - `#tag-generation-status` says " is not downloaded. Download it from the - legacy Auto-tag page." (the server's own sentence), so no job starts only to fail. + `#tag-generation-status` says " is not downloaded. Download it in System › + Auto-tagging." (the server's own sentence), so no job starts only to fail. Start validates the fields as Generate does (native validity; nothing is sent). The concept comparison (`images-concept.spec.mjs`) uses the concept page itself as the oracle, with Add to collection and Remove links… removed and a copy of its Add @@ -2587,7 +2587,7 @@ the connection's `total_changes` is 0 at close and the data reads back unchanged `-shm` bytes are not compared: every WAL reader rewrites its read marks there, so no reader can keep it byte-identical (`test_health_check_never_writes`). -Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–D114; D121–D131 assigned): +Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–D114; D121–D134 assigned): - **D92** (shell): no "Sample configuration" chip. The Save status is truthful: "No unsaved changes.", "Unsaved changes.", "Saved in this browser." / "Saved to the server @@ -2621,7 +2621,8 @@ Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111– output folder is used." when that is the case. - **D94** (all sections): a section not delivered yet shows the hand-off panel "Your current tools are ready" with a link to each legacy tool that still owns it (the - plan's legacy coverage table); retired section by section, gone after S12. + plan's legacy coverage table); retired section by section, gone after S12. Since S11 + (#328) only Diagnostics & logs hands off (`SECTION_HANDOFFS`, to the admin page). - **D95** (health report): the concept's `health-check` dialog titled "Health report", with no sample note. It lists every check in the server's order: Integrity check, Foreign keys & pair references, Canonical hashes, Pair uniqueness (added), Missing @@ -2997,6 +2998,59 @@ Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111– dialog on that scope. Pinned by the thumbnail dialog comparison (mapped) and its mutation check (system-concept.spec) and `the thumbnail row counts outdated previews` (workspace-system.spec). +- **D132** (Auto-tagging, S11, #328, OQ1/OQ5/OQ6/OQ10): "Tagging model" (`#tag-model`) + lists the four real engines in the concept's order, each `