diff --git a/TESTING.md b/TESTING.md index ffcfde2..1bb77e9 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1683,6 +1683,9 @@ targets of the longest legal path (1024 code points, even escaped as `\uXXXX`) f unchanged, `-` skipped and `!` error. - The progress line is the scan's: "Processing file N of M · about X min left", the estimate (D56's `_time_left`) shown once two files are done, each file one batch. + An ended job never keeps it: its status line becomes " + · N files processed" (or the storage-restored sentence), + pinned by `test_an_ended_job_shows_no_time_left` (#379 live check). - Download report saves `prompt-manager-image-tags-.json`. - The start toasts are "Tagging started." / "Auto-tagging started." (D85); the end toasts follow the design's pattern with the job's counts, as the scan's: " @@ -3149,25 +3152,39 @@ the per-section hand-off panel, was retired by S12 (#329): every section is nati the persisted count the Tag image dialog's Suggestions field (D91) starts from. Pinned by the comparison's mapping (it inserts the row) and `defaults save through Save preferences and the Tag image dialog uses them`. -- **D134** (Auto-tagging, S11, OQ2/OQ3/OQ7/OQ8): an `.operation` "Tag output images" - with "Review" follows `.row-actions`: the legacy prompt-tag batch. Review is +- **D134** (Auto-tagging, S11, OQ2/OQ3/OQ7/OQ8; amended by #379): an `.operation` "Tag + output images" with "Review" follows `.row-actions`: the per-image batch (#379), whose + description reads "Tag each output image from its own pixels with the model above. + Prompts are not changed. Review the scope first.". Review is `aria-disabled` (the press explains) while the chosen engine is not downloaded or a value it uses is invalid; otherwise it posts `/autotag/batch/preview` with the drafted engine and policy (read-only) and opens "Review tag output images" in the concept's record-changing review (its `maintenancePlan` shell): `h3` "Tag output images", the - `dl.meta-grid` Model, Existing tags, Images linked, Prompts, Already tagged, Prompts - to tag ("up to N"), Settings ("General <g> · Character <c>" or "Template of N - characters") and Backup, the upper-bound sentence in the `p.small`, the D55 backup box - and Cancel / "Start tagging", which starts the cancellable `autotag` job and opens it. - Replace after review, a model not on disk, a short disk, no linked image and nothing to - tag each show the rows with only Cancel and the reason ("Replacing tags needs a review - of each image. Use Preview tagging review, or choose Skip or Merge for a batch run.", - "Every prompt with an output image is already tagged.", "Not enough free disk space: - the safety backup needs about X MB and Y MB is free. Nothing is written until there is - room.") in the close-only review Integrations shares (`openCloseOnlyReview`). Pinned by + `dl.meta-grid` Model, Existing tags, Images, Already auto-tagged, Images to tag + ("about N"), Time, Settings ("General <g> · Character <c>" or "Template of N + characters", then " · up to L tags per image" with the preview's saved + `suggestion_limit`) and Backup, the estimate sentence ("Nothing is written until you + start. The number to tag is an estimate: images changed since they were tagged are + tagged again, and missing files are skipped. Activity reports what was tagged.") in the + `p.small`, the D55 backup box and Cancel / "Start tagging", which starts the + cancellable `autotag` job and opens it. Time (DEC11) is `m = max(1, + ceil(estimate_seconds / 60))` read "about m min" under 90 and "about H h M min" (or + "about H h") from 90, followed by " on a processor; a graphics card can be faster." + or, when the preview says WD14 runs on CUDA, " on the graphics card."; the language + models read "Language-model engines take longer per image. Activity shows the time + left once tagging starts.". The option label "Skip already tagged" is the concept's; + for this batch it skips images already auto-tagged (hand tags do not count). Replace + after review, a model not on disk, a short disk, no image and nothing to tag each show + the rows with only Cancel and the reason ("Replacing tags needs a review of each + image. Use Preview tagging review, or choose Skip or Merge for a batch run.", "The + output folders contain no images to tag.", "Every output image is already + auto-tagged.", "Not enough free disk space: the safety backup needs about X MB and Y + MB is free. Nothing is written until there is room.") in the close-only review + Integrations shares (`openCloseOnlyReview`). Pinned by `tag output images review matches the concept`, its mutation check and `the Cancel-only review matches the concept review shell` (system-tagging-concept.spec) and - the batch tests of workspace-system-tagging.spec. + the batch tests of workspace-system-tagging.spec; the mutation check renames each grid + label in turn. - **D109** (Logs, S12, #329): the concept's Logs surface on the live server. The toolbar keeps "Level", Refresh and the download (renamed "Download log": the active log, `/logs/download/prompt_manager.log`, through the api client), with all six minimum @@ -3775,19 +3792,45 @@ images", "Tag output images (legacy page)", "Tagging one image (legacy page)", `409 busy`; the legacy SSE route attaches to the same job and keeps its frames. - **Unload**: `POST /prompt_manager/autotag/unload` answers `{success, message}` or `409 model-busy`. -- **Tag output images**: `POST /prompt_manager/autotag/batch/preview {model_type, - existing}` writes nothing and opens no file; it answers `images`, `prompts`, - `already_tagged`, `to_tag` (an upper bound), `model {type, label, size_label, - downloaded}`, `existing`, `review_required`, `backup {policy}` and `disk`. `POST +- **Tag output images** (per image since #379; prompt tags and prompt text are never + written): its targets are every still image the output folders list (the Images + listing of each root, linked to a prompt or not; videos are not targets; a file under + two overlapping roots is one target). `POST /prompt_manager/autotag/batch/preview + {model_type, existing}` writes nothing and opens or stats no image file (it may wait + for the shared listing walk); it answers `images`, `already_tagged` (images with any + auto-tagged byte record), `to_tag` (`images − already_tagged` under Skip, `images` + under Merge; both estimates), `estimate_seconds` (WD14: `to_tag` × 1.6 s on the + processor or × `WD14_GPU_SECONDS_PER_IMAGE` when ONNX Runtime offers CUDA, see + `py/tagging_defaults.py`; `null` for GGUF and HF), `estimate_device` (`cpu`, `gpu`, + or `null`), `suggestion_limit` (the saved count the run uses), `model {type, label, + size_label, downloaded}`, `existing`, `review_required`, `backup {policy}` and + `disk`. `POST /prompt_manager/jobs/autotag {model_type, existing, options}` (WD14: both thresholds; GGUF/HF: `instructions`, the template) starts the cancellable `autotag` job (202), title "Tag output images", admission `records`, a safety backup by the policy first. + The job runs the `image-tags` loop (`_tag_each_file`) on each image: a capture, at + most `suggestion_limit` model tags, a re-hash, then a union into that image's tags + that also marks its bytes auto-tagged (`media_image_autotags`, one row per byte + record; an empty result and an unsupported still are marked too; File changed, File + not found, Too large to verify, read, inference and tag-limit errors are not). Phases, + each checked for Cancel: backup, model load, listing ("Listing output images…"), a + pre-pass that stats each file without opening it (gone: "File not found", over 256 + MiB: "Too large to verify", under Skip a marked current record: "Already + auto-tagged"; counted with `record_counted`, rows up to the cap, no log line each), + then the loop ("Processing file {k + n} of {T}", the time left from the model's own + rate, D56 unchanged). It ends "Tagged N images", "Cancelled. Tags already added were + kept.", or failed with "Prompt storage was restored; start the job again.", "The + output folders could not be listed. Try again." or the model's refusal. Refusals: `400 review-required` (Replace after review never runs as a batch), `409 model-unavailable` ("<label> is not downloaded. Download it in System › Auto-tagging."), `507 insufficient-space`, `409 busy`, `409 model-busy`. The section toasts `model-busy` with its operation, the other three with the server's sentence, `busy` with the running job (D57), anything else "<what> could not start. Try again.". Outcomes: Tagged / Unchanged / Skipped / Errors (the `image-tags` descriptor, D111). + `tests/test_autotag_batch.py` pins it, including two images of one prompt getting + different tags with no `prompts` or `prompt_tags` row changed; the marks' schema group + is pinned by `tests/test_autotag_schema.py` (forward-only: one empty `WITHOUT ROWID` + table, optional in old backups, refused without the annotation group). No browser test downloads or runs a model: the injected `TestModel` answers, its `download_model` raises unless `/__test__/tagging/download {mode: hold | succeed | @@ -3923,7 +3966,7 @@ list stay unscoped: any row's control proves them. | Auto-tag model download | `admin.html:1413–1438`; `admin.js:3680–3747`; `gallery.js:3412` | S11 Auto-tagging | covered | `tagging` · button · "Download" | Download starts the job and Activity shows it without Cancel (workspace-system-tagging.spec.mjs) | | Unload model | `gallery.js:3480` | S11 Auto-tagging | covered | `tagging` · button · "Unload" | Unload is enabled only while a model is loaded, and a held gate's refusal names the operation (workspace-system-tagging.spec.mjs) | | Tag template, WD14 thresholds, already-tagged mode | `admin.html:1150–1198`; `admin.js:3748` | S11 Auto-tagging | covered | `tagging` · combobox · "Existing tags" | defaults save through Save preferences and the Tag image dialog uses them (workspace-system-tagging.spec.mjs) | -| Batch auto-tag output images | `admin.html:191,1210,1279–1318`; `admin.js:3748–3840` | S11 Auto-tagging | covered | `tagging` · button · "Review" | Tag output images reviews the scope and starts a batch job (workspace-system-tagging.spec.mjs) | +| Batch auto-tag output images | `admin.html:191,1210,1279–1318`; `admin.js:3748–3840` | S11 Auto-tagging | covered | `tagging` · button · "Review" | Tag output images reviews the images, tags each one and never a prompt (#379) (workspace-system-tagging.spec.mjs) | | Batch review loop | `admin.html:1206,1218–1364`; `admin.js:3841–4172` | S11 + image-tagging plan | covered | `tagging` · button · "Preview tagging review" | Preview tagging review opens the Tag image dialog for the selected image (workspace-system-tagging.spec.mjs) | | Gallery per-image auto-tag / review | `gallery.html:601,605`; `gallery.js:3506,3674–3753` | Image-tagging plan (selection bar) | covered | — | Auto-tag selected with Apply automatically runs a job with the chosen model, thresholds and suggestion count (workspace-tagging-bulk.spec.mjs) | | Gallery thumbnail generation | `gallery.js:1852,1967` | S5/S7 | covered | `sources` · button · "Review options" · in "Thumbnail generation" | Start closes the dialog, toasts, and opens the job in Activity (workspace-system-sources.spec.mjs) | diff --git a/database/annotations.py b/database/annotations.py index 219a503..6c0ee6b 100644 --- a/database/annotations.py +++ b/database/annotations.py @@ -19,6 +19,11 @@ links (``PromptDatabase._move_tag_links``). Cheap reads (spec amendment S1) treat a byte record as *current* only while the file's present ``(size, mtime_ns)`` equals the record's verified fingerprint; a file without a current record shows no image tags. + +A write made from a tagging model's result passes ``mark_autotagged`` (#379): +the byte record is then marked auto-tagged (``media_image_autotags``) in the +same transaction, so a mark never exists without the write it records. Hand +tags never mark. """ import re @@ -34,6 +39,7 @@ from .annotation_schema import ( MAX_LOCATOR_LENGTH, METADATA_TABLE, ) +from .autotag_schema import AUTOTAG_TABLE from .library_commands import _UUID, LibraryCommandError, _validate_epoch from .models import utc_now_iso from .vocabulary import validate_vocabulary_tag_name @@ -96,6 +102,14 @@ _TAGS_OF = ( "JOIN tags t ON t.id=mit.tag_id WHERE mit.byte_id IN ({marks}) " "ORDER BY mit.byte_id,mit.position" ) +_MARK = f"INSERT OR IGNORE INTO {AUTOTAG_TABLE}(byte_id) VALUES(?)" # nosec B608 +_MARKED_OF = ( + f"SELECT byte_id FROM {AUTOTAG_TABLE} WHERE byte_id IN ({{marks}})" # nosec B608 +) +_MARKED_LOCATORS = ( + f"SELECT DISTINCT f.locator FROM {AUTOTAG_TABLE} m " # nosec B608 + f"JOIN {BYTES_TABLE} b ON b.id=m.byte_id JOIN {FILES_TABLE} f ON f.id=b.file_id" +) _RECORDS = ( "SELECT f.locator,b.id,b.size,b.mtime_ns,t.name " # nosec B608 f"FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b ON b.file_id=f.id " @@ -301,14 +315,17 @@ class AnnotationOperations: expected_epoch: str, expected_tag_revision: int, tags: list, + mark_autotagged: bool = False, ) -> dict: """Replace the set of these exact bytes, guarded by epoch and revision. Creates the owner and byte record when the bytes are new (an explicit empty set is a real save). The identical set at the current revision, or one behind it (a lost-response retry), is ``replayed`` and writes - nothing; any other revision mismatch is ``tag-conflict`` with the - current set and no writes. + no tag; any other revision mismatch is ``tag-conflict`` with the + current set and no writes. ``mark_autotagged`` marks the bytes + auto-tagged on a saved or replayed receipt (the Tag image dialog's + save after Generate suggestions). """ _validate_target(locator, sha256, size, mtime_ns) _validate_epoch(expected_epoch) @@ -329,6 +346,8 @@ class AnnotationOperations: ): if refingerprinted: conn.execute(_BUMP_INDEX) + if mark_autotagged: + conn.execute(_MARK, (byte_id,)) return _receipt(epoch, sha256, current, revision, True) if expected_tag_revision != revision: raise AnnotationCommandError( @@ -337,6 +356,8 @@ class AnnotationOperations: {"sha256": sha256, "tags": current, "tag_revision": revision}, ) revision = self._write_image_tags(conn, byte_id, tags) + if mark_autotagged: + conn.execute(_MARK, (byte_id,)) return _receipt(epoch, sha256, tags, revision, False) def add_image_tags( @@ -348,16 +369,22 @@ class AnnotationOperations: mtime_ns: int, expected_epoch: str, tags: list, + mark_autotagged: bool = False, ) -> dict: """Union ``tags`` into the set of these exact bytes, keeping stored spelling. ``added`` counts the names that were new (case-insensitively). A union over ``MAX_IMAGE_TAGS`` is ``tag-limit`` with the current set and no writes; a union that adds nothing to an existing record is ``replayed``. + ``mark_autotagged`` (a tagging model's result) also marks the bytes + auto-tagged, replayed or not; only such a call may pass no tags (the + model found none: the image is marked and its set is unchanged). """ _validate_target(locator, sha256, size, mtime_ns) _validate_epoch(expected_epoch) tags = normalize_image_tags(tags) + if not tags and not mark_autotagged: + raise ValueError("tags must name at least one tag") with self._pair_transaction() as conn: epoch = self._library_epoch(conn, expected_epoch) now = utc_now_iso() @@ -375,6 +402,8 @@ class AnnotationOperations: epoch, {"sha256": sha256, "tags": current, "tag_revision": revision}, ) + if mark_autotagged: + conn.execute(_MARK, (byte_id,)) if not created and not added: if refingerprinted: conn.execute(_BUMP_INDEX) @@ -403,6 +432,41 @@ class AnnotationOperations: conn.execute(_BUMP_INDEX) return cursor.rowcount > 0 + @staticmethod + def _validated_stats(stats: object) -> dict: + if not isinstance(stats, dict): + raise ValueError( + "stats must map locators to (size, mtime_ns) integer pairs" + ) + for locator, fingerprint in stats.items(): + _validate_fingerprint(locator, fingerprint) + return stats + + @staticmethod + def _current_byte_records(conn: sqlite3.Connection, stats: dict) -> dict: + """``{locator: byte_id}`` of each locator's current byte record. + + The one rule every cheap read uses: a record is current while its + verified ``(size, mtime_ns)`` equals the file's present fingerprint + in ``stats``; the most recently verified wins (ties on ``(verified_at, + byte_id)``). A locator without a current record is left out. Lookups + go in chunks of 500 locators. + """ + locators = list(stats) + chosen: dict[str, tuple[tuple[str, str], str]] = {} + for start in range(0, len(locators), _LOOKUP_CHUNK): + chunk = locators[start : start + _LOOKUP_CHUNK] + marks = ",".join("?" * len(chunk)) + for locator, byte_id, size, mtime_ns, verified_at in conn.execute( + _FINGERPRINTS.format(marks=marks), chunk + ): + if (size, mtime_ns) != stats[locator]: + continue + key = (verified_at, byte_id) + if locator not in chosen or key > chosen[locator][0]: + chosen[locator] = (key, byte_id) + return {locator: byte_id for locator, (_key, byte_id) in chosen.items()} + def current_image_tags(self, stats: dict) -> dict: """Image tags of the files whose current ``(size, mtime_ns)`` is verified. @@ -411,29 +475,11 @@ class AnnotationOperations: exactly that verified fingerprint; the most recently verified record wins when several do. Lookups go in chunks of 500 locators. """ - if not isinstance(stats, dict): - raise ValueError( - "stats must map locators to (size, mtime_ns) integer pairs" - ) - for locator, fingerprint in stats.items(): - _validate_fingerprint(locator, fingerprint) - locators = list(stats) - chosen: dict[str, tuple[tuple[str, str], str]] = {} - result: dict[str, list[str]] = {} + stats = self._validated_stats(stats) with self._annotation_read() as conn: - for start in range(0, len(locators), _LOOKUP_CHUNK): - chunk = locators[start : start + _LOOKUP_CHUNK] - marks = ",".join("?" * len(chunk)) - for locator, byte_id, size, mtime_ns, verified_at in conn.execute( - _FINGERPRINTS.format(marks=marks), chunk - ): - if (size, mtime_ns) != stats[locator]: - continue - key = (verified_at, byte_id) - if locator not in chosen or key > chosen[locator][0]: - chosen[locator] = (key, byte_id) - owners = {byte_id: locator for locator, (_key, byte_id) in chosen.items()} - result = {locator: [] for locator in chosen} + current = self._current_byte_records(conn, stats) + owners = {byte_id: locator for locator, byte_id in current.items()} + result: dict[str, list[str]] = {locator: [] for locator in current} byte_ids = list(owners) for start in range(0, len(byte_ids), _LOOKUP_CHUNK): chunk = byte_ids[start : start + _LOOKUP_CHUNK] @@ -442,6 +488,34 @@ class AnnotationOperations: result[owners[byte_id]].append(name) return result + def autotag_state(self, stats: dict) -> dict: + """``{locator: bool}``: is each file's current byte record auto-tagged? + + ``stats`` as for :meth:`current_image_tags`, whose rule picks the + current record; a locator without one is False. One read snapshot. + """ + stats = self._validated_stats(stats) + with self._annotation_read() as conn: + current = self._current_byte_records(conn, stats) + byte_ids = list(dict.fromkeys(current.values())) + marked: set[str] = set() + for start in range(0, len(byte_ids), _LOOKUP_CHUNK): + chunk = byte_ids[start : start + _LOOKUP_CHUNK] + marks = ",".join("?" * len(chunk)) + marked.update( + row[0] + for row in conn.execute(_MARKED_OF.format(marks=marks), chunk) + ) + return {locator: current.get(locator) in marked for locator in stats} + + def autotagged_locators(self) -> set: + """Locators with at least one auto-tagged byte record, current or not. + + No fingerprint is checked (the Tag output images review's estimate). + """ + with self._annotation_read() as conn: + return {row[0] for row in conn.execute(_MARKED_LOCATORS)} + def iter_image_tag_records(self) -> Iterator[tuple[str, int, int, list[str]]]: """Stream every byte record as ``(locator, size, mtime_ns, tags)``. diff --git a/database/autotag_schema.py b/database/autotag_schema.py new file mode 100644 index 0000000..29e2db1 --- /dev/null +++ b/database/autotag_schema.py @@ -0,0 +1,140 @@ +"""Auto-tagged marks (#379): which image bytes a tagging model has handled. + +One table, ``media_image_autotags``: one row per byte record of the image +annotation group (``media_annotation_bytes``) means "a tagging model looked at +these exact bytes and its result was stored", or that the bytes were found to +be untaggable. "Skip already tagged" skips an image whose current byte record +is marked; hand tags never mark. Rows are written only by the image tag store +(``database/annotations.py``), in the transaction of the tag write they +record. + +The group installs like ``image_favorites``: optional in old backups, +required once a database is current, and never backfilled. Like the Needs +review group it is valid only beside a valid image annotation group, because +every mark cites one byte record (``ON DELETE CASCADE``). Validation reads +only: the catalog compared after whitespace compaction, the columns and the +foreign key through the PRAGMAs, then one bounded value pass and the foreign +key check. +""" + +import re +import sqlite3 +from typing import Optional + +from .annotation_schema import BYTES_TABLE +from .annotation_schema import TABLES as ANNOTATION_TABLES +from .annotation_schema import _catalog as _annotation_catalog +from .annotation_schema import _uuid, annotation_schema_error + +AUTOTAG_TABLE = "media_image_autotags" +_DDL = ( + f"CREATE TABLE {AUTOTAG_TABLE} (byte_id TEXT PRIMARY KEY " + f"REFERENCES {BYTES_TABLE}(id) ON DELETE CASCADE) WITHOUT ROWID" +) +# (cid, name, declared type, not null, default, primary key position): the +# key of a WITHOUT ROWID table is implicitly NOT NULL. +_COLUMNS = [(0, "byte_id", "TEXT", 1, None, 1)] +# (column, parent table, parent column, on delete) +_FOREIGN_KEYS = {("byte_id", BYTES_TABLE, "id", "CASCADE")} + +NEEDS_ANNOTATIONS = "Image auto-tag marks need the image annotation tables" +INVALID_SCHEMA = "Image auto-tag marks have an invalid schema" +INVALID_VALUES = "Image auto-tag marks contain invalid values" +MISSING = "Missing image auto-tag marks schema" + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql).lower() + + +def _catalog(conn: sqlite3.Connection) -> list: + return conn.execute( + "SELECT type,name,sql FROM sqlite_master WHERE name=? OR tbl_name=?", + (AUTOTAG_TABLE, AUTOTAG_TABLE), + ).fetchall() + + +def autotag_without_annotations_error(conn: sqlite3.Connection) -> Optional[str]: + """Catalog only: the marks table present without its annotation tables. + + Startup asks this before the annotation group installs, which would make + a marks table copied without it look valid. It reads no rows. + """ + if not _catalog(conn): + return None + present = { + name for kind, name, _sql in _annotation_catalog(conn) if kind == "table" + } + if present != set(ANNOTATION_TABLES): + return NEEDS_ANNOTATIONS + return None + + +def autotag_schema_error( + conn: sqlite3.Connection, + *, + allow_missing: bool = True, + annotations_valid: bool = False, +) -> Optional[str]: + """Reject an unsupported marks table or invalid marks; reads only. + + A database without the table is acceptable for backups that predate it + unless ``allow_missing`` is false. A present table requires a present, + valid image annotation group; ``annotations_valid`` says the caller has + just validated it (startup's annotation install), so only its presence is + checked again. + """ + objects = _catalog(conn) + if not objects: + return None if allow_missing else MISSING + if annotations_valid: + if autotag_without_annotations_error(conn): + return NEEDS_ANNOTATIONS + elif annotation_schema_error(conn, allow_missing=False): + return NEEDS_ANNOTATIONS + actual = {(kind, name, _compact(sql or "")) for kind, name, sql in objects} + if actual != {("table", AUTOTAG_TABLE, _compact(_DDL))}: + return INVALID_SCHEMA + columns = [ + (row[0], row[1], row[2].upper(), row[3], row[4], row[5]) + for row in conn.execute(f"PRAGMA table_info({AUTOTAG_TABLE})") + ] + if columns != _COLUMNS: + return INVALID_SCHEMA + foreign_keys = { + (row[3], row[2], row[4], row[6]) + for row in conn.execute(f"PRAGMA foreign_key_list({AUTOTAG_TABLE})") + } + if foreign_keys != _FOREIGN_KEYS: + return INVALID_SCHEMA + # Only this module's constants are interpolated into the SQL below. + if conn.execute( + f"SELECT 1 FROM {AUTOTAG_TABLE} " # nosec B608 + f"WHERE NOT ({_uuid('byte_id')}) LIMIT 1" + ).fetchone(): + return INVALID_VALUES + if conn.execute(f"PRAGMA foreign_key_check({AUTOTAG_TABLE})").fetchone(): + return INVALID_VALUES + return None + + +def migrate_autotag_state(conn: sqlite3.Connection) -> None: + """Install the table atomically; validate before and after creating it. + + Runs after ``migrate_annotation_state`` (which has just validated the + annotation group). Creates one empty table: no row is read or copied. + """ + conn.execute("SAVEPOINT media_autotag_schema") + try: + error = autotag_schema_error(conn, annotations_valid=True) + if error: + raise sqlite3.DatabaseError(error) + conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + error = autotag_schema_error(conn, allow_missing=False, annotations_valid=True) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE media_autotag_schema") + except BaseException: + conn.execute("ROLLBACK TO media_autotag_schema") + conn.execute("RELEASE media_autotag_schema") + raise diff --git a/database/models.py b/database/models.py index e49ff94..dad9cf0 100644 --- a/database/models.py +++ b/database/models.py @@ -23,6 +23,11 @@ from . import restore_stage from .library_schema import library_schema_error, migrate_library_state from .image_favorites import image_favorites_schema_error, migrate_image_favorites from .annotation_schema import annotation_schema_error, migrate_annotation_state +from .autotag_schema import ( + autotag_schema_error, + autotag_without_annotations_error, + migrate_autotag_state, +) from .review_schema import ( migrate_review_state, review_schema_error, @@ -338,8 +343,13 @@ class PromptModel: review_error = review_without_annotations_error(conn) if review_error: raise sqlite3.DatabaseError(review_error) + # The auto-tag marks cite annotation bytes the same way (#379). + autotag_error = autotag_without_annotations_error(conn) + if autotag_error: + raise sqlite3.DatabaseError(autotag_error) migrate_annotation_state(conn) migrate_review_state(conn) + migrate_autotag_state(conn) # Last: the older generated_images rebuilds run only before file_path. migrate_workflow_documents(conn) @@ -1225,6 +1235,11 @@ class PromptModel: ) if review_error: return False, review_error + autotag_error = autotag_schema_error( + conn, allow_missing=not require_library_state + ) + if autotag_error: + return False, autotag_error documents_error = workflow_documents_schema_error( conn, allow_missing=not require_library_state, documents=documents ) diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index b69c521..5d0b7fe 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -15,12 +15,11 @@ from urllib.parse import quote from .job_routes import busy_response 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 .maintenance import disk_plan from .prompts import thumbnail_url_for from . import thumbnail_settings from .. import model_gate @@ -29,7 +28,7 @@ from aiohttp import web _INTERNAL_ERROR = "An internal error occurred. Check server logs for details." -# Tag output images (#328): the native prompt-tag batch. +# Tag output images (#328, per image since #379). AUTOTAG_JOB_TITLE = "Tag output images" AUTOTAG_PREVIEW_LEASE = "Tag output images review" AUTOTAG_OWNER_KIND = "autotag-job" @@ -200,6 +199,18 @@ async def _refusal(title): yield {"type": "error", "status": "error", "message": message} +def wd14_device(service) -> str: + """Worker thread: ``"gpu"`` when WD14 runs (or would run) on CUDA, else + ``"cpu"``; a service without the probe is ``"cpu"``.""" + probe = getattr(service, "wd14_device", None) + if not callable(probe): + return "cpu" + try: + return "gpu" if probe() == "gpu" else "cpu" + except Exception: + return "cpu" + + class AutotagRoutesMixin: """Mixin providing auto-tagging API endpoints.""" @@ -802,6 +813,8 @@ class AutotagRoutesMixin: async def _preview_autotag_batch(self, request): from ..autotag import get_autotag_service + from ..config import AutotagConfig + from ..tagging_defaults import WD14_MODEL_TYPES, wd14_seconds_per_image data, error = await self._batch_body( request, ("model_type", "existing"), BATCH_PREVIEW_BODY_MAX @@ -813,11 +826,14 @@ class AutotagRoutesMixin: return self._bad_request(message) model_type, existing = data["model_type"], data["existing"] policy = self._safety_backup_policy() + service = get_autotag_service() try: - candidates = await self._run_in_executor(self.db.autotag_batch_candidates) - status = await self._run_in_executor( - get_autotag_service().get_models_status - ) + targets = await self._autotag_targets() + marked = await self._run_in_executor(self.db.autotagged_locators) + status = await self._run_in_executor(service.get_models_status) + device = None + if model_type in WD14_MODEL_TYPES: + device = await self._run_in_executor(wd14_device, service) disk, _short, _size = await self._run_in_executor( functools.partial( disk_plan, @@ -831,16 +847,21 @@ class AutotagRoutesMixin: 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) + images = len(targets) + tagged = sum(1 for item in targets if item.locator in marked) + to_tag = images - tagged if existing == "skip" else images + estimate = None + if device is not None: + estimate = round(to_tag * wd14_seconds_per_image(device), 3) return web.json_response( { "success": True, - "images": len(candidates), - "prompts": len(prompts), + "images": images, "already_tagged": tagged, "to_tag": to_tag, + "estimate_seconds": estimate, + "estimate_device": device, + "suggestion_limit": AutotagConfig.suggestion_limit(), "model": { "type": model_type, "label": config["label"], @@ -862,9 +883,9 @@ class AutotagRoutesMixin: 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. + ``model-busy``. No await follows the busy check. The job tags each + output image (``_run_autotag_batch``), never a prompt (#379). """ - from ..autotag import get_autotag_service from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig from .annotation_suggest import OptionError, validate_options @@ -873,7 +894,7 @@ class AutotagRoutesMixin: ) if error is not None: return error - config, message = self._batch_choice(data) + _config, message = self._batch_choice(data) if message: return self._bad_request(message) model_type, existing, options = ( @@ -897,144 +918,24 @@ class AutotagRoutesMixin: {"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 - ) + job, refusal = await self._start_autotag( + model_type=model_type, + existing=existing, + options=options, + keep_in_memory=True, + title=AUTOTAG_JOB_TITLE, + owner_kind=AUTOTAG_OWNER_KIND, ) 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)) + if refusal is not None: + return refusal.response 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: diff --git a/py/api/image_tag_jobs.py b/py/api/image_tag_jobs.py index 38fa9b6..09aeb87 100644 --- a/py/api/image_tag_jobs.py +++ b/py/api/image_tag_jobs.py @@ -18,13 +18,21 @@ the suggestion worker and the storage call ``add_image_tags``. Files are handled one at a time and the job stops at a file boundary when cancelled, keeping what it already wrote. Rows and log lines carry root-relative paths and public root labels only. + +Auto-tag marks each image it ran the model on as auto-tagged (#379), with the +union of its tags. The ``autotag`` job ("Tag output images", started by +``POST /jobs/autotag``) runs the same per-file loop over every still image the +output folders list (:meth:`ImageTagJobsMixin._run_autotag_batch`): image +tags only, never a prompt tag or prompt text. """ import asyncio +import dataclasses import functools import os import time from dataclasses import dataclass, field +from typing import NamedTuple from aiohttp import web @@ -45,10 +53,17 @@ from .annotations import ( _submit, _target, ) -from .autotag_routes import model_busy_response -from .images import _public_path +from .autotag_routes import BATCH_CANCELLED, MODEL_BUSY_ERROR, model_busy_response +from .images import _current_fingerprint, _listing_posix, _public_path from .job_routes import busy_response -from .jobs import IMAGE_TAG_GROUP_NAMES, IMAGE_TAG_JOB_KIND, JobBusy +from .jobs import ( + AUTOTAG_BATCH_JOB_KIND, + AUTOTAG_GROUP_NAMES, + IMAGE_TAG_GROUP_NAMES, + IMAGE_TAG_JOB_KIND, + JobBusy, +) +from .maintenance import disk_plan, short_space_message from .prompts import bad_request try: @@ -107,6 +122,14 @@ READ_FAILED = "The file could not be read." SOURCE_CHANGED = "The source changed. Refresh the image collection." STORAGE_RESTORED = "Prompt storage was restored; start the job again." BUSY_RETRY = "The file could not be read now. Run the job again for it." +# Tag output images (#379): the per-image batch. +ALREADY_AUTOTAGGED = "Already auto-tagged" +LISTING_FAILED = "The output folders could not be listed. Try again." +LISTING_STATUS = "Listing output images…" +_INTERNAL_ERROR = "An internal error occurred. Check server logs for details." +# The pre-pass stats and looks up this many targets at a time, checking for a +# cancel between chunks. +PREPASS_CHUNK = 500 _monotonic = time.monotonic # the job's clock (a test can fake it) @@ -246,6 +269,92 @@ async def _capacity_retried(job, attempt): return await attempt() +class Refusal(NamedTuple): + """Why a Tag output images start was refused. + + ``response`` is the native route's JSON answer; ``message`` the sentence + a legacy stream frame carries. + """ + + status: int + code: str + message: str + response: web.Response + + +@dataclass(frozen=True) +class AutotagRun: + """What one Tag output images job runs with (validated at its start).""" + + model_type: str + existing: str + options: dict + keep_in_memory: bool + backup_on: bool + owner: object + + +@dataclass(frozen=True) +class _BatchTarget: + """One still image of the batch: the job target and its private keys.""" + + target: tuple # (root_index, root_scope, relative): the image-tags target + path: str # the listed file, for the pre-pass stat (never sent out) + locator: str # the image tag key: normcase(join(realpath(root), relative)) + label: str # the root's public label + + +def _still_targets(api, index, root, files, seen: set) -> list: + """Worker thread: one root's listed still images as _BatchTargets. + + Newest first (the listing's order). A locator an earlier root already + produced (overlapping roots) is dropped; ``seen`` collects them. + """ + scope = api._gallery_root_scope(root, index) + label = _public_path(root) + real = os.path.realpath(root) + prefix = os.path.join(str(root), "") + found = [] + for path, _mtime in files: + relative = _listing_posix(root, prefix, path) + if is_video_path(relative): + continue + locator = os.path.normcase(os.path.join(real, *relative.split("/"))) + if locator in seen: + continue + seen.add(locator) + found.append(_BatchTarget((index, scope, relative), str(path), locator, label)) + return found + + +def _prepass_chunk(db, chunk, skip_marked: bool, max_file: int): + """Worker thread: ``(remaining, skips)`` for a chunk of _BatchTargets. + + Stats each file (no open, no hash): gone is "File not found", over the + capture limit "Too large to verify"; under "Skip already tagged" a file + whose current bytes are marked is "Already auto-tagged". ``skips`` are + ``(target, detail)`` pairs in target order. + """ + stats, skips, alive = {}, [], [] + for item in chunk: + fingerprint = _current_fingerprint(item.path) + if fingerprint is None: + skips.append((item, FILE_NOT_FOUND)) + elif fingerprint[0] > max_file: + skips.append((item, TOO_LARGE)) + else: + stats[item.locator] = fingerprint + alive.append(item) + marked = db.autotag_state(stats) if skip_marked and stats else {} + remaining = [] + for item in alive: + if marked.get(item.locator): + skips.append((item, ALREADY_AUTOTAGGED)) + else: + remaining.append(item) + return remaining, skips + + def _load_model(service, model_type: str, use_gpu: bool) -> None: """Worker thread: load (or swap to) *model_type*; never downloads.""" if service.is_model_loaded() and service.get_loaded_model_type() == model_type: @@ -254,6 +363,25 @@ def _load_model(service, model_type: str, use_gpu: bool) -> None: raise RuntimeError(f"{model_type} could not be loaded") +def _refusal(status: int, code: str, message: str) -> Refusal: + response = web.json_response( + {"success": False, "code": code, "error": message}, status=status + ) + return Refusal(status, code, message, response) + + +def _busy_refusal(busy: JobBusy) -> Refusal: + title = getattr(busy.running, "title", "") or "Another job" + return Refusal(409, "busy", f"{title} already running", busy_response(busy)) + + +def _batch_cancelled(job) -> str: + """Say the batch was cancelled; its end state (the runner ends the job).""" + job.update(status_line=BATCH_CANCELLED) + job.log_line(BATCH_CANCELLED) + return "cancelled" + + class ImageTagJobsMixin: """The ``image-tags`` job kind: start route and runner.""" @@ -301,6 +429,216 @@ class ImageTagJobsMixin: job.task.add_done_callback(lambda _task: model_gate.release(owner)) return web.json_response({"success": True, "job": job.summary()}, status=202) + # -- Tag output images (#379) --------------------------------------------- + + async def _autotag_targets(self) -> list: + """Every still image the output folders list, as _BatchTargets. + + Root by root in configured order, each newest first (the shared + listing, ``_get_gallery_files``, used as it is even when stale). Videos + are not targets; a file under two overlapping roots is one target. + """ + seen: set = set() + found = [] + for index, root in enumerate(list(self._get_all_output_dirs())): + files = await self._get_gallery_files(root) + found.extend( + await self._run_in_executor( + _still_targets, self, index, root, files, seen + ) + ) + return found + + async def _start_autotag( + self, *, model_type, existing, options, keep_in_memory, title, owner_kind + ): + """Register a Tag output images job: ``(job, None)`` or ``(None, Refusal)``. + + ``options`` are already validated (``validate_options``). The checks + run in this order: 409 ``model-unavailable``; 507 + ``insufficient-space`` (only when the safety backup is on); 409 + ``busy`` (another records job); 409 ``model-busy``. No await follows + the busy check. The gate is taken before the job registers and + released when the task ends, also before its runner ever starts. + """ + from ..autotag import MODELS, get_autotag_service + + backup_on = self._safety_backup_policy() == "on" + 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=backup_on, vacuum=False + ) + ) + label = MODELS[model_type]["label"] + if not (status.get(model_type) or {}).get("downloaded"): + message = ( + f"{label} is not downloaded. Download it in System › Auto-tagging." + ) + return None, _refusal(409, "model-unavailable", message) + if backup_on and short is not None: + return None, _refusal(507, "insufficient-space", short_space_message(short)) + busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records") + if busy is not None: + return None, _busy_refusal(JobBusy(busy)) + owner = model_gate.new_owner(owner_kind) + if not model_gate.try_acquire(owner): + response = model_busy_response() + return None, Refusal(409, "model-busy", MODEL_BUSY_ERROR, response) + run = AutotagRun( + model_type, existing, options, keep_in_memory, backup_on, owner + ) + try: + job = self._jobs.start( + AUTOTAG_BATCH_JOB_KIND, + title, + admission="records", + cancellable=True, + scope={"model_type": model_type, "existing": existing, "label": label}, + groups=AUTOTAG_GROUP_NAMES, + runner=functools.partial(self._run_autotag_batch, run), + ) + except BaseException as refused: + model_gate.release(owner) + if isinstance(refused, JobBusy): + return None, _busy_refusal(refused) + raise + job.task.add_done_callback(lambda _task: model_gate.release(owner)) + return job, None + + async def _run_autotag_batch(self, run: AutotagRun, job) -> None: + """The ``autotag`` job: backup, model, listing, pre-pass, per-image loop. + + ``job.cancelled`` is checked between phases (and per pre-pass chunk + and per image). Each phase publishes ``{"type": "phase", "phase"}``. + With ``keep_in_memory`` false the model is unloaded under the gate. + The job ends only after that cleanup and the gate release, so a job + that reads as ended never still holds the gate. + """ + state = "failed" + try: + state = await self._autotag_batch_phases(run, job) + except Exception: + self.logger.exception("Tag output images failed") + job.update(status_line=_INTERNAL_ERROR) + job.log_line(f"Tagging failed · {_INTERNAL_ERROR}") + finally: + try: + if not run.keep_in_memory: + await self._unload_batch_model(run.owner) + if run.backup_on: + await self._prune_backup_pool(job, self.db) + finally: + model_gate.release(run.owner) + job.publish({"type": state, "progress": job.progress}) + job.finish(state) + + async def _unload_batch_model(self, owner) -> None: + from ..autotag import get_autotag_service + + service = get_autotag_service() + try: + if service.is_model_loaded(): + await model_gate.run(owner, service.unload_model) + except Exception: + self.logger.exception("The tagging model could not be unloaded") + + async def _autotag_batch_phases(self, run: AutotagRun, job) -> str: + """Run the phases; the job's end state (the caller ends the job).""" + job.publish({"type": "phase", "phase": "backup"}) + ok, _backup = await self._safety_backup_step(job, self.db, run.backup_on) + if not ok: + return "failed" + if job.cancelled: + return _batch_cancelled(job) + job.publish({"type": "phase", "phase": "model"}) + job.update(status_line="Loading the tagging model") + spec = _Spec("auto", (), model_type=run.model_type, options=run.options) + service = await self._prepare_tagging_model(job, spec, run.owner) + if service is None: + return "failed" + if job.cancelled: + return _batch_cancelled(job) + job.publish({"type": "phase", "phase": "listing"}) + job.update(status_line=LISTING_STATUS) + epoch = await self._run_in_executor(self.db.get_execution_epoch) + try: + targets = await self._autotag_targets() + except Exception: + self.logger.exception("Tag output images could not list the outputs") + job.update(status_line=LISTING_FAILED) + job.log_line(f"! {LISTING_FAILED}") + return "failed" + total = len(targets) + job.scope["images"] = total + job.update(total=total) + policy = ( + "skip images already auto-tagged" if run.existing == "skip" else "merge" + ) + job.log_line( + f"Tagging started · {job.scope.get('label') or run.model_type} · " + f"{_plural(total, 'image')} · {policy}" + ) + if job.cancelled: + return _batch_cancelled(job) + remaining = await self._autotag_prepass(job, targets, run.existing == "skip") + if remaining is None: + return _batch_cancelled(job) + done = total - len(remaining) + spec = dataclasses.replace( + spec, targets=tuple(item.target for item in remaining) + ) + loop_run = _Run(job, spec, epoch, run.owner, service) + if not await self._tag_each_file(loop_run, done_before=done): + job.update(status_line=STORAGE_RESTORED) + job.log_line(f"Tagging stopped · {STORAGE_RESTORED}") + return "failed" + if job.cancelled: + return _batch_cancelled(job) + finished = f"Tagged {_plural(job.counters.get('tagged', 0), 'image')}" + job.update(status_line=finished, progress=100) + job.log_line(finished) + return "finished" + + async def _autotag_prepass(self, job, targets, skip_marked: bool): + """The targets left for the model, or None when cancelled meanwhile. + + Opens and hashes nothing: one stat per file and one indexed read per + chunk. Skips are counted with ``record_counted`` (exact counters, rows + up to the cap, no log line each). + """ + max_file = self._capture_store().max_file + remaining, skipped = [], 0 + total = len(targets) + for start in range(0, total, PREPASS_CHUNK): + if job.cancelled: + return None + chunk = targets[start : start + PREPASS_CHUNK] + left, skips = await self._run_in_executor( + _prepass_chunk, self.db, chunk, skip_marked, max_file + ) + remaining.extend(left) + if skips: + job.record_counted( + "skipped", + len(skips), + [ + { + "path": item.target[2], + "root": item.label, + "root_index": item.target[0], + "detail": detail, + } + for item, detail in skips + ], + ) + skipped += len(skips) + job.update(processed=skipped, progress=skipped * 100 // total) + job.publish({"type": "progress", "progress": job.progress}) + if skipped: + job.log_line(f"Skipped {_plural(skipped, 'image')} before tagging") + return remaining + # -- runner ---------------------------------------------------------------- def _job_roots(self, targets) -> dict: @@ -331,28 +669,39 @@ class ImageTagJobsMixin: job.finish("failed") return if not await self._tag_each_file(run): + job.update(status_line=STORAGE_RESTORED) job.log_line(f"{verb} stopped · {STORAGE_RESTORED}") job.finish("failed") return state = "cancelled" if job.cancelled else "finished" - job.log_line(f"{verb} {state} · {_plural(job.processed, 'file')} processed") + # An ended job never keeps the loop's "about X min left" (#379). + ended = f"{verb} {state} · {_plural(job.processed, 'file')} processed" + job.update(status_line=ended) + job.log_line(ended) job.publish({"type": state, "progress": job.progress}) job.finish(state) finally: if owner is not None: model_gate.release(owner) - async def _tag_each_file(self, run: _Run) -> bool: - """Tag the targets in order; False when the job had to stop.""" - job, total = run.job, len(run.spec.targets) + async def _tag_each_file(self, run: _Run, *, done_before: int = 0) -> bool: + """Tag the targets in order; False when the job had to stop. + + ``done_before`` files were already counted (the batch's pre-pass): file + n of this loop is "file done_before + n" of the job, and the time left + uses this loop's rate only. + """ + job, count = run.job, len(run.spec.targets) + total = done_before + count run.roots = await self._run_in_executor(self._job_roots, run.spec.targets) started = _monotonic() for number, target in enumerate(run.spec.targets, 1): if job.cancelled: break # The scan's estimate (D56), with each finished file as one batch. - left = _time_left(number - 1, total, number - 1, _monotonic() - started) - job.update(status_line=f"Processing file {number:,} of {total:,}{left}") + left = _time_left(number - 1, count, number - 1, _monotonic() - started) + shown = done_before + number + job.update(status_line=f"Processing file {shown:,} of {total:,}{left}") job.publish({"type": "progress", "progress": job.progress}) stopped = False try: @@ -363,7 +712,7 @@ class ImageTagJobsMixin: self.logger.exception("Image tagging failed for one file") group, detail = "errors", SAVE_FAILED self._record_tag_outcome(run, target, group, detail) - job.update(processed=number, progress=number * 100 // total) + job.update(processed=shown, progress=shown * 100 // total) if stopped: return False return True @@ -439,9 +788,17 @@ class ImageTagJobsMixin: store.release(capture.token) async def _auto_tag_capture(self, run: _Run, store, capture, root, relative): - """Suggest on the capture, then union only onto the same current bytes.""" + """Suggest on the capture, then union only onto the same current bytes. + + The bytes are marked auto-tagged with the union, also when the model + returned no names; captured bytes no model can read are marked too + (and stay skipped), so a "Skip already tagged" run passes them by. + """ if not capture.inferable: - return "skipped", NOT_A_STILL + group, detail = await self._union_image_tags( + capture, [], run.epoch, mark_autotagged=True + ) + return (group, detail) if group == "errors" else ("skipped", NOT_A_STILL) if store.acquire(capture.token) is None: raise _JobStopped # every capture was invalidated (restore) spec = run.spec @@ -467,21 +824,30 @@ class ImageTagJobsMixin: if facts.locator != capture.locator or facts.sha256 != capture.sha256: return "skipped", FILE_CHANGED names = [str(name) for name, _score in pairs] - if not names: - return "unchanged", NO_NEW_TAGS - return await self._union_image_tags(facts, names, run.epoch) + return await self._union_image_tags( + facts, names, run.epoch, mark_autotagged=True + ) - async def _union_image_tags(self, facts, tags: list, epoch: str): - """``add_image_tags`` on the verified bytes, as a (group, detail).""" + async def _union_image_tags( + self, facts, tags: list, epoch: str, *, mark_autotagged: bool = False + ): + """``add_image_tags`` on the verified bytes, as a (group, detail). + + ``mark_autotagged`` (a model's result) marks the bytes auto-tagged in + the same write; then ``tags`` may be empty. + """ try: receipt = await self._run_in_executor( - self.db.add_image_tags, - locator=facts.locator, - sha256=facts.sha256, - size=facts.size, - mtime_ns=facts.mtime_ns, - expected_epoch=epoch, - tags=tags, + functools.partial( + self.db.add_image_tags, + locator=facts.locator, + sha256=facts.sha256, + size=facts.size, + mtime_ns=facts.mtime_ns, + expected_epoch=epoch, + tags=tags, + mark_autotagged=mark_autotagged, + ) ) except LibraryCommandError as exc: if exc.code == "storage-changed": diff --git a/py/autotag.py b/py/autotag.py index 9a6d55b..8938ed5 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -756,6 +756,24 @@ class AutoTagService: self.logger.info("Model unloaded, memory freed") + def wd14_device(self) -> str: + """``"gpu"`` when WD14 runs (or would run) on CUDA, else ``"cpu"``. + + A loaded WD14 session answers from its own first provider; otherwise + ONNX Runtime's available providers say whether CUDA would be offered + (a session can still fall back to the processor). Loads no model. + """ + tagger = self._tagger + if tagger is not None and tagger[0] == "wd14": + providers = tagger[1][0].get_providers() + return "gpu" if providers[:1] == ["CUDAExecutionProvider"] else "cpu" + try: + import onnxruntime as ort + except ImportError: + return "cpu" + available = ort.get_available_providers() + return "gpu" if "CUDAExecutionProvider" in available else "cpu" + def is_model_loaded(self) -> bool: """Check if a model is currently loaded.""" return self._tagger is not None diff --git a/py/tagging_defaults.py b/py/tagging_defaults.py index be75c9e..575a65a 100644 --- a/py/tagging_defaults.py +++ b/py/tagging_defaults.py @@ -13,7 +13,23 @@ 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). +# Tag output images' review estimate (#379): seconds per image for the WD14 +# engines, by where ONNX Runtime runs them. The graphics card figure was +# measured in PR 1's live check (RTX 4090, CUDA: Auto-tag selected on 100 +# images in 16.7 s with the model load, 0.167 s/image). The processor figure +# is an upper bound from the maintainer's cancelled CPU run (19 min / 697 +# images, backup, model load and skips included), kept as the fallback. The +# language-model engines have no figure. +WD14_SECONDS_PER_IMAGE = 1.6 +WD14_GPU_SECONDS_PER_IMAGE = 0.17 + + +def wd14_seconds_per_image(device: str) -> float: + """The estimate's rate for ``"gpu"`` or ``"cpu"`` (anything else).""" + return WD14_GPU_SECONDS_PER_IMAGE if device == "gpu" else WD14_SECONDS_PER_IMAGE + + +# What a batch does with an image that already has tags (#328, #379). EXISTING_POLICIES = ("skip", "merge", "replace_after_review") # The language-model tag template every page starts from until one is saved. diff --git a/tests/browser/system-tagging-concept.spec.mjs b/tests/browser/system-tagging-concept.spec.mjs index 61b1971..b02f5a8 100644 --- a/tests/browser/system-tagging-concept.spec.mjs +++ b/tests/browser/system-tagging-concept.spec.mjs @@ -33,7 +33,7 @@ test.afterEach(async ({ page, request, browserName }) => { const PREVIEW_URL = '/prompt_manager/autotag/batch/preview'; const MODEL_OPTIONS = ['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB']; const SUGGESTIONS = 'From 1 to 100: the most suggestions Generate shows for one image.'; -const BATCH = 'Add tags to each prompt from its linked output images with the model above. Review the scope first.'; +const BATCH = 'Tag each output image from its own pixels with the model above. Prompts are not changed. Review the scope first.'; /** Runs in the page: Auto-tagging's surface under `mainSelector`: rows, field, actions, operation, geometry and tokens. */ function measureTagging(mainSelector) { @@ -272,11 +272,19 @@ function measureReview() { } const number = value => value.toLocaleString('en-US'); +/** DEC11 (#379): the Time row the preview's estimate reads as (WD14 on the test server's processor). */ +function timeRow(plan) { + const minutes = Math.max(1, Math.ceil(plan.estimate_seconds / 60)); + const text = minutes < 90 ? `about ${minutes} min` + : `about ${Math.floor(minutes / 60)} h${minutes % 60 ? ` ${minutes % 60} min` : ''}`; + return `${text} on a processor; a graphics card can be faster.`; +} /** - * D134: the server's batch preview in the concept's review: "Review tag output images" with "Tag output images" as its - * `h3`, the model, policy, counts and settings in the concept's `dl.meta-grid` (Backup in the concept's wording), the - * upper-bound sentence in the concept's `p.small`, the D55 backup box and Cancel / "Start tagging". + * D134 (#379): the server's batch preview in the concept's review: "Review tag output images" with "Tag output images" + * as its `h3`, the model, policy, image counts, time estimate and settings in the concept's `dl.meta-grid` (Backup in + * the concept's wording), the estimate sentence in the concept's `p.small`, the D55 backup box and Cancel / "Start + * tagging". */ function mapBatchConcept(reference, plan) { const [cancel, start] = reference.actions.buttons; @@ -286,13 +294,14 @@ function mapBatchConcept(reference, plan) { tree: ['h3', 'dl.meta-grid', 'p.small', 'div.notice', 'div.dialog-actions'], heading: { ...reference.heading, text: 'Tag output images' }, grid: { ...reference.grid, rows: [['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Merge with existing'], - ['Images linked', number(plan.images)], ['Prompts', number(plan.prompts)], ['Already tagged', number(plan.already_tagged)], - ['Prompts to tag', `up to ${number(plan.to_tag)}`], ['Settings', 'General 0.35 · Character 0.85'], + ['Images', number(plan.images)], ['Already auto-tagged', number(plan.already_tagged)], + ['Images to tag', `about ${number(plan.to_tag)}`], ['Time', timeRow(plan)], + ['Settings', `General 0.35 · Character 0.85 · up to ${number(plan.suggestion_limit)} tags per image`], ['Backup', 'Create and verify before changes']] }, backup: { tree: ['div', 'div.row-actions'], text: [['h3', 'Back up the database first (recommended)'], ['p', 'Downloads a copy of the prompts database before this operation changes any records.']], buttons: ['Back up now'] }, - note: { ...reference.note, text: 'Nothing is written until you start. Missing files and files outside the output folders ' - + 'are counted, so the prompts to tag are an upper bound; Activity reports what was tagged.' }, + note: { ...reference.note, text: 'Nothing is written until you start. The number to tag is an estimate: images changed ' + + 'since they were tagged are tagged again, and missing files are skipped. Activity reports what was tagged.' }, actions: { ...reference.actions, buttons: [cancel, { ...start, text: 'Start tagging' }] }, }; } @@ -308,8 +317,8 @@ async function conceptReview(page, theme) { async function productionReview(page, theme) { await productionTagging(page, theme); - // Every concept-fixture prompt with an image is already tagged, so under the saved Skip there is nothing to start - // (the Cancel-only review). Merge reviews a batch with work: the full review the concept's shell is compared with. + // Merge reviews every listed still image (#379): a batch with work, the full review the concept's shell is compared + // with. await page.locator('#workspace-main #tag-existing').selectOption('merge'); const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL)); await page.locator('#workspace-main .operation').getByRole('button', { name: 'Review', exact: true }).click(); @@ -326,8 +335,9 @@ for (const theme of THEMES) for (const [width, height] of SIZES) { expect(concept.tree).toEqual(['p.demo-note', 'h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions']); await page.screenshot({ path: testInfo.outputPath(`concept-plan-review-${theme}-${width}.png`) }); const { plan, production } = await productionReview(page, theme); - // Under Merge every prompt with a linked image is tagged again: the batch has work to review. + // Under Merge every listed still image is tagged again: the batch has work to review. expect(plan.to_tag).toBeGreaterThan(0); + expect(plan.estimate_device).toBe('cpu'); await page.screenshot({ path: testInfo.outputPath(`production-tagging-review-${theme}-${width}.png`) }); expect(diffInventory(mapBatchConcept(concept, plan), production)).toEqual([]); }); @@ -372,9 +382,17 @@ test('tag output images review comparison detects mutations', async ({ page }) = const expected = mapBatchConcept(concept, plan); const paths = async () => diffInventory(expected, await page.evaluate(measureReview)).map(item => item.path); expect(await paths()).toEqual([]); - await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'To tag'; }); - expect(await paths()).toEqual(['grid.rows.5.0']); - await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'Prompts to tag'; }); + // Every label of the grid is compared (#379's rows included): a changed one is named by its row. + for (const [index, [term]] of expected.grid.rows.entries()) { + await page.evaluate(([at, label]) => { + document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[at].textContent = label; + }, [index, `${term}!`]); + expect(await paths()).toEqual([`grid.rows.${index}.0`]); + await page.evaluate(([at, label]) => { + document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[at].textContent = label; + }, [index, term]); + } + expect(await paths()).toEqual([]); const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' }); expect(await paths()).toEqual(expect.arrayContaining(['grid.columnGap'])); await undoGrid(); diff --git a/tests/browser/workflows.spec.mjs b/tests/browser/workflows.spec.mjs index 6ae6d97..62d8645 100644 --- a/tests/browser/workflows.spec.mjs +++ b/tests/browser/workflows.spec.mjs @@ -42,8 +42,13 @@ test('actual gallery form saves a prompt that the admin page can edit and reload await page.locator(`[data-action="edit"][data-prompt-id="${id}"]`).click(); const editor = page.locator(`.prompt-text[data-id="${id}"]`); await editor.fill('Browser edited portrait'); + // The page shows the edit at once and saves it after: reload only once the save answered (a reload that wins the + // race abandons the PUT, as seen on WebKit in CI). + const savedEdit = page.waitForResponse(item => item.request().method() === 'PUT' + && item.url().endsWith(`/prompt_manager/prompts/${id}`)); await editor.press('Enter'); await expect(page.locator('.prompt-text').filter({ hasText: 'Browser edited portrait' })).toBeVisible(); + expect((await savedEdit).ok()).toBe(true); await page.reload(); await expect(page.locator('.prompt-text').filter({ hasText: 'Browser edited portrait' })).toBeVisible(); const saved = await request.post('/prompt_manager/save', { data: { text: 'Browser saved landscape' } }); diff --git a/tests/browser/workspace-system-tagging.spec.mjs b/tests/browser/workspace-system-tagging.spec.mjs index a2720b1..0d9fa58 100644 --- a/tests/browser/workspace-system-tagging.spec.mjs +++ b/tests/browser/workspace-system-tagging.spec.mjs @@ -315,22 +315,28 @@ test('Preview tagging review without a selection toasts the D115 sentence', asyn await expect(tagDialog(page)).toHaveCount(0); }); -test('Tag output images reviews the scope and starts a batch job', async ({ page, request }) => { +test('Tag output images reviews the images, tags each one and never a prompt (#379)', async ({ page, request }) => { const before = await serverJson(request, '/__test__/tagging/snapshot'); await openTagging(page); - // The fixture's one prompt is already tagged: under Merge it is tagged again with what its images give. - await page.locator('#tag-existing').selectOption('merge'); + // The saved policy is Skip: no image is auto-tagged yet, so every listed still image is to tag. const previewed = page.waitForResponse(item => new URL(item.url()).pathname === PREVIEW_URL); await reviewButton(page).click(); const answer = await previewed; - expect(answer.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'merge' }); + expect(answer.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'skip' }); const plan = await answer.json(); + // Four still images (two linked, a copy, an unlinked one); the video is not a target. + expect([plan.images, plan.already_tagged, plan.to_tag]).toEqual([4, 0, 4]); + expect(plan.prompts).toBeUndefined(); await expect(dialog(page).locator('#modal-title')).toHaveText('Review tag output images'); await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Tag output images'); - expect(await gridOf(page)).toEqual([['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Merge with existing'], - ['Images linked', String(plan.images)], ['Prompts', String(plan.prompts)], ['Already tagged', String(plan.already_tagged)], - ['Prompts to tag', `up to ${plan.to_tag}`], ['Settings', 'General 0.35 · Character 0.85'], + expect(await gridOf(page)).toEqual([['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Skip already tagged'], + ['Images', '4'], ['Already auto-tagged', '0'], ['Images to tag', 'about 4'], + ['Time', 'about 1 min on a processor; a graphics card can be faster.'], + ['Settings', `General 0.35 · Character 0.85 · up to ${plan.suggestion_limit} tags per image`], ['Backup', 'Create and verify before changes']]); + await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Nothing is written until you start. The number to ' + + 'tag is an estimate: images changed since they were tagged are tagged again, and missing files are skipped. ' + + 'Activity reports what was tagged.'); await expect(dialog(page).locator('#modal-body > .notice h3')).toHaveText('Back up the database first (recommended)'); // The review wrote nothing. expect(await serverJson(request, '/__test__/tagging/snapshot')).toEqual(before); @@ -338,15 +344,23 @@ test('Tag output images reviews the scope and starts a batch job', async ({ page const started = page.waitForResponse(item => new URL(item.url()).pathname === '/prompt_manager/jobs/autotag'); await dialog(page).getByRole('button', { name: 'Start tagging', exact: true }).click(); const start = await started; - expect(start.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'merge', + expect(start.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'skip', options: { general_threshold: 0.35, character_threshold: 0.85 } }); const { job } = await start.json(); await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`); await expect(page.locator('#workspace-main h1')).toHaveText('Tag output images'); const finished = await ended(request, job.id); expect(finished.state).toBe('finished'); - const after = await serverJson(request, '/__test__/tagging/snapshot'); - expect(after.prompt_tags.length).toBeGreaterThan(before.prompt_tags.length); + expect(finished.status_line).toBe('Tagged 4 images'); + // Each image got its own tags; no prompt, pair or link row changed. + expect(await serverJson(request, '/__test__/tagging/snapshot')).toEqual(before); + + // Under Skip every image is now auto-tagged: nothing to start. + await openTagging(page); + await reviewButton(page).click(); + await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Every output image is already auto-tagged.'); + expect((await gridOf(page)).slice(2, 5)).toEqual([['Images', '4'], ['Already auto-tagged', '4'], ['Images to tag', 'about 0']]); + await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel']); }); test('Replace after review shows only Cancel with the reason', async ({ page, request }) => { @@ -355,18 +369,13 @@ test('Replace after review shows only Cancel with the reason', async ({ page, re await page.locator('#tag-existing').selectOption('replace_after_review'); await reviewButton(page).click(); await expect(dialog(page).locator('#modal-title')).toHaveText('Review tag output images'); - expect((await gridOf(page)).map(([term]) => term)).toEqual(['Model', 'Existing tags', 'Images linked', 'Prompts', 'Already tagged', - 'Prompts to tag', 'Settings', 'Backup']); + expect((await gridOf(page)).map(([term]) => term)).toEqual(['Model', 'Existing tags', 'Images', 'Already auto-tagged', + 'Images to tag', 'Time', 'Settings', 'Backup']); await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Replacing tags needs a review of each image. ' + 'Use Preview tagging review, or choose Skip or Merge for a batch run.'); await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel']); await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click(); await expect(dialog(page)).toHaveCount(0); - // Under Skip the fixture's only prompt is already tagged: nothing to start either. - await page.locator('#tag-existing').selectOption('skip'); - await reviewButton(page).click(); - await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Every prompt with an output image is already tagged.'); - await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel']); expect(await serverJson(request, '/__test__/tagging/snapshot')).toEqual(before); expect((await serverJson(request, '/prompt_manager/jobs')).jobs.filter(item => item.kind === 'autotag')).toEqual([]); }); diff --git a/tests/js/workspace-system-tagging.test.js b/tests/js/workspace-system-tagging.test.js index a8a9c6d..95c03b4 100644 --- a/tests/js/workspace-system-tagging.test.js +++ b/tests/js/workspace-system-tagging.test.js @@ -360,19 +360,23 @@ test('Preview tagging review opens the Tag image dialog with the section’s val } finally { scope.dispose(); dom.restore(); } }); -const PLAN = { images: 3, prompts: 2, already_tagged: 1, to_tag: 1, +const PLAN = { images: 3, already_tagged: 1, to_tag: 2, estimate_seconds: 3.2, estimate_device: 'cpu', suggestion_limit: 10, model: { type: 'wd14-swinv2', label: 'WD14 SwinV2', size_label: '~400 MB', downloaded: true }, existing: 'skip', review_required: false, backup: { policy: 'on' }, disk: { free_bytes: 10 * 2 ** 30, required_bytes: 2 ** 20, sufficient: true } }; -test('the batch review grid, body and Cancel-only reasons (D134, CUJ 4–5)', async () => { - const { batchGrid, batchBody, batchCloseNote, existingLabel, REPLACE_NOTE, NOTHING_TO_TAG, NO_LINKED_IMAGES } = await tagging; +test('the batch review grid, body and Cancel-only reasons (D134, #379, CUJ 4–5)', async () => { + const { batchGrid, batchBody, batchCloseNote, existingLabel, REPLACE_NOTE, NOTHING_TO_TAG, NO_IMAGES } = await tagging; assert.deepEqual(batchGrid(PLAN, HOST), [['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Skip already tagged'], - ['Images linked', '3'], ['Prompts', '2'], ['Already tagged', '1'], ['Prompts to tag', 'up to 1'], - ['Settings', 'General 0.35 · Character 0.85'], ['Backup', 'Create and verify before changes']]); + ['Images', '3'], ['Already auto-tagged', '1'], ['Images to tag', 'about 2'], + ['Time', 'about 1 min on a processor; a graphics card can be faster.'], + ['Settings', 'General 0.35 · Character 0.85 · up to 10 tags per image'], ['Backup', 'Create and verify before changes']]); const llm = { ...HOST, autotag_model_type: 'gguf' }; - assert.deepEqual(batchGrid({ ...PLAN, model: { type: 'gguf' }, backup: { policy: 'off' } }, llm).slice(-2), - [['Settings', 'Template of 27 characters'], ['Backup', 'No safety backup (the policy is off)']]); + assert.deepEqual(batchGrid({ ...PLAN, model: { type: 'gguf' }, backup: { policy: 'off' }, estimate_seconds: null, + estimate_device: null, suggestion_limit: 4 }, llm).slice(-3), + [['Time', 'Language-model engines take longer per image. Activity shows the time left once tagging starts.'], + ['Settings', 'Template of 27 characters · up to 4 tags per image'], ['Backup', 'No safety backup (the policy is off)']]); assert.equal(batchGrid({}, HOST)[2][1], '0'); + assert.equal(batchGrid({}, HOST)[6][1], 'General 0.35 · Character 0.85'); assert.equal(existingLabel('odd'), 'odd'); assert.deepEqual(batchBody(HOST), { model_type: 'wd14-swinv2', existing: 'skip', options: { general_threshold: 0.35, character_threshold: 0.85 } }); @@ -385,9 +389,31 @@ test('the batch review grid, body and Cancel-only reasons (D134, CUJ 4–5)', as assert.equal(batchCloseNote({ ...PLAN, model: { ...PLAN.model, downloaded: false } }), 'WD14 SwinV2 is not downloaded. Download it first.'); assert.equal(batchCloseNote({ ...PLAN, disk: { free_bytes: 2 ** 20, required_bytes: 2 ** 21, sufficient: false } }), 'Not enough free disk space: the safety backup needs about 2.0 MB and 1.0 MB is free. Nothing is written until there is room.'); - assert.equal(batchCloseNote({ ...PLAN, prompts: 0, to_tag: 0 }), NO_LINKED_IMAGES); + assert.equal(batchCloseNote({ ...PLAN, images: 0, to_tag: 0 }), NO_IMAGES); + assert.equal(batchCloseNote({ ...PLAN, images: undefined }), NO_IMAGES); + assert.equal(NO_IMAGES, 'The output folders contain no images to tag.'); assert.equal(batchCloseNote({ ...PLAN, to_tag: 0 }), NOTHING_TO_TAG); - assert.equal(NOTHING_TO_TAG, 'Every prompt with an output image is already tagged.'); + assert.equal(NOTHING_TO_TAG, 'Every output image is already auto-tagged.'); +}); + +test('the review Time row (DEC11, #379): minutes, hours and the engine it runs on', async () => { + const { durationText, timeText, TIME_UNKNOWN } = await tagging; + assert.equal(durationText(1), 'about 1 min'); + assert.equal(durationText(0), 'about 1 min'); + assert.equal(durationText(89), 'about 89 min'); + assert.equal(durationText(90), 'about 1 h 30 min'); + assert.equal(durationText(120), 'about 2 h'); + assert.equal(durationText(631), 'about 10 h 31 min'); + assert.equal(durationText(Number.NaN), 'about 1 min'); + // 23,645 images at 1.6 s on a processor (the design's CUJ 1). + assert.equal(timeText({ estimate_seconds: 37832, estimate_device: 'cpu' }), + 'about 10 h 31 min on a processor; a graphics card can be faster.'); + assert.equal(timeText({ estimate_seconds: 59, estimate_device: 'gpu' }), 'about 1 min on the graphics card.'); + assert.equal(timeText({ estimate_seconds: 0, estimate_device: 'cpu' }), 'about 1 min on a processor; a graphics card can be faster.'); + for (const plan of [{ estimate_seconds: null }, {}, { estimate_seconds: -1 }, { estimate_seconds: Infinity }, null]) { + assert.equal(timeText(plan), TIME_UNKNOWN); + } + assert.equal(TIME_UNKNOWN, 'Language-model engines take longer per image. Activity shows the time left once tagging starts.'); }); test('showBatchPlan opens the review with the backup box, or the Cancel-only review', async () => { diff --git a/tests/test_annotation_schema.py b/tests/test_annotation_schema.py index 4961117..ce5c502 100644 --- a/tests/test_annotation_schema.py +++ b/tests/test_annotation_schema.py @@ -35,6 +35,8 @@ REVIEW_TABLES = ( "media_review_items", "media_review_metadata", ) +# The auto-tag marks (#379) cite annotation bytes too; an old backup lacks them. +CITING_TABLES = REVIEW_TABLES + ("media_image_autotags",) # The exact DDL the plan pins; the module must produce precisely these objects. DDL = { METADATA: ( @@ -214,7 +216,7 @@ class AnnotationSchemaTests(unittest.TestCase): """ with closing(sqlite3.connect(path)) as raw, raw: raw.execute("PRAGMA foreign_keys=OFF") - for table in REVIEW_TABLES + tuple(reversed(TABLES)): + for table in CITING_TABLES + tuple(reversed(TABLES)): raw.execute(f"DROP TABLE {table}") def test_fresh_install_and_reopen_is_idempotent(self): diff --git a/tests/test_annotations.py b/tests/test_annotations.py index a5aa600..a70e219 100644 --- a/tests/test_annotations.py +++ b/tests/test_annotations.py @@ -561,6 +561,142 @@ class CurrentImageTagsTests(AnnotationTestCase): ) +class AutotagMarkTests(AnnotationTestCase): + """The auto-tagged mark (#379): per byte record, written with the tags.""" + + def marked(self): + conn = self.db.model.get_connection() + return sorted( + row[0] for row in conn.execute("SELECT * FROM media_image_autotags") + ) + + def byte_id(self, loc, digest=DIGEST_A): + conn = self.db.model.get_connection() + return conn.execute( + "SELECT b.id FROM media_annotation_bytes b JOIN media_annotation_files f " + "ON f.id=b.file_id WHERE f.locator=? AND b.sha256=?", + (loc, digest), + ).fetchone()[0] + + def test_hand_tags_never_mark(self): + loc = self.locator("a.png") + self.replace(loc, tags=["keeper"]) + self.add(loc, tags=["more"]) + self.assertEqual(self.marked(), []) + self.assertEqual(self.db.autotag_state({loc: (SIZE, MTIME)}), {loc: False}) + self.assertEqual(self.db.autotagged_locators(), set()) + + def test_model_union_marks_including_replayed_and_empty_results(self): + loc = self.locator("a.png") + receipt = self.add(loc, tags=["bird"], mark_autotagged=True) + self.assertEqual((receipt["added"], receipt["replayed"]), (1, False)) + self.assertEqual(self.marked(), [self.byte_id(loc)]) + # A replayed union and an empty result still mark (idempotently). + other = self.locator("b.png") + self.replace(other, DIGEST_B, ["keeper"]) + replayed = self.add(other, DIGEST_B, ["KEEPER"], mark_autotagged=True) + self.assertTrue(replayed["replayed"]) + empty = self.locator("c.png") + nothing = self.add(empty, tags=[], mark_autotagged=True) + self.assertEqual((nothing["tags"], nothing["added"]), ([], 0)) + self.add(empty, tags=[], mark_autotagged=True) + self.assertEqual( + self.marked(), + sorted( + [self.byte_id(loc), self.byte_id(other, DIGEST_B), self.byte_id(empty)] + ), + ) + self.assertEqual( + self.db.autotag_state( + {loc: (SIZE, MTIME), other: (SIZE, MTIME), empty: (SIZE, MTIME)} + ), + {loc: True, other: True, empty: True}, + ) + self.schema_clean() + + def test_empty_union_without_mark_is_refused(self): + with self.assertRaises(ValueError): + self.add(self.locator("a.png"), tags=[]) + self.assertEqual(self.counts()["media_annotation_files"], 0) + + def test_tag_limit_union_writes_no_tag_and_no_mark(self): + loc = self.locator("a.png") + self.replace(loc, tags=[f"t{i}" for i in range(100)]) + before = self.rows(*ANNOTATION_TABLES, "tags") + with self.assertRaises(AnnotationCommandError) as ctx: + self.add(loc, tags=["one more"], mark_autotagged=True) + self.assertEqual(ctx.exception.code, "tag-limit") + self.assertEqual(self.rows(*ANNOTATION_TABLES, "tags"), before) + self.assertEqual(self.marked(), []) + + def test_replace_marks_on_saved_and_replayed_receipts_not_on_conflict(self): + loc = self.locator("a.png") + saved = self.replace(loc, tags=["a"], mark_autotagged=True) + self.assertFalse(saved["replayed"]) + self.assertEqual(self.marked(), [self.byte_id(loc)]) + other = self.locator("b.png") + self.replace(other, DIGEST_B, ["x"]) + replayed = self.replace( + other, DIGEST_B, ["x"], revision=1, mark_autotagged=True + ) + self.assertTrue(replayed["replayed"]) + third = self.locator("c.png") + self.replace(third, tags=["y"]) + with self.assertRaises(AnnotationCommandError): + self.replace(third, tags=["z"], revision=5, mark_autotagged=True) + self.assertEqual( + self.marked(), sorted([self.byte_id(loc), self.byte_id(other, DIGEST_B)]) + ) + + def test_mark_follows_the_bytes_not_the_file(self): + loc = self.locator("a.png") + self.add(loc, tags=["bird"], mark_autotagged=True) + # The file's bytes changed: the new record is current and unmarked. + self.replace(loc, DIGEST_B, ["hand"], mtime_ns=MTIME + 1) + self.assertEqual(self.db.autotag_state({loc: (SIZE, MTIME + 1)}), {loc: False}) + self.assertEqual(self.db.autotag_state({loc: (SIZE, MTIME)}), {loc: True}) + # The review's estimate counts any marked record of the locator. + self.assertEqual(self.db.autotagged_locators(), {loc}) + missing = self.locator("none.png") + self.assertEqual(self.db.autotag_state({missing: (1, 1)}), {missing: False}) + + def test_state_and_tags_choose_the_same_current_record(self): + loc = self.locator("a.png") + # Two records with the same verified fingerprint: the newest wins. + self.replace(loc, DIGEST_A, ["old bytes"]) + with patch.object( + annotations_module, "utc_now_iso", return_value="2099-01-01T00:00:00" + ): + self.add(loc, DIGEST_B, ["new bytes"], mark_autotagged=True) + stats = {loc: (SIZE, MTIME)} + self.assertEqual(self.db.current_image_tags(stats), {loc: ["new bytes"]}) + self.assertEqual(self.db.autotag_state(stats), {loc: True}) + with patch.object( + annotations_module, "utc_now_iso", return_value="2199-01-01T00:00:00" + ): + self.add(loc, DIGEST_A, ["older bytes again"]) + self.assertEqual( + self.db.current_image_tags(stats), {loc: ["old bytes", "older bytes again"]} + ) + self.assertEqual(self.db.autotag_state(stats), {loc: False}) + + def test_marks_never_touch_prompt_rows(self): + conn = self.db.model.get_connection() + prompt_id = self.db.save_prompt("a beach", tags=["beach"]) + before = [ + [tuple(row) for row in conn.execute(f"SELECT * FROM {table} ORDER BY 1")] + for table in ("prompts", "prompt_tags") + ] + self.add(self.locator("a.png"), tags=["bird"], mark_autotagged=True) + self.replace(self.locator("b.png"), DIGEST_B, ["waves"], mark_autotagged=True) + after = [ + [tuple(row) for row in conn.execute(f"SELECT * FROM {table} ORDER BY 1")] + for table in ("prompts", "prompt_tags") + ] + self.assertIsNotNone(prompt_id) + self.assertEqual(after, before) + + class AnnotationIsolationTests(AnnotationTestCase): def test_cascade_from_vocabulary_delete_keeps_byte_records(self): loc = self.locator("a.png") diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 444594c..da42059 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -1893,13 +1893,15 @@ class TestTagOutputImagesPreview(BatchJobTestCase): self.assertEqual(self.service.generate_calls, []) self.assertEqual(self.service.load_calls, []) - async def test_preview_counts_skip_and_merge(self): + async def test_preview_counts_listed_still_images_not_prompts(self): self._seed() skip = await (await self._preview(existing="skip")).json() + # t, f1, f2 and m are listed; gone.png is missing, o.png is outside. self.assertEqual( - {k: skip[k] for k in ("images", "prompts", "already_tagged", "to_tag")}, - {"images": 6, "prompts": 5, "already_tagged": 1, "to_tag": 4}, + {k: skip[k] for k in ("images", "already_tagged", "to_tag")}, + {"images": 4, "already_tagged": 0, "to_tag": 4}, ) + self.assertNotIn("prompts", skip) self.assertEqual( skip["model"], { @@ -1916,7 +1918,7 @@ class TestTagOutputImagesPreview(BatchJobTestCase): set(skip["disk"]), {"free_bytes", "required_bytes", "sufficient"} ) merge = await (await self._preview(existing="merge")).json() - self.assertEqual(merge["to_tag"], 5) + self.assertEqual(merge["to_tag"], 4) leaked = [s for s in _strings(skip) if _absolute(s)] self.assertEqual(leaked, []) @@ -1924,7 +1926,7 @@ class TestTagOutputImagesPreview(BatchJobTestCase): self._seed() data = await (await self._preview(existing="replace_after_review")).json() self.assertTrue(data["review_required"]) - self.assertEqual(data["prompts"], 5) + self.assertEqual(data["images"], 4) async def test_preview_refuses_bad_bodies(self): for body in ( @@ -1979,119 +1981,6 @@ class TestTagOutputImagesJob(BatchJobTestCase): "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" @@ -2103,41 +1992,6 @@ class TestTagOutputImagesJob(BatchJobTestCase): 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" @@ -2185,7 +2039,7 @@ class TestTagOutputImagesJob(BatchJobTestCase): ("database folder", "the safety backup", 10, 1), 10, ) - with patch("py.api.autotag_routes.disk_plan", return_value=short): + with patch("py.api.image_tag_jobs.disk_plan", return_value=short): resp = await self._start_job() data = await resp.json() self.assertEqual(resp.status, 507) @@ -2203,14 +2057,6 @@ class TestTagOutputImagesJob(BatchJobTestCase): 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). @@ -2342,7 +2188,7 @@ class TestTagOutputImagesEdges(BatchJobTestCase): 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")] + self.api, "_autotag_prepass", side_effect=RuntimeError("/srv/x") ): job = await self._run_job(existing="merge") self.assertEqual(job.state, "failed") diff --git a/tests/test_autotag_batch.py b/tests/test_autotag_batch.py new file mode 100644 index 0000000..c6703ec --- /dev/null +++ b/tests/test_autotag_batch.py @@ -0,0 +1,696 @@ +"""Tag output images tags each image, never its prompt (#379, PR 1). + +``POST /prompt_manager/jobs/autotag`` runs the ``image-tags`` loop over every +still image the output folders list, linked to a prompt or not: each image +gets its own model tags in its image tag set and is marked auto-tagged. +"Skip already tagged" skips images whose current bytes are marked, counted +before the loop without opening them. Prompt tags and prompt text never +change. No real model is loaded: the engine is a fake that answers per image. +""" + +import asyncio +import json +import ntpath +import os +import posixpath +import re +import sqlite3 +import threading +from pathlib import Path +from unittest.mock import patch + +try: + from test_api_images import make_png + from test_image_tag_jobs import HookedTagger, ImageTagJobTestCase +except ImportError: + from tests.test_api_images import make_png + from tests.test_image_tag_jobs import HookedTagger, ImageTagJobTestCase + +from database.execution_history_schema import rotate_history_epoch +from py import model_gate +from py.api import admin as admin_module +from py.api import image_tag_jobs as jobs_module +from py.config import PromptManagerConfig +from utils.hashing import generate_prompt_hash + +JOB = "/prompt_manager/jobs/autotag" +PREVIEW = "/prompt_manager/autotag/batch/preview" +WAIT = 30.0 +WD14 = {"general_threshold": 0.35, "character_threshold": 0.85} +CANCELLED = "Cancelled. Tags already added were kept." +_DRIVE = re.compile(r"^[A-Za-z]:[\\/]") + + +class PerImageTagger(HookedTagger): + """Answers each image from its own bytes: ``by_bytes`` maps exact file + contents to the names the "model" sees; any other image gets SCORED.""" + + def __init__(self): + super().__init__() + self.by_bytes = {} + + def generate_tag_suggestions(self, image_path, **kwargs): + result = super().generate_tag_suggestions(image_path, **kwargs) + names = self.by_bytes.get(self.read_bytes[-1]) + if names is None: + return result + return [(name, 0.9) for name in names][: kwargs["limit"]] + + +def _strings(value): + if isinstance(value, str): + yield value + elif isinstance(value, dict): + for item in value.values(): + yield from _strings(item) + elif isinstance(value, (list, tuple)): + for item in value: + yield from _strings(item) + + +def _absolute(text): + return posixpath.isabs(text) or ntpath.isabs(text) or bool(_DRIVE.match(text)) + + +class BatchTestCase(ImageTagJobTestCase): + async def get_application(self): + app = await super().get_application() + self.tagger = PerImageTagger() + patcher = patch("py.autotag.get_autotag_service", return_value=self.tagger) + patcher.start() + self.addCleanup(patcher.stop) + backup = PromptManagerConfig.SAFETY_BACKUP + self.addCleanup(setattr, PromptManagerConfig, "SAFETY_BACKUP", backup) + PromptManagerConfig.SAFETY_BACKUP = False + return app + + # -- helpers --------------------------------------------------------- + + def image(self, name, color, sees=None): + path = make_png(self.output_dir / name, size=(8, 6), color=color) + if sees is not None: + self.tagger.by_bytes[path.read_bytes()] = list(sees) + return path + + async def relist(self): + """Walk the output folder again (a cached listing is used as it is).""" + await self.api._listing(self.output_dir, fresh=True) + + def prompt(self, text, tags=()): + return self.api.db.save_prompt( + text=text, + category=None, + tags=list(tags), + rating=None, + prompt_hash=generate_prompt_hash(text), + ) + + def link(self, prompt_id, path): + return self.api.db.link_image_to_prompt(prompt_id, str(path)) + + def prompt_rows(self): + conn = sqlite3.connect(self._temp_db.name) + try: + return { + table: conn.execute(f"SELECT * FROM {table} ORDER BY 1,2").fetchall() + for table in ("prompts", "prompt_tags") + } + finally: + conn.close() + + def marked(self): + conn = sqlite3.connect(self._temp_db.name) + try: + return { + row[0] + for row in conn.execute( + "SELECT f.locator FROM media_image_autotags m JOIN " + "media_annotation_bytes b ON b.id=m.byte_id JOIN " + "media_annotation_files f ON f.id=b.file_id" + ) + } + finally: + conn.close() + + async def preview(self, existing="skip", model_type="wd14-vit"): + response = await self.client.post( + PREVIEW, json={"model_type": model_type, "existing": existing} + ) + return response.status, await self.json_of(response) + + async def start_batch(self, existing="skip", model_type="wd14-vit", options=None): + if options is None: + options = ( + dict(WD14) + if model_type.startswith("wd14") + else {"instructions": "Tags please."} + ) + response = await self.client.post( + JOB, + json={"model_type": model_type, "existing": existing, "options": options}, + ) + return response.status, await self.json_of(response) + + async def run_batch(self, **kwargs): + status, data = await self.start_batch(**kwargs) + self.assertEqual(status, 202, data) + job = self.api._jobs.get(data["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + return job + + +class TestPerImageWrites(BatchTestCase): + async def test_two_images_of_one_prompt_get_their_own_tags_and_no_prompt_changes( + self, + ): + prompt_id = self.prompt("a girl at the beach", tags=["beach"]) + a = self.image("a.png", (10, 20, 30), sees=["1girl", "bird"]) + b = self.image("b.png", (40, 50, 60), sees=["1girl", "waves"]) + self.link(prompt_id, a) + self.link(prompt_id, b) + before = self.prompt_rows() + + job = await self.run_batch(existing="skip") + + self.assertEqual(job.state, "finished", job.status_line) + self.assertEqual(self.stored("a.png")["tags"], ["1girl", "bird"]) + self.assertEqual(self.stored("b.png")["tags"], ["1girl", "waves"]) + self.assertEqual(self.prompt_rows(), before) + self.assertEqual(self.api.db.get_prompt_by_id(prompt_id)["tags"], ["beach"]) + self.assertEqual(self.marked(), {self.locator("a.png"), self.locator("b.png")}) + self.assertEqual( + job.counters, {"tagged": 2, "unchanged": 0, "skipped": 0, "errors": 0} + ) + self.assertEqual(job.status_line, "Tagged 2 images") + self.assertEqual((job.processed, job.total, job.progress), (2, 2, 100)) + self.assertEqual(job.scope["images"], 2) + # Inference read the private captures, one per image. + self.assertEqual(len(self.tagger.calls), 2) + + async def test_unlinked_images_are_tagged_and_videos_are_not_targets(self): + self.image("free.png", (1, 2, 3), sees=["tree"]) + (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 8) + + job = await self.run_batch() + + self.assertEqual(self.stored("free.png")["tags"], ["tree"]) + self.assertEqual(job.total, 1) + self.assertEqual(job.status_line, "Tagged 1 image") + rows = [row["path"] for rows in job.outcomes.values() for row in rows] + self.assertNotIn("clip.mp4", rows) + + async def test_hand_tags_do_not_block_a_skip_run(self): + path = self.image("c.png", (7, 7, 7), sees=["bird"]) + self.saved_set("c.png", path.read_bytes(), ["keeper"]) + + job = await self.run_batch(existing="skip") + + self.assertEqual(self.stored("c.png")["tags"], ["keeper", "bird"]) + self.assertEqual(job.counters["tagged"], 1) + + async def test_skip_counts_marked_images_without_opening_them_and_merge_retags( + self, + ): + self.image("a.png", (1, 1, 1), sees=["bird"]) + self.image("b.png", (2, 2, 2), sees=["cat"]) + first = await self.run_batch() + self.assertEqual(first.counters["tagged"], 2) + self.tagger.calls.clear() + + captures = [] + store = self.api._capture_store() + real_capture = store.capture + + def spy(*args, **kwargs): + captures.append(args) + return real_capture(*args, **kwargs) + + with patch.object(store, "capture", side_effect=spy): + skipped = await self.run_batch(existing="skip") + self.assertEqual(captures, []) + self.assertEqual(self.tagger.calls, []) + self.assertEqual( + skipped.counters, {"tagged": 0, "unchanged": 0, "skipped": 2, "errors": 0} + ) + self.assertEqual( + sorted(self.rows(skipped, "skipped")), + [("a.png", "Already auto-tagged"), ("b.png", "Already auto-tagged")], + ) + self.assertFalse(any("Already auto-tagged" in line for line in skipped.log)) + self.assertEqual((skipped.processed, skipped.total), (2, 2)) + self.assertEqual(skipped.status_line, "Tagged 0 images") + + self.tagger.by_bytes[(self.output_dir / "a.png").read_bytes()] = ["bird", "sky"] + merged = await self.run_batch(existing="merge") + self.assertEqual(len(self.tagger.calls), 2) + self.assertEqual(self.stored("a.png")["tags"], ["bird", "sky"]) + self.assertEqual(self.rows(merged, "tagged"), [("a.png", "Added 1 tag")]) + self.assertEqual(self.rows(merged, "unchanged"), [("b.png", "No new tags")]) + + async def test_changed_bytes_are_tagged_again(self): + self.image("a.png", (1, 1, 1), sees=["bird"]) + await self.run_batch() + self.image("a.png", (9, 9, 9), sees=["cat"]) + job = await self.run_batch(existing="skip") + self.assertEqual(self.rows(job, "tagged"), [("a.png", "Added 1 tag")]) + self.assertEqual(self.stored("a.png")["tags"], ["cat"]) + + async def test_an_image_with_no_model_tags_is_marked_and_skipped_next_time(self): + self.image("plain.png", (3, 3, 3), sees=[]) + job = await self.run_batch() + self.assertEqual(self.rows(job, "unchanged"), [("plain.png", "No new tags")]) + self.assertEqual(self.marked(), {self.locator("plain.png")}) + again = await self.run_batch() + self.assertEqual( + self.rows(again, "skipped"), [("plain.png", "Already auto-tagged")] + ) + + async def test_unsupported_still_is_marked_and_skipped_by_the_next_run(self): + (self.output_dir / "broken.png").write_bytes(b"not an image") + job = await self.run_batch() + self.assertEqual( + self.rows(job, "skipped"), [("broken.png", "Not a supported still image")] + ) + self.assertEqual(self.marked(), {self.locator("broken.png")}) + again = await self.run_batch() + self.assertEqual( + self.rows(again, "skipped"), [("broken.png", "Already auto-tagged")] + ) + self.assertEqual(self.tagger.calls, []) + + async def test_too_large_is_skipped_before_any_copy(self): + self.image("big.png", (5, 5, 5)) + self.api._capture_store().max_file = 16 + store = self.api._capture_store() + with patch.object(store, "capture", side_effect=AssertionError("copied")): + job = await self.run_batch() + self.assertEqual( + self.rows(job, "skipped"), [("big.png", "Too large to verify")] + ) + self.assertEqual(self.marked(), set()) + + async def test_overlapping_roots_list_a_file_once(self): + self.image("sub/a.png", (4, 4, 4), sees=["bird"]) + sub = self.output_dir / "sub" + self.api._get_all_output_dirs = lambda: [self.output_dir, sub] + job = await self.run_batch() + self.assertEqual(job.total, 1) + self.assertEqual(len(self.tagger.calls), 1) + self.assertEqual(self.rows(job, "tagged"), [("sub/a.png", "Added 1 tag")]) + status, data = await self.preview() + self.assertEqual((status, data["images"]), (200, 1)) + + async def test_rows_log_and_status_carry_no_absolute_paths(self): + self.image("a.png", (1, 1, 1)) + (self.output_dir / "broken.png").write_bytes(b"not an image") + job = await self.run_batch() + leaked = [s for s in _strings(job.report()) if _absolute(s)] + self.assertEqual(leaked, []) + self.assertNotIn(str(self.output_dir), json.dumps(job.report())) + + +class TestBatchLifecycle(BatchTestCase): + async def test_cancel_after_one_image_keeps_it_and_the_rerun_skips_it(self): + for n in range(3): + self.image(f"{n}.png", (n, n, n), sees=[f"t{n}"]) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + status, data = await self.start_batch() + self.assertEqual(status, 202, data) + job = self.api._jobs.get(data["job"]["id"]) + try: + await self.wait_for(entered) + cancel = await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel") + self.assertEqual(cancel.status, 200) + finally: + self.tagger.block = None + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.status_line, CANCELLED) + self.assertEqual(job.counters["tagged"], 1) + self.assertEqual(len(self.marked()), 1) + self.assertIsNone(model_gate.holder()) + + lines = [] + real_update = type(job).update + + def spy(self_job, **kwargs): + if kwargs.get("status_line"): + lines.append(kwargs["status_line"]) + return real_update(self_job, **kwargs) + + with patch.object(type(job), "update", spy): + rerun = await self.run_batch() + self.assertEqual(rerun.counters["skipped"], 1) + self.assertEqual(rerun.counters["tagged"], 2) + self.assertIn("Processing file 2 of 3", lines) + self.assertIn("Processing file 3 of 3", lines) + self.assertNotIn("Processing file 1 of 3", lines) + self.assertEqual((rerun.processed, rerun.total), (3, 3)) + + async def test_cancel_in_each_phase_writes_nothing(self): + self.image("a.png", (1, 1, 1)) + for phase in ( + "_safety_backup_step", + "_prepare_tagging_model", + "_autotag_targets", + ): + with self.subTest(phase): + real = getattr(self.api, phase) + + async def cancelling(*args, _real=real, **kwargs): + result = await _real(*args, **kwargs) + self.api._jobs.active("autotag").cancelled = True + return result + + with patch.object(self.api, phase, side_effect=cancelling): + job = await self.run_batch() + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.status_line, CANCELLED) + self.assertEqual(self.marked(), set()) + self.assertEqual(self.tagger.calls, []) + self.assertIsNone(model_gate.holder()) + + async def test_cancel_between_prepass_chunks_keeps_counts(self): + for n in range(3): + self.image(f"{n}.png", (n, n, n)) + await self.run_batch() + seen = [] + real = jobs_module._prepass_chunk + + def chunk(db, items, skip_marked, max_file): + result = real(db, items, skip_marked, max_file) + job = self.api._jobs.active("autotag") + seen.append(job.processed) + job.cancelled = True # cancel after the first chunk + return result + + with ( + patch.object(jobs_module, "PREPASS_CHUNK", 1), + patch.object(jobs_module, "_prepass_chunk", side_effect=chunk), + ): + job = await self.run_batch() + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.counters["skipped"], 1) + self.assertEqual((job.processed, job.total), (1, 3)) + self.assertEqual(seen, [0]) + + async def test_the_job_ends_only_after_its_cleanup(self): + self.image("a.png", (1, 1, 1)) + PromptManagerConfig.SAFETY_BACKUP = True + states = [] + + async def prune(job, _db): + states.append((job.state, model_gate.holder() is not None)) + + with ( + patch("py.api.maintenance.take_safety_backup", return_value="db.backup_x"), + patch.object(self.api, "_prune_backup_pool", side_effect=prune), + ): + job = await self.run_batch() + self.assertEqual(job.state, "finished") + self.assertEqual(states, [("running", True)]) + self.assertIsNone(model_gate.holder()) + + async def test_prepass_publishes_its_progress(self): + for n in range(2): + self.image(f"{n}.png", (n, n, n)) + await self.run_batch() + published = [] + real = type(self.api._jobs).start + + def start(registry, *args, **kwargs): + job = real(registry, *args, **kwargs) + original = job.publish + job.publish = lambda payload: ( + published.append(dict(payload)), + original(payload), + ) + return job + + with patch.object(type(self.api._jobs), "start", start): + await self.run_batch() + self.assertIn({"type": "progress", "progress": 100}, published) + self.assertEqual(published[-1], {"type": "finished", "progress": 100}) + + async def test_restore_mid_run_fails_the_job(self): + self.image("a.png", (1, 1, 1)) + self.image("b.png", (2, 2, 2)) + + def restore(_path): + conn = sqlite3.connect(self._temp_db.name) + try: + rotate_history_epoch(conn) + finally: + conn.close() + + self.tagger.on_generate = restore + job = await self.run_batch() + self.assertEqual(job.state, "failed") + self.assertEqual( + job.status_line, "Prompt storage was restored; start the job again." + ) + + async def test_listing_error_fails_the_job(self): + self.image("a.png", (1, 1, 1)) + with patch.object( + self.api, "_get_gallery_files", side_effect=OSError("/secret/out") + ): + job = await self.run_batch() + self.assertEqual(job.state, "failed") + self.assertEqual( + job.status_line, "The output folders could not be listed. Try again." + ) + self.assertNotIn("/secret", json.dumps(job.report())) + + async def test_model_that_cannot_load_fails_the_job(self): + self.image("a.png", (1, 1, 1)) + self.tagger.load_error = RuntimeError("no CUDA") + job = await self.run_batch() + self.assertEqual(job.state, "failed") + self.assertEqual( + job.status_line, "Auto-tagging failed. Check server logs for details." + ) + self.assertIsNone(model_gate.holder()) + + async def test_zero_targets_finish_with_tagged_zero(self): + job = await self.run_batch() + self.assertEqual(job.state, "finished") + self.assertEqual(job.status_line, "Tagged 0 images") + self.assertEqual((job.processed, job.total), (0, 0)) + + async def test_runner_unloads_after_a_cancel_when_not_kept_in_memory(self): + self.image("a.png", (1, 1, 1)) + owner = model_gate.new_owner("legacy-batch") + self.assertTrue(model_gate.try_acquire(owner)) + run = jobs_module.AutotagRun( + model_type="wd14-vit", + existing="skip", + options={**WD14, "use_gpu": True, "suggestion_limit": 10}, + keep_in_memory=False, + backup_on=False, + owner=owner, + ) + + from py.api.jobs import AUTOTAG_GROUP_NAMES, Job as RealJob + + job = RealJob( + "autotag", + "Tag output images", + admission="records", + groups=AUTOTAG_GROUP_NAMES, + ) + real_targets = self.api._autotag_targets + + async def cancel_after_listing(): + targets = await real_targets() + job.cancelled = True + return targets + + with patch.object( + self.api, "_autotag_targets", side_effect=cancel_after_listing + ): + await self.api._run_autotag_batch(run, job) + self.assertEqual(job.state, "cancelled") + self.assertEqual(self.tagger.unload_calls, 1) + self.assertIsNone(model_gate.holder()) + + async def test_time_left_uses_the_models_rate_only(self): + for n in range(4): + self.image(f"{n}.png", (n, n, n)) + await self.run_batch() + for n in range(4, 8): + self.image(f"{n}.png", (n, n, n)) + await self.relist() + calls = [] + real = jobs_module._time_left + + def spy(*args): + calls.append(args) + return real(*args) + + with patch.object(jobs_module, "_time_left", side_effect=spy): + job = await self.run_batch() + self.assertEqual(job.counters["skipped"], 4) + # (done, total, batches, elapsed): over the 4 images the model runs on. + self.assertEqual( + [(c[0], c[1], c[2]) for c in calls], [(n, 4, n) for n in range(4)] + ) + # The shared D56 text is unchanged. + self.assertEqual(admin_module._time_left(5, 10, 5, 300), " · about 5 min left") + + +class TestStartRefusals(BatchTestCase): + async def test_unavailable_model_and_bad_bodies(self): + self.tagger.downloaded["wd14-vit"] = False + status, data = await self.start_batch() + self.assertEqual((status, data["code"]), (409, "model-unavailable")) + status, data = await self.start_batch(existing="replace_after_review") + self.assertEqual((status, data["code"]), (400, "review-required")) + status, _data = await self.start_batch(options={**WD14, "suggestion_limit": 3}) + self.assertEqual(status, 400) + self.assertEqual(self.api._jobs.list(), []) + self.assertIsNone(model_gate.holder()) + + async def test_model_busy_when_the_gate_is_held(self): + owner = model_gate.new_owner("image-annotations") + self.assertTrue(model_gate.try_acquire(owner)) + try: + status, data = await self.start_batch() + finally: + model_gate.release(owner) + self.assertEqual((status, data["code"]), (409, "model-busy")) + + +class TestBatchBeside(BatchTestCase): + async def test_legacy_batch_and_unload_are_refused_while_the_job_runs(self): + self.image("a.png", (1, 1, 1)) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + status, data = await self.start_batch() + self.assertEqual(status, 202, data) + job = self.api._jobs.get(data["job"]["id"]) + try: + await self.wait_for(entered) + self.assertTrue(model_gate.holder().startswith("autotag-job:")) + legacy = await self.client.post("/prompt_manager/autotag/start") + self.assertEqual(legacy.status, 200) + frames = [ + json.loads(line[len("data: ") :]) + for line in (await legacy.text()).splitlines() + if line.startswith("data: ") + ] + self.assertEqual( + frames, + [ + { + "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: + self.tagger.block = None + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(job.state, "finished") + + async def test_language_model_batch_passes_instructions_per_call(self): + self.image("a.png", (1, 1, 1)) + job = await self.run_batch(model_type="gguf") + self.assertEqual(job.state, "finished", job.status_line) + self.assertEqual(self.tagger.calls[0]["prompt"], "Tags please.") + self.assertEqual(self.tagger.custom_prompt, "legacy prompt") + + async def test_inference_errors_are_fixed_strings_and_not_marked(self): + self.image("a.png", (1, 1, 1)) + + def explode(_path): + raise RuntimeError("/srv/secret/model exploded") + + self.tagger.on_generate = explode + job = await self.run_batch() + self.assertEqual( + self.rows(job, "errors"), + [("a.png", "Auto-tagging failed. Check server logs for details.")], + ) + self.assertNotIn("/srv/secret", json.dumps(job.report())) + self.assertEqual(self.marked(), set()) + + +class TestPreview(BatchTestCase): + async def test_preview_counts_images_and_estimates_time(self): + prompt_id = self.prompt("p") + a = self.image("a.png", (1, 1, 1)) + self.image("b.png", (2, 2, 2)) + self.image("c.png", (3, 3, 3)) + self.link(prompt_id, a) + (self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 8) + status, fresh = await self.preview() + self.assertEqual(status, 200, fresh) + self.assertEqual( + {k: fresh[k] for k in ("images", "already_tagged", "to_tag")}, + {"images": 3, "already_tagged": 0, "to_tag": 3}, + ) + self.assertNotIn("prompts", fresh) + self.assertEqual(fresh["suggestion_limit"], 10) + self.assertEqual(fresh["estimate_device"], "cpu") + self.assertAlmostEqual(fresh["estimate_seconds"], 3 * 1.6) + + self.tagger.block = None + await self.run_batch() + self.image("d.png", (4, 4, 4)) + await self.relist() + skip = (await self.preview("skip"))[1] + self.assertEqual( + {k: skip[k] for k in ("images", "already_tagged", "to_tag")}, + {"images": 4, "already_tagged": 3, "to_tag": 1}, + ) + merge = (await self.preview("merge"))[1] + self.assertEqual(merge["to_tag"], 4) + for model_type in ("gguf", "hf"): + llm = (await self.preview(model_type=model_type))[1] + self.assertIsNone(llm["estimate_seconds"]) + self.assertIsNone(llm["estimate_device"]) + self.assertEqual([s for s in _strings(skip) if _absolute(s)], []) + + async def test_preview_uses_the_graphics_card_rate_when_wd14_runs_on_cuda(self): + self.image("a.png", (1, 1, 1)) + self.tagger.wd14_device = lambda: "gpu" + data = (await self.preview())[1] + self.assertEqual(data["estimate_device"], "gpu") + from py.tagging_defaults import WD14_GPU_SECONDS_PER_IMAGE + + self.assertAlmostEqual(data["estimate_seconds"], WD14_GPU_SECONDS_PER_IMAGE) + + async def test_preview_writes_nothing_and_opens_no_media_file(self): + self.image("a.png", (1, 1, 1)) + await self.run_batch() + self.image("b.png", (2, 2, 2)) + await self.relist() # the listing, primed + files = [self._temp_db.name + suffix for suffix in ("", "-wal")] + before = [Path(f).read_bytes() if os.path.exists(f) else None for f in files] + media = str(self.output_dir) + real_open, real_stat = open, os.stat + + def guarded_open(file, *args, **kwargs): + if str(file).startswith(media): + raise AssertionError(f"opened {file}") + return real_open(file, *args, **kwargs) + + def guarded_stat(path, *args, **kwargs): + if str(path).startswith(media) and str(path) != media: + raise AssertionError(f"stat {path}") + return real_stat(path, *args, **kwargs) + + with patch("builtins.open", guarded_open), patch("os.stat", guarded_stat): + status, data = await self.preview() + self.assertEqual(status, 200, data) + after = [Path(f).read_bytes() if os.path.exists(f) else None for f in files] + self.assertEqual(after, before) + self.assertEqual(data["to_tag"], 1) + self.assertEqual(self.tagger.calls[1:], []) diff --git a/tests/test_autotag_engine.py b/tests/test_autotag_engine.py index 53e0bca..7f7faa3 100644 --- a/tests/test_autotag_engine.py +++ b/tests/test_autotag_engine.py @@ -884,3 +884,29 @@ class TestImageConversion(EngineTestCase): if __name__ == "__main__": unittest.main() + + +class TestWd14Device(EngineTestCase): + """Where WD14 runs, for the Tag output images estimate (#379).""" + + def test_loaded_session_answers_from_its_first_provider(self): + cpu = FakeOnnxSession([0.5] * len(WD14_TAGS)) + service = self._service(("wd14", (cpu, WD14_TAGS))) + _touch_model_files(self.models_dir, "wd14-vit") + service.load_model("wd14-vit") + self.assertEqual(service.wd14_device(), "cpu") + cpu.get_providers = lambda: ["CUDAExecutionProvider", "CPUExecutionProvider"] + self.assertEqual(service.wd14_device(), "gpu") + + def test_without_a_session_the_available_providers_answer(self): + service = self._service(None) + fake = types.SimpleNamespace( + get_available_providers=lambda: ["CUDAExecutionProvider"] + ) + with patch.dict(sys.modules, {"onnxruntime": fake}): + self.assertEqual(service.wd14_device(), "gpu") + fake.get_available_providers = lambda: ["CPUExecutionProvider"] + with patch.dict(sys.modules, {"onnxruntime": fake}): + self.assertEqual(service.wd14_device(), "cpu") + with patch.dict(sys.modules, {"onnxruntime": None}): + self.assertEqual(service.wd14_device(), "cpu") diff --git a/tests/test_autotag_schema.py b/tests/test_autotag_schema.py new file mode 100644 index 0000000..c7f3c2c --- /dev/null +++ b/tests/test_autotag_schema.py @@ -0,0 +1,318 @@ +"""Auto-tagged marks schema group (#379): install, validation and restore.""" + +import hashlib +import re +import shutil +import sqlite3 +import tempfile +import unittest +from contextlib import closing +from pathlib import Path + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.models import PromptModel + +try: # guarded so a missing module fails assertions, not test loading + from database.autotag_schema import ( + AUTOTAG_TABLE, + autotag_schema_error, + autotag_without_annotations_error, + migrate_autotag_state, + ) +except ImportError: + AUTOTAG_TABLE = None + autotag_schema_error = None + autotag_without_annotations_error = None + migrate_autotag_state = None + +TABLE = "media_image_autotags" +DDL = ( + f"CREATE TABLE {TABLE} (byte_id TEXT PRIMARY KEY " + "REFERENCES media_annotation_bytes(id) ON DELETE CASCADE) WITHOUT ROWID" +) +ANNOTATION_TABLES = ( + "media_annotation_metadata", + "media_annotation_files", + "media_annotation_bytes", + "media_image_tags", +) +REVIEW_TABLES = ( + "media_review_metadata", + "media_review_items", + "media_review_decisions", +) +FILE_ID = "0f8e4b2a-1c3d-4e5f-8a9b-0c1d2e3f4a5b" +BYTE_ID = "6a1b2c3d-4e5f-4a6b-8c9d-0e1f2a3b4c5d" +OTHER_ID = "7b2c3d4e-5f6a-4b7c-9d0e-1f2a3b4c5d6e" +DIGEST = hashlib.sha256(b"img").hexdigest() +NOW = "2026-10-27T00:00:00.000000+00:00" +NEEDS_ANNOTATIONS = "Image auto-tag marks need the image annotation tables" +INVALID_VALUES = "Image auto-tag marks contain invalid values" +INVALID_SCHEMA = "Image auto-tag marks have an invalid schema" + + +def _compact(sql): + return re.sub(r"\s+", "", sql or "").lower() + + +def catalog(conn): + return { + (kind, name, _compact(sql)) + for kind, name, sql in conn.execute( + "SELECT type,name,sql FROM sqlite_master WHERE name=? OR tbl_name=?", + (TABLE, TABLE), + ) + } + + +def marks(conn): + return [tuple(row) for row in conn.execute(f"SELECT byte_id FROM {TABLE}")] + + +def insert_byte(conn, file_id=FILE_ID, byte_id=BYTE_ID, locator="/out/img.png"): + conn.execute( + "INSERT INTO media_annotation_files VALUES(?,?,?)", (file_id, locator, NOW) + ) + conn.execute( + "INSERT INTO media_annotation_bytes VALUES(?,?,?,?,?,?,?,0)", + (byte_id, file_id, DIGEST, 3, 1_700_000_000_000_000_000, NOW, NOW), + ) + + +class AutotagSchemaTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.addCleanup(assert_closed, self, self.temp.name) + self.addCleanup(PromptModel.reset_schema_cache) + self.models = [] + self.addCleanup(self.close_models) + self.path = str(Path(self.temp.name) / "live.db") + + def close_models(self): + for model in self.models: + model.close_all() + + def model(self, path=None): + model = PromptModel(path or self.path) + self.models.append(model) + return model + + def require_module(self): + self.assertIsNotNone( + autotag_schema_error, "database.autotag_schema is not implemented" + ) + + def raw(self, path, *statements, foreign_keys=False): + with closing(sqlite3.connect(path)) as raw, raw: + raw.execute(f"PRAGMA foreign_keys={'ON' if foreign_keys else 'OFF'}") + for statement in statements: + if isinstance(statement, tuple): + raw.execute(*statement) + else: + raw.execute(statement) + + def closed_current(self, name="current.db", *, with_mark=True): + """A closed, current database with one byte record (and its mark).""" + path = str(Path(self.temp.name) / name) + model = self.model(path) + with model.get_connection() as conn: + insert_byte(conn) + if with_mark: + conn.execute(f"INSERT INTO {TABLE} VALUES(?)", (BYTE_ID,)) + model.close_all() + return path + + def test_fresh_install_has_exactly_the_pinned_table(self): + self.require_module() + self.assertEqual(AUTOTAG_TABLE, TABLE) + model = self.model() + with model.get_connection() as conn: + # WITHOUT ROWID with a primary key: no sqlite_autoindex row. + self.assertEqual(catalog(conn), {("table", TABLE, _compact(DDL))}) + self.assertEqual( + [tuple(row) for row in conn.execute(f"PRAGMA table_info({TABLE})")], + [(0, "byte_id", "TEXT", 1, None, 1)], + ) + self.assertIsNone(autotag_schema_error(conn, allow_missing=False)) + self.assertEqual(marks(conn), []) + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (True, "ok"), + ) + + def test_reopen_and_rerun_are_idempotent(self): + self.require_module() + path = self.closed_current() + for _ in range(2): + PromptModel.reset_schema_cache() + model = self.model(path) + with model.get_connection() as conn: + migrate_autotag_state(conn) + self.assertEqual(marks(conn), [(BYTE_ID,)]) + self.assertEqual(catalog(conn), {("table", TABLE, _compact(DDL))}) + model.close_all() + + def test_old_database_without_the_table_gains_it_and_keeps_its_rows(self): + self.require_module() + path = self.closed_current(with_mark=False) + self.raw(path, f"DROP TABLE {TABLE}") + with closing(sqlite3.connect(path)) as raw: + before = raw.execute("SELECT * FROM media_annotation_bytes").fetchall() + self.assertIsNone(autotag_schema_error(raw)) + self.assertEqual( + autotag_schema_error(raw, allow_missing=False), + "Missing image auto-tag marks schema", + ) + self.assertTrue(PromptModel.verify_database_file(path)[0]) + self.assertEqual( + PromptModel.verify_database_file(path, require_library_state=True), + (False, "Missing image auto-tag marks schema"), + ) + PromptModel.reset_schema_cache() + model = self.model(path) + with model.get_connection() as conn: + self.assertEqual(catalog(conn), {("table", TABLE, _compact(DDL))}) + self.assertEqual( + [ + tuple(row) + for row in conn.execute("SELECT * FROM media_annotation_bytes") + ], + before, + ) + model.close_all() + + def test_table_without_annotation_group_is_refused(self): + self.require_module() + path = self.closed_current(with_mark=False) + self.raw( + path, + *( + f"DROP TABLE {t}" + for t in REVIEW_TABLES + tuple(reversed(ANNOTATION_TABLES)) + ), + ) + with closing(sqlite3.connect(path)) as raw: + self.assertEqual(autotag_without_annotations_error(raw), NEEDS_ANNOTATIONS) + for allow_missing in (True, False): + self.assertEqual( + autotag_schema_error(raw, allow_missing=allow_missing), + NEEDS_ANNOTATIONS, + ) + self.assertEqual( + PromptModel.verify_database_file(path), (False, NEEDS_ANNOTATIONS) + ) + PromptModel.reset_schema_cache() + with self.assertRaises(sqlite3.DatabaseError) as raised: + self.model(path) + self.assertIn(NEEDS_ANNOTATIONS, str(raised.exception)) + + def test_invalid_values_are_refused_without_repair(self): + self.require_module() + cases = { + "not-a-uuid": (f"INSERT INTO {TABLE} VALUES('not-a-uuid')",), + "upper-uuid": ((f"INSERT INTO {TABLE} VALUES(?)", (OTHER_ID.upper(),)),), + "dangling": ((f"INSERT INTO {TABLE} VALUES(?)", (OTHER_ID,)),), + } + for name, statements in cases.items(): + with self.subTest(name): + path = self.closed_current(f"{name}.db", with_mark=False) + self.raw(path, *statements) + with closing(sqlite3.connect(path)) as raw: + self.assertEqual(autotag_schema_error(raw), INVALID_VALUES) + before = marks(raw) + self.assertEqual( + PromptModel.verify_database_file(path), (False, INVALID_VALUES) + ) + PromptModel.reset_schema_cache() + with self.assertRaises(sqlite3.DatabaseError): + self.model(path) + with closing(sqlite3.connect(path)) as raw: + self.assertEqual(marks(raw), before) + + def test_altered_definitions_are_refused(self): + self.require_module() + cases = { + "rowid": f"CREATE TABLE {TABLE} (byte_id TEXT PRIMARY KEY " + "REFERENCES media_annotation_bytes(id) ON DELETE CASCADE)", + "restrict": f"CREATE TABLE {TABLE} (byte_id TEXT PRIMARY KEY " + "REFERENCES media_annotation_bytes(id) ON DELETE RESTRICT) WITHOUT ROWID", + "extra column": f"CREATE TABLE {TABLE} (byte_id TEXT PRIMARY KEY " + "REFERENCES media_annotation_bytes(id) ON DELETE CASCADE, engine TEXT) " + "WITHOUT ROWID", + } + for name, ddl in cases.items(): + with self.subTest(name): + path = self.closed_current(f"{name}.db", with_mark=False) + self.raw(path, f"DROP TABLE {TABLE}", ddl) + with closing(sqlite3.connect(path)) as raw: + self.assertEqual(autotag_schema_error(raw), INVALID_SCHEMA) + self.assertFalse(PromptModel.verify_database_file(path)[0]) + path = self.closed_current("index.db", with_mark=False) + self.raw(path, f"CREATE INDEX idx_extra ON {TABLE}(byte_id)") + with closing(sqlite3.connect(path)) as raw: + self.assertEqual(autotag_schema_error(raw), INVALID_SCHEMA) + + def test_deleting_the_byte_record_cascades_to_its_mark(self): + self.require_module() + model = self.model() + with model.get_connection() as conn: + insert_byte(conn) + conn.execute(f"INSERT INTO {TABLE} VALUES(?)", (BYTE_ID,)) + conn.execute(f"INSERT OR IGNORE INTO {TABLE} VALUES(?)", (BYTE_ID,)) + self.assertEqual(marks(conn), [(BYTE_ID,)]) + conn.execute("DELETE FROM media_annotation_files") + self.assertEqual(marks(conn), []) + + def test_restore_of_a_pre_change_backup_migrates(self): + self.require_module() + src = self.closed_current("backup.db", with_mark=False) + self.raw(src, f"DROP TABLE {TABLE}") + live = self.model() + live.restore_from_file(src) + with closing(sqlite3.connect(self.path)) as raw: + self.assertEqual(catalog(raw), {("table", TABLE, _compact(DDL))}) + self.assertEqual( + raw.execute("SELECT id FROM media_annotation_bytes").fetchall(), + [(BYTE_ID,)], + ) + self.assertEqual( + PromptModel.verify_database_file(self.path, require_library_state=True), + (True, "ok"), + ) + with closing(sqlite3.connect(src)) as raw: # the source is untouched + self.assertEqual(catalog(raw), set()) + + def test_install_failure_rolls_back(self): + self.require_module() + src = self.closed_current("copy.db", with_mark=False) + self.raw(src, f"DROP TABLE {TABLE}") + path = str(Path(self.temp.name) / "rollback.db") + shutil.copyfile(src, path) + with closing(sqlite3.connect(path)) as raw: + raw.execute("BEGIN") + raw.execute("CREATE TABLE marker(x)") + # An invalid group already present: nothing created, outer work kept. + raw.execute(f"CREATE TABLE {TABLE}(byte_id TEXT)") + with self.assertRaises(sqlite3.DatabaseError): + migrate_autotag_state(raw) + self.assertEqual( + raw.execute( + f"SELECT sql FROM sqlite_master WHERE name='{TABLE}'" + ).fetchone()[0], + f"CREATE TABLE {TABLE}(byte_id TEXT)", + ) + self.assertIsNotNone( + raw.execute( + "SELECT 1 FROM sqlite_master WHERE name='marker'" + ).fetchone() + ) + raw.rollback() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index e4aa9a8..2437220 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -468,6 +468,7 @@ class TestRestoreFromFile(BackupTestCase): def test_restore_old_backup_without_annotations(self): tables = ( + "media_image_autotags", "media_review_decisions", "media_review_items", "media_review_metadata", diff --git a/tests/test_image_tag_jobs.py b/tests/test_image_tag_jobs.py index 04d0dde..25943b8 100644 --- a/tests/test_image_tag_jobs.py +++ b/tests/test_image_tag_jobs.py @@ -690,6 +690,32 @@ class TestFixRound(ImageTagJobTestCase): ], ) + async def test_an_ended_job_shows_no_time_left(self): + """#379 live check: a finished or cancelled job never says "min left".""" + for index in range(3): + self.png(f"f{index}.png", color=(index, 7, 7)) + clock = iter([0.0, 0.0, 120.0, 240.0]) + with patch( + "py.api.image_tag_jobs._monotonic", side_effect=lambda: next(clock, 240.0) + ): + job = await self.run_job(self.auto_body([f"f{i}.png" for i in range(3)])) + self.assertEqual(job.state, "finished") + self.assertEqual(job.status_line, "Auto-tag finished · 3 files processed") + + self.png("g0.png", color=(1, 1, 1)) + self.png("g1.png", color=(2, 2, 2)) + entered, proceed = threading.Event(), threading.Event() + self.tagger.block = (entered, proceed) + job = await self.start_ok(self.auto_body(["g0.png", "g1.png"])) + await self.wait_for(entered) + await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel") + self.tagger.block = None + proceed.set() + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.status_line, "Auto-tag cancelled · 1 file processed") + self.assertNotIn("left", job.status_line) + def maximal_target(self, number): """A distinct target whose path is 1024 non-ASCII code points long.""" head = f"{number:03d}/" + ("界" * 200 + "/") * 5 diff --git a/tests/test_review_schema.py b/tests/test_review_schema.py index 1ee935e..b97780a 100644 --- a/tests/test_review_schema.py +++ b/tests/test_review_schema.py @@ -371,7 +371,10 @@ class ReviewSchemaTests(unittest.TestCase): self.require_module() for label, dropped in ( ("annotations only", TABLES), - ("predates annotations", ANNOTATION_TABLES + TABLES), + ( + "predates annotations", + ("media_image_autotags",) + ANNOTATION_TABLES + TABLES, + ), ): with self.subTest(label): path = str(Path(self.temp.name) / f"{label.replace(' ', '-')}.db") diff --git a/web/js/workspace/system-tagging.mjs b/web/js/workspace/system-tagging.mjs index aaea542..d9b369f 100644 --- a/web/js/workspace/system-tagging.mjs +++ b/web/js/workspace/system-tagging.mjs @@ -15,10 +15,11 @@ * value (the toast names it). D132: the concept's two simulated job buttons are "Download" (or * "Download again" for a model on disk), which starts the `model_download` Activity job and opens it, and "Unload", * a direct `POST /autotag/unload` (enabled while a model is loaded; a busy model's refusal names the operation). - * - D134: the `.operation` "Tag output images" › "Review" previews the prompt-tag batch - * (`POST /autotag/batch/preview`, read-only) and shows it in the shared review dialog with the D55 backup box; "Start - * tagging" runs the cancellable `autotag` Activity job. A review that cannot start (Replace after review, a short - * disk, nothing to tag) offers only Cancel and says why. + * - D134: the `.operation` "Tag output images" › "Review" previews the per-image batch (#379, + * `POST /autotag/batch/preview`, read-only: image counts and a time estimate) and shows it in the shared review dialog + * with the D55 backup box; "Start tagging" runs the cancellable `autotag` Activity job, which tags each output image + * from its own pixels and never a prompt. A review that cannot start (Replace after review, a short disk, no image, + * nothing to tag) offers only Cancel and says why. * * Every value is a host draft saved by the one Save row (system.mjs); every listener sits on the section's scope. */ @@ -44,7 +45,7 @@ export const REVIEW_TITLE = 'Review tag output images'; export const BATCH_TITLE = 'Tag output images'; export const START_LABEL = 'Start tagging'; export const INTRO = 'One tagging configuration for images and prompt-library actions. Every result can be reviewed before applying.'; -export const BATCH_DESCRIPTION = 'Add tags to each prompt from its linked output images with the model above. Review the scope first.'; +export const BATCH_DESCRIPTION = 'Tag each output image from its own pixels with the model above. Prompts are not changed. Review the scope first.'; export const SUGGESTIONS_DESCRIPTION = 'From 1 to 100: the most suggestions Generate shows for one image.'; /** D132: the engines in the concept's order, with the labels and sizes the server states when it cannot be read. */ @@ -97,8 +98,12 @@ export const SETTING_PROBLEMS = Object.freeze({ export const characters = text => [...String(text ?? '')].length; export const REPLACE_NOTE = 'Replacing tags needs a review of each image. Use Preview tagging review, or choose Skip or Merge for a batch run.'; -export const NOTHING_TO_TAG = 'Every prompt with an output image is already tagged.'; -export const NO_LINKED_IMAGES = 'No output image is linked to a prompt yet, so there is nothing to tag.'; +export const NOTHING_TO_TAG = 'Every output image is already auto-tagged.'; +export const NO_IMAGES = 'The output folders contain no images to tag.'; +/** The Time row's tails (#379): where WD14 runs, and the language models, which have no estimate. */ +export const TIME_ON_CPU = ' on a processor; a graphics card can be faster.'; +export const TIME_ON_GPU = ' on the graphics card.'; +export const TIME_UNKNOWN = 'Language-model engines take longer per image. Activity shows the time left once tagging starts.'; export const STATUS_UNREAD = 'The model status is still being read. Try again in a moment.'; export const SETTINGS_UNREAD = 'The settings are still being read. Try again in a moment.'; export const UNLOAD_FAILED = 'The tagging model could not be unloaded. Try again.'; @@ -254,19 +259,40 @@ export async function startJob({ activity, scope, kind, body, what, notify = sho /** The existing-tags policy's label. */ export const existingLabel = value => EXISTING_OPTIONS.find(([id]) => id === value)?.[1] ?? String(value ?? ''); -/** D134: the review's rows from the server's preview and the drafts it was made with. */ +/** + * DEC11 (#379): a whole number of minutes as the review's Time text: "about m min" under 90, then "about H h M min", + * or "about H h" on the hour. + */ +export function durationText(minutes) { + const m = Math.max(1, Math.ceil(Number.isFinite(minutes) ? minutes : 1)); + if (m < 90) return `about ${m} min`; + const hours = Math.floor(m / 60); + const rest = m % 60; + return rest === 0 ? `about ${hours} h` : `about ${hours} h ${rest} min`; +} + +/** The review's Time row: the preview's `estimate_seconds` where WD14 runs (`estimate_device`), else the sentence. */ +export function timeText(plan) { + const seconds = plan?.estimate_seconds; + if (typeof seconds !== 'number' || !Number.isFinite(seconds) || seconds < 0) return TIME_UNKNOWN; + return durationText(Math.ceil(seconds / 60)) + (plan.estimate_device === 'gpu' ? TIME_ON_GPU : TIME_ON_CPU); +} + +/** D134: the review's rows from the server's preview and the drafts it was made with (image counts, #379). */ export function batchGrid(plan, values) { const id = plan.model?.type ?? values.autotag_model_type; - const settings = CLASSIFIERS.has(id) + let settings = CLASSIFIERS.has(id) ? `General ${values.autotag_general_threshold} · Character ${values.autotag_character_threshold}` : `Template of ${number(characters(values.autotag_template))} characters`; + // The saved count the run uses (the preview's), not an unsaved draft. + if (Number.isSafeInteger(plan.suggestion_limit)) settings += ` · up to ${number(plan.suggestion_limit)} tags per image`; return [ ['Model', `${modelLabel(id, plan.model)} · ${modelSize(id, plan.model)}`], ['Existing tags', existingLabel(plan.existing ?? values.autotag_existing)], - ['Images linked', number(plan.images)], - ['Prompts', number(plan.prompts)], - ['Already tagged', number(plan.already_tagged)], - ['Prompts to tag', `up to ${number(plan.to_tag)}`], + ['Images', number(plan.images)], + ['Already auto-tagged', number(plan.already_tagged)], + ['Images to tag', `about ${number(plan.to_tag)}`], + ['Time', timeText(plan)], ['Settings', settings], ['Backup', backupText(plan.backup)], ]; @@ -291,13 +317,13 @@ export function batchCloseNote(plan) { if (plan.review_required === true) return REPLACE_NOTE; if (plan.model?.downloaded === false) return notDownloaded(modelLabel(plan.model.type, plan.model)); if (plan.disk?.sufficient === false) return diskShortNote(plan.disk); - if (!Number.isSafeInteger(plan.prompts) || plan.prompts === 0) return NO_LINKED_IMAGES; + if (!Number.isSafeInteger(plan.images) || plan.images === 0) return NO_IMAGES; if (plan.to_tag === 0) return NOTHING_TO_TAG; return null; } -export const BATCH_SENTENCE = 'Nothing is written until you start. Missing files and files outside the output folders are ' - + 'counted, so the prompts to tag are an upper bound; Activity reports what was tagged.'; +export const BATCH_SENTENCE = 'Nothing is written until you start. The number to tag is an estimate: images changed since ' + + 'they were tagged are tagged again, and missing files are skipped. Activity reports what was tagged.'; /** The Cancel-only review: the counts and why the batch cannot start (CUJ 5), in the shared close-only shell. */ export function openBatchCloseOnly({ scope, trigger, grid, note }) {