* fix(tagging): Tag output images tags each image, not its prompt (#379, PR 1) The native "Tag output images" job now runs the image-tags loop over every still image the output folders list, linked or not: each image gets its own model tags (at most "Suggestions per image") in its image tag set, and its bytes are marked auto-tagged in a new media_image_autotags table. No prompt tag or prompt text is written. - database/autotag_schema.py: one forward-only WITHOUT ROWID table, optional in old backups, refused without the annotation group. - add_image_tags/replace_image_tags take mark_autotagged; autotag_state and autotagged_locators read the marks through the one current-record rule. - "Skip already tagged" skips images whose current bytes are marked, counted in a stat-only pre-pass; hand tags do not count. - The review counts images and estimates the time (WD14 on CPU or CUDA). - System > Auto-tagging review rows and copy, D134 and the Auto-tagging contract follow. Part of #379 * fix(tagging): end the batch job after its cleanup and report pre-pass progress (#379) The job read as ended while it still held the model gate, so a quick restart could be refused with model-busy. Pre-pass chunks now publish progress, and a cancel between chunks keeps exact counts. Part of #379 * fix(tagging): measured GPU rate, and an ended image-tags job shows no time left (#379) The live check measured WD14 on CUDA at 0.167 s per image (100 images in 16.7 s with the model load), so the review estimate uses 0.17 s. A finished or cancelled Auto-tag/Add tags job kept "Processing file N of N · about X min left"; it now ends with "<verb> <state> · N files processed". Part of #379 * fix(tagging): legacy pages tag each image, not its prompt (#379, PR 2) The admin and gallery batch starts the shared Tag output images job ("Tag output images (legacy page)") and relays its progress; closing the window cancels it. Review asks per image with Skip first and applies by image and sha256; apply refuses a prompt_id body. The Tag image dialog marks an image auto-tagged only after Generate suggestions on that capture. A second legacy batch (also one that wins the start race) answers JSON 409, and both pages now show that sentence. * test(tagging): gallery review units apply by image; record each progress change The gallery review unit tests still pinned the removed prompt_id apply. The legacy batch spec read the progress text once per observer callback, so a phase frame and the first image frame landing in one task lost the phase (WebKit); it now records every mutation record.
This commit is contained in:
Vito
+54
-1
@@ -3826,7 +3826,9 @@ threshold off that grid is refused ("General-tag confidence must be a multiple o
|
||||
0.05.", `validate_saved_threshold`; `draftProblem` holds back Preview tagging review and
|
||||
the batch Review with the same sentence) and falls back on load; an accepted one is
|
||||
stored exactly on the grid (#369). A per-request threshold keeps its own check: `/image-annotations/suggest` and
|
||||
`/jobs/autotag` take 0–1; the legacy `/autotag/single` and `/autotag/start` only require a number. `GET /autotag/models` reports the
|
||||
`/jobs/autotag` take 0–1; the legacy `/autotag/start` passes its options through
|
||||
`validate_options` and starts the `autotag` job titled "Tag output images (legacy page)";
|
||||
the legacy `/autotag/single` only requires a number. `GET /autotag/models` reports the
|
||||
saved values (`wd14_general_threshold`, `wd14_character_threshold`, `suggestion_limit`,
|
||||
`default_prompt`, `model_type`, `existing`), so the Tag image and Auto-tag selected
|
||||
dialogs start from them; a preset (the bulk review, Preview tagging review) wins.
|
||||
@@ -3892,6 +3894,57 @@ images", "Tag output images (legacy page)", "Tagging one image (legacy page)",
|
||||
is pinned by `tests/test_autotag_schema.py` (forward-only: one empty `WITHOUT ROWID`
|
||||
table, optional in old backups, refused without the annotation group).
|
||||
|
||||
- **The legacy pages (#379, PR 2)** write image tags, never prompt tags.
|
||||
`POST /autotag/start` (the admin and gallery batch) reads `model_type`, `prompt`,
|
||||
`skip_tagged` (Skip or Merge), `keep_in_memory` and the thresholds (else the saved
|
||||
ones), checks them with `validate_options` (WD14: both thresholds; GGUF/HF:
|
||||
`instructions` = the non-blank `prompt`, else the saved template; the saved
|
||||
"Suggestions per image") and starts the `autotag` job "Tag output images (legacy
|
||||
page)" (gate owner `legacy-batch`) through the native start, so the safety backup and
|
||||
the disk check apply. A second legacy batch answers JSON 409 `{success: false, error:
|
||||
"Auto-tagging is already running"}` (also when another legacy start wins the race past
|
||||
that check); `SseStream` hands a non-2xx JSON `error` to `onerror` as `serverMessage`,
|
||||
which both pages show instead of "Connection lost during auto tagging". Any other refusal (a job in the way: "<title>
|
||||
already running", `model-unavailable`, `insufficient-space`, `model-busy`, a bad
|
||||
option) is one `{type: "error", status: "error", message}` frame. The stream relays
|
||||
the job: the phases as `{type: "progress", progress: 0 | 5 | 8, status: "Taking a
|
||||
safety backup…" | "Loading model…" | "Listing output images…"}`, each image as
|
||||
`{type: "progress", progress: 10 + progress × 85 / 100, status, processed, tagged,
|
||||
skipped}` (at most every 0.5 s; `processed` = Tagged + Unchanged + Errors), and the end
|
||||
as `complete` (`processed`, `tagged`, `skipped`, `errors`, `skipped_outside_roots: 0`,
|
||||
`status` "Complete" or "No still images found", `model_status`, `model_loaded`),
|
||||
`{type: "cancelled", message: "Cancelled. Tags already added were kept."}` or an
|
||||
`error` frame with the job's status line. A job quiet for 5 s repeats its last frame,
|
||||
and a closed stream cancels the job (the legacy pages have no other stop).
|
||||
`POST /autotag/single` takes `existing` (`"skip"` or `"merge"`, default merge). Under
|
||||
Skip, an image whose current bytes are auto-tagged answers `{success, already_tagged:
|
||||
true, image_tags, prompt_id, filename}` after one stat, without the model gate.
|
||||
Otherwise the image is captured and the model runs on the capture under the gate with
|
||||
the saved cap; the answer adds `tags`, `sha256` (the captured bytes) and `image_tags`,
|
||||
and `service.custom_prompt` is never set. Refusals: `409 model-unavailable`, `413
|
||||
too-large`, `404 file-not-found`, `409 file-changed`, `503 busy`, `415 not-a-still`
|
||||
(the image is marked), `409 model-busy`. `POST /autotag/apply` takes exactly
|
||||
`{image_id, sha256, tags}` or `{path, root_index?, sha256, tags}` (64 KiB), re-hashes
|
||||
the file and unions the tags into that image (`add_image_tags`, marked): `{success,
|
||||
added_tags, total_tags}`; an empty list only marks the image. A body with `prompt_id`
|
||||
is `400` "This page is out of date. Reload it."; a changed file `409 file-changed`
|
||||
("The image changed since its tags were generated."); a full set `409 tag-limit`; a
|
||||
restore `409 storage-changed`. The admin review asks each image with Skip first: an
|
||||
auto-tagged one shows "Image Already Auto-tagged" / "A tagging model already tagged
|
||||
this image. Generate new tags?" with its image tags (Skip, Skip All Tagged, Re-tag asks
|
||||
again with Merge); the gallery review always merges. Apply & Next sends the image and
|
||||
its `sha256` (linked or not). The batch copy: "Skip images already auto-tagged",
|
||||
"Auto tagging complete! Tagged N images." ("1 image"), and under both progress bars
|
||||
"Closing this window stops tagging. Tags already added are kept. Prompt Manager's
|
||||
Activity page shows the time left.". Pinned by `tests/test_autotag_legacy.py`,
|
||||
`tests/js/image-helpers.test.js` and `legacy-autotag.spec`.
|
||||
- **Tag image dialog marks (DEC4).** A save after Generate suggestions ran on that
|
||||
capture marks the image auto-tagged (also a replayed save, and a saved retry after
|
||||
`tag-conflict`); a save without Generate does not. The API keeps the suggested capture
|
||||
tokens; each new one drops the tokens whose capture is gone (`expires_in`, which never
|
||||
resets a capture's idle clock), a release or a saved receipt drops its own, and
|
||||
closing the captures clears them (`TestDialogMarks`).
|
||||
|
||||
No browser test downloads or runs a model: the injected `TestModel` answers, its
|
||||
`download_model` raises unless `/__test__/tagging/download {mode: hold | succeed |
|
||||
fail}` installs a fake that waits on a fixture event, and every tagging reset releases
|
||||
|
||||
@@ -1808,37 +1808,6 @@ class PromptDatabase(
|
||||
reader = DocumentReader(conn)
|
||||
return [self._image_row_to_dict(row, reader) for row in cursor]
|
||||
|
||||
def autotag_batch_candidates(self) -> List[Dict[str, Any]]:
|
||||
"""The images a Tag output images batch looks at (#328).
|
||||
|
||||
``{id, image_path, prompt_id, prompt_tags}`` for every image with a
|
||||
stored path that is linked to a prompt, newest first (the
|
||||
``get_all_images`` order, ties newest row first). Never reads
|
||||
``workflow_data`` or the other documents: a batch needs none of them.
|
||||
"""
|
||||
sql = (
|
||||
"SELECT gi.id, gi.image_path, gi.prompt_id, "
|
||||
"(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt "
|
||||
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = gi.prompt_id) "
|
||||
"AS prompt_tags "
|
||||
"FROM generated_images gi "
|
||||
"INNER JOIN prompts p ON gi.prompt_id = p.id "
|
||||
"WHERE gi.image_path IS NOT NULL AND gi.image_path != '' "
|
||||
"ORDER BY gi.generation_time DESC, gi.id DESC"
|
||||
)
|
||||
with self.model.get_connection() as conn:
|
||||
return [
|
||||
{
|
||||
"id": row["id"],
|
||||
"image_path": row["image_path"],
|
||||
"prompt_id": row["prompt_id"],
|
||||
"prompt_tags": (
|
||||
row["prompt_tags"].split("|||") if row["prompt_tags"] else []
|
||||
),
|
||||
}
|
||||
for row in conn.execute(sql)
|
||||
]
|
||||
|
||||
def get_all_images(
|
||||
self,
|
||||
limit: int = 0,
|
||||
|
||||
@@ -363,6 +363,9 @@ class PromptManagerAPI(
|
||||
# Private media captures of the annotation routes (py/api/annotations.py),
|
||||
# created on first use with their restore subscription.
|
||||
self._captures = None
|
||||
# Capture tokens Generate suggestions ran on: a save of one marks the
|
||||
# image auto-tagged (#379).
|
||||
self._suggested_tokens = set()
|
||||
self._captures_unsubscribe = None
|
||||
# Set by close_annotation_captures (shutdown): no store is recreated.
|
||||
self._captures_closed = False
|
||||
|
||||
@@ -296,6 +296,7 @@ class AnnotationRoutesMixin:
|
||||
self._captures_closed = True
|
||||
unsubscribe, self._captures_unsubscribe = self._captures_unsubscribe, None
|
||||
store, self._captures = self._captures, None
|
||||
self._suggested_tokens = set()
|
||||
if unsubscribe is not None:
|
||||
unsubscribe()
|
||||
if store is not None:
|
||||
@@ -515,6 +516,7 @@ class AnnotationRoutesMixin:
|
||||
store = self._captures
|
||||
if store is not None and _TOKEN.fullmatch(token) is not None:
|
||||
store.release(token)
|
||||
self._suggested_tokens.discard(token)
|
||||
return web.json_response(
|
||||
{"success": True}, headers={"Cache-Control": "no-store"}
|
||||
)
|
||||
@@ -564,6 +566,9 @@ class AnnotationRoutesMixin:
|
||||
return _coded("capture-expired", _OTHER_FILE)
|
||||
if facts.sha256 != capture.sha256:
|
||||
return _coded("file-changed")
|
||||
# A save after Generate suggestions ran on this capture marks the
|
||||
# image auto-tagged (#379), so a Skip batch leaves the review alone.
|
||||
suggested = token in self._suggested_tokens
|
||||
try:
|
||||
receipt = await self._run_in_executor(
|
||||
self.db.replace_image_tags,
|
||||
@@ -574,6 +579,7 @@ class AnnotationRoutesMixin:
|
||||
expected_epoch=expected_epoch,
|
||||
expected_tag_revision=revision,
|
||||
tags=tags,
|
||||
mark_autotagged=suggested,
|
||||
)
|
||||
except AnnotationCommandError as exc:
|
||||
current = exc.current or {}
|
||||
@@ -590,10 +596,23 @@ class AnnotationRoutesMixin:
|
||||
except Exception:
|
||||
self.logger.exception("Image tag save failed")
|
||||
return _error(500, "save-failed", "The tags could not be saved.")
|
||||
# Dropped only after a saved or replayed receipt: a retry after a
|
||||
# tag-conflict still marks.
|
||||
self._suggested_tokens.discard(token)
|
||||
return web.json_response(
|
||||
{"success": True, **receipt}, headers={"Cache-Control": "no-store"}
|
||||
)
|
||||
|
||||
def _remember_suggested(self, store, token: str) -> None:
|
||||
"""Note that Generate ran on *token*; forget tokens whose capture is
|
||||
gone. ``expires_in`` does not reset a capture's idle clock."""
|
||||
self._suggested_tokens = {
|
||||
kept
|
||||
for kept in self._suggested_tokens
|
||||
if store.expires_in(kept) is not None
|
||||
}
|
||||
self._suggested_tokens.add(token)
|
||||
|
||||
# -- suggest ---------------------------------------------------------------
|
||||
|
||||
async def suggest_image_tags(self, request: web.Request) -> web.Response:
|
||||
@@ -675,6 +694,7 @@ class AnnotationRoutesMixin:
|
||||
finally:
|
||||
if not started:
|
||||
store.release_use(token)
|
||||
self._remember_suggested(store, token)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
|
||||
+520
-409
File diff suppressed because it is too large
Load Diff
@@ -95,8 +95,14 @@ async function watchProgress(page) {
|
||||
await page.evaluate(() => {
|
||||
window.progressSeen = [];
|
||||
const node = document.getElementById('autoTagProgressPercent');
|
||||
new MutationObserver(() => window.progressSeen.push(node.textContent))
|
||||
.observe(node, { childList: true, characterData: true, subtree: true });
|
||||
// Each change from its own record: frames that land in one task share one
|
||||
// callback, and the node's text by then is only the last of them.
|
||||
new MutationObserver(records => {
|
||||
for (const record of records) {
|
||||
if (record.type === 'characterData') window.progressSeen.push(record.target.data);
|
||||
for (const added of record.addedNodes) window.progressSeen.push(added.textContent);
|
||||
}
|
||||
}).observe(node, { childList: true, characterData: true, subtree: true });
|
||||
});
|
||||
}
|
||||
|
||||
@@ -108,25 +114,58 @@ async function runGalleryBatch(page) {
|
||||
await page.locator('#startAutoTagBtn').click();
|
||||
await done;
|
||||
await expect(page.locator('#autoTagProgressModal')).toBeHidden();
|
||||
await expect(page.getByText(/^Auto tagging complete! Applied tags to \d+ prompts\.$/)).toBeVisible();
|
||||
// #379: the batch counts images, never prompts.
|
||||
await expect(page.getByText(/^Auto tagging complete! Tagged \d+ images?\.$/)).toBeVisible();
|
||||
}
|
||||
|
||||
const CLOSE_NOTE = 'Closing this window stops tagging. Tags already added are kept. Prompt Manager\'s Activity page shows the '
|
||||
+ 'time left.';
|
||||
|
||||
async function serverJobs(request) {
|
||||
const response = await request.get('/prompt_manager/jobs');
|
||||
await expect(response).toBeOK();
|
||||
return (await response.json()).jobs;
|
||||
}
|
||||
|
||||
async function runNativeBatch(request) {
|
||||
const started = await request.post('/prompt_manager/jobs/autotag', { data: { model_type: 'wd14-swinv2', existing: 'skip',
|
||||
options: { general_threshold: 0.35, character_threshold: 0.85 } } });
|
||||
expect(started.status()).toBe(202);
|
||||
const { job } = await started.json();
|
||||
await expect.poll(async () => (await serverJobs(request)).find(item => item.id === job.id)?.state).toBe('finished');
|
||||
}
|
||||
|
||||
test('gallery batch modal counts progress and closes on complete', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await runGalleryBatch(page);
|
||||
const seen = await page.evaluate(() => window.progressSeen);
|
||||
// The fixture's prompts are already tagged, so the batch skips them; its progress frames still move the bar.
|
||||
expect(seen).toContain('95%');
|
||||
// The phase frames (model 5%, listing 8%) then the per-image frames (from 10%) move the bar.
|
||||
expect(seen).toContain('8%');
|
||||
expect(seen.some(value => parseInt(value, 10) >= 10)).toBe(true);
|
||||
await expect(page.locator('#autoTagCloseNote')).toHaveText(CLOSE_NOTE);
|
||||
// The batch ran as the autotag job of the legacy page.
|
||||
expect((await serverJobs(request)).filter(job => job.kind === 'autotag').map(job => [job.title, job.state]))
|
||||
.toEqual([['Tag output images (legacy page)', 'finished']]);
|
||||
expect((await downloads(request)).length).toBe(0);
|
||||
});
|
||||
|
||||
test('admin batch copy names images, not prompts (#379)', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await expect(page.locator('input[name="autoTagMode"][value="skip"] + span')).toHaveText('Skip images already auto-tagged');
|
||||
await expect(page.getByText('"auto-scanned" is not counted as a real tag')).toHaveCount(0);
|
||||
await expect(page.locator('#autoTagCloseNote')).toHaveText(CLOSE_NOTE);
|
||||
await expect(page.locator('#retagConfirmModal h3')).toHaveText('Image Already Auto-tagged');
|
||||
await expect(page.locator('#retagConfirmModal')).toContainText('A tagging model already tagged this image. Generate new tags?');
|
||||
});
|
||||
|
||||
test('admin review opens the first image', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await page.locator('#startReviewBtn').click();
|
||||
// The first image's prompt is already tagged: confirm a new pass, then the review opens on that image.
|
||||
await page.locator('#retagConfirmBtn').click();
|
||||
// #379: no image is auto-tagged yet (prompt tags no longer count), so the review opens on the first image.
|
||||
await expect(page.locator('#retagConfirmModal')).toBeHidden();
|
||||
await expect(page.locator('#autoTagReviewModal')).toBeVisible();
|
||||
await expect(page.locator('#reviewImageName')).not.toHaveText(/^(image_name\.png|Unknown|)$/);
|
||||
await expect(page.getByText('Failed to start review mode')).toHaveCount(0);
|
||||
@@ -150,3 +189,70 @@ test('gallery unload works and is refused while the gate is held', async ({ page
|
||||
await unload.click();
|
||||
await expect(page.getByText('WD14-SWINV2 model unloaded successfully')).toBeVisible();
|
||||
});
|
||||
|
||||
/** The admin review lists every row, videos too: Skip any image that is not one until the confirmation shows. */
|
||||
async function untilAutoTagged(page, confirm) {
|
||||
// Each image's skip check answers before its confirmation or its error shows.
|
||||
const failed = page.locator('#reviewTagsVisible .text-pm-error');
|
||||
for (let tries = 0; tries < 5; tries++) {
|
||||
await expect.poll(async () => (await confirm.isVisible()) || (await failed.count()) > 0).toBe(true);
|
||||
if (await confirm.isVisible()) return;
|
||||
const answered = page.waitForResponse(item => item.url().endsWith('/prompt_manager/autotag/single'));
|
||||
await page.locator('#skipReviewBtn').click();
|
||||
await answered;
|
||||
}
|
||||
await expect(confirm).toBeVisible();
|
||||
}
|
||||
|
||||
// #379: the admin review asks before re-tagging an image a model already tagged, and Apply & Next writes that image's
|
||||
// tags (its bytes, never its prompt).
|
||||
test('admin review: an auto-tagged image asks first; Re-tag then Apply & Next tags the image', async ({ page, request }) => {
|
||||
await runNativeBatch(request);
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await page.locator('#startReviewBtn').click();
|
||||
const confirm = page.locator('#retagConfirmModal');
|
||||
await untilAutoTagged(page, confirm);
|
||||
await expect(confirm.locator('#retagExistingTags span').first()).toHaveText('sandstone');
|
||||
await page.locator('#retagConfirmBtn').click();
|
||||
await expect(page.locator('#autoTagReviewModal')).toBeVisible();
|
||||
await expect(page.locator('#reviewTagsVisible .tag-chip').first()).toContainText('sandstone');
|
||||
const applied = page.waitForRequest(item => item.url().endsWith('/prompt_manager/autotag/apply'));
|
||||
const answer = page.waitForResponse(item => item.url().endsWith('/prompt_manager/autotag/apply'));
|
||||
await page.locator('#applyReviewBtn').click();
|
||||
const body = (await applied).postDataJSON();
|
||||
expect(body.prompt_id).toBeUndefined();
|
||||
expect(body.sha256).toMatch(/^[0-9a-f]{64}$/);
|
||||
expect((await answer).status()).toBe(200);
|
||||
});
|
||||
|
||||
test('admin review: Skip All Tagged passes every auto-tagged image', async ({ page, request }) => {
|
||||
await runNativeBatch(request);
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await page.locator('#startReviewBtn').click();
|
||||
const confirm = page.locator('#retagConfirmModal');
|
||||
await untilAutoTagged(page, confirm);
|
||||
await page.locator('#retagSkipAllBtn').click();
|
||||
await expect(page.locator('#retagConfirmModal')).toBeHidden();
|
||||
await expect(page.getByText('Review complete!')).toBeVisible();
|
||||
});
|
||||
|
||||
test('gallery review: Apply & Next tags the image by its path and bytes', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
const suggested = page.waitForRequest(item => item.url().endsWith('/prompt_manager/autotag/single'));
|
||||
await page.locator('#startReviewBtn').click();
|
||||
expect((await suggested).postDataJSON().existing).toBe('merge');
|
||||
await expect(page.locator('#reviewTagsContainer .tag-chip').first()).toContainText('sandstone');
|
||||
const applied = page.waitForRequest(item => item.url().endsWith('/prompt_manager/autotag/apply'));
|
||||
const answer = page.waitForResponse(item => item.url().endsWith('/prompt_manager/autotag/apply'));
|
||||
await page.locator('#applyReviewBtn').click();
|
||||
const body = (await applied).postDataJSON();
|
||||
expect(Object.keys(body).sort()).toEqual(['path', 'root_index', 'sha256', 'tags'].filter(key => key in body).sort());
|
||||
expect(body.prompt_id).toBeUndefined();
|
||||
expect((await answer).status()).toBe(200);
|
||||
});
|
||||
|
||||
@@ -132,7 +132,8 @@ test('the gallery cache clear keeps preferences and other tools caches', async (
|
||||
}))).toEqual({ theme: 'light', other: 'keep', prefixed: 'keep', notice: '3.2.0', caches: ['other-tool-gallery-thumbnails'] });
|
||||
});
|
||||
|
||||
test('actual gallery review generates and persists tags for the selected root', async ({ page, request }) => {
|
||||
// #379: the reviewed tags are that image's own tags; no prompt changes.
|
||||
test('actual gallery review generates and persists tags for the selected root\'s image', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
|
||||
await page.locator('#autoTagBtn').click();
|
||||
@@ -140,13 +141,15 @@ test('actual gallery review generates and persists tags for the selected root',
|
||||
await page.locator('#startReviewBtn').click();
|
||||
const generation = await generated;
|
||||
expect(generation.request().postDataJSON().root_index).toBe(0);
|
||||
await expect(page.locator('#reviewTagsContainer')).toContainText('generated');
|
||||
await expect(page.locator('#reviewTagsContainer')).toContainText('sandstone');
|
||||
const result = page.waitForResponse(response => response.url().endsWith('/prompt_manager/autotag/apply'));
|
||||
await page.locator('#applyReviewBtn').click();
|
||||
expect((await (await result).json()).success).toBe(true);
|
||||
const data = await (await request.get('/prompt_manager/recent?limit=20')).json();
|
||||
expect(data.results.find(prompt => prompt.text === 'Cat portrait').tags).toContain('generated');
|
||||
expect(data.results.find(prompt => prompt.text === 'Dog landscape').tags).not.toContain('generated');
|
||||
expect(data.results.find(prompt => prompt.text === 'Cat portrait').tags).toEqual(['cats']);
|
||||
expect(data.results.find(prompt => prompt.text === 'Dog landscape').tags).toEqual(['dogs']);
|
||||
const tagged = await (await request.get('/prompt_manager/images/output?mode=workspace&tags=sandstone')).json();
|
||||
expect(tagged.images.map(image => `${image.root_index}:${image.path}`)).toEqual(['0:shared/same.png']);
|
||||
});
|
||||
|
||||
test('actual Tags page reloads the selected filter and opens its image viewer', async ({ page }) => {
|
||||
|
||||
@@ -36,33 +36,48 @@ function fixture() {
|
||||
const reply = (index, data) => requests[index].resolve({ json: async () => data });
|
||||
return { manager, requests, reply, notifications };
|
||||
}
|
||||
test('gallery review retains root identity and applies tags with the returned prompt ID', async () => {
|
||||
const SHA = 'a'.repeat(64);
|
||||
test('gallery review retains root identity and applies the tags to the image by its sha256 (#379)', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const load = manager.loadNextReviewImage();
|
||||
assert.equal(requests[0].body.root_index, 1);
|
||||
reply(0, { success: true, prompt_id: 42, tags: ['smile'] });
|
||||
assert.equal(requests[0].body.existing, 'merge');
|
||||
reply(0, { success: true, prompt_id: 42, tags: ['smile'], sha256: SHA, image_tags: [] });
|
||||
await load;
|
||||
manager.loadNextReviewImage = () => {};
|
||||
const apply = manager.applyReviewTags();
|
||||
assert.deepEqual(requests[1].body, { prompt_id: 42, tags: ['smile'] });
|
||||
assert.equal(requests[1].url, '/prompt_manager/autotag/apply');
|
||||
assert.deepEqual(requests[1].body, ImageHelpers.autotagApplyBody(
|
||||
{ id: 'digest', path: 'same.png', root_index: 1 }, SHA, ['smile']));
|
||||
assert.equal('prompt_id' in requests[1].body, false);
|
||||
reply(1, { success: true });
|
||||
await apply;
|
||||
});
|
||||
test('unlinked images do not submit an invalid apply request or advance', async () => {
|
||||
test('an unlinked image applies to the image like a linked one', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const load = manager.loadNextReviewImage();
|
||||
reply(0, { success: true, prompt_id: null, tags: ['smile'] });
|
||||
reply(0, { success: true, prompt_id: null, tags: ['smile'], sha256: SHA, image_tags: [] });
|
||||
await load;
|
||||
manager.loadNextReviewImage = () => {};
|
||||
const apply = manager.applyReviewTags();
|
||||
if (requests[1]) reply(1, { success: false, error: "Unexpected apply" });
|
||||
assert.equal(requests[1].body.sha256, SHA);
|
||||
assert.deepEqual(requests[1].body.tags, ['smile']);
|
||||
reply(1, { success: true });
|
||||
await apply;
|
||||
assert.equal(requests.length, 1);
|
||||
assert.equal(manager.autoTagState.reviewIndex, 0);
|
||||
});
|
||||
test('an answer without a captured sha256 never sends an apply', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const load = manager.loadNextReviewImage();
|
||||
reply(0, { success: true, prompt_id: 42, tags: ['smile'] });
|
||||
await load;
|
||||
manager.loadNextReviewImage = () => {};
|
||||
await manager.applyReviewTags();
|
||||
assert.equal(requests.filter(r => r.url === '/prompt_manager/autotag/apply').length, 0);
|
||||
});
|
||||
test('a failed apply keeps the current image available for retry', async () => {
|
||||
const { manager, requests, reply } = fixture();
|
||||
const load = manager.loadNextReviewImage();
|
||||
reply(0, { success: true, prompt_id: 42, tags: ['smile'] });
|
||||
reply(0, { success: true, prompt_id: 42, tags: ['smile'], sha256: SHA, image_tags: [] });
|
||||
await load;
|
||||
const apply = manager.applyReviewTags();
|
||||
reply(1, { success: false, error: 'Try again' });
|
||||
|
||||
@@ -5,6 +5,8 @@ const {
|
||||
formatImageCaption,
|
||||
workflowDocumentHtml,
|
||||
autotagSingleBody,
|
||||
autotagApplyBody,
|
||||
autotagCompleteMessage,
|
||||
fallbackImageSource,
|
||||
} = require("../../web/js/image-helpers.js");
|
||||
|
||||
@@ -121,3 +123,27 @@ test('metadata lookup preserves encoded filename characters and works behind a p
|
||||
assert.equal(metadataLookupUrl('/prompt_manager/images/serve/a.png?v=123'), '/prompt_manager/images/prompt/a.png');
|
||||
assert.equal(metadataLookupUrl('/elsewhere/a.png'), null);
|
||||
});
|
||||
|
||||
test("autotagSingleBody asks to skip auto-tagged images only when told (#379)", () => {
|
||||
assert.equal(autotagSingleBody({ id: 1 }, WD14).existing, "merge");
|
||||
assert.equal(autotagSingleBody({ id: 1 }, { ...WD14, existing: "skip" }).existing, "skip");
|
||||
assert.equal(autotagSingleBody({ id: 1 }, { ...WD14, existing: "odd" }).existing, "merge");
|
||||
});
|
||||
|
||||
test("autotagApplyBody names the image, its bytes and the tags, never a prompt (#379)", () => {
|
||||
const sha = "a".repeat(64);
|
||||
assert.deepEqual(autotagApplyBody({ id: 4, prompt_id: 7 }, sha, ["bird"]), { image_id: 4, sha256: sha, tags: ["bird"] });
|
||||
assert.deepEqual(autotagApplyBody({ id: "digest", path: "sub/a.png", root_index: 1 }, sha, []),
|
||||
{ path: "sub/a.png", root_index: 1, sha256: sha, tags: [] });
|
||||
assert.deepEqual(autotagApplyBody({ id: 2 }, undefined, null), { image_id: 2, sha256: "", tags: [] });
|
||||
const tags = ["x"];
|
||||
autotagApplyBody({ id: 1 }, sha, tags).tags.push("y");
|
||||
assert.deepEqual(tags, ["x"]);
|
||||
});
|
||||
|
||||
test("autotagCompleteMessage counts images (#379)", () => {
|
||||
assert.equal(autotagCompleteMessage(2), "Auto tagging complete! Tagged 2 images.");
|
||||
assert.equal(autotagCompleteMessage(1), "Auto tagging complete! Tagged 1 image.");
|
||||
assert.equal(autotagCompleteMessage(0), "Auto tagging complete! Tagged 0 images.");
|
||||
assert.equal(autotagCompleteMessage(undefined), "Auto tagging complete! Tagged 0 images.");
|
||||
});
|
||||
|
||||
@@ -97,6 +97,30 @@ test("connect reports a non-2xx response through onerror and never onmessage", a
|
||||
assert.equal(source.readyState, source.CLOSED);
|
||||
});
|
||||
|
||||
test("a non-2xx JSON answer's error sentence reaches onerror as serverMessage", async () => {
|
||||
let error = null;
|
||||
const fetchImpl = async () => new Response(
|
||||
JSON.stringify({ success: false, error: "Auto-tagging is already running" }),
|
||||
{ status: 409, headers: { "Content-Type": "application/json" } },
|
||||
);
|
||||
const source = connect("/x", { fetchImpl });
|
||||
source.onerror = (err) => { error = err; };
|
||||
await until(() => error !== null);
|
||||
assert.match(error.message, /409/);
|
||||
assert.equal(error.status, 409);
|
||||
assert.equal(error.serverMessage, "Auto-tagging is already running");
|
||||
});
|
||||
|
||||
test("a non-2xx answer without a JSON error has no serverMessage", async () => {
|
||||
let error = null;
|
||||
const fetchImpl = async () => new Response("nope", { status: 500 });
|
||||
const source = connect("/x", { fetchImpl });
|
||||
source.onerror = (err) => { error = err; };
|
||||
await until(() => error !== null);
|
||||
assert.equal(error.status, 500);
|
||||
assert.equal(error.serverMessage, undefined);
|
||||
});
|
||||
|
||||
test("connect reports a network failure through onerror", async () => {
|
||||
let error = null;
|
||||
const source = connect("/x", { fetchImpl: async () => { throw new TypeError("Failed to fetch"); } });
|
||||
|
||||
@@ -363,23 +363,20 @@ class TestModelGate(SuggestTestCase):
|
||||
prompt_id = self._save_prompt("batch prompt")
|
||||
self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png"))
|
||||
started, proceed = threading.Event(), threading.Event()
|
||||
|
||||
def slow_generate(image_path, **kwargs):
|
||||
started.set()
|
||||
proceed.wait(WAIT)
|
||||
return ["batch"]
|
||||
|
||||
self.tagger.generate_tags = slow_generate
|
||||
# The legacy batch is the autotag job (#379): it suggests per image.
|
||||
self.tagger.block = (started, proceed)
|
||||
batch = asyncio.ensure_future(
|
||||
self.client.post("/prompt_manager/autotag/start?model_type=wd14-vit")
|
||||
)
|
||||
try:
|
||||
await asyncio.get_running_loop().run_in_executor(None, started.wait, WAIT)
|
||||
calls = len(self.tagger.calls)
|
||||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||||
self.assertEqual((status, body["code"]), (409, "model-busy"))
|
||||
self.assertEqual(self.tagger.calls, [])
|
||||
self.assertEqual(len(self.tagger.calls), calls)
|
||||
self.assertEqual(self.tagger.loaded, "wd14-vit") # not swapped
|
||||
finally:
|
||||
self.tagger.block = None
|
||||
proceed.set()
|
||||
response = await batch
|
||||
await response.read()
|
||||
|
||||
+69
-411
@@ -5,6 +5,7 @@ no network request is made.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
@@ -34,6 +35,7 @@ from aiohttp.test_utils import AioHTTPTestCase
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from py.api import PromptManagerAPI
|
||||
from py.api.autotag_routes import MODEL_BUSY_ERROR
|
||||
from py.api.jobs import JobBusy
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
@@ -70,6 +72,7 @@ class FakeAutotagService:
|
||||
self.status_calls = 0
|
||||
self.status_threads = []
|
||||
self.status_paths = {}
|
||||
self.suggest_calls = 0
|
||||
|
||||
def get_models_status(self):
|
||||
self.status_calls += 1
|
||||
@@ -115,6 +118,25 @@ class FakeAutotagService:
|
||||
raise self.generate_error
|
||||
return list(self.tags)
|
||||
|
||||
def generate_tag_suggestions(
|
||||
self,
|
||||
image_path,
|
||||
*,
|
||||
limit,
|
||||
prompt=None,
|
||||
general_threshold=None,
|
||||
character_threshold=None,
|
||||
):
|
||||
"""The suggestion entry point (#379): the fixed tags, at most ``limit``."""
|
||||
self.suggest_calls += 1
|
||||
tags = self.generate_tags(
|
||||
image_path,
|
||||
prompt=prompt,
|
||||
general_threshold=general_threshold,
|
||||
character_threshold=character_threshold,
|
||||
)
|
||||
return [(name, None) for name in tags[:limit]]
|
||||
|
||||
|
||||
class AutotagAPITestCase(AioHTTPTestCase):
|
||||
"""Stands up the PromptManager routes with a temp DB, temp output dir and a
|
||||
@@ -175,8 +197,17 @@ class AutotagAPITestCase(AioHTTPTestCase):
|
||||
)
|
||||
|
||||
def _make_image(self, directory, name="img.png"):
|
||||
"""A small, valid image (the legacy review captures and probes it)."""
|
||||
from PIL import Image
|
||||
|
||||
path = directory / name
|
||||
path.write_bytes(b"\x89PNG\r\n\x1a\n")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shade = sum(name.encode()) % 256
|
||||
fmt = {".jpg": "JPEG", ".jpeg": "JPEG", ".webp": "WEBP"}.get(
|
||||
path.suffix.lower(), "PNG"
|
||||
)
|
||||
with Image.new("RGB", (4, 4), (shade, 64, 128)) as image:
|
||||
image.save(path, fmt)
|
||||
return path
|
||||
|
||||
def _link_image(self, prompt_id, path):
|
||||
@@ -216,10 +247,9 @@ class TestAutotagSingleContainment(AutotagAPITestCase):
|
||||
self.assertEqual(data["tags"], ["1girl", "smile"])
|
||||
self.assertEqual(data["prompt_id"], prompt_id)
|
||||
self.assertEqual(len(self.service.generate_calls), 1)
|
||||
self.assertEqual(
|
||||
os.path.realpath(self.service.generate_calls[0][0]),
|
||||
os.path.realpath(str(image)),
|
||||
)
|
||||
# #379: the model reads a private capture of the image, never the file.
|
||||
self.assertNotIn(str(self.output_dir), self.service.generate_calls[0][0])
|
||||
self.assertEqual(data["sha256"], hashlib.sha256(image.read_bytes()).hexdigest())
|
||||
|
||||
async def test_path_inside_output_dir_is_tagged(self):
|
||||
image = self._make_image(self.output_dir)
|
||||
@@ -449,238 +479,6 @@ class TestAutotagSideEffectRoutesArePostOnly(AutotagAPITestCase):
|
||||
self.assertTrue((await resp.json())["success"])
|
||||
|
||||
|
||||
class TestAutotagBatch(AutotagAPITestCase):
|
||||
"""Batch tagging over the DB image list with SSE progress reporting."""
|
||||
|
||||
def _seed(self):
|
||||
"""Three prompts: A (one image), B (two images), C (missing file)."""
|
||||
a = self._save_prompt("prompt a")
|
||||
b = self._save_prompt("prompt b")
|
||||
c = self._save_prompt("prompt c")
|
||||
self._link_image(a, self._make_image(self.output_dir, "a1.png"))
|
||||
self._link_image(b, self._make_image(self.output_dir, "b1.png"))
|
||||
self._link_image(b, self._make_image(self.output_dir, "b2.png"))
|
||||
self._link_image(c, self.output_dir / "missing.png")
|
||||
return a, b, c
|
||||
|
||||
async def _start(self, query="model_type=gguf", body=None):
|
||||
if body is None:
|
||||
resp = await self.client.request(
|
||||
"POST", f"/prompt_manager/autotag/start?{query}"
|
||||
)
|
||||
else:
|
||||
resp = await self._post_json(f"/prompt_manager/autotag/start?{query}", body)
|
||||
return resp, self._sse_events(await resp.text())
|
||||
|
||||
async def test_batch_tags_each_prompt_once_and_reports_counts(self):
|
||||
a, b, c = self._seed()
|
||||
|
||||
resp, events = await self._start()
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
done = events[-1]
|
||||
self.assertEqual(done["type"], "complete")
|
||||
self.assertEqual(done["processed"], 2)
|
||||
self.assertEqual(done["tagged"], 2)
|
||||
self.assertEqual(done["skipped"], 2)
|
||||
self.assertEqual(done["errors"], 0)
|
||||
self.assertEqual(len(self.service.generate_calls), 2)
|
||||
self.assertEqual(self.api.db.get_prompt_by_id(a)["tags"], ["1girl", "smile"])
|
||||
self.assertEqual(self.api.db.get_prompt_by_id(b)["tags"], ["1girl", "smile"])
|
||||
self.assertEqual(self.api.db.get_prompt_by_id(c)["tags"], [])
|
||||
progress = [e for e in events if e["type"] == "progress"]
|
||||
self.assertTrue(all(0 <= e["progress"] <= 100 for e in progress))
|
||||
self.assertTrue(any("processed" in e for e in progress))
|
||||
|
||||
async def test_rows_outside_output_dirs_are_skipped_and_counted(self):
|
||||
inside = self._save_prompt("inside")
|
||||
outside = self._save_prompt("outside")
|
||||
self._link_image(inside, self._make_image(self.output_dir, "in.png"))
|
||||
secret = self._make_image(self.secret_dir, "password.png")
|
||||
self._link_image(outside, secret)
|
||||
|
||||
_, events = await self._start()
|
||||
|
||||
done = events[-1]
|
||||
self.assertEqual(done["type"], "complete")
|
||||
self.assertEqual(done["processed"], 1)
|
||||
self.assertEqual(done["tagged"], 1)
|
||||
self.assertEqual(done["skipped"], 1)
|
||||
self.assertEqual(done["skipped_outside_roots"], 1)
|
||||
self.assertEqual(len(self.service.generate_calls), 1)
|
||||
self.assertEqual(
|
||||
os.path.realpath(self.service.generate_calls[0][0]),
|
||||
os.path.realpath(str(self.output_dir / "in.png")),
|
||||
)
|
||||
self.assertEqual(self.api.db.get_prompt_by_id(outside)["tags"], [])
|
||||
progress = [e for e in events if e["type"] == "progress" and "skipped" in e]
|
||||
self.assertTrue(progress)
|
||||
self.assertTrue(all("skipped_outside_roots" in e for e in progress))
|
||||
body = json.dumps(events)
|
||||
self.assertNotIn(str(self.secret_dir), body)
|
||||
|
||||
async def test_second_start_while_running_is_409(self):
|
||||
self._seed()
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def slow_generate(image_path, **kwargs):
|
||||
started.set()
|
||||
release.wait(5)
|
||||
return ["slow"]
|
||||
|
||||
self.service.generate_tags = slow_generate
|
||||
first = asyncio.ensure_future(self._start())
|
||||
try:
|
||||
await asyncio.get_event_loop().run_in_executor(None, started.wait, 5)
|
||||
self.assertTrue(started.is_set())
|
||||
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start?model_type=gguf"
|
||||
)
|
||||
|
||||
self.assertEqual(resp.status, 409)
|
||||
self.assertEqual(
|
||||
await resp.json(),
|
||||
{"success": False, "error": "Auto-tagging is already running"},
|
||||
)
|
||||
finally:
|
||||
release.set()
|
||||
resp1, events = await first
|
||||
self.assertEqual(resp1.status, 200)
|
||||
self.assertEqual(events[-1]["type"], "complete")
|
||||
|
||||
async def test_lock_is_released_after_a_run(self):
|
||||
self._seed()
|
||||
|
||||
resp1, _ = await self._start()
|
||||
resp2, events = await self._start()
|
||||
|
||||
self.assertEqual(resp1.status, 200)
|
||||
self.assertEqual(resp2.status, 200)
|
||||
self.assertEqual(events[-1]["type"], "complete")
|
||||
|
||||
async def test_lock_is_released_when_the_run_fails(self):
|
||||
self._seed()
|
||||
self.service.load_error = RuntimeError("no model")
|
||||
|
||||
_, events = await self._start()
|
||||
self.assertEqual(events[-1]["type"], "error")
|
||||
|
||||
self.service.load_error = None
|
||||
resp, events = await self._start()
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(events[-1]["type"], "complete")
|
||||
|
||||
async def test_already_tagged_prompts_are_skipped_unless_asked(self):
|
||||
p = self._save_prompt("tagged", tags=["cat"])
|
||||
self._link_image(p, self._make_image(self.output_dir, "t.png"))
|
||||
|
||||
_, events = await self._start("model_type=gguf&skip_tagged=true")
|
||||
self.assertEqual(events[-1]["skipped"], 1)
|
||||
self.assertEqual(self.service.generate_calls, [])
|
||||
|
||||
_, events = await self._start("model_type=gguf&skip_tagged=false")
|
||||
self.assertEqual(events[-1]["tagged"], 1)
|
||||
self.assertEqual(
|
||||
self.api.db.get_prompt_by_id(p)["tags"], ["cat", "1girl", "smile"]
|
||||
)
|
||||
|
||||
async def test_prompt_already_holding_the_tags_counts_as_skipped(self):
|
||||
p = self._save_prompt("same", tags=["1girl", "smile"])
|
||||
self._link_image(p, self._make_image(self.output_dir, "s.png"))
|
||||
|
||||
_, events = await self._start("model_type=gguf&skip_tagged=false")
|
||||
|
||||
self.assertEqual(events[-1]["skipped"], 1)
|
||||
self.assertEqual(events[-1]["tagged"], 0)
|
||||
|
||||
async def test_generation_errors_are_counted_and_do_not_abort(self):
|
||||
self._seed()
|
||||
self.service.generate_error = RuntimeError("boom")
|
||||
|
||||
_, events = await self._start()
|
||||
|
||||
self.assertEqual(events[-1]["type"], "complete")
|
||||
# a1, b1 and b2 each fail: a failed prompt is retried on its next image
|
||||
self.assertEqual(events[-1]["errors"], 3)
|
||||
self.assertEqual(events[-1]["tagged"], 0)
|
||||
|
||||
async def test_json_body_parameters_are_accepted(self):
|
||||
self._seed()
|
||||
|
||||
_, events = await self._start(
|
||||
"", {"model_type": "wd14-vit", "general_threshold": 0.5, "prompt": "p"}
|
||||
)
|
||||
|
||||
self.assertEqual(events[-1]["type"], "complete")
|
||||
self.assertEqual(self.service.load_calls, [("wd14-vit", True)])
|
||||
self.assertEqual(self.service.generate_calls[0][2], 0.5)
|
||||
self.assertEqual(self.service.custom_prompt, "p")
|
||||
|
||||
async def test_keep_in_memory_false_unloads_after_batch(self):
|
||||
self._seed()
|
||||
|
||||
_, events = await self._start("model_type=gguf&keep_in_memory=false")
|
||||
|
||||
self.assertEqual(events[-1]["model_loaded"], False)
|
||||
self.assertEqual(self.service.unload_calls, 1)
|
||||
|
||||
async def test_model_not_downloaded_is_an_error_event(self):
|
||||
self.service.downloaded = False
|
||||
|
||||
_, events = await self._start()
|
||||
|
||||
self.assertEqual(events[-1]["type"], "error")
|
||||
self.assertEqual(self.service.load_calls, [])
|
||||
|
||||
async def test_load_failure_is_an_error_event_without_details(self):
|
||||
self.service.load_error = RuntimeError("/srv/models/secret.gguf is corrupt")
|
||||
|
||||
_, events = await self._start()
|
||||
|
||||
self.assertEqual(events[-1]["type"], "error")
|
||||
self.assertNotIn("/srv/models", events[-1]["message"])
|
||||
|
||||
async def test_non_numeric_threshold_is_400(self):
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start?general_threshold=abc"
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_invalid_json_body_is_400(self):
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/autotag/start",
|
||||
data="{not json",
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_client_disconnect_stops_the_batch_and_unloads(self):
|
||||
self._seed()
|
||||
real_write = web.StreamResponse.write
|
||||
state = {"writes": 0}
|
||||
|
||||
async def flaky_write(response, data):
|
||||
state["writes"] += 1
|
||||
if b"Processing " in data:
|
||||
raise ConnectionResetError("Cannot write to closing transport")
|
||||
return await real_write(response, data)
|
||||
|
||||
with patch.object(web.StreamResponse, "write", flaky_write):
|
||||
try:
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start?keep_in_memory=false"
|
||||
)
|
||||
await resp.read()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.assertEqual(len(self.service.generate_calls), 1)
|
||||
self.assertEqual(self.service.unload_calls, 1)
|
||||
|
||||
|
||||
class TestAutotagModelsAndDownloadErrors(AutotagAPITestCase):
|
||||
|
||||
async def test_models_error_is_500(self):
|
||||
@@ -834,7 +632,15 @@ class TestAutotagSingleEdgeCases(AutotagAPITestCase):
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(self.service.load_calls, [("wd14-vit", True)])
|
||||
self.assertEqual(self.service.custom_prompt, "terse")
|
||||
# #379: the request's prompt is for this call only.
|
||||
self.assertEqual(self.service.custom_prompt, "")
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/single",
|
||||
{"path": str(image), "model_type": "gguf", "prompt": "terse"},
|
||||
)
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(self.service.generate_calls[-1][1], "terse")
|
||||
self.assertEqual(self.service.custom_prompt, "")
|
||||
|
||||
async def test_image_id_without_stored_path_is_404(self):
|
||||
prompt_id = self._save_prompt()
|
||||
@@ -850,64 +656,6 @@ class TestAutotagSingleEdgeCases(AutotagAPITestCase):
|
||||
self.assertEqual(resp.status, 404)
|
||||
|
||||
|
||||
class TestApplyAutotag(AutotagAPITestCase):
|
||||
|
||||
async def test_missing_prompt_id_is_400(self):
|
||||
resp = await self._post_json("/prompt_manager/autotag/apply", {"tags": ["a"]})
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_no_tags_is_a_noop(self):
|
||||
resp = await self._post_json("/prompt_manager/autotag/apply", {"prompt_id": 1})
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
self.assertEqual(data["message"], "No tags to apply")
|
||||
|
||||
async def test_unknown_prompt_is_404(self):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply", {"prompt_id": 999, "tags": ["a"]}
|
||||
)
|
||||
self.assertEqual(resp.status, 404)
|
||||
|
||||
async def test_tags_are_merged_without_duplicates(self):
|
||||
prompt_id = self._save_prompt(tags=["cat"])
|
||||
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply",
|
||||
{"prompt_id": prompt_id, "tags": ["cat", "dog"]},
|
||||
)
|
||||
|
||||
data = await resp.json()
|
||||
self.assertEqual(data["added_tags"], ["dog"])
|
||||
self.assertEqual(data["total_tags"], 2)
|
||||
self.assertEqual(
|
||||
self.api.db.get_prompt_by_id(prompt_id)["tags"], ["cat", "dog"]
|
||||
)
|
||||
|
||||
async def test_string_tags_from_db_are_split(self):
|
||||
prompt_id = self._save_prompt()
|
||||
with patch.object(
|
||||
self.api.db,
|
||||
"get_prompt_by_id",
|
||||
return_value={"id": prompt_id, "tags": "cat, dog"},
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply",
|
||||
{"prompt_id": prompt_id, "tags": ["dog", "owl"]},
|
||||
)
|
||||
data = await resp.json()
|
||||
self.assertEqual(data["added_tags"], ["owl"])
|
||||
self.assertEqual(data["total_tags"], 3)
|
||||
|
||||
async def test_db_error_is_500(self):
|
||||
with patch.object(
|
||||
self.api.db, "get_prompt_by_id", side_effect=RuntimeError("db")
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply", {"prompt_id": 1, "tags": ["a"]}
|
||||
)
|
||||
self.assertEqual(resp.status, 500)
|
||||
|
||||
|
||||
class TestAutotagErrorBodiesHideServerPaths(AutotagAPITestCase):
|
||||
"""Every autotag handler reports exceptions through the shared public helper."""
|
||||
|
||||
@@ -932,9 +680,12 @@ class TestAutotagErrorBodiesHideServerPaths(AutotagAPITestCase):
|
||||
await self._assert_hidden(resp)
|
||||
|
||||
async def test_apply_error(self):
|
||||
with patch.object(self.api.db, "get_prompt_by_id", side_effect=self._locked()):
|
||||
image = self._make_image(self.output_dir)
|
||||
digest = hashlib.sha256(image.read_bytes()).hexdigest()
|
||||
with patch.object(self.api.db, "add_image_tags", side_effect=self._locked()):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/autotag/apply", {"prompt_id": 1, "tags": ["a"]}
|
||||
"/prompt_manager/autotag/apply",
|
||||
{"path": str(image), "sha256": digest, "tags": ["a"]},
|
||||
)
|
||||
await self._assert_hidden(resp)
|
||||
|
||||
@@ -1008,10 +759,13 @@ class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase):
|
||||
"POST", "/prompt_manager/autotag/download/gguf"
|
||||
)
|
||||
)
|
||||
await self._busy(
|
||||
await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start?model_type=gguf"
|
||||
)
|
||||
# The legacy batch is read as a stream: one error frame (#379).
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start?model_type=gguf"
|
||||
)
|
||||
self.assertEqual(
|
||||
self._sse_events(await resp.text()),
|
||||
[{"type": "error", "status": "error", "message": MODEL_BUSY_ERROR}],
|
||||
)
|
||||
await self._busy(
|
||||
await self.client.request("POST", "/prompt_manager/autotag/unload")
|
||||
@@ -1026,7 +780,7 @@ class TestLegacyRoutesShareTheModelGate(AutotagAPITestCase):
|
||||
self.assertEqual(self.service.unload_calls, 0)
|
||||
self.assertEqual(self.service.get_loaded_model_type(), "wd14-vit")
|
||||
|
||||
# A refused start released the batch lock: the next start runs.
|
||||
# A refused start left no job behind: the next start runs.
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start?model_type=gguf"
|
||||
)
|
||||
@@ -1238,7 +992,7 @@ if __name__ == "__main__":
|
||||
|
||||
|
||||
class TestGalleryReviewContract(AutotagAPITestCase):
|
||||
async def test_root_relative_gallery_targets_generate_and_apply_to_their_prompt(
|
||||
async def test_root_relative_gallery_targets_generate_and_apply_to_their_image(
|
||||
self,
|
||||
):
|
||||
second = self.tmp_root / "second-output"
|
||||
@@ -1255,16 +1009,22 @@ class TestGalleryReviewContract(AutotagAPITestCase):
|
||||
self.assertEqual(response.status, 200)
|
||||
data = await response.json()
|
||||
self.assertEqual(data["prompt_id"], pid)
|
||||
self.assertEqual(
|
||||
self.service.generate_calls[-1][0], str((root / "same.png").resolve())
|
||||
)
|
||||
response = await self._post_json(
|
||||
"/prompt_manager/autotag/apply",
|
||||
{"prompt_id": data["prompt_id"], "tags": data["tags"]},
|
||||
{
|
||||
"path": "same.png",
|
||||
"root_index": index,
|
||||
"sha256": data["sha256"],
|
||||
"tags": data["tags"],
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status, 200)
|
||||
# The image's own tags; the prompt is never written (#379).
|
||||
self.assertEqual(self.api.db.get_prompt_by_id(pid)["tags"], [])
|
||||
locator = os.path.normcase(os.path.realpath(str(root / "same.png")))
|
||||
self.assertEqual(
|
||||
set(self.api.db.get_prompt_by_id(pid)["tags"]), set(data["tags"])
|
||||
self.api.db.read_image_tags(locator, data["sha256"])["tags"],
|
||||
data["tags"],
|
||||
)
|
||||
|
||||
async def test_gallery_unlinked_image_returns_no_prompt_id(self):
|
||||
@@ -2058,108 +1818,6 @@ class TestTagOutputImagesJob(BatchJobTestCase):
|
||||
self.assertTrue(inside)
|
||||
|
||||
|
||||
# The legacy batch's per-image frames over BatchJobTestCase's library, newest
|
||||
# first, with every throttle check passing (pinned from main's generator).
|
||||
GOLDEN_LEGACY_ROWS = [
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 24,
|
||||
"status": "Processing 1/6...",
|
||||
"processed": 1,
|
||||
"tagged": 1,
|
||||
"skipped": 0,
|
||||
"skipped_outside_roots": 0,
|
||||
},
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 38,
|
||||
"status": "Skipping 2/6 (outside allowed directories)...",
|
||||
"processed": 1,
|
||||
"tagged": 1,
|
||||
"skipped": 1,
|
||||
"skipped_outside_roots": 1,
|
||||
},
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 52,
|
||||
"status": "Skipping 3/6 (file missing)...",
|
||||
"processed": 1,
|
||||
"tagged": 1,
|
||||
"skipped": 2,
|
||||
"skipped_outside_roots": 1,
|
||||
},
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 66,
|
||||
"status": "Processing 4/6...",
|
||||
"processed": 2,
|
||||
"tagged": 2,
|
||||
"skipped": 2,
|
||||
"skipped_outside_roots": 1,
|
||||
},
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 80,
|
||||
"status": "Skipping 5/6 (prompt already processed)...",
|
||||
"processed": 2,
|
||||
"tagged": 2,
|
||||
"skipped": 3,
|
||||
"skipped_outside_roots": 1,
|
||||
},
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 95,
|
||||
"status": "Skipping 6/6 (already tagged)...",
|
||||
"processed": 2,
|
||||
"tagged": 2,
|
||||
"skipped": 4,
|
||||
"skipped_outside_roots": 1,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class TestLegacyBatchFrames(BatchJobTestCase):
|
||||
|
||||
async def test_legacy_batch_events_unchanged(self):
|
||||
"""Golden frames of the legacy batch over the seeded library."""
|
||||
self._seed()
|
||||
ticks = iter(range(1000)) # every throttle check sees a second pass
|
||||
with patch("py.api.autotag_routes._monotonic", side_effect=lambda: next(ticks)):
|
||||
resp = await self.client.post(
|
||||
"/prompt_manager/autotag/start?model_type=wd14-vit&skip_tagged=true"
|
||||
)
|
||||
frames = self._sse_events(await resp.text()) # the stream runs here
|
||||
self.assertEqual(
|
||||
frames,
|
||||
[
|
||||
{"type": "progress", "progress": 0, "status": "Loading model..."},
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 5,
|
||||
"status": "Model loaded. Fetching all images from database...",
|
||||
},
|
||||
{
|
||||
"type": "progress",
|
||||
"progress": 10,
|
||||
"status": "Found 6 images. Processing...",
|
||||
},
|
||||
*GOLDEN_LEGACY_ROWS,
|
||||
{
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
"processed": 2,
|
||||
"tagged": 2,
|
||||
"skipped": 4,
|
||||
"errors": 0,
|
||||
"skipped_outside_roots": 1,
|
||||
"status": "Complete",
|
||||
"model_status": "Model kept in memory",
|
||||
"model_loaded": True,
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestTagOutputImagesEdges(BatchJobTestCase):
|
||||
|
||||
async def test_cancel_during_backup_loads_no_model(self):
|
||||
|
||||
@@ -0,0 +1,617 @@
|
||||
"""The legacy pages and the Tag image dialog tag images, never prompts (#379, PR 2).
|
||||
|
||||
- ``POST /autotag/start`` (legacy admin and gallery batch) starts the
|
||||
``autotag`` job titled "Tag output images (legacy page)" and relays it as
|
||||
the legacy frames; closing the stream cancels it.
|
||||
- ``POST /autotag/single`` suggests on a capture with the saved cap and
|
||||
answers the captured ``sha256``; ``existing: "skip"`` answers
|
||||
``already_tagged`` for an auto-tagged image without the model gate.
|
||||
- ``POST /autotag/apply`` unions into that image's tags and marks it; a
|
||||
``prompt_id`` body is refused.
|
||||
- A Tag image dialog save after Generate suggestions marks the image.
|
||||
|
||||
No real model runs: the per-image fake from test_autotag_batch answers.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
try:
|
||||
from test_autotag_batch import BatchTestCase, _absolute, _strings
|
||||
except ImportError:
|
||||
from tests.test_autotag_batch import BatchTestCase, _absolute, _strings
|
||||
|
||||
from database.execution_history_schema import rotate_history_epoch
|
||||
from py import model_gate
|
||||
from py.api import autotag_routes, image_tag_jobs
|
||||
from py.config import AutotagConfig
|
||||
|
||||
START = "/prompt_manager/autotag/start"
|
||||
SINGLE = "/prompt_manager/autotag/single"
|
||||
APPLY = "/prompt_manager/autotag/apply"
|
||||
SUGGEST = "/prompt_manager/image-annotations/suggest"
|
||||
OBSERVE = "/prompt_manager/image-annotations/observe"
|
||||
TAGS = "/prompt_manager/image-annotations/tags"
|
||||
WAIT = 30.0
|
||||
LEGACY_TITLE = "Tag output images (legacy page)"
|
||||
|
||||
|
||||
def frames_of(text):
|
||||
return [
|
||||
json.loads(line[len("data: ") :])
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data: ")
|
||||
]
|
||||
|
||||
|
||||
class LegacyTestCase(BatchTestCase):
|
||||
async def legacy(self, query="model_type=wd14-vit&skip_tagged=true"):
|
||||
response = await self.client.post(f"{START}?{query}")
|
||||
return response, frames_of(await response.text())
|
||||
|
||||
def digest(self, path):
|
||||
return hashlib.sha256((self.output_dir / path).read_bytes()).hexdigest()
|
||||
|
||||
async def single(self, body):
|
||||
response = await self.client.post(SINGLE, json=body)
|
||||
return response.status, await response.json()
|
||||
|
||||
async def apply(self, body):
|
||||
response = await self.client.post(APPLY, json=body)
|
||||
return response.status, await response.json()
|
||||
|
||||
def snapshot_prompts(self):
|
||||
conn = sqlite3.connect(self._temp_db.name)
|
||||
try:
|
||||
return [
|
||||
conn.execute(f"SELECT * FROM {t} ORDER BY 1,2").fetchall()
|
||||
for t in ("prompts", "prompt_tags")
|
||||
]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class TestLegacyBatch(LegacyTestCase):
|
||||
async def test_the_legacy_batch_tags_each_image_and_relays_frames(self):
|
||||
prompt = self.prompt("one prompt", tags=["beach"])
|
||||
a = self.image("a.png", (1, 2, 3), sees=["bird"])
|
||||
b = self.image("b.png", (4, 5, 6), sees=["waves"])
|
||||
self.link(prompt, a)
|
||||
self.link(prompt, b)
|
||||
before = self.snapshot_prompts()
|
||||
ticks = iter(range(1000))
|
||||
with patch.object(
|
||||
autotag_routes, "_monotonic", side_effect=lambda: next(ticks)
|
||||
):
|
||||
response, frames = await self.legacy()
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.stored("a.png")["tags"], ["bird"])
|
||||
self.assertEqual(self.stored("b.png")["tags"], ["waves"])
|
||||
self.assertEqual(self.snapshot_prompts(), before)
|
||||
phases = [f["status"] for f in frames if f.get("progress") in (0, 5, 8)]
|
||||
self.assertEqual(
|
||||
phases,
|
||||
sorted(
|
||||
phases,
|
||||
key=[
|
||||
"Taking a safety backup…",
|
||||
"Loading model…",
|
||||
"Listing output images…",
|
||||
].index,
|
||||
),
|
||||
)
|
||||
self.assertIn("Listing output images…", phases)
|
||||
rows = [f for f in frames if f["type"] == "progress" and "processed" in f]
|
||||
self.assertTrue(rows)
|
||||
for row in rows:
|
||||
self.assertEqual(
|
||||
set(row),
|
||||
{"type", "progress", "status", "processed", "tagged", "skipped"},
|
||||
)
|
||||
self.assertTrue(10 <= row["progress"] <= 95, row)
|
||||
self.assertEqual(
|
||||
frames[-1],
|
||||
{
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
"processed": 2,
|
||||
"tagged": 2,
|
||||
"skipped": 0,
|
||||
"errors": 0,
|
||||
"skipped_outside_roots": 0,
|
||||
"status": "Complete",
|
||||
"model_status": "Model kept in memory",
|
||||
"model_loaded": True,
|
||||
},
|
||||
)
|
||||
job = self.api._jobs.list()[0]
|
||||
self.assertEqual(
|
||||
(job.title, job.kind, job.state), (LEGACY_TITLE, "autotag", "finished")
|
||||
)
|
||||
self.assertEqual([s for s in _strings(frames) if _absolute(s)], [])
|
||||
|
||||
# Skip: the marked images are skipped before the model runs.
|
||||
self.tagger.calls.clear()
|
||||
_, again = await self.legacy()
|
||||
self.assertEqual(again[-1]["skipped"], 2)
|
||||
self.assertEqual(again[-1]["processed"], 0)
|
||||
self.assertEqual(self.tagger.calls, [])
|
||||
|
||||
async def test_no_still_images_and_keep_in_memory_false(self):
|
||||
_, frames = await self.legacy("model_type=wd14-vit&keep_in_memory=false")
|
||||
self.assertEqual(frames[-1]["status"], "No still images found")
|
||||
self.assertEqual(frames[-1]["model_loaded"], False)
|
||||
self.assertEqual(frames[-1]["model_status"], "Model unloaded")
|
||||
self.assertEqual(self.tagger.unload_calls, 1)
|
||||
|
||||
async def test_language_model_options_come_from_the_page_or_the_template(self):
|
||||
self.image("a.png", (1, 1, 1))
|
||||
await self.legacy("model_type=gguf&prompt=Tags%20only")
|
||||
self.assertEqual(self.tagger.calls[-1]["prompt"], "Tags only")
|
||||
self.assertEqual(self.tagger.calls[-1]["limit"], 10)
|
||||
self.image("b.png", (2, 2, 2))
|
||||
await self.relist()
|
||||
await self.legacy("model_type=gguf&prompt=%20")
|
||||
self.assertEqual(self.tagger.calls[-1]["prompt"], AutotagConfig.TEMPLATE)
|
||||
self.assertEqual(self.tagger.custom_prompt, "legacy prompt")
|
||||
|
||||
async def test_each_refusal_is_one_error_frame(self):
|
||||
self.tagger.downloaded["wd14-vit"] = False
|
||||
_, frames = await self.legacy()
|
||||
self.assertEqual(len(frames), 1)
|
||||
self.assertEqual(frames[0]["type"], "error")
|
||||
self.assertEqual(frames[0]["status"], "error")
|
||||
self.assertIn("WD14 ViT is not downloaded", frames[0]["message"])
|
||||
self.tagger.downloaded["wd14-vit"] = True
|
||||
|
||||
_, frames = await self.legacy("model_type=wd14-vit&general_threshold=2")
|
||||
self.assertEqual(
|
||||
frames,
|
||||
[
|
||||
{
|
||||
"type": "error",
|
||||
"status": "error",
|
||||
"message": "general_threshold must be between 0 and 1",
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
owner = model_gate.new_owner("image-annotations")
|
||||
self.assertTrue(model_gate.try_acquire(owner))
|
||||
try:
|
||||
_, frames = await self.legacy()
|
||||
finally:
|
||||
model_gate.release(owner)
|
||||
self.assertEqual(frames[0]["message"], autotag_routes.MODEL_BUSY_ERROR)
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
response = await self.client.post(f"{START}?general_threshold=abc")
|
||||
self.assertEqual(response.status, 400)
|
||||
|
||||
async def test_busy_jobs_refuse_the_legacy_batch(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)
|
||||
_, frames = await self.legacy()
|
||||
self.assertEqual(
|
||||
frames,
|
||||
[
|
||||
{
|
||||
"type": "error",
|
||||
"status": "error",
|
||||
"message": "Tag output images already running",
|
||||
}
|
||||
],
|
||||
)
|
||||
finally:
|
||||
self.tagger.block = None
|
||||
proceed.set()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
|
||||
# A second legacy batch keeps its JSON 409.
|
||||
self.image("b.png", (2, 2, 2))
|
||||
await self.relist()
|
||||
entered, proceed = threading.Event(), threading.Event()
|
||||
self.tagger.block = (entered, proceed)
|
||||
first = asyncio.ensure_future(
|
||||
self.legacy("model_type=wd14-vit&skip_tagged=false")
|
||||
)
|
||||
try:
|
||||
await self.wait_for(entered)
|
||||
second = await self.client.post(f"{START}?model_type=wd14-vit")
|
||||
self.assertEqual(second.status, 409)
|
||||
self.assertEqual(
|
||||
await second.json(),
|
||||
{"success": False, "error": "Auto-tagging is already running"},
|
||||
)
|
||||
finally:
|
||||
self.tagger.block = None
|
||||
proceed.set()
|
||||
_, frames = await first
|
||||
self.assertEqual(frames[-1]["type"], "complete")
|
||||
|
||||
async def test_a_legacy_start_that_loses_the_race_gets_the_json_409(self):
|
||||
# Two starts can both pass the title check before either registers its
|
||||
# job; the loser's busy refusal names the legacy batch and answers the
|
||||
# same JSON 409 as a second legacy batch.
|
||||
refusal = image_tag_jobs.Refusal(
|
||||
409,
|
||||
"busy",
|
||||
f"{autotag_routes.LEGACY_BATCH_TITLE} already running",
|
||||
web.json_response({}, status=409),
|
||||
)
|
||||
|
||||
async def refused(**_kwargs):
|
||||
return None, refusal
|
||||
|
||||
with patch.object(self.api, "_start_autotag", side_effect=refused):
|
||||
response = await self.client.post(f"{START}?model_type=wd14-vit")
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(
|
||||
await response.json(),
|
||||
{"success": False, "error": "Auto-tagging is already running"},
|
||||
)
|
||||
|
||||
async def test_closing_the_stream_cancels_the_job(self):
|
||||
for n in range(3):
|
||||
self.image(f"{n}.png", (n, n, n))
|
||||
real_write = web.StreamResponse.write
|
||||
|
||||
async def flaky_write(response, data):
|
||||
if b'"processed"' in data:
|
||||
raise ConnectionResetError("Cannot write to closing transport")
|
||||
return await real_write(response, data)
|
||||
|
||||
with patch.object(web.StreamResponse, "write", flaky_write):
|
||||
try:
|
||||
response = await self.client.post(f"{START}?model_type=wd14-vit")
|
||||
await response.read()
|
||||
except Exception: # nosec B110 - the client sees the dropped stream
|
||||
pass
|
||||
job = self.api._jobs.list()[0]
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
self.assertEqual(job.state, "cancelled")
|
||||
self.assertLess(job.counters["tagged"], 3)
|
||||
|
||||
async def test_a_quiet_job_repeats_its_last_frame(self):
|
||||
self.image("a.png", (1, 1, 1))
|
||||
entered, proceed = threading.Event(), threading.Event()
|
||||
self.tagger.block = (entered, proceed)
|
||||
with patch.object(autotag_routes, "LEGACY_REPEAT_AFTER", 0.05):
|
||||
run = asyncio.ensure_future(self.legacy())
|
||||
await self.wait_for(entered)
|
||||
await asyncio.sleep(0.3)
|
||||
self.tagger.block = None
|
||||
proceed.set()
|
||||
_, frames = await run
|
||||
statuses = [f.get("status") for f in frames]
|
||||
repeated = [s for s in statuses if s and statuses.count(s) > 1]
|
||||
self.assertTrue(repeated, statuses)
|
||||
self.assertEqual(frames[-1]["type"], "complete")
|
||||
|
||||
async def test_a_cancel_from_activity_ends_with_a_cancelled_frame(self):
|
||||
for n in range(2):
|
||||
self.image(f"{n}.png", (n, n, n))
|
||||
entered, proceed = threading.Event(), threading.Event()
|
||||
self.tagger.block = (entered, proceed)
|
||||
run = asyncio.ensure_future(self.legacy())
|
||||
try:
|
||||
await self.wait_for(entered)
|
||||
job = self.api._jobs.active("autotag")
|
||||
await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel")
|
||||
finally:
|
||||
self.tagger.block = None
|
||||
proceed.set()
|
||||
_, frames = await run
|
||||
self.assertEqual(
|
||||
frames[-1],
|
||||
{
|
||||
"type": "cancelled",
|
||||
"message": "Cancelled. Tags already added were kept.",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class TestLegacySingleAndApply(LegacyTestCase):
|
||||
async def test_single_suggests_on_a_capture_and_apply_tags_only_that_image(self):
|
||||
prompt = self.prompt("shared prompt", tags=["keep"])
|
||||
a = self.image("a.png", (1, 2, 3), sees=["bird", "sky"])
|
||||
b = self.image("b.png", (4, 5, 6), sees=["cat"])
|
||||
image_a = self.link(prompt, a)
|
||||
self.link(prompt, b)
|
||||
self.saved_set("b.png", b.read_bytes(), ["hand"])
|
||||
before = self.snapshot_prompts()
|
||||
status, data = await self.single(
|
||||
{"image_id": image_a, "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(data["tags"], ["bird", "sky"])
|
||||
self.assertEqual(data["sha256"], self.digest("a.png"))
|
||||
self.assertEqual(data["image_tags"], [])
|
||||
self.assertEqual(data["prompt_id"], prompt)
|
||||
self.assertEqual(self.tagger.custom_prompt, "legacy prompt")
|
||||
# Inference read the private capture, not the file.
|
||||
self.assertNotIn(str(self.output_dir), self.tagger.read_paths[-1])
|
||||
status, applied = await self.apply(
|
||||
{"image_id": image_a, "sha256": data["sha256"], "tags": ["bird"]}
|
||||
)
|
||||
self.assertEqual(
|
||||
(status, applied),
|
||||
(200, {"success": True, "added_tags": ["bird"], "total_tags": 1}),
|
||||
)
|
||||
self.assertEqual(self.stored("a.png")["tags"], ["bird"])
|
||||
self.assertEqual(self.stored("b.png")["tags"], ["hand"])
|
||||
self.assertEqual(self.snapshot_prompts(), before)
|
||||
self.assertEqual(self.marked(), {self.locator("a.png")})
|
||||
|
||||
async def test_skip_answers_already_tagged_without_the_gate(self):
|
||||
a = self.image("a.png", (1, 2, 3), sees=["bird"])
|
||||
await self.run_batch()
|
||||
self.tagger.calls.clear()
|
||||
owner = model_gate.new_owner("image-tag-job")
|
||||
self.assertTrue(model_gate.try_acquire(owner))
|
||||
try:
|
||||
status, data = await self.single(
|
||||
{
|
||||
"path": "a.png",
|
||||
"root_index": 0,
|
||||
"existing": "skip",
|
||||
"model_type": "wd14-vit",
|
||||
}
|
||||
)
|
||||
finally:
|
||||
model_gate.release(owner)
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(data["already_tagged"], True)
|
||||
self.assertEqual(data["image_tags"], ["bird"])
|
||||
self.assertEqual(self.tagger.calls, [])
|
||||
# Hand tags alone do not count; merge always suggests.
|
||||
b = self.image("b.png", (9, 9, 9))
|
||||
self.saved_set("b.png", b.read_bytes(), ["hand"])
|
||||
status, data = await self.single(
|
||||
{"path": str(b), "existing": "skip", "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertNotIn("already_tagged", data)
|
||||
self.assertEqual(data["image_tags"], ["hand"])
|
||||
status, data = await self.single(
|
||||
{"path": str(a), "existing": "merge", "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertNotIn("already_tagged", data)
|
||||
|
||||
async def test_single_uses_the_saved_cap_and_the_refusal_table(self):
|
||||
self.image("a.png", (1, 2, 3))
|
||||
AutotagConfig.SUGGESTION_LIMIT = 3
|
||||
status, data = await self.single(
|
||||
{"path": "a.png", "root_index": 0, "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertEqual((status, len(data["tags"])), (200, 3))
|
||||
(self.output_dir / "broken.png").write_bytes(b"not an image")
|
||||
await self.relist()
|
||||
status, data = await self.single(
|
||||
{"path": "broken.png", "root_index": 0, "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertEqual((status, data["code"]), (415, "not-a-still"))
|
||||
self.assertIn(self.locator("broken.png"), self.marked())
|
||||
self.tagger.downloaded["gguf"] = False
|
||||
status, data = await self.single(
|
||||
{"path": "a.png", "root_index": 0, "model_type": "gguf"}
|
||||
)
|
||||
self.assertEqual((status, data["code"]), (409, "model-unavailable"))
|
||||
self.api._capture_store().max_file = 16
|
||||
status, data = await self.single(
|
||||
{"path": "a.png", "root_index": 0, "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertEqual((status, data["code"]), (413, "too-large"))
|
||||
self.api._capture_store().max_file = 2**20
|
||||
owner = model_gate.new_owner("image-annotations")
|
||||
self.assertTrue(model_gate.try_acquire(owner))
|
||||
try:
|
||||
status, data = await self.single(
|
||||
{"path": "a.png", "root_index": 0, "model_type": "wd14-vit"}
|
||||
)
|
||||
finally:
|
||||
model_gate.release(owner)
|
||||
self.assertEqual((status, data["code"]), (409, "model-busy"))
|
||||
|
||||
async def test_apply_body_and_answer_table(self):
|
||||
a = self.image("a.png", (1, 2, 3))
|
||||
digest = self.digest("a.png")
|
||||
status, data = await self.apply({"prompt_id": 1, "tags": ["a"]})
|
||||
self.assertEqual(
|
||||
(status, data["error"]), (400, "This page is out of date. Reload it.")
|
||||
)
|
||||
status, _ = await self.apply(
|
||||
{"path": "a.png", "root_index": 0, "sha256": "x", "tags": ["a"]}
|
||||
)
|
||||
self.assertEqual(status, 400)
|
||||
status, _ = await self.apply(
|
||||
{"path": "a.png", "root_index": 0, "sha256": digest, "tags": ["bad,name"]}
|
||||
)
|
||||
self.assertEqual(status, 400)
|
||||
status, _ = await self.apply({"path": "a.png", "sha256": digest, "tags": []})
|
||||
self.assertEqual(status, 403) # a bare relative path is not inside a root
|
||||
# An empty reviewed list only marks the image.
|
||||
status, data = await self.apply({"path": str(a), "sha256": digest, "tags": []})
|
||||
self.assertEqual((status, data["added_tags"]), (200, []))
|
||||
self.assertEqual(self.marked(), {self.locator("a.png")})
|
||||
# The file changed since its tags were generated.
|
||||
status, data = await self.apply(
|
||||
{"path": str(a), "sha256": "0" * 64, "tags": ["x"]}
|
||||
)
|
||||
self.assertEqual((status, data["code"]), (409, "file-changed"))
|
||||
# Too many tags.
|
||||
full = self.image("full.png", (7, 7, 7))
|
||||
self.saved_set("full.png", full.read_bytes(), [f"t{i}" for i in range(100)])
|
||||
status, data = await self.apply(
|
||||
{"path": str(full), "sha256": self.digest("full.png"), "tags": ["one_more"]}
|
||||
)
|
||||
self.assertEqual((status, data["code"]), (409, "tag-limit"))
|
||||
# Gone.
|
||||
status, data = await self.apply(
|
||||
{"path": "gone.png", "root_index": 0, "sha256": digest, "tags": ["x"]}
|
||||
)
|
||||
self.assertIn(status, (403, 404))
|
||||
|
||||
async def test_apply_after_a_restore_writes_nothing(self):
|
||||
a = self.image("a.png", (1, 2, 3))
|
||||
digest = self.digest("a.png")
|
||||
real = self.api.db.get_execution_epoch
|
||||
|
||||
def stale_epoch():
|
||||
epoch = real()
|
||||
conn = sqlite3.connect(self._temp_db.name)
|
||||
try:
|
||||
rotate_history_epoch(conn)
|
||||
finally:
|
||||
conn.close()
|
||||
return epoch
|
||||
|
||||
with patch.object(self.api.db, "get_execution_epoch", side_effect=stale_epoch):
|
||||
status, data = await self.apply(
|
||||
{"path": str(a), "sha256": digest, "tags": ["x"]}
|
||||
)
|
||||
self.assertEqual((status, data["code"]), (409, "storage-changed"))
|
||||
self.assertEqual(self.stored("a.png")["tags"], [])
|
||||
|
||||
async def test_image_id_and_path_bodies_reach_the_same_bytes(self):
|
||||
prompt = self.prompt("p")
|
||||
a = self.image("sub/a.png", (1, 2, 3), sees=["x"])
|
||||
image_id = self.link(prompt, a)
|
||||
_, by_id = await self.single({"image_id": image_id, "model_type": "wd14-vit"})
|
||||
_, by_path = await self.single(
|
||||
{"path": "sub/a.png", "root_index": 0, "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertEqual(by_id["sha256"], by_path["sha256"])
|
||||
status, _ = await self.apply(
|
||||
{"image_id": image_id, "sha256": by_id["sha256"], "tags": ["x"]}
|
||||
)
|
||||
self.assertEqual(status, 200)
|
||||
status, data = await self.apply(
|
||||
{
|
||||
"path": "sub/a.png",
|
||||
"root_index": 0,
|
||||
"sha256": by_path["sha256"],
|
||||
"tags": ["y"],
|
||||
}
|
||||
)
|
||||
self.assertEqual((status, data["total_tags"]), (200, 2))
|
||||
|
||||
|
||||
class TestDialogMarks(LegacyTestCase):
|
||||
"""DEC4: a Tag image dialog save after Generate suggestions marks the image."""
|
||||
|
||||
def target(self, path):
|
||||
return {
|
||||
"root_index": 0,
|
||||
"root_scope": self.api._gallery_root_scope(self.output_dir, 0),
|
||||
"path": path,
|
||||
}
|
||||
|
||||
async def observe(self, path):
|
||||
response = await self.client.post(OBSERVE, json={"target": self.target(path)})
|
||||
data = await response.json()
|
||||
self.assertEqual(response.status, 200, data)
|
||||
return data
|
||||
|
||||
async def suggest(self, observed, path):
|
||||
response = await self.client.post(
|
||||
SUGGEST,
|
||||
json={
|
||||
"target": self.target(path),
|
||||
"capture_token": observed["capture"]["token"],
|
||||
"expected_epoch": observed["library_epoch"],
|
||||
"model_type": "wd14-vit",
|
||||
"options": {"general_threshold": 0.35, "character_threshold": 0.85},
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status, 200, await response.text())
|
||||
|
||||
async def save(self, observed, path, tags, revision=None):
|
||||
response = await self.client.put(
|
||||
TAGS,
|
||||
json={
|
||||
"target": self.target(path),
|
||||
"capture_token": observed["capture"]["token"],
|
||||
"sha256": observed["capture"]["sha256"],
|
||||
"expected_epoch": observed["library_epoch"],
|
||||
"expected_tag_revision": (
|
||||
observed["tag_revision"] if revision is None else revision
|
||||
),
|
||||
"tags": tags,
|
||||
},
|
||||
)
|
||||
return response.status, await response.json()
|
||||
|
||||
async def test_save_after_generate_marks_and_a_skip_batch_leaves_it(self):
|
||||
self.image("d.png", (1, 2, 3), sees=["bird", "sky"])
|
||||
observed = await self.observe("d.png")
|
||||
await self.suggest(observed, "d.png")
|
||||
status, _ = await self.save(observed, "d.png", ["sky"]) # bird removed
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(self.marked(), {self.locator("d.png")})
|
||||
job = await self.run_batch(existing="skip")
|
||||
self.assertEqual(self.rows(job, "skipped"), [("d.png", "Already auto-tagged")])
|
||||
self.assertEqual(self.stored("d.png")["tags"], ["sky"])
|
||||
|
||||
async def test_save_without_generate_does_not_mark(self):
|
||||
self.image("e.png", (1, 2, 3))
|
||||
observed = await self.observe("e.png")
|
||||
status, _ = await self.save(observed, "e.png", ["hand"])
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(self.marked(), set())
|
||||
|
||||
async def test_a_conflict_then_a_saved_retry_marks(self):
|
||||
f = self.image("f.png", (1, 2, 3))
|
||||
observed = await self.observe("f.png")
|
||||
await self.suggest(observed, "f.png")
|
||||
self.saved_set("f.png", f.read_bytes(), ["other"]) # another writer
|
||||
status, data = await self.save(observed, "f.png", ["mine"])
|
||||
self.assertEqual((status, data["code"]), (409, "tag-conflict"))
|
||||
self.assertEqual(self.marked(), set())
|
||||
status, _ = await self.save(
|
||||
observed, "f.png", ["mine"], revision=data["tag_revision"]
|
||||
)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(self.marked(), {self.locator("f.png")})
|
||||
|
||||
async def test_another_capture_does_not_inherit_the_mark(self):
|
||||
self.image("g.png", (1, 2, 3))
|
||||
self.image("h.png", (4, 5, 6))
|
||||
first = await self.observe("g.png")
|
||||
await self.suggest(first, "g.png")
|
||||
second = await self.observe("h.png")
|
||||
status, _ = await self.save(second, "h.png", ["x"])
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(self.marked(), set())
|
||||
self.assertIn(first["capture"]["token"], self.api._suggested_tokens)
|
||||
|
||||
async def test_gone_tokens_are_dropped_on_the_next_insert_without_touching_clocks(
|
||||
self,
|
||||
):
|
||||
self.image("i.png", (1, 2, 3))
|
||||
self.image("j.png", (4, 5, 6))
|
||||
first = await self.observe("i.png")
|
||||
await self.suggest(first, "i.png")
|
||||
store = self.api._capture_store()
|
||||
store.release(first["capture"]["token"])
|
||||
second = await self.observe("j.png")
|
||||
seen = []
|
||||
real_get = store.get
|
||||
with patch.object(
|
||||
store, "get", side_effect=lambda t: seen.append(t) or real_get(t)
|
||||
):
|
||||
await self.suggest(second, "j.png")
|
||||
self.assertEqual(seen, [])
|
||||
self.assertEqual(self.api._suggested_tokens, {second["capture"]["token"]})
|
||||
@@ -1398,50 +1398,6 @@ class TestModelHousekeeping(DatabaseTestCase):
|
||||
self.assertIsNot(self.db.model.get_connection(), failing)
|
||||
|
||||
|
||||
class TestAutotagBatchCandidates(DatabaseTestCase):
|
||||
"""#328: the batch's candidate read never loads workflow documents."""
|
||||
|
||||
def test_autotag_candidates_skip_workflow_data(self):
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
first = self.db.save_prompt(
|
||||
text="a", tags=["one", "two"], prompt_hash=generate_prompt_hash("a")
|
||||
)
|
||||
second = self.db.save_prompt(text="b", prompt_hash=generate_prompt_hash("b"))
|
||||
old = self.db.link_image_to_prompt(
|
||||
first, "/out/a.png", {"workflow": {"big": "x" * 1000}}
|
||||
)
|
||||
new = self.db.link_image_to_prompt(second, "/out/b.png")
|
||||
statements = []
|
||||
conn = self.db.model.get_connection() # this thread's own connection
|
||||
conn.set_trace_callback(statements.append)
|
||||
try:
|
||||
rows = self.db.autotag_batch_candidates()
|
||||
finally:
|
||||
conn.set_trace_callback(None)
|
||||
self.assertEqual(
|
||||
rows[0],
|
||||
{
|
||||
"id": new,
|
||||
"image_path": "/out/b.png",
|
||||
"prompt_id": second,
|
||||
"prompt_tags": [],
|
||||
},
|
||||
)
|
||||
self.assertEqual(
|
||||
{**rows[1], "prompt_tags": sorted(rows[1]["prompt_tags"])},
|
||||
{
|
||||
"id": old,
|
||||
"image_path": "/out/a.png",
|
||||
"prompt_id": first,
|
||||
"prompt_tags": ["one", "two"],
|
||||
},
|
||||
)
|
||||
self.assertEqual(len(rows), 2)
|
||||
self.assertTrue(statements)
|
||||
self.assertFalse(any("workflow_data" in sql for sql in statements))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
@@ -535,7 +535,9 @@ class TestHeldJobStates(_LegacyHoldCase):
|
||||
self.assertIn("event: error", text)
|
||||
self.assertEqual(await self._listed(), [("thumbnails", "failed")])
|
||||
|
||||
async def test_legacy_autotag_that_ends_in_error_is_failed(self):
|
||||
async def test_legacy_autotag_refused_at_start_leaves_no_job(self):
|
||||
# #379: the legacy batch is the autotag job; a model that is not
|
||||
# downloaded refuses the start (one error frame) and registers nothing.
|
||||
service = MagicMock()
|
||||
service.get_models_status.return_value = {} # model not downloaded
|
||||
with patch("py.autotag.get_autotag_service", return_value=service):
|
||||
@@ -545,7 +547,7 @@ class TestHeldJobStates(_LegacyHoldCase):
|
||||
text = await resp.text()
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertIn('"type": "error"', text)
|
||||
self.assertEqual(await self._listed(), [("autotag", "failed")])
|
||||
self.assertEqual(await self._listed(), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+4
-4
@@ -1187,14 +1187,13 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
|
||||
<div class="space-y-1.5">
|
||||
<label class="flex items-center p-2.5 bg-pm-input rounded-pm-md cursor-pointer hover:bg-pm-hover border border-pm">
|
||||
<input type="radio" name="autoTagMode" value="skip" checked class="w-3.5 h-3.5 accent-[var(--pm-accent)]">
|
||||
<span class="ml-2.5 text-pm text-xs">Skip images with existing tags</span>
|
||||
<span class="ml-2.5 text-pm text-xs">Skip images already auto-tagged</span>
|
||||
</label>
|
||||
<label class="flex items-center p-2.5 bg-pm-input rounded-pm-md cursor-pointer hover:bg-pm-hover border border-pm">
|
||||
<input type="radio" name="autoTagMode" value="retag" class="w-3.5 h-3.5 accent-[var(--pm-accent)]">
|
||||
<span class="ml-2.5 text-pm text-xs">Re-tag all images</span>
|
||||
</label>
|
||||
</div>
|
||||
<p class="text-pm-muted text-[11px] mt-1">Note: "auto-scanned" is not counted as a real tag</p>
|
||||
</div>
|
||||
|
||||
<!-- Action Buttons -->
|
||||
@@ -1293,6 +1292,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
|
||||
<div class="w-full bg-pm-input rounded-full h-1.5">
|
||||
<div id="autoTagProgressBar" class="bg-pm-accent h-1.5 rounded-full transition-all duration-300" style="width: 0%"></div>
|
||||
</div>
|
||||
<p id="autoTagCloseNote" class="text-pm-muted text-[11px] mt-1.5">Closing this window stops tagging. Tags already added are kept. Prompt Manager's Activity page shows the time left.</p>
|
||||
</div>
|
||||
|
||||
<div class="grid grid-cols-3 gap-3 mb-4">
|
||||
@@ -1332,7 +1332,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
|
||||
<!-- Re-tag Confirmation Modal -->
|
||||
<div id="retagConfirmModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50">
|
||||
<div class="bg-pm-surface rounded-pm-lg p-4 max-w-lg w-full mx-4 border border-pm shadow-pm">
|
||||
<h3 class="text-sm font-semibold text-pm mb-3">Image Already Tagged</h3>
|
||||
<h3 class="text-sm font-semibold text-pm mb-3">Image Already Auto-tagged</h3>
|
||||
|
||||
<div class="mb-3">
|
||||
<img id="retagPreviewImage" src="" alt="Preview" class="w-full h-40 object-contain bg-pm-primary rounded-pm-md">
|
||||
@@ -1345,7 +1345,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<p class="text-pm-secondary text-xs mb-4">This image already has tags. Would you like to generate new tags?</p>
|
||||
<p class="text-pm-secondary text-xs mb-4">A tagging model already tagged this image. Generate new tags?</p>
|
||||
|
||||
<div class="flex gap-2 justify-end">
|
||||
<button id="retagSkipBtn"
|
||||
|
||||
@@ -682,6 +682,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
|
||||
<div class="w-full bg-pm-input rounded-full h-2">
|
||||
<div id="autoTagProgressBar" class="bg-pm-accent h-2 rounded-full transition-all duration-300" style="width: 0%"></div>
|
||||
</div>
|
||||
<p id="autoTagCloseNote" class="text-pm-muted text-xs mt-1.5">Closing this window stops tagging. Tags already added are kept. Prompt Manager's Activity page shows the time left.</p>
|
||||
</div>
|
||||
|
||||
<!-- Stats -->
|
||||
|
||||
+52
-65
@@ -3801,7 +3801,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
} else if (data.type === 'complete') {
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification(`Auto tagging complete! Applied tags to ${data.tagged || 0} prompts.`, 'success');
|
||||
this.showNotification(ImageHelpers.autotagCompleteMessage(data.tagged), 'success');
|
||||
this.refreshList(); // Refresh the prompt list
|
||||
} else if (data.type === 'error') {
|
||||
this.autoTagState.eventSource.close();
|
||||
@@ -3819,7 +3819,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
console.error('EventSource readyState:', this.autoTagState.eventSource.readyState);
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification('Connection lost during auto tagging', 'error');
|
||||
this.showNotification((error && error.serverMessage) || 'Connection lost during auto tagging', 'error');
|
||||
};
|
||||
} catch (error) {
|
||||
console.error('Auto tag error:', error);
|
||||
@@ -3894,10 +3894,14 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
}
|
||||
}
|
||||
|
||||
// Check if tags array has "real" tags (excluding auto-scanned and prepend:* metadata)
|
||||
getRealTags(tags) {
|
||||
if (!tags || !Array.isArray(tags)) return [];
|
||||
return tags.filter(tag => tag !== 'auto-scanned' && !tag.startsWith('prepend:') && !tag.startsWith('append:'));
|
||||
/** POST /autotag/single for `image` (#379: `existing` "skip" answers already_tagged for an auto-tagged image). */
|
||||
async requestReviewTags(image, existing) {
|
||||
const response = await ApiClient.fetchUrl('/prompt_manager/autotag/single', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify(ImageHelpers.autotagSingleBody(image, { ...this.autoTagState, existing }))
|
||||
});
|
||||
return response.json();
|
||||
}
|
||||
|
||||
// Show the re-tag confirmation modal
|
||||
@@ -3949,26 +3953,48 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
// always provides a servable url; never rebuild it from the path.
|
||||
const imageUrl = this.getImageUrl(image);
|
||||
|
||||
// Check if image already has real tags (excluding auto-scanned)
|
||||
const realTags = this.getRealTags(image.prompt_tags);
|
||||
this.autoTagState.currentImage = image;
|
||||
this.autoTagState.currentSha256 = null;
|
||||
this.autoTagState.tagsExpanded = false; // Reset accordion for new image
|
||||
|
||||
if (realTags.length > 0) {
|
||||
// Skip if "Skip All Tagged" was selected
|
||||
if (this.autoTagState.skipAllTagged) {
|
||||
this.autoTagState.reviewIndex++;
|
||||
return this.loadNextReviewImage();
|
||||
let data;
|
||||
try {
|
||||
// Each image is first asked with "skip": an image a tagging model already tagged answers
|
||||
// already_tagged (its image tags, never its prompt's) without running the model (#379).
|
||||
data = await this.requestReviewTags(image, 'skip');
|
||||
if (data.success && data.already_tagged) {
|
||||
if (this.autoTagState.skipAllTagged) {
|
||||
this.autoTagState.reviewIndex++;
|
||||
return this.loadNextReviewImage();
|
||||
}
|
||||
const choice = await this.showRetagConfirmation(image, imageUrl, data.image_tags || []);
|
||||
if (choice === 'skip' || choice === 'skipAll') {
|
||||
this.autoTagState.reviewIndex++;
|
||||
return this.loadNextReviewImage();
|
||||
}
|
||||
// choice === 'retag': generate new tags for it
|
||||
this.showReviewImage(image, imageUrl);
|
||||
data = await this.requestReviewTags(image, 'merge');
|
||||
} else {
|
||||
this.showReviewImage(image, imageUrl);
|
||||
}
|
||||
|
||||
// Show confirmation modal
|
||||
const choice = await this.showRetagConfirmation(image, imageUrl, realTags);
|
||||
|
||||
if (choice === 'skip' || choice === 'skipAll') {
|
||||
this.autoTagState.reviewIndex++;
|
||||
return this.loadNextReviewImage();
|
||||
if (data.success) {
|
||||
this.autoTagState.currentTags = data.tags;
|
||||
this.autoTagState.currentSha256 = data.sha256 || null;
|
||||
this.renderReviewTags();
|
||||
} else {
|
||||
document.getElementById('reviewTagsVisible').innerHTML =
|
||||
`<div class="text-pm-error">Error: ${escapeHtml(data.error)}</div>`;
|
||||
}
|
||||
// choice === 'retag' - continue with tagging
|
||||
} catch (error) {
|
||||
console.error('Error generating tags:', error);
|
||||
this.showReviewImage(image, imageUrl);
|
||||
document.getElementById('reviewTagsVisible').innerHTML =
|
||||
'<div class="text-pm-error">Failed to generate tags</div>';
|
||||
}
|
||||
}
|
||||
|
||||
showReviewImage(image, imageUrl) {
|
||||
document.getElementById('reviewImage').src = imageUrl;
|
||||
// The record's file name, else the last segment of its path (#328: `filename` was never declared).
|
||||
const filename = image.filename || String(image.image_path || image.path || '').split(/[\\/]/).pop();
|
||||
@@ -3977,37 +4003,6 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
document.getElementById('reviewTagsHidden').innerHTML = '';
|
||||
document.getElementById('reviewTagsAccordion').classList.add('hidden');
|
||||
document.getElementById('reviewTagsToggle').classList.add('hidden');
|
||||
|
||||
// Store prompt_id from database for applying tags later
|
||||
this.autoTagState.currentPromptId = image.prompt_id;
|
||||
this.autoTagState.tagsExpanded = false; // Reset accordion for new image
|
||||
|
||||
try {
|
||||
// Review rows come from /images/all (database ids), so this resolves to image_id
|
||||
const requestBody = ImageHelpers.autotagSingleBody(image, this.autoTagState);
|
||||
|
||||
const response = await ApiClient.fetchUrl('/prompt_manager/autotag/single', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify(requestBody)
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
if (data.success) {
|
||||
this.autoTagState.currentTags = data.tags;
|
||||
// Use prompt_id from database, fallback to API response
|
||||
this.autoTagState.currentPromptId = image.prompt_id || data.prompt_id;
|
||||
this.renderReviewTags();
|
||||
} else {
|
||||
document.getElementById('reviewTagsVisible').innerHTML =
|
||||
`<div class="text-pm-error">Error: ${escapeHtml(data.error)}</div>`;
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Error generating tags:', error);
|
||||
document.getElementById('reviewTagsVisible').innerHTML =
|
||||
'<div class="text-pm-error">Failed to generate tags</div>';
|
||||
}
|
||||
}
|
||||
|
||||
/** /prompt_manager/images/all is capped per call; collect every page. */
|
||||
@@ -4126,15 +4121,9 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
}
|
||||
|
||||
async applyReviewTags() {
|
||||
if (this.autoTagState.currentTags.length === 0) {
|
||||
this.skipReviewImage();
|
||||
return;
|
||||
}
|
||||
|
||||
const promptId = this.autoTagState.currentPromptId;
|
||||
|
||||
if (!promptId) {
|
||||
this.showNotification('No linked prompt found for this image', 'warning');
|
||||
// #379: the reviewed tags go to this image (linked or not), never its prompt; an empty
|
||||
// list still marks the image as reviewed by a tagging model.
|
||||
if (!this.autoTagState.currentSha256) {
|
||||
this.skipReviewImage();
|
||||
return;
|
||||
}
|
||||
@@ -4143,10 +4132,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
const response = await ApiClient.fetchUrl('/prompt_manager/autotag/apply', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
prompt_id: promptId,
|
||||
tags: this.autoTagState.currentTags
|
||||
})
|
||||
body: JSON.stringify(ImageHelpers.autotagApplyBody(
|
||||
this.autoTagState.currentImage, this.autoTagState.currentSha256, this.autoTagState.currentTags))
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
+20
-16
@@ -3560,19 +3560,24 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
} else if (data.type === 'complete') {
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification(`Auto tagging complete! Applied tags to ${data.tagged || 0} prompts.`, 'success');
|
||||
this.showNotification(ImageHelpers.autotagCompleteMessage(data.tagged), 'success');
|
||||
this.loadImages(); // Refresh the gallery
|
||||
} else if (data.type === 'error') {
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification(`Auto tag error: ${data.message}`, 'error');
|
||||
} else if (data.type === 'cancelled') {
|
||||
// Cancelled in Prompt Manager's Activity (#379): as its own cancel.
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification('Auto tagging cancelled', 'info');
|
||||
}
|
||||
};
|
||||
|
||||
this.autoTagState.eventSource.onerror = () => {
|
||||
this.autoTagState.eventSource.onerror = (error) => {
|
||||
this.autoTagState.eventSource.close();
|
||||
this.hideModal("autoTagProgressModal");
|
||||
this.showNotification('Connection lost during auto tagging', 'error');
|
||||
this.showNotification((error && error.serverMessage) || 'Connection lost during auto tagging', 'error');
|
||||
};
|
||||
} catch (error) {
|
||||
console.error('Auto tag error:', error);
|
||||
@@ -3629,6 +3634,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
this.reviewRequest = request;
|
||||
this.autoTagState.currentTags = [];
|
||||
this.autoTagState.reviewPromptId = null;
|
||||
this.autoTagState.reviewImage = null;
|
||||
this.autoTagState.reviewSha256 = null;
|
||||
if (this.autoTagState.reviewIndex >= this.autoTagState.reviewImages.length) {
|
||||
this.hideModal("autoTagReviewModal");
|
||||
this.showNotification('Review complete!', 'success');
|
||||
@@ -3643,8 +3650,9 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
document.getElementById('reviewTagsContainer').innerHTML = '<div class="text-pm-secondary">Generating tags...</div>';
|
||||
|
||||
try {
|
||||
// Gallery entries come from /images/output: id is a path digest, so address by path
|
||||
const requestBody = ImageHelpers.autotagSingleBody(image, this.autoTagState);
|
||||
// Gallery entries come from /images/output: id is a path digest, so address by path.
|
||||
// The gallery review always suggests (it never skipped tagged images, #379).
|
||||
const requestBody = ImageHelpers.autotagSingleBody(image, { ...this.autoTagState, existing: 'merge' });
|
||||
|
||||
const response = await ApiClient.fetchUrl('/prompt_manager/autotag/single', {
|
||||
method: 'POST',
|
||||
@@ -3657,6 +3665,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
if (this.reviewRequest !== request) return;
|
||||
if (data.success) {
|
||||
this.autoTagState.reviewPromptId = data.prompt_id ?? null;
|
||||
this.autoTagState.reviewImage = image;
|
||||
this.autoTagState.reviewSha256 = data.sha256 || null;
|
||||
this.autoTagState.currentTags = data.tags;
|
||||
this.renderReviewTags();
|
||||
} else {
|
||||
@@ -3697,16 +3707,12 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
}
|
||||
|
||||
async applyReviewTags() {
|
||||
if (this.autoTagState.currentTags.length === 0) {
|
||||
// #379: the reviewed tags go to this image (linked or not), never its prompt; an empty
|
||||
// list still marks the image as reviewed by a tagging model.
|
||||
if (!this.autoTagState.reviewSha256) {
|
||||
this.skipReviewImage();
|
||||
return;
|
||||
}
|
||||
|
||||
const promptId = this.autoTagState.reviewPromptId;
|
||||
if (!promptId) {
|
||||
this.showNotification('This image has no linked prompt. Scan the output folder first.', 'warning');
|
||||
return;
|
||||
}
|
||||
const request = this.reviewRequest;
|
||||
if (this.reviewApplying === request) return;
|
||||
this.reviewApplying = request;
|
||||
@@ -3715,10 +3721,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
|
||||
const response = await ApiClient.fetchUrl('/prompt_manager/autotag/apply', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
prompt_id: promptId,
|
||||
tags: this.autoTagState.currentTags
|
||||
})
|
||||
body: JSON.stringify(ImageHelpers.autotagApplyBody(
|
||||
this.autoTagState.reviewImage, this.autoTagState.reviewSha256, this.autoTagState.currentTags))
|
||||
});
|
||||
|
||||
const data = await response.json();
|
||||
|
||||
+44
-10
@@ -70,10 +70,28 @@
|
||||
* path digest as id, which would 404, so those are addressed by path.
|
||||
*
|
||||
* @param {{id?: unknown, path?: unknown}} image
|
||||
* @param {{modelType: string, generalThreshold?: number, characterThreshold?: number, prompt?: string}} settings
|
||||
* `existing` (#379): "skip" answers `already_tagged` for an image a model
|
||||
* already tagged; anything else is "merge" (always suggest).
|
||||
*
|
||||
* @param {{id?: unknown, path?: unknown}} image
|
||||
* @param {{modelType: string, generalThreshold?: number, characterThreshold?: number, prompt?: string, existing?: string}} settings
|
||||
* @returns {object} JSON-serialisable request body
|
||||
*/
|
||||
function autotagSingleBody(image, settings) {
|
||||
const target = autotagTarget(image);
|
||||
const existing = settings.existing === "skip" ? "skip" : "merge";
|
||||
const modelType = String(settings.modelType || "");
|
||||
const params = modelType.startsWith("wd14")
|
||||
? {
|
||||
general_threshold: settings.generalThreshold,
|
||||
character_threshold: settings.characterThreshold,
|
||||
}
|
||||
: { prompt: settings.prompt };
|
||||
return { ...target, model_type: modelType, existing, ...params };
|
||||
}
|
||||
|
||||
/** The image a legacy single/apply body names: its row id, or its path (and root). */
|
||||
function autotagTarget(image) {
|
||||
const entry = image || {};
|
||||
const target =
|
||||
typeof entry.path === "string" && entry.path && !Number.isInteger(entry.id)
|
||||
@@ -82,14 +100,22 @@
|
||||
if ("path" in target && Number.isInteger(entry.root_index)) {
|
||||
target.root_index = entry.root_index;
|
||||
}
|
||||
const modelType = String(settings.modelType || "");
|
||||
const params = modelType.startsWith("wd14")
|
||||
? {
|
||||
general_threshold: settings.generalThreshold,
|
||||
character_threshold: settings.characterThreshold,
|
||||
}
|
||||
: { prompt: settings.prompt };
|
||||
return { ...target, model_type: modelType, ...params };
|
||||
return target;
|
||||
}
|
||||
|
||||
/**
|
||||
* Body for POST /prompt_manager/autotag/apply (#379): the image, the
|
||||
* `sha256` /autotag/single answered and the reviewed tags. The tags go to
|
||||
* that image only, never to its prompt; an empty list marks it reviewed.
|
||||
*/
|
||||
function autotagApplyBody(image, sha256, tags) {
|
||||
return { ...autotagTarget(image), sha256: String(sha256 || ""), tags: Array.isArray(tags) ? [...tags] : [] };
|
||||
}
|
||||
|
||||
/** The legacy batch's completion toast (#379): images, never prompts. */
|
||||
function autotagCompleteMessage(tagged) {
|
||||
const count = Number.isSafeInteger(tagged) && tagged > 0 ? tagged : 0;
|
||||
return `Auto tagging complete! Tagged ${count} ${count === 1 ? "image" : "images"}.`;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -122,7 +148,15 @@
|
||||
}
|
||||
}
|
||||
|
||||
const api = { formatImageCaption, workflowDocumentHtml, autotagSingleBody, fallbackImageSource, metadataLookupUrl };
|
||||
const api = {
|
||||
formatImageCaption,
|
||||
workflowDocumentHtml,
|
||||
autotagSingleBody,
|
||||
autotagApplyBody,
|
||||
autotagCompleteMessage,
|
||||
fallbackImageSource,
|
||||
metadataLookupUrl,
|
||||
};
|
||||
if (typeof module !== "undefined" && module.exports) module.exports = api;
|
||||
else root.ImageHelpers = api;
|
||||
})(typeof window !== "undefined" ? window : globalThis);
|
||||
|
||||
+15
-1
@@ -59,6 +59,20 @@
|
||||
* @returns {{onmessage: Function|null, onerror: Function|null, onclose: Function|null,
|
||||
* readyState: number, close: Function, CONNECTING: number, OPEN: number, CLOSED: number}}
|
||||
*/
|
||||
// A non-2xx answer as an Error carrying its status and, when the body is
|
||||
// the API's JSON {error}, that sentence as serverMessage for the page.
|
||||
async function httpError(response) {
|
||||
const error = new Error(`HTTP ${response.status}`);
|
||||
error.status = response.status;
|
||||
try {
|
||||
const data = await response.json();
|
||||
if (data && typeof data.error === "string" && data.error) error.serverMessage = data.error;
|
||||
} catch (_) {
|
||||
// Not JSON: the status alone.
|
||||
}
|
||||
return error;
|
||||
}
|
||||
|
||||
function connect(url, { method = "GET", body, fetchImpl } = {}) {
|
||||
const doFetch = fetchImpl || ((...args) => globalThis.fetch(...args));
|
||||
const controller = new AbortController();
|
||||
@@ -103,7 +117,7 @@
|
||||
if (body !== undefined) init.body = body;
|
||||
const requestUrl = root.ApiClient ? root.ApiClient.resolveUrl(url) : url;
|
||||
const response = await doFetch(requestUrl, init);
|
||||
if (!response.ok) throw new Error(`HTTP ${response.status}`);
|
||||
if (!response.ok) throw await httpError(response);
|
||||
if (source.readyState === CLOSED) return;
|
||||
source.readyState = OPEN;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user