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

The download continues in Activity.

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

The download continues in Activity.

diff --git a/web/js/admin.js b/web/js/admin.js index b02fb90..be78250 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -3689,6 +3689,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; document.getElementById('downloadStatus').textContent = 'Preparing...'; document.getElementById('downloadProgressPercent').textContent = '0%'; document.getElementById('downloadProgressBar').style.width = '0%'; + document.getElementById('downloadProgressBar').classList.remove('animate-pulse'); this.hideModal("autoTagModal"); this.showModal("autoTagDownloadModal"); @@ -3699,18 +3700,21 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; this.autoTagState.downloadEventSource.onmessage = (event) => { const data = JSON.parse(event.data); - if (data.status === 'downloading') { - const percent = Math.round(data.progress * 100); - document.getElementById('downloadStatus').textContent = data.message || 'Downloading...'; - document.getElementById('downloadProgressPercent').textContent = `${percent}%`; - document.getElementById('downloadProgressBar').style.width = `${percent}%`; - } else if (data.status === 'complete') { + // The server sends `type` (progress / complete / error); the download is an + // Activity job with no byte progress, so the bar is indeterminate (#328). + if (data.type === 'progress') { + document.getElementById('downloadStatus').textContent = 'Downloading\u2026'; + document.getElementById('downloadProgressPercent').textContent = ''; + const bar = document.getElementById('downloadProgressBar'); + bar.style.width = '100%'; + bar.classList.add('animate-pulse'); + } else if (data.type === 'complete') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal"); this.checkAutoTagModels(); this.showNotification('Model downloaded successfully!', 'success'); - } else if (data.status === 'error') { + } else if (data.type === 'error') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal"); diff --git a/web/js/gallery.js b/web/js/gallery.js index 97834b0..daf45ec 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -3421,6 +3421,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; document.getElementById('downloadStatus').textContent = 'Preparing...'; document.getElementById('downloadProgressPercent').textContent = '0%'; document.getElementById('downloadProgressBar').style.width = '0%'; + document.getElementById('downloadProgressBar').classList.remove('animate-pulse'); this.hideModal("autoTagModal"); this.showModal("autoTagDownloadModal"); @@ -3431,18 +3432,21 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; this.autoTagState.downloadEventSource.onmessage = (event) => { const data = JSON.parse(event.data); - if (data.status === 'downloading') { - const percent = Math.round(data.progress * 100); - document.getElementById('downloadStatus').textContent = data.message || 'Downloading...'; - document.getElementById('downloadProgressPercent').textContent = `${percent}%`; - document.getElementById('downloadProgressBar').style.width = `${percent}%`; - } else if (data.status === 'complete') { + // The server sends `type` (progress / complete / error); the download is an + // Activity job with no byte progress, so the bar is indeterminate (#328). + if (data.type === 'progress') { + document.getElementById('downloadStatus').textContent = 'Downloading\u2026'; + document.getElementById('downloadProgressPercent').textContent = ''; + const bar = document.getElementById('downloadProgressBar'); + bar.style.width = '100%'; + bar.classList.add('animate-pulse'); + } else if (data.type === 'complete') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal"); this.checkAutoTagModels(); this.showNotification('Model downloaded successfully!', 'success'); - } else if (data.status === 'error') { + } else if (data.type === 'error') { this.autoTagState.downloadEventSource.close(); this.hideModal("autoTagDownloadModal"); this.showModal("autoTagModal");