From 41027aa1b18f2eb818e32f0aeee9bc74e5b1325c Mon Sep 17 00:00:00 2001 From: Vito Date: Wed, 28 Oct 2026 19:58:05 +0000 Subject: [PATCH] fix(autotag): report a missing backend package instead of offering the model as ready (#426) Model status now says, per backend, which Python packages cannot be imported (importlib.util.find_spec, nothing imported): llama-cpp-python for GGUF, torch/transformers/accelerate/bitsandbytes for HuggingFace, onnxruntime for WD14, with the pip install command and a note (the CUDA CMAKE_ARGS for llama-cpp-python, onnxruntime-gpu for WD14). Nothing is installed. load_model refuses such a model before the loader with MissingBackendPackages (one WARNING line, no traceback); a ModuleNotFoundError from a loader becomes the same refusal. The suggest route, the legacy single and batch routes and both tagging jobs answer 409 model-unavailable with that sentence. System > Auto-tagging reads "Downloaded. Needs llama-cpp-python: pip install llama-cpp-python. ..." instead of "Downloaded. Not loaded.", Review explains instead of reviewing, and the bulk tagging dialog keeps Start disabled with the sentence. "downloaded" keeps its files-only meaning, so the legacy pages are unchanged. Closes #424 --- py/api/annotation_suggest.py | 17 ++ py/api/annotations.py | 14 +- py/api/autotag_routes.py | 24 +- py/api/image_tag_jobs.py | 21 +- py/autotag.py | 128 ++++++++- tests/browser/server.py | 5 + .../browser/workspace-system-tagging.spec.mjs | 26 ++ tests/browser/workspace-tagging-bulk.spec.mjs | 20 ++ tests/browser/workspace-tagging.spec.mjs | 12 + tests/browser/workspace_tagging_fixture.py | 25 +- .../js/workspace-images-tagging-bulk.test.js | 35 +++ tests/js/workspace-system-tagging.test.js | 54 ++++ tests/test_api_annotation_suggest.py | 45 +++- tests/test_api_autotag.py | 24 ++ tests/test_autotag_batch.py | 43 +++ tests/test_autotag_engine.py | 6 +- tests/test_autotag_legacy.py | 33 +++ tests/test_autotag_packages.py | 248 ++++++++++++++++++ tests/test_image_tag_jobs.py | 13 + web/js/workspace/image-tagging-state.mjs | 39 ++- web/js/workspace/system-tagging.mjs | 23 +- 21 files changed, 815 insertions(+), 40 deletions(-) create mode 100644 tests/test_autotag_packages.py diff --git a/py/api/annotation_suggest.py b/py/api/annotation_suggest.py index befef53..9411640 100644 --- a/py/api/annotation_suggest.py +++ b/py/api/annotation_suggest.py @@ -86,6 +86,23 @@ def unavailable_message(model_name: str) -> str: return f"{model_name} is not downloaded. {UNAVAILABLE_HINT}" +def model_refusal(status: dict, model_type: str, model_name: str): + """Why *model_type* cannot run, from ``get_models_status()``, else None. + + Not downloaded first (``unavailable_message``), then a backend package + that cannot be imported (#424). Both are a 409 ``model-unavailable``. + """ + entry = status.get(model_type) or {} + if not entry.get("downloaded"): + return unavailable_message(model_name) + missing = entry.get("missing_packages") or [] + if missing: + from ..autotag import missing_packages_message + + return missing_packages_message(model_type, [str(name) for name in missing]) + return None + + def run_suggestions(service, capture_path: str, model_type: str, options: dict): """Blocking: load (or swap to) *model_type*, then suggest for the capture. diff --git a/py/api/annotations.py b/py/api/annotations.py index d9a27a0..7f3ab94 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -21,9 +21,9 @@ from aiohttp import web from .annotation_suggest import ( SUGGEST_FIELDS, OptionError, + model_refusal, run_suggestions, suggestion_rows, - unavailable_message, validate_options, ) from .autotag_routes import model_busy_response @@ -595,7 +595,7 @@ class AnnotationRoutesMixin: it, including a legacy batch), never downloads (``model-unavailable``) and never writes. """ - from ..autotag import MODELS, get_autotag_service + from ..autotag import MODELS, MissingBackendPackages, get_autotag_service from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig data, error = await self._annotation_body(request, SUGGEST_FIELDS) @@ -633,9 +633,10 @@ class AnnotationRoutesMixin: return error service = get_autotag_service() status = await self._run_in_executor(service.get_models_status) - if not status.get(model_type, {}).get("downloaded"): - name = MODELS.get(model_type, {}).get("name", model_type) - return _coded("model-unavailable", unavailable_message(name)) + name = MODELS.get(model_type, {}).get("name", model_type) + refusal = model_refusal(status, model_type, name) + if refusal is not None: + return _coded("model-unavailable", refusal) owner = model_gate.new_owner("image-annotations") if not model_gate.try_acquire(owner): return model_busy_response() @@ -654,6 +655,9 @@ class AnnotationRoutesMixin: options, on_skipped=functools.partial(store.release_use, token), ) + except MissingBackendPackages as exc: + # A package gone since the status read (load_model logged it). + return _coded("model-unavailable", str(exc)) except Exception: self.logger.exception("Tag suggestion failed") return _error( diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 7314f0f..edc8ca8 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -102,6 +102,17 @@ def size_label(size: str) -> str: return f"{match.group(1)} {match.group(2)}" if match else str(size or "") +def package_entry(flags: dict) -> dict: + """The public package fields of one status entry (#424): the pip names + that cannot be imported, the install command and note (strings or null).""" + missing = flags.get("missing_packages") or [] + text = { + key: flags[key] if isinstance(flags.get(key), str) else None + for key in ("install_command", "install_note") + } + return {"missing_packages": [str(name) for name in missing], **text} + + 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).""" @@ -122,6 +133,7 @@ def public_models(status: dict, loaded_type: Optional[str]) -> dict: for key in ("model_exists", "mmproj_exists"): if key in flags: entry[key] = bool(flags[key]) + entry.update(package_entry(flags)) models[model_type] = entry return models @@ -750,6 +762,7 @@ class AutotagRoutesMixin: "downloaded": bool( (status.get(model_type) or {}).get("downloaded") ), + **package_entry(status.get(model_type) or {}), }, "existing": existing, "review_required": existing == "replace_after_review", @@ -1073,8 +1086,8 @@ class AutotagRoutesMixin: return self.db.current_image_tags({locator: fingerprint}).get(locator, []) async def _suggest_single(self, target, model_type, options, base): - from ..autotag import MODELS, get_autotag_service - from .annotation_suggest import unavailable_message + from ..autotag import MODELS, MissingBackendPackages, get_autotag_service + from .annotation_suggest import model_refusal from .annotations import _pinned_suggestions, _settled, _submit try: @@ -1084,8 +1097,8 @@ class AutotagRoutesMixin: service = get_autotag_service() status = await self._run_in_executor(service.get_models_status) - if not (status.get(model_type) or {}).get("downloaded"): - message = unavailable_message(MODELS[model_type]["name"]) + message = model_refusal(status, model_type, MODELS[model_type]["name"]) + if message is not None: return self._coded_error(409, "model-unavailable", message) owner = model_gate.new_owner("legacy-single") if not model_gate.try_acquire(owner): @@ -1131,6 +1144,9 @@ class AutotagRoutesMixin: options, on_skipped=functools.partial(store.release_use, capture.token), ) + except MissingBackendPackages as exc: + # A package gone since the status read (load_model logged it). + return self._coded_error(409, "model-unavailable", str(exc)) finally: store.release(capture.token) finally: diff --git a/py/api/image_tag_jobs.py b/py/api/image_tag_jobs.py index 61c7942..5d1f954 100644 --- a/py/api/image_tag_jobs.py +++ b/py/api/image_tag_jobs.py @@ -42,7 +42,7 @@ from .. import model_gate from .admin import _plural, _time_left from .annotation_suggest import ( OptionError, - unavailable_message, + model_refusal, validate_options, ) from .annotations import ( @@ -512,10 +512,8 @@ class ImageTagJobsMixin: ) ) label = MODELS[model_type]["label"] - if not (status.get(model_type) or {}).get("downloaded"): - message = ( - f"{label} is not downloaded. Download it in System › Auto-tagging." - ) + message = model_refusal(status, model_type, label) + if message is not None: return None, _refusal(409, "model-unavailable", message) if backup_on and short is not None: return None, _refusal(507, "insufficient-space", short_space_message(short)) @@ -792,13 +790,13 @@ class ImageTagJobsMixin: async def _prepare_tagging_model(self, job, spec: _Spec, owner): """The loaded service, or None (logged) when the job cannot run.""" - from ..autotag import MODELS, get_autotag_service + from ..autotag import MODELS, MissingBackendPackages, get_autotag_service service = get_autotag_service() status = await self._run_in_executor(service.get_models_status) - if not status.get(spec.model_type, {}).get("downloaded"): - name = MODELS.get(spec.model_type, {}).get("name", spec.model_type) - message = unavailable_message(name) + name = MODELS.get(spec.model_type, {}).get("name", spec.model_type) + message = model_refusal(status, spec.model_type, name) + if message is not None: job.update(status_line=message) job.log_line(f"! {message}") return None @@ -810,6 +808,11 @@ class ImageTagJobsMixin: spec.model_type, spec.options["use_gpu"], ) + except MissingBackendPackages as exc: + # A package gone since the status read (load_model logged it). + job.update(status_line=str(exc)) + job.log_line(f"! {exc}") + return None except Exception: self.logger.exception("The auto-tag job could not load its model") job.update(status_line=INFERENCE_FAILED) diff --git a/py/autotag.py b/py/autotag.py index d002f62..727af79 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -7,6 +7,7 @@ Refactored from standalone_tagger.py for integration with PromptManager API. import errno import gc +import importlib.util import inspect import os import sys @@ -155,6 +156,91 @@ MODELS = { }, } +# The Python modules each backend imports when it loads, with their pip names +# (#424). None of them is a PromptManager requirement: they are optional, and +# llama-cpp-python needs a build for the user's GPU. Nothing is ever installed. +BACKEND_PACKAGES: Dict[str, Tuple[Tuple[str, str], ...]] = { + "gguf": (("llama_cpp", "llama-cpp-python"),), + # The default 8-bit load needs bitsandbytes, and device_map needs accelerate. + "hf": ( + ("torch", "torch"), + ("transformers", "transformers"), + ("accelerate", "accelerate"), + ("bitsandbytes", "bitsandbytes"), + ), + "wd14-swinv2": (("onnxruntime", "onnxruntime"),), + "wd14-vit": (("onnxruntime", "onnxruntime"),), +} +_INSTALL_NOTES = { + "gguf": ( + 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python' + ), + "wd14-swinv2": "Use onnxruntime-gpu for CUDA.", + "wd14-vit": "Use onnxruntime-gpu for CUDA.", +} +_PIP_NAMES = { + module: pip for packages in BACKEND_PACKAGES.values() for module, pip in packages +} + + +def _importable(module: str) -> bool: + """Whether *module* (a top-level name) can be imported; imports nothing.""" + if sys.modules.get(module) is not None: + return True + try: + return importlib.util.find_spec(module) is not None + except (ImportError, ValueError, OSError): + return False + + +def missing_packages(model_type: str) -> List[str]: + """The pip names of the packages *model_type* needs and cannot import.""" + return [ + pip + for module, pip in BACKEND_PACKAGES.get(model_type, ()) + if not _importable(module) + ] + + +def package_fields(model_type: str, missing: List[str]) -> Dict[str, Any]: + """The status fields of *missing*: pip names, install command and note.""" + missing = list(missing) + return { + "missing_packages": missing, + "install_command": "pip install " + " ".join(missing) if missing else None, + "install_note": _INSTALL_NOTES.get(model_type) if missing else None, + } + + +def _joined(names: List[str]) -> str: + if len(names) <= 1: + return "".join(names) + return ", ".join(names[:-1]) + " and " + names[-1] + + +def missing_packages_message(model_type: str, missing: List[str]) -> str: + """The one sentence every refusal of a model with missing packages uses.""" + label = MODELS.get(model_type, {}).get("label", model_type) + fields = package_fields(model_type, missing) + pronoun = "it" if len(missing) == 1 else "them" + message = ( + f"{label} needs {_joined(missing)}. Install {pronoun} into ComfyUI's " + f"Python: {fields['install_command']}." + ) + if fields["install_note"]: + message += " " + fields["install_note"] + return message + + +class MissingBackendPackages(RuntimeError): + """A model whose backend packages cannot be imported (#424).""" + + def __init__(self, model_type: str, missing: List[str]): + self.model_type = model_type + self.missing = list(missing) + super().__init__(missing_packages_message(model_type, self.missing)) + + # Default prompts DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE @@ -167,6 +253,10 @@ SYSTEM_PROMPT = ( ) +def _none(_model_type: str) -> List[str]: + return [] + + def _load_rgb(image_path: str) -> Image.Image: """Decode the image as RGB and close the file before inference. @@ -203,6 +293,7 @@ class AutoTagService: self, models_dir: Optional[Path] = None, model_loader: Optional[ModelLoader] = None, + package_probe: Optional[Callable[[str], List[str]]] = None, ): """Initialize the AutoTag service. @@ -213,9 +304,16 @@ class AutoTagService: ``(kind, backend)`` tuple used by ``generate_tags``. Defaults to ``_load_model``, which imports the ML backends lazily. Tests inject a fake here. + package_probe: Callable ``(model_type)`` returning the pip names + of the backend packages that cannot be imported. + Defaults to ``missing_packages`` with the default + loader; an injected loader brings its own backend. """ self.logger = get_logger("autotag.service") self._model_loader: ModelLoader = model_loader or self._load_model + if package_probe is None: + package_probe = missing_packages if model_loader is None else _none + self._package_probe = package_probe # Determine models directory if models_dir: @@ -320,6 +418,9 @@ class AutoTagService: # Get the actual path (local or cache) model_status["model_path"] = str(self._get_hf_model_path()) + model_status.update( + package_fields(model_type, self._package_probe(model_type)) + ) status[model_type] = model_status return status @@ -455,13 +556,9 @@ class AutoTagService: def _load_wd14_tagger(self, model_type: str, use_gpu: bool = True): """Load WD14 ONNX model and tag list.""" - try: - import onnxruntime as ort - except ImportError: - raise RuntimeError( - "onnxruntime is required for WD14 models. " - "Install with: pip install onnxruntime (or onnxruntime-gpu for CUDA)" - ) + # A missing onnxruntime is refused by load_model before this runs + # (#424); a ModuleNotFoundError here becomes the same refusal. + import onnxruntime as ort import csv config = MODELS[model_type] @@ -694,6 +791,7 @@ class AutoTagService: Raises: ValueError: If model_type is invalid RuntimeError: If model not downloaded or loading fails + MissingBackendPackages: If its backend packages cannot be imported """ if model_type not in MODELS: raise ValueError(f"Invalid model type: {model_type}") @@ -705,6 +803,11 @@ class AutoTagService: status = self.get_models_status() if not status[model_type]["downloaded"]: raise RuntimeError(f"Model {model_type} not downloaded") + missing = status[model_type].get("missing_packages") or [] + if missing: + refusal = MissingBackendPackages(model_type, missing) + self.logger.warning(str(refusal)) + raise refusal try: self._tagger = self._model_loader(model_type, use_gpu) @@ -712,6 +815,17 @@ class AutoTagService: self.logger.info(f"Model {model_type} loaded successfully") return True + except ModuleNotFoundError as e: + # A package find_spec found but that does not import (a broken + # install), or a dependency of one: the same refusal, no traceback. + self._tagger = None + self._current_model_type = None + module = (e.name or "").split(".")[0] or str(e) + refusal = MissingBackendPackages( + model_type, [_PIP_NAMES.get(module, module)] + ) + self.logger.warning(str(refusal)) + raise refusal from None except Exception as e: self.logger.error(f"Failed to load model {model_type}: {e}") self._tagger = None diff --git a/tests/browser/server.py b/tests/browser/server.py index dab1801..ef7e24c 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -75,6 +75,8 @@ class TestModel: loaded = None # Models the Tag image dialog workflows mark as not downloaded. unavailable: set = set() + # Pip names per model whose backend packages "cannot be imported" (#424). + missing_packages: dict = {} # Scored names the concept comparisons answer instead of the fixed list # (``/__test__/tagging/suggestions``); None keeps the fixed list. suggestions: list | None = None @@ -82,10 +84,13 @@ class TestModel: models_config: dict = {} def get_models_status(self) -> dict: + from py.autotag import package_fields + return { name: { "downloaded": name not in self.unavailable, "loaded": self.loaded == name, + **package_fields(name, self.missing_packages.get(name, [])), } for name in self.models_config } diff --git a/tests/browser/workspace-system-tagging.spec.mjs b/tests/browser/workspace-system-tagging.spec.mjs index 0d9fa58..5a5a640 100644 --- a/tests/browser/workspace-system-tagging.spec.mjs +++ b/tests/browser/workspace-system-tagging.spec.mjs @@ -123,6 +123,32 @@ test('model options are the four real engines with size and the description show await expect(actionOf(page, 'Unload')).toBeDisabled(); }); +const NEEDS_LLAMA = 'GGUF JoyCaption needs llama-cpp-python. Install it into ComfyUI\'s Python: pip install llama-cpp-python. ' + + 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python'; + +test('a downloaded model whose backend package is missing says so and is not offered as ready (#424)', async ({ page, request }) => { + await expect(await request.post('/__test__/tagging/models', { data: { unavailable: ['wd14-vit'], + missing_packages: { gguf: ['llama-cpp-python'], 'wd14-vit': ['onnxruntime'] } } })).toBeOK(); + await openTagging(page); + const line = modelRow(page).locator('p').first(); + await page.locator('#tag-model').selectOption('gguf'); + await expect(line).toHaveText('Downloaded. Needs llama-cpp-python: pip install llama-cpp-python. ' + + 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python'); + // Files and packages are separate: Download (again) stays offered. + await expect(actionOf(page, 'Download again')).not.toHaveAttribute('aria-disabled', 'true'); + // Review explains instead of reviewing; nothing is previewed. + await expect(reviewButton(page)).toHaveAttribute('aria-disabled', 'true'); + await reviewButton(page).dispatchEvent('click'); + await expect(toast(page)).toHaveText(NEEDS_LLAMA); + expect(page.modelCalls.filter(call => call.endsWith('/autotag/batch/preview'))).toEqual([]); + // Not downloaded reads differently and still names what it will need. + await page.locator('#tag-model').selectOption('wd14-vit'); + await expect(line).toHaveText('Not downloaded. Download size ~400 MB. Also needs onnxruntime.'); + await page.locator('#tag-model').selectOption('wd14-swinv2'); + await expect(line).toHaveText('Downloaded. Not loaded.'); + await expect(reviewButton(page)).not.toHaveAttribute('aria-disabled', 'true'); +}); + test('entering the section makes no model call besides status', async ({ page }) => { await openTagging(page); await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Integrations', exact: true }).click(); diff --git a/tests/browser/workspace-tagging-bulk.spec.mjs b/tests/browser/workspace-tagging-bulk.spec.mjs index 01bbda1..6f3d114 100644 --- a/tests/browser/workspace-tagging-bulk.spec.mjs +++ b/tests/browser/workspace-tagging-bulk.spec.mjs @@ -186,6 +186,26 @@ test('a model that is not downloaded keeps Start disabled and says why', async ( expect(posted).toBe(0); }); +const NEEDS_LLAMA = 'GGUF JoyCaption needs llama-cpp-python. Install it into ComfyUI\'s Python: pip install llama-cpp-python. ' + + 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python'; + +test('a model missing its backend package keeps Start disabled and names the install command (#424)', async ({ page, request }) => { + await expect(await request.post('/__test__/tagging/models', { data: { unavailable: [], + missing_packages: { gguf: ['llama-cpp-python'] } } })).toBeOK(); + await select(page, FIRST); + await barButton(page, 'Auto-tag selected').click(); + const status = modal(page).locator('#tag-generation-status'); + const start = actionButtons(page).last(); + await modal(page).getByRole('combobox', { name: 'Tagging model', exact: true }).selectOption('gguf'); + await expect(status).toHaveText(NEEDS_LLAMA); + await expect(start).toHaveAttribute('aria-disabled', 'true'); + await modal(page).getByRole('combobox', { name: 'Tagging model', exact: true }).selectOption('wd14-vit'); + await expect(status).toHaveText(''); + await expect(start).not.toHaveAttribute('aria-disabled', 'true'); + await actionButtons(page).first().click(); + await expect(modal(page)).toHaveCount(0); +}); + test('Suggestions field prefills from the setting and caps chips at the entered count', async ({ page, request }) => { await expect(await request.post('/__test__/tagging/defaults', { data: { general: 0.5, character: 0.9, limit: 7 } })).toBeOK(); await select(page, FIRST); diff --git a/tests/browser/workspace-tagging.spec.mjs b/tests/browser/workspace-tagging.spec.mjs index eaefe7b..0c4cf8e 100644 --- a/tests/browser/workspace-tagging.spec.mjs +++ b/tests/browser/workspace-tagging.spec.mjs @@ -260,6 +260,18 @@ test('unavailable model shows inline status and never downloads', async ({ page, await expect(dialog(page).locator('.dialog-actions .ca-primary')).toBeFocused(); }); +test('a missing backend package is the server\'s install sentence, no traceback page (#424)', async ({ page, request }) => { + await expect(await request.post('/__test__/tagging/models', { data: { unavailable: [], + missing_packages: { 'wd14-swinv2': ['onnxruntime'] } } })).toBeOK(); + await openTag(page); + const answer = await generated(page); + expect(answer.status()).toBe(409); + expect(await answer.json()).toMatchObject({ success: false, code: 'model-unavailable' }); + await expect(status(page)).toHaveText('WD14 SwinV2 needs onnxruntime. Install it into ComfyUI\'s Python: ' + + 'pip install onnxruntime. Use onnxruntime-gpu for CUDA.'); + await expect(chips(page)).toHaveCount(0); +}); + test('video tag is manual only', async ({ page, request }) => { await openTag(page, VIDEO); const preview = dialog(page).locator('video.tag-image-preview'); diff --git a/tests/browser/workspace_tagging_fixture.py b/tests/browser/workspace_tagging_fixture.py index e125076..6bc62a8 100644 --- a/tests/browser/workspace_tagging_fixture.py +++ b/tests/browser/workspace_tagging_fixture.py @@ -41,6 +41,18 @@ FILES = ( ("loose/unlinked.png", "unlinked"), ) MODELS = ("wd14-swinv2", "wd14-vit", "gguf", "hf") + + +def _valid_missing(value) -> bool: + """``{"": ["", ...]}`` with known models and string names.""" + return isinstance(value, dict) and all( + name in MODELS + and isinstance(names, list) + and all(isinstance(pip, str) and pip for pip in names) + for name, names in value.items() + ) + + # Unlinked copies enough to virtualize the Images grid (it renders all up to 120). MAX_EXTRA_IMAGES = 160 _TABLES = ( @@ -217,6 +229,7 @@ def register_tagging_commands( if not 0 <= extra <= MAX_EXTRA_IMAGES: raise web.HTTPBadRequest() model.unavailable = set() + model.missing_packages = {} model.suggestions = None model.loaded = None restore_download() @@ -235,17 +248,25 @@ def register_tagging_commands( @routes.post("/__test__/tagging/models") async def models(request: web.Request) -> web.Response: - """Mark fake models as not downloaded (``{"unavailable": [...]}``).""" + """Mark fake models as not downloaded (``{"unavailable": [...]}``) and, + optionally, as missing backend packages (``"missing_packages": + {"": ["", ...]}``, #424).""" data = await request.json() if request.can_read_body else None if ( request.query or not isinstance(data, dict) - or set(data) != {"unavailable"} + or "unavailable" not in data + or set(data) - {"unavailable", "missing_packages"} or not isinstance(data["unavailable"], list) or any(name not in MODELS for name in data["unavailable"]) + or not _valid_missing(data.get("missing_packages", {})) ): raise web.HTTPBadRequest() model.unavailable = set(data["unavailable"]) + model.missing_packages = { + name: list(names) + for name, names in data.get("missing_packages", {}).items() + } return web.json_response({"success": True}) @routes.post("/__test__/tagging/loaded") diff --git a/tests/js/workspace-images-tagging-bulk.test.js b/tests/js/workspace-images-tagging-bulk.test.js index ac6dbef..9ffec9a 100644 --- a/tests/js/workspace-images-tagging-bulk.test.js +++ b/tests/js/workspace-images-tagging-bulk.test.js @@ -135,6 +135,41 @@ test('a model that is not downloaded is named with the server sentence', () => { assert.throws(() => modelUnavailable({}, 'nope'), TypeError); }); +test('a downloaded model whose backend package is missing names the package and the command (#424)', () => { + const { modelUnavailable, needsPackages } = state; + const gguf = { downloaded: true, name: 'GGUF (Recommended)', label: 'GGUF JoyCaption', missing_packages: ['llama-cpp-python'], + install_command: 'pip install llama-cpp-python', + install_note: 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python' }; + const sentence = 'GGUF JoyCaption needs llama-cpp-python. Install it into ComfyUI\'s Python: pip install llama-cpp-python. ' + + 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python'; + assert.equal(modelUnavailable({ gguf }, 'gguf'), sentence); + // Not downloaded is named first, in its own sentence. + assert.equal(modelUnavailable({ gguf: { ...gguf, downloaded: false } }, 'gguf'), + 'GGUF (Recommended) is not downloaded. Download it in System › Auto-tagging.'); + assert.equal(modelUnavailable({ gguf: { ...gguf, missing_packages: [] } }, 'gguf'), ''); + assert.equal(needsPackages('HuggingFace JoyCaption', { missing_packages: ['accelerate', 'bitsandbytes'] }), + 'HuggingFace JoyCaption needs accelerate and bitsandbytes. Install them into ComfyUI\'s Python: pip install accelerate bitsandbytes.'); + assert.equal(needsPackages('X', { missing_packages: ['a', 'b', 'c'], install_command: 'pip install a b c' }), + 'X needs a, b and c. Install them into ComfyUI\'s Python: pip install a b c.'); + assert.equal(needsPackages('X', null), ''); + assert.equal(needsPackages('X', { missing_packages: [1, ''] }), ''); + // Without the server's label the option's own name is used. + assert.ok(modelUnavailable({ 'wd14-vit': { downloaded: true, missing_packages: ['onnxruntime'] } }, 'wd14-vit') + .startsWith('WD14 ViT needs onnxruntime.')); +}); + +test('the bulk dialog keeps Start disabled while the chosen model misses a package (#424)', async () => { + const { createTaggingOptionsController } = state; + const models = { 'wd14-vit': { downloaded: true, label: 'WD14 ViT', missing_packages: ['onnxruntime'], + install_command: 'pip install onnxruntime' } }; + const options = createTaggingOptionsController({ api: { get: async () => ({ success: true, models, model_type: 'wd14-vit' }) } }); + await options.open(); + assert.equal(options.state.model, 'wd14-vit'); + assert.equal(options.state.canStart, false); + assert.equal(options.state.unavailable, 'WD14 ViT needs onnxruntime. Install it into ComfyUI\'s Python: pip install onnxruntime.'); + options.dispose(); +}); + test('the bulk options start from the server defaults and wait for them before Start', async () => { const { createTaggingOptionsController } = state; let answer; diff --git a/tests/js/workspace-system-tagging.test.js b/tests/js/workspace-system-tagging.test.js index 95c03b4..317b914 100644 --- a/tests/js/workspace-system-tagging.test.js +++ b/tests/js/workspace-system-tagging.test.js @@ -482,6 +482,60 @@ test('Tag output images › Review previews the drafts, else says why; a refused } finally { scope.dispose(); dom.restore(); } }); +// #424: GGUF on disk without llama-cpp-python. +const NEEDS_LLAMA = { missing_packages: ['llama-cpp-python'], install_command: 'pip install llama-cpp-python', + install_note: 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python' }; +const NEEDS_SENTENCE = 'GGUF JoyCaption needs llama-cpp-python. Install it into ComfyUI\'s Python: pip install llama-cpp-python. ' + + 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python'; + +test('a model on disk whose backend package is missing never reads as ready (#424)', async () => { + const { modelDescription, readStatus, batchCloseNote, STATUS_COPY } = await tagging; + const status = readStatus({ ...MODELS, models: { ...MODELS.models, + gguf: entry('GGUF JoyCaption', '~4.5 GB', { ...NEEDS_LLAMA }), + 'wd14-vit': entry('WD14 ViT', '~400 MB', { downloaded: false, missing_packages: ['onnxruntime'], + install_command: 'pip install onnxruntime', install_note: 'Use onnxruntime-gpu for CUDA.' }), + hf: entry('HuggingFace JoyCaption', '~16 GB', { downloaded: false, model_exists: undefined, + missing_packages: ['accelerate', 'bitsandbytes'], install_command: 'pip install accelerate bitsandbytes' }) } }); + const ready = modelDescription(status, 'gguf'); + assert.notEqual(ready, STATUS_COPY.downloaded); + assert.equal(ready, 'Downloaded. Needs llama-cpp-python: pip install llama-cpp-python. ' + + 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python'); + // Not downloaded reads differently, and still says what it will need. + assert.equal(modelDescription(status, 'wd14-vit'), 'Not downloaded. Download size ~400 MB. Also needs onnxruntime.'); + assert.equal(modelDescription(status, 'hf'), 'Not downloaded. Download size ~16 GB. Also needs accelerate and bitsandbytes.'); + const partial = readStatus({ ...MODELS, models: { ...MODELS.models, gguf: entry('GGUF JoyCaption', '~4.5 GB', + { downloaded: false, model_exists: true, mmproj_exists: false, ...NEEDS_LLAMA }) } }); + assert.equal(modelDescription(partial, 'gguf'), 'Partly downloaded. Download again to finish. Also needs llama-cpp-python.'); + // Untrusted shapes are ignored; a missing command is built from the names. + const odd = readStatus({ ...MODELS, models: { ...MODELS.models, gguf: entry('GGUF JoyCaption', '~4.5 GB', + { missing_packages: ['llama-cpp-python', 3, ''], install_command: 7, install_note: {} }) } }); + assert.equal(modelDescription(odd, 'gguf'), 'Downloaded. Needs llama-cpp-python: pip install llama-cpp-python.'); + const junk = readStatus({ ...MODELS, models: { ...MODELS.models, gguf: entry('GGUF JoyCaption', '~4.5 GB', + { missing_packages: 'llama-cpp-python' }) } }); + assert.equal(modelDescription(junk, 'gguf'), STATUS_COPY.downloaded); + // The batch review says why it cannot start, in the server's sentence. + assert.equal(batchCloseNote({ ...PLAN, model: { type: 'gguf', label: 'GGUF JoyCaption', downloaded: true, ...NEEDS_LLAMA } }), + NEEDS_SENTENCE); + assert.equal(batchCloseNote({ ...PLAN, model: { ...PLAN.model, downloaded: false, ...NEEDS_LLAMA } }), + 'WD14 SwinV2 is not downloaded. Download it first.'); +}); + +test('Review explains a missing backend package instead of reviewing (#424)', async () => { + const models = { ...MODELS, models: { ...MODELS.models, gguf: entry('GGUF JoyCaption', '~4.5 GB', { ...NEEDS_LLAMA }) } }; + const { dom, scope, button, byId, calls } = await mounted({ models }); + try { + byId('tag-model').value = 'gguf'; await byId('tag-model').dispatch('change'); + const review = button('Review'); + assert.equal(review.getAttribute('aria-disabled'), 'true'); + assert.equal(review.dataset.reason, NEEDS_SENTENCE); + await review.dispatch('click'); + assert.equal(calls.toasts.at(-1), NEEDS_SENTENCE); + assert.deepEqual(calls.post, []); + // Download stays offered: files and packages are separate. + assert.equal(button('Download again').getAttribute('aria-disabled'), null); + } finally { scope.dispose(); dom.restore(); } +}); + test('settingRefusal marks only an Auto-tagging save\'s plain 400 as invalid-setting (CUJ 2a)', async () => { const { settingRefusal } = await import('../../web/js/workspace/system.mjs'); const refusal = (extra = {}) => Object.assign(new Error('General-tag confidence must be between 0 and 1.'), diff --git a/tests/test_api_annotation_suggest.py b/tests/test_api_annotation_suggest.py index de097d6..7d1b353 100644 --- a/tests/test_api_annotation_suggest.py +++ b/tests/test_api_annotation_suggest.py @@ -29,6 +29,7 @@ except ImportError: model_gate = None from py import config as config_module +from py.autotag import MissingBackendPackages, package_fields AutotagConfig = getattr(config_module, "AutotagConfig", None) @@ -37,6 +38,11 @@ WAIT = 30.0 WD14 = {"general_threshold": 0.35, "character_threshold": 0.85} LLM = {"instructions": "Generate tags."} UNAVAILABLE = "Download it in System › Auto-tagging." +NEEDS_LLAMA = ( + "GGUF JoyCaption needs llama-cpp-python. Install it into ComfyUI's Python: " + 'pip install llama-cpp-python. For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" ' + "pip install llama-cpp-python" +) SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [ (f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150) @@ -59,6 +65,8 @@ class FakeTagger: "hf": {}, } self.downloaded = {name: True for name in self.models_config} + # pip names of the backend packages that cannot be imported (#424) + self.missing = {name: [] for name in self.models_config} self.loaded = None self.load_calls = [] self.unload_calls = 0 @@ -71,7 +79,11 @@ class FakeTagger: def get_models_status(self): return { - name: {"downloaded": self.downloaded[name], "model_path": None} + name: { + "downloaded": self.downloaded[name], + "model_path": None, + **package_fields(name, self.missing[name]), + } for name in self.models_config } @@ -315,6 +327,37 @@ class TestModelAvailability(SuggestTestCase): self.assertEqual(self.tagger.calls, []) self.assertIsNone(model_gate.holder()) + async def test_missing_backend_package_is_refused_before_the_model(self): + observed = await self.observed() + self.tagger.missing["gguf"] = ["llama-cpp-python"] + + status, body = await self.suggest(self.suggest_body(observed, "gguf")) + + self.assertEqual((status, body["code"]), (409, "model-unavailable")) + self.assertEqual(body["success"], False) + self.assertEqual(body["error"], NEEDS_LLAMA) + self.assertEqual(self.tagger.load_calls, []) + self.assertIsNone(model_gate.holder()) + # Not downloaded still reads as before, whatever the packages. + self.tagger.downloaded["gguf"] = False + status, body = await self.suggest(self.suggest_body(observed, "gguf")) + self.assertIn("is not downloaded", body["error"]) + + async def test_a_package_lost_after_the_check_is_the_same_409_without_a_traceback( + self, + ): + observed = await self.observed() + self.tagger.load_error = MissingBackendPackages("gguf", ["llama-cpp-python"]) + + with self.assertNoLogs(self.api.logger, level="ERROR"): + status, body = await self.suggest(self.suggest_body(observed, "gguf")) + + self.assertEqual((status, body["code"]), (409, "model-unavailable")) + self.assertEqual(body["error"], NEEDS_LLAMA) + self.assertIsNone(model_gate.holder()) + entry = self.api._captures._entries[observed["capture"]["token"]] + self.assertEqual(entry.uses, 0) + async def test_load_failure_is_a_generic_error_and_frees_everything(self): observed = await self.observed() self.tagger.load_error = RuntimeError(f"cannot load {self.output_dir}") diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index cfe3f47..4b009ba 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -73,6 +73,7 @@ class FakeAutotagService: self.status_calls = 0 self.status_threads = [] self.status_paths = {} + self.status_extra = {} self.suggest_calls = 0 def get_models_status(self): @@ -82,6 +83,7 @@ class FakeAutotagService: name: { "downloaded": self.downloaded, "model_path": self.status_paths.get(name), + **self.status_extra.get(name, {}), } for name in self.models_config } @@ -1111,6 +1113,24 @@ class TestModelStatus(AutotagAPITestCase): leaked = [text for text in _strings(data) if _absolute(text)] self.assertEqual(leaked, []) + async def test_models_status_names_missing_backend_packages(self): + fields = { + "missing_packages": ["llama-cpp-python"], + "install_command": "pip install llama-cpp-python", + "install_note": "For a CUDA build: ...", + } + self.service.status_extra = {"gguf": fields} + models = (await self._models())["models"] + self.assertTrue(models["gguf"]["downloaded"]) + self.assertEqual( + {key: models["gguf"][key] for key in fields}, + fields, + ) + self.assertEqual( + {key: models["wd14-vit"][key] for key in fields}, + {"missing_packages": [], "install_command": None, "install_note": None}, + ) + async def test_models_status_adds_label_size_label_and_loaded(self): self.service._loaded_type = "wd14-vit" data = await self._models() @@ -1669,6 +1689,10 @@ class TestTagOutputImagesPreview(BatchJobTestCase): "label": "WD14 ViT", "size_label": "~400 MB", "downloaded": True, + # #424: the backend packages it cannot import, none here. + "missing_packages": [], + "install_command": None, + "install_note": None, }, ) self.assertEqual(skip["existing"], "skip") diff --git a/tests/test_autotag_batch.py b/tests/test_autotag_batch.py index c6703ec..4b63d2f 100644 --- a/tests/test_autotag_batch.py +++ b/tests/test_autotag_batch.py @@ -553,6 +553,35 @@ class TestStartRefusals(BatchTestCase): self.assertEqual(self.api._jobs.list(), []) self.assertIsNone(model_gate.holder()) + async def test_a_missing_backend_package_refuses_the_start(self): + self.tagger.missing["gguf"] = ["llama-cpp-python"] + status, data = await self.start_batch(model_type="gguf") + self.assertEqual((status, data["code"]), (409, "model-unavailable")) + self.assertEqual(data["success"], False) + self.assertTrue( + data["error"].startswith("GGUF JoyCaption needs llama-cpp-python."), data + ) + self.assertIn("pip install llama-cpp-python", data["error"]) + self.assertEqual(self.api._jobs.list(), []) + self.assertEqual(self.tagger.load_calls, []) + self.assertIsNone(model_gate.holder()) + + async def test_a_package_lost_after_the_start_fails_the_job_without_a_traceback( + self, + ): + from py.autotag import MissingBackendPackages + + self.image("a.png", (1, 1, 1)) + self.tagger.load_error = MissingBackendPackages("gguf", ["llama-cpp-python"]) + with self.assertNoLogs(self.api.logger, level="ERROR"): + job = await self.run_batch(model_type="gguf") + self.assertEqual(job.state, "failed") + self.assertTrue( + job.status_line.startswith("GGUF JoyCaption needs"), job.status_line + ) + self.assertTrue(any("pip install llama-cpp-python" in line for line in job.log)) + self.assertIsNone(model_gate.holder()) + async def test_model_busy_when_the_gate_is_held(self): owner = model_gate.new_owner("image-annotations") self.assertTrue(model_gate.try_acquire(owner)) @@ -658,6 +687,20 @@ class TestPreview(BatchTestCase): self.assertIsNone(llm["estimate_device"]) self.assertEqual([s for s in _strings(skip) if _absolute(s)], []) + async def test_preview_names_missing_backend_packages(self): + self.image("a.png", (1, 1, 1)) + data = (await self.preview(model_type="gguf"))[1] + self.assertEqual(data["model"]["missing_packages"], []) + self.assertIsNone(data["model"]["install_command"]) + self.tagger.missing["gguf"] = ["llama-cpp-python"] + data = (await self.preview(model_type="gguf"))[1] + self.assertTrue(data["model"]["downloaded"]) + self.assertEqual(data["model"]["missing_packages"], ["llama-cpp-python"]) + self.assertEqual( + data["model"]["install_command"], "pip install llama-cpp-python" + ) + self.assertIn("DGGML_CUDA", data["model"]["install_note"]) + async def test_preview_uses_the_graphics_card_rate_when_wd14_runs_on_cuda(self): self.image("a.png", (1, 1, 1)) self.tagger.wd14_device = lambda: "gpu" diff --git a/tests/test_autotag_engine.py b/tests/test_autotag_engine.py index 7f7faa3..ec3e468 100644 --- a/tests/test_autotag_engine.py +++ b/tests/test_autotag_engine.py @@ -800,8 +800,12 @@ class TestDefaultLoaders(EngineTestCase): transformers, model, processor, inputs = self._fake_transformers() pixel = inputs["pixel_values"] service = AutoTagService(models_dir=self.models_dir) + # The 8-bit load also needs these two (#424); stand-ins satisfy the check. + stubs = { + name: types.ModuleType(name) for name in ("accelerate", "bitsandbytes") + } - with _patch_modules({"torch": torch, "transformers": transformers}): + with _patch_modules({"torch": torch, "transformers": transformers, **stubs}): service.load_model("hf") tags = service.generate_tags(str(self.image_path)) diff --git a/tests/test_autotag_legacy.py b/tests/test_autotag_legacy.py index 126ef67..f8a379a 100644 --- a/tests/test_autotag_legacy.py +++ b/tests/test_autotag_legacy.py @@ -161,6 +161,39 @@ class TestLegacyBatch(LegacyTestCase): self.assertEqual(self.tagger.calls[-1]["prompt"], AutotagConfig.TEMPLATE) self.assertEqual(self.tagger.custom_prompt, "legacy prompt") + async def test_a_missing_backend_package_is_one_error_frame(self): + self.tagger.missing["gguf"] = ["llama-cpp-python"] + _, frames = await self.legacy("model_type=gguf") + self.assertEqual(len(frames), 1) + self.assertEqual(frames[0]["type"], "error") + self.assertTrue( + frames[0]["message"].startswith("GGUF JoyCaption needs llama-cpp-python.") + ) + self.assertEqual(self.tagger.load_calls, []) + + async def test_legacy_single_refuses_a_missing_backend_package(self): + from py.autotag import MissingBackendPackages + + self.image("a.png", (1, 1, 1)) + await self.relist() + self.tagger.missing["gguf"] = ["llama-cpp-python"] + status, data = await self.single( + {"path": "a.png", "root_index": 0, "model_type": "gguf"} + ) + self.assertEqual((status, data["code"]), (409, "model-unavailable")) + self.assertIn("pip install llama-cpp-python", data["error"]) + self.assertEqual(self.tagger.load_calls, []) + # Lost after the check: the same 409, no ERROR record, nothing held. + self.tagger.missing["gguf"] = [] + self.tagger.load_error = MissingBackendPackages("gguf", ["llama-cpp-python"]) + with self.assertNoLogs(self.api.logger, level="ERROR"): + status, data = await self.single( + {"path": "a.png", "root_index": 0, "model_type": "gguf"} + ) + self.assertEqual((status, data["code"]), (409, "model-unavailable")) + self.assertIn("pip install llama-cpp-python", data["error"]) + self.assertIsNone(model_gate.holder()) + async def test_each_refusal_is_one_error_frame(self): self.tagger.downloaded["wd14-vit"] = False _, frames = await self.legacy() diff --git a/tests/test_autotag_packages.py b/tests/test_autotag_packages.py new file mode 100644 index 0000000..fb910ce --- /dev/null +++ b/tests/test_autotag_packages.py @@ -0,0 +1,248 @@ +"""#424: each tagging backend reports the Python packages it cannot import. + +``find_spec`` is patched throughout: no ML backend is installed or imported. +""" + +import os +import sys +import types +import unittest +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +try: + from test_autotag_engine import ( + EngineTestCase, + FakeGgufModel, + _patch_modules, + _touch_model_files, + ) +except ImportError: # pragma: no cover - run as a package + from tests.test_autotag_engine import ( + EngineTestCase, + FakeGgufModel, + _patch_modules, + _touch_model_files, + ) + +from py import autotag # noqa: E402 +from py.autotag import ( # noqa: E402 + BACKEND_PACKAGES, + MODELS, + AutoTagService, + MissingBackendPackages, + missing_packages, + package_fields, +) + +GGUF_CUDA = 'For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python' + + +def _find_spec_without(*absent): + """A ``find_spec`` that knows every module except *absent*.""" + + def find_spec(name, *args): + return None if name in absent else object() + + return find_spec + + +def _missing_from(*absent): + return patch.object( + autotag.importlib.util, "find_spec", side_effect=_find_spec_without(*absent) + ) + + +class TestMissingPackages(unittest.TestCase): + + def test_every_backend_names_its_modules_and_pip_names(self): + self.assertEqual(BACKEND_PACKAGES["gguf"], (("llama_cpp", "llama-cpp-python"),)) + self.assertEqual( + [pip for _module, pip in BACKEND_PACKAGES["hf"]], + ["torch", "transformers", "accelerate", "bitsandbytes"], + ) + for model in ("wd14-swinv2", "wd14-vit"): + self.assertEqual(BACKEND_PACKAGES[model], (("onnxruntime", "onnxruntime"),)) + self.assertEqual(set(BACKEND_PACKAGES), set(MODELS)) + + def test_reports_pip_names_of_modules_find_spec_cannot_find(self): + with _missing_from("llama_cpp", "accelerate", "bitsandbytes"): + self.assertEqual(missing_packages("gguf"), ["llama-cpp-python"]) + self.assertEqual(missing_packages("hf"), ["accelerate", "bitsandbytes"]) + self.assertEqual(missing_packages("wd14-vit"), []) + + def test_find_spec_errors_count_as_missing(self): + for error in (ValueError("no spec"), ImportError("broken"), OSError("io")): + with self.subTest(error=type(error).__name__): + with patch.object( + autotag.importlib.util, "find_spec", side_effect=error + ): + self.assertEqual(missing_packages("wd14-vit"), ["onnxruntime"]) + + def test_a_module_already_imported_is_present_without_a_lookup(self): + with _patch_modules({"llama_cpp": types.ModuleType("llama_cpp")}): + with patch.object( + autotag.importlib.util, "find_spec", side_effect=AssertionError + ): + self.assertEqual(missing_packages("gguf"), []) + + def test_the_check_imports_nothing(self): + before = set(sys.modules) + missing_packages("gguf") + missing_packages("hf") + missing_packages("wd14-vit") + imported = set(sys.modules) - before + heavy = {"llama_cpp", "torch", "transformers", "onnxruntime"} + self.assertFalse(imported & heavy, imported) + + def test_unknown_model_has_no_packages(self): + self.assertEqual(missing_packages("nope"), []) + + def test_package_fields_and_the_one_sentence(self): + self.assertEqual( + package_fields("gguf", []), + {"missing_packages": [], "install_command": None, "install_note": None}, + ) + fields = package_fields("gguf", ["llama-cpp-python"]) + self.assertEqual(fields["missing_packages"], ["llama-cpp-python"]) + self.assertEqual(fields["install_command"], "pip install llama-cpp-python") + self.assertEqual(fields["install_note"], GGUF_CUDA) + self.assertEqual( + package_fields("hf", ["torch", "bitsandbytes"])["install_command"], + "pip install torch bitsandbytes", + ) + self.assertIsNone(package_fields("hf", ["torch"])["install_note"]) + self.assertIn( + "onnxruntime-gpu", + package_fields("wd14-vit", ["onnxruntime"])["install_note"], + ) + error = MissingBackendPackages("gguf", ["llama-cpp-python"]) + self.assertIsInstance(error, RuntimeError) + self.assertEqual( + str(error), + "GGUF JoyCaption needs llama-cpp-python. Install it into ComfyUI's " + "Python: pip install llama-cpp-python. " + GGUF_CUDA, + ) + self.assertEqual( + str(MissingBackendPackages("hf", ["accelerate", "bitsandbytes"])), + "HuggingFace JoyCaption needs accelerate and bitsandbytes. Install them " + "into ComfyUI's Python: pip install accelerate bitsandbytes.", + ) + + +class TestStatusReportsPackages(EngineTestCase): + + def test_downloaded_keeps_its_files_only_meaning(self): + _touch_model_files(self.models_dir, "gguf") + _touch_model_files(self.models_dir, "wd14-vit") + service = AutoTagService(models_dir=self.models_dir) + + with _missing_from("llama_cpp"): + status = service.get_models_status() + + self.assertTrue(status["gguf"]["downloaded"]) + self.assertEqual(status["gguf"]["missing_packages"], ["llama-cpp-python"]) + self.assertEqual( + status["gguf"]["install_command"], "pip install llama-cpp-python" + ) + self.assertEqual(status["gguf"]["install_note"], GGUF_CUDA) + self.assertEqual(status["wd14-vit"]["missing_packages"], []) + self.assertIsNone(status["wd14-vit"]["install_command"]) + # A model not on disk still says which packages it will need. + self.assertFalse(status["wd14-swinv2"]["downloaded"]) + self.assertIn("missing_packages", status["hf"]) + + def test_an_injected_loader_brings_its_own_backend(self): + service = self._service(("gguf", FakeGgufModel("x"))) + with _missing_from("llama_cpp"): + self.assertEqual( + service.get_models_status()["gguf"]["missing_packages"], [] + ) + + def test_an_injected_probe_is_used_for_status_and_load(self): + _touch_model_files(self.models_dir, "gguf") + service = AutoTagService( + models_dir=self.models_dir, + model_loader=lambda *_: ("gguf", FakeGgufModel("x")), + package_probe=lambda model_type: ["llama-cpp-python"], + ) + self.assertEqual( + service.get_models_status()["gguf"]["missing_packages"], + ["llama-cpp-python"], + ) + with self.assertRaises(MissingBackendPackages): + service.load_model("gguf") + + +class TestLoadFailsFast(EngineTestCase): + + def test_missing_package_refuses_before_the_loader_without_an_error_log(self): + _touch_model_files(self.models_dir, "gguf") + service = AutoTagService(models_dir=self.models_dir) + calls = [] + service._model_loader = lambda *args: calls.append(args) + + with ( + _missing_from("llama_cpp"), + self.assertLogs(service.logger, level="DEBUG") as logs, + ): + with self.assertRaises(MissingBackendPackages) as ctx: + service.load_model("gguf") + + self.assertEqual(calls, []) + self.assertEqual(ctx.exception.model_type, "gguf") + self.assertEqual(ctx.exception.missing, ["llama-cpp-python"]) + self.assertIn("pip install llama-cpp-python", str(ctx.exception)) + self.assertFalse(service.is_model_loaded()) + levels = {record.levelname for record in logs.records} + self.assertNotIn("ERROR", levels) + self.assertIn("WARNING", levels) + self.assertTrue(all(record.exc_info is None for record in logs.records)) + + def test_not_downloaded_is_reported_before_packages(self): + service = AutoTagService(models_dir=self.models_dir) + with _missing_from("llama_cpp"): + with self.assertRaises(RuntimeError) as ctx: + service.load_model("gguf") + self.assertNotIsInstance(ctx.exception, MissingBackendPackages) + + def test_a_module_not_found_from_the_loader_becomes_the_same_refusal(self): + _touch_model_files(self.models_dir, "gguf") + + def loader(model_type, use_gpu): + raise ModuleNotFoundError("No module named 'llama_cpp'", name="llama_cpp") + + service = AutoTagService( + models_dir=self.models_dir, + model_loader=loader, + package_probe=lambda _model: [], + ) + with self.assertLogs(service.logger, level="DEBUG") as logs: + with self.assertRaises(MissingBackendPackages) as ctx: + service.load_model("gguf") + self.assertEqual(ctx.exception.missing, ["llama-cpp-python"]) + self.assertNotIn("ERROR", {record.levelname for record in logs.records}) + self.assertFalse(service.is_model_loaded()) + + def test_an_unknown_module_not_found_names_the_module(self): + def loader(model_type, use_gpu): + raise ModuleNotFoundError( + "No module named 'sentencepiece'", name="sentencepiece" + ) + + hf_dir = self.models_dir / MODELS["hf"]["subdir"] + hf_dir.mkdir(parents=True, exist_ok=True) + (hf_dir / "config.json").write_text("{}") + service = AutoTagService( + models_dir=self.models_dir, + model_loader=loader, + package_probe=lambda _model: [], + ) + with self.assertRaises(MissingBackendPackages) as ctx: + service.load_model("hf") + self.assertEqual(ctx.exception.missing, ["sentencepiece"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_image_tag_jobs.py b/tests/test_image_tag_jobs.py index f593aaf..b555c79 100644 --- a/tests/test_image_tag_jobs.py +++ b/tests/test_image_tag_jobs.py @@ -299,6 +299,19 @@ class TestAutoTag(ImageTagJobTestCase): self.tagger.download_model.assert_not_called() self.assertIsNone(model_gate.holder()) + async def test_a_missing_backend_package_fails_the_job_before_the_first_file(self): + self.png("a.png") + self.tagger.missing["wd14-vit"] = ["onnxruntime"] + + with self.assertNoLogs(self.api.logger, level="ERROR"): + job = await self.run_job(self.auto_body(["a.png"])) + + self.assertEqual(job.state, "failed") + self.assertEqual(job.processed, 0) + self.assertTrue(job.status_line.startswith("WD14 ViT needs onnxruntime.")) + self.assertIn("pip install onnxruntime", job.status_line) + self.assertEqual(self.tagger.load_calls, []) + async def test_model_unavailable_fails_the_job_before_the_first_file(self): self.png("a.png") self.tagger.downloaded["wd14-vit"] = False diff --git a/web/js/workspace/image-tagging-state.mjs b/web/js/workspace/image-tagging-state.mjs index 18b2e88..6bb7413 100644 --- a/web/js/workspace/image-tagging-state.mjs +++ b/web/js/workspace/image-tagging-state.mjs @@ -203,17 +203,48 @@ async function readDefaults(api, signal) { } export const UNAVAILABLE_HINT = 'Download it in System › Auto-tagging.'; +const nonEmpty = value => typeof value === 'string' && value.trim() !== ''; +const joined = names => (names.length > 1 ? `${names.slice(0, -1).join(', ')} and ${names.at(-1)}` : names.join('')); + +/** The pip names a `/autotag/models` entry says its backend cannot import (#424); only non-empty strings count. */ +export const missingPackages = entry => (object(entry) && Array.isArray(entry.missing_packages) + ? entry.missing_packages.filter(nonEmpty) : []); + +/** `pip install …` for `entry`'s missing packages: the server's command, else built from the names; '' when none. */ +export function installCommand(entry) { + const missing = missingPackages(entry); + if (!missing.length) return ''; + return nonEmpty(entry.install_command) ? entry.install_command : `pip install ${missing.join(' ')}`; +} + +/** + * The server's missing-package sentence (`missing_packages_message`, #424) for `name`, else '' (nothing missing): + * " needs llama-cpp-python. Install it into ComfyUI's Python: pip install llama-cpp-python. ". + */ +export function needsPackages(name, entry) { + const missing = missingPackages(entry); + if (!missing.length) return ''; + const note = nonEmpty(entry.install_note) ? ` ${entry.install_note}` : ''; + return `${name} needs ${joined(missing)}. Install ${missing.length === 1 ? 'it' : 'them'} into ComfyUI's Python: ` + + `${installCommand(entry)}.${note}`; +} + +/** "a, b and c": the missing package names as the sentences list them. */ +export const packageList = entry => joined(missingPackages(entry)); + /** * Why `model` cannot run, from the `/autotag/models` status (the server's own `model-unavailable` sentence), else '' - * (downloaded, or the status is unknown: the job then reports a missing model itself). + * (ready, or the status is unknown: the job then reports a missing model itself). Not downloaded is named first, then + * a backend package that cannot be imported (#424). */ export function modelUnavailable(models, model) { const option = modelOf(model); requireValue(option, 'Unknown tagging model'); const status = object(models) && object(models[model]) ? models[model] : null; - if (!status || status.downloaded !== false) return ''; - const name = typeof status.name === 'string' && status.name.trim() ? status.name : option.label.split(' · ')[0]; - return `${name} is not downloaded. ${UNAVAILABLE_HINT}`; + if (!status) return ''; + const fallback = option.label.split(' · ')[0]; + if (status.downloaded === false) return `${nonEmpty(status.name) ? status.name : fallback} is not downloaded. ${UNAVAILABLE_HINT}`; + return needsPackages(nonEmpty(status.label) ? status.label : fallback.replace(' (Recommended)', ''), status); } function initialState(source, preset) { diff --git a/web/js/workspace/system-tagging.mjs b/web/js/workspace/system-tagging.mjs index d9b369f..6a96f17 100644 --- a/web/js/workspace/system-tagging.mjs +++ b/web/js/workspace/system-tagging.mjs @@ -29,7 +29,7 @@ import { HOST_UNAVAILABLE } from './system-appearance.mjs'; import { openCloseOnlyReview, openReviewDialog } from './workspace-dialogs.mjs'; import { BACKUP_COPY, backupText, holding, reviewRefusal } from './maintenance-review.mjs'; import { busyToast, downloadDatabaseBackup } from './activity-scan.mjs'; -import { MODEL_OPTIONS } from './image-tagging-state.mjs'; +import { MODEL_OPTIONS, installCommand, missingPackages, needsPackages, packageList } from './image-tagging-state.mjs'; import { openImageTagging } from './image-tagging.mjs'; import { selectedImage } from './images.mjs'; import { NO_SELECTION } from './vocabulary.mjs'; @@ -80,6 +80,7 @@ export const STATUS_COPY = Object.freeze({ loaded: 'Downloaded and loaded.', partial: 'Partly downloaded. Download again to finish.', missing: size => `Not downloaded. Download size ${size}.`, + needs: (names, command, note) => `Downloaded. Needs ${names}: ${command}.${typeof note === 'string' && note.trim() ? ` ${note}` : ''}`, }); /** The messages that name a setting the section cannot use (the server's own wording). */ @@ -140,15 +141,21 @@ export function modelOptions(models) { /** * D132: the Tagging model description for `status` (`{phase: loading|loaded|failed, models}`) and the chosen `model`. - * GGUF with only one of its two files on disk is partly downloaded. + * GGUF with only one of its two files on disk is partly downloaded. A backend package the server cannot import (#424) + * turns "Downloaded" into "Downloaded. Needs : pip install ." and adds "Also needs ." otherwise. */ export function modelDescription(status, model) { if (status?.phase === 'loading') return STATUS_COPY.loading; const entry = status?.models?.[model]; if (status?.phase !== 'loaded' || !object(entry)) return STATUS_COPY.failed; - if (entry.downloaded === true) return entry.loaded === true ? STATUS_COPY.loaded : STATUS_COPY.downloaded; - if (partlyDownloaded(model, entry)) return STATUS_COPY.partial; - return STATUS_COPY.missing(modelSize(model, entry)); + // #424: a backend package that cannot be imported; a model on disk without it never reads as ready. + const needs = missingPackages(entry).length ? packageList(entry) : ''; + if (entry.downloaded === true) { + if (needs) return STATUS_COPY.needs(needs, installCommand(entry), entry.install_note); + return entry.loaded === true ? STATUS_COPY.loaded : STATUS_COPY.downloaded; + } + const files = partlyDownloaded(model, entry) ? STATUS_COPY.partial : STATUS_COPY.missing(modelSize(model, entry)); + return needs ? `${files} Also needs ${needs}.` : files; } /** GGUF with only one of its two files on disk (a download that did not finish). */ @@ -316,6 +323,7 @@ export function diskShortNote(disk) { export function batchCloseNote(plan) { if (plan.review_required === true) return REPLACE_NOTE; if (plan.model?.downloaded === false) return notDownloaded(modelLabel(plan.model.type, plan.model)); + if (missingPackages(plan.model).length) return needsPackages(modelLabel(plan.model.type, plan.model), plan.model); if (plan.disk?.sufficient === false) return diskShortNote(plan.disk); if (!Number.isSafeInteger(plan.images) || plan.images === 0) return NO_IMAGES; if (plan.to_tag === 0) return NOTHING_TO_TAG; @@ -484,8 +492,9 @@ export function mountTagging(container, { api, activity, scope, drafts, notify = unload.disabled = !(status.phase === 'loaded' && status.loadedType); if (review.textContent === 'Reviewing…') return; const batchProblem = known() ? draftProblem(values, { limit: false }) : unknown; - setAvailable(review, entry?.downloaded === false ? notDownloaded(modelLabel(values.autotag_model_type, entry)) - : batchProblem?.message); + const label = modelLabel(values.autotag_model_type, entry); + setAvailable(review, entry?.downloaded === false ? notDownloaded(label) + : needsPackages(label, entry) || batchProblem?.message); }; let reads = 0;