From a6d4facfbfefb2716d999babcb4e4c6e781cda77 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 18:43:10 -0700 Subject: [PATCH] refactor: autotag model loading is injectable --- py/api/autotag_routes.py | 489 ++++++++++++++++++----------------- py/autotag.py | 33 ++- tests/test_api_autotag.py | 152 +++++++++++ tests/test_autotag_engine.py | 331 ++++++++++++++++++++++++ 4 files changed, 756 insertions(+), 249 deletions(-) create mode 100644 tests/test_autotag_engine.py diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 80e1615..af24af2 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -91,41 +91,13 @@ class AutotagRoutesMixin: self.logger.error(f"Get autotag models error: {e}") return web.json_response({"success": False, "error": str(e)}, status=500) - async def download_autotag_model(self, request): - """Download an AutoTag model with streaming progress.""" - model_type = request.match_info.get("model_type") - - async def stream_response(): - try: - from ..autotag import get_autotag_service - - service = get_autotag_service() - - if model_type not in service.models_config: - yield f"data: {json.dumps({'type': 'error', 'message': f'Invalid model type: {model_type}'})}\n\n" - return - - yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Starting download...'})}\n\n" - - progress_data = {"last_progress": 0} - - def progress_callback(status: str, progress: float): - progress_data["last_progress"] = progress - - loop = asyncio.get_event_loop() - success = await loop.run_in_executor( - None, lambda: service.download_model(model_type, progress_callback) - ) - - if success: - yield f"data: {json.dumps({'type': 'complete', 'progress': 100, 'status': 'Download complete'})}\n\n" - else: - yield f"data: {json.dumps({'type': 'error', 'message': 'Download failed'})}\n\n" - - except Exception as e: - self.logger.exception("Download model error") - yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\n\n" + async def _stream_sse(self, request, events): + """Write an async iterator of dict events as server-sent events. + A client disconnect raises ConnectionResetError from ``write``; the + event generator is then closed so its cleanup runs before re-raising + (aiohttp logs premature disconnects at debug level). + """ response = web.StreamResponse( status=200, reason="OK", @@ -135,15 +107,62 @@ class AutotagRoutesMixin: "Connection": "keep-alive", }, ) - await response.prepare(request) - - async for chunk in stream_response(): - await response.write(chunk.encode("utf-8")) - + try: + async for event in events: + await response.write(f"data: {json.dumps(event)}\n\n".encode("utf-8")) + except (ConnectionResetError, asyncio.CancelledError): + self.logger.info("Client disconnected; stopping autotag stream") + await events.aclose() + raise await response.write_eof() return response + async def download_autotag_model(self, request): + """Download an AutoTag model with streaming progress (POST only).""" + model_type = request.match_info.get("model_type") + + 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 self._run_in_executor( + 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": "An internal error occurred. Check server logs for details.", + } + + return await self._stream_sse(request, events()) + async def _read_start_params(self, request) -> Optional[dict]: """Merge query-string and optional JSON-body parameters for autotag/start. @@ -160,221 +179,211 @@ class AutotagRoutesMixin: params.update(body) return params + @staticmethod + def _real_tags(tags) -> list: + """Normalise a tag field to a list without bookkeeping markers.""" + if isinstance(tags, str): + tags = [t.strip() for t in tags.split(",") if t.strip()] + return [ + t + for t in (tags or []) + if t != "auto-scanned" + and not t.startswith("prepend:") + and not t.startswith("append:") + ] + + async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool: + """Merge ``tags`` into the prompt; True when something new was added.""" + existing_prompt = await self._run_in_executor( + self.db.get_prompt_by_id, prompt_id + ) + if not existing_prompt: + return False + existing_tags = existing_prompt.get("tags", []) + if isinstance(existing_tags, str): + existing_tags = [t.strip() for t in existing_tags.split(",") if t.strip()] + new_tags = [t for t in tags if t not in existing_tags] + if not new_tags: + return False + await self._run_in_executor( + self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags + ) + return True + + async def _autotag_batch_events(self, service, opts: dict): + """Yield SSE progress dicts while tagging every image linked to a prompt.""" + import time as _time + + model_type = opts["model_type"] + keep_in_memory = opts["keep_in_memory"] + status = service.get_models_status() + if not status.get(model_type, {}).get("downloaded"): + yield {"type": "error", "message": f"Model {model_type} not downloaded"} + return + + yield {"type": "progress", "progress": 0, "status": "Loading model..."} + try: + await self._run_in_executor(service.load_model, model_type, True) + except Exception: + self.logger.exception("Failed to load model") + yield { + "type": "error", + "message": "Failed to load model. Check server logs for details.", + } + return + + if opts["custom_prompt"]: + service.custom_prompt = opts["custom_prompt"] + + yield { + "type": "progress", + "progress": 5, + "status": "Model loaded. Fetching all images from database...", + } + images = await self._run_in_executor(self.db.get_all_images) + total = len(images) + if total == 0: + yield { + "type": "complete", + "processed": 0, + "tagged": 0, + "skipped": 0, + "status": "No images with linked prompts found", + } + service.unload_model() + return + + yield { + "type": "progress", + "progress": 10, + "status": f"Found {total} images. Processing...", + } + + counts = {"processed": 0, "tagged": 0, "skipped": 0, "errors": 0} + tagged_prompt_ids = set() + last_update = _time.monotonic() + finished = False + + def progress_event(i, message): + return { + "type": "progress", + "progress": 10 + int((i + 1) / total * 85), + "status": message, + "processed": counts["processed"], + "tagged": counts["tagged"], + "skipped": counts["skipped"], + } + + try: + for i, image_data in enumerate(images): + image_path = image_data.get("image_path") + prompt_id = image_data.get("prompt_id") + skip_reason = None + if not image_path or not prompt_id: + skip_reason = "unlinked" + elif prompt_id in tagged_prompt_ids: + skip_reason = "prompt already processed" + elif not Path(image_path).exists(): + skip_reason = "file missing" + elif opts["skip_tagged"] and self._real_tags( + image_data.get("prompt_tags", []) + ): + tagged_prompt_ids.add(prompt_id) + skip_reason = "already tagged" + + if skip_reason is not None: + counts["skipped"] += 1 + now = _time.monotonic() + if (now - last_update) >= 0.5 or i == total - 1: + yield progress_event( + i, f"Skipping {i + 1}/{total} ({skip_reason})..." + ) + await asyncio.sleep(0.01) + last_update = now + continue + + try: + tags = await self._run_in_executor( + service.generate_tags, + str(image_path), + general_threshold=opts["general_threshold"], + character_threshold=opts["character_threshold"], + ) + counts["processed"] += 1 + if tags and await self._apply_tags_to_prompt(prompt_id, tags): + counts["tagged"] += 1 + else: + counts["skipped"] += 1 + tagged_prompt_ids.add(prompt_id) + except Exception as img_err: + self.logger.error( + f"Error processing {os.path.basename(str(image_path))}: " + f"{img_err}" + ) + counts["errors"] += 1 + counts["processed"] += 1 + + yield progress_event(i, f"Processing {i + 1}/{total}...") + await asyncio.sleep(0.01) + last_update = _time.monotonic() + + finished = True + if not keep_in_memory: + service.unload_model() + yield { + "type": "complete", + "progress": 100, + **counts, + "status": "Complete", + "model_status": ( + "Model kept in memory" if keep_in_memory else "Model unloaded" + ), + "model_loaded": keep_in_memory, + } + finally: + if not finished and not keep_in_memory: + service.unload_model() + async def start_autotag(self, request): """Start batch auto-tagging with streaming progress (POST only).""" params = await self._read_start_params(request) if params is None: - return web.json_response( - {"success": False, "error": "Request body must be a JSON object"}, - status=400, - ) - model_type = params.get("model_type", "gguf") - custom_prompt = params.get("prompt", "") - skip_tagged = str(params.get("skip_tagged", "true")).lower() == "true" - keep_in_memory = str(params.get("keep_in_memory", "true")).lower() == "true" - use_gpu = True + return self._bad_request("Request body must be a JSON object") - # WD14 threshold params - general_threshold = params.get("general_threshold") - character_threshold = params.get("character_threshold") + opts = { + "model_type": params.get("model_type", "gguf"), + "custom_prompt": params.get("prompt", ""), + "skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true", + "keep_in_memory": str(params.get("keep_in_memory", "true")).lower() + == "true", + "general_threshold": params.get("general_threshold"), + "character_threshold": params.get("character_threshold"), + } try: - if general_threshold is not None: - general_threshold = float(general_threshold) - if character_threshold is not None: - character_threshold = float(character_threshold) + for key in ("general_threshold", "character_threshold"): + if opts[key] is not None: + opts[key] = float(opts[key]) except (ValueError, TypeError): - return web.json_response( - { - "success": False, - "error": "general_threshold and character_threshold must be numeric", - }, - status=400, + return self._bad_request( + "general_threshold and character_threshold must be numeric" ) - async def stream_response(): + async def events(): try: from ..autotag import get_autotag_service - import time as _time - - service = get_autotag_service() - - status = service.get_models_status() - if not status.get(model_type, {}).get("downloaded"): - yield f"data: {json.dumps({'type': 'error', 'message': f'Model {model_type} not downloaded'})}\n\n" - return - - yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Loading model...'})}\n\n" - - loop = asyncio.get_event_loop() - try: - await loop.run_in_executor( - None, lambda: service.load_model(model_type, use_gpu) - ) - except Exception as e: - self.logger.exception("Failed to load model") - yield f"data: {json.dumps({'type': 'error', 'message': 'Failed to load model. Check server logs for details.'})}\n\n" - return - - if custom_prompt: - service.custom_prompt = custom_prompt - - yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': 'Model loaded. Fetching all images from database...'})}\n\n" - - images = await self._run_in_executor(self.db.get_all_images) - - total_files = len(images) - if total_files == 0: - yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'tagged': 0, 'skipped': 0, 'status': 'No images with linked prompts found'})}\n\n" - service.unload_model() - return - - yield f"data: {json.dumps({'type': 'progress', 'progress': 10, 'status': f'Found {total_files} images. Processing...'})}\n\n" - - processed = 0 - tagged = 0 - skipped = 0 - errors = 0 - tagged_prompt_ids = set() - - last_update_time = _time.monotonic() - - for i, image_data in enumerate(images): - image_path = image_data.get("image_path") - prompt_id = image_data.get("prompt_id") - - if not image_path or not prompt_id: - skipped += 1 - continue - - if prompt_id in tagged_prompt_ids: - skipped += 1 - now = _time.monotonic() - if (now - last_update_time) >= 0.5 or i == total_files - 1: - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (prompt already processed)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = now - continue - - if not Path(image_path).exists(): - skipped += 1 - now = _time.monotonic() - if (now - last_update_time) >= 0.5 or i == total_files - 1: - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (file missing)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = now - continue - - if skip_tagged: - prompt_tags = image_data.get("prompt_tags", []) - if isinstance(prompt_tags, str): - prompt_tags = [ - t.strip() for t in prompt_tags.split(",") if t.strip() - ] - real_tags = [ - t - for t in prompt_tags - if t != "auto-scanned" - and not t.startswith("prepend:") - and not t.startswith("append:") - ] - if real_tags: - tagged_prompt_ids.add(prompt_id) - skipped += 1 - now = _time.monotonic() - if (now - last_update_time) >= 0.5 or i == total_files - 1: - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (already tagged)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = now - continue - - try: - tags = await loop.run_in_executor( - None, - lambda p=str(image_path): service.generate_tags( - p, - general_threshold=general_threshold, - character_threshold=character_threshold, - ), - ) - - processed += 1 - - if tags: - existing_prompt = await self._run_in_executor( - self.db.get_prompt_by_id, prompt_id - ) - if existing_prompt: - existing_tags = existing_prompt.get("tags", []) - if isinstance(existing_tags, str): - existing_tags = [ - t.strip() - for t in existing_tags.split(",") - if t.strip() - ] - - new_tags = [t for t in tags if t not in existing_tags] - if new_tags: - all_tags = existing_tags + new_tags - await self._run_in_executor( - self.db.update_prompt_metadata, - prompt_id, - tags=all_tags, - ) - tagged += 1 - else: - skipped += 1 - else: - skipped += 1 - else: - skipped += 1 - - tagged_prompt_ids.add(prompt_id) - - progress = 10 + int((i + 1) / total_files * 85) - yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Processing {i+1}/{total_files}...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n" - await asyncio.sleep(0.01) - last_update_time = _time.monotonic() - - except Exception as img_err: - self.logger.error(f"Error processing {image_path}: {img_err}") - errors += 1 - processed += 1 - - if not keep_in_memory: - service.unload_model() - model_status = "Model unloaded" - else: - model_status = "Model kept in memory" - - yield f"data: {json.dumps({'type': 'complete', 'progress': 100, 'processed': processed, 'tagged': tagged, 'skipped': skipped, 'errors': errors, 'status': 'Complete', 'model_status': model_status, 'model_loaded': keep_in_memory})}\n\n" - - except Exception as e: - self.logger.error(f"AutoTag error: {e}") - import traceback + async for event in self._autotag_batch_events( + get_autotag_service(), opts + ): + yield event + except Exception: self.logger.exception("AutoTag error") - yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\n\n" + yield { + "type": "error", + "message": "An internal error occurred. Check server logs for details.", + } - response = web.StreamResponse( - status=200, - reason="OK", - headers={ - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Connection": "keep-alive", - }, - ) - - await response.prepare(request) - - async for chunk in stream_response(): - await response.write(chunk.encode("utf-8")) - - await response.write_eof() - return response + return await self._stream_sse(request, events()) def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]: """Return the real path when it is a file inside an output directory.""" diff --git a/py/autotag.py b/py/autotag.py index 7197675..ab426af 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -5,12 +5,16 @@ Provides JoyCaption-based automatic tagging for images using LLM models. Refactored from standalone_tagger.py for integration with PromptManager API. """ +import errno import gc import os from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple from PIL import Image +# A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend) +ModelLoader = Callable[[str, bool], Tuple[str, Any]] + # Try to import logging from utils, fallback to standard logging try: from ..utils.logging_config import get_logger @@ -98,14 +102,23 @@ class AutoTagService: service.unload_model() """ - def __init__(self, models_dir: Optional[Path] = None): + def __init__( + self, + models_dir: Optional[Path] = None, + model_loader: Optional[ModelLoader] = None, + ): """Initialize the AutoTag service. Args: models_dir: Directory for storing models. If None, uses ComfyUI's folder_paths.models_dir / "LLM" path. + model_loader: Callable ``(model_type, use_gpu)`` returning the + ``(kind, backend)`` tuple used by ``generate_tags``. + Defaults to ``_load_model``, which imports the ML + backends lazily. Tests inject a fake here. """ self.logger = get_logger("autotag.service") + self._model_loader: ModelLoader = model_loader or self._load_model # Determine models directory if models_dir: @@ -591,13 +604,7 @@ class AutoTagService: raise RuntimeError(f"Model {model_type} not downloaded") try: - if model_type == "gguf": - self._tagger = self._load_gguf_tagger(use_gpu) - elif model_type.startswith("wd14"): - self._tagger = self._load_wd14_tagger(model_type, use_gpu) - else: - self._tagger = self._load_hf_tagger() - + self._tagger = self._model_loader(model_type, use_gpu) self._current_model_type = model_type self.logger.info(f"Model {model_type} loaded successfully") return True @@ -608,6 +615,14 @@ class AutoTagService: self._current_model_type = None raise RuntimeError(f"Failed to load model: {e}") + def _load_model(self, model_type: str, use_gpu: bool = True) -> Tuple[str, Any]: + """Default loader: dispatch to a backend, importing it only now.""" + if model_type == "gguf": + return self._load_gguf_tagger(use_gpu) + if model_type.startswith("wd14"): + return self._load_wd14_tagger(model_type, use_gpu) + return self._load_hf_tagger() + def _load_gguf_tagger(self, use_gpu: bool = True): """Load GGUF-based tagger.""" from llama_cpp import Llama @@ -735,7 +750,7 @@ class AutoTagService: raise RuntimeError("No model loaded. Call load_model() first.") if not os.path.exists(image_path): - raise FileNotFoundError(f"Image not found: {image_path}") + raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path)) use_prompt = prompt or self._custom_prompt diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index dc77133..fb05f54 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -380,5 +380,157 @@ class TestAutotagSideEffectRoutesArePostOnly(AutotagAPITestCase): self.assertTrue((await resp.json())["success"]) +class TestAutotagBatch(AutotagAPITestCase): + """Batch tagging over the DB image list with SSE progress reporting.""" + + def _seed(self): + """Three prompts: A (one image), B (two images), C (missing file).""" + a = self._save_prompt("prompt a") + b = self._save_prompt("prompt b") + c = self._save_prompt("prompt c") + self._link_image(a, self._make_image(self.output_dir, "a1.png")) + self._link_image(b, self._make_image(self.output_dir, "b1.png")) + self._link_image(b, self._make_image(self.output_dir, "b2.png")) + self._link_image(c, self.output_dir / "missing.png") + return a, b, c + + async def _start(self, query="model_type=gguf", body=None): + if body is None: + resp = await self.client.request( + "POST", f"/prompt_manager/autotag/start?{query}" + ) + else: + resp = await self._post_json(f"/prompt_manager/autotag/start?{query}", body) + return resp, self._sse_events(await resp.text()) + + async def test_batch_tags_each_prompt_once_and_reports_counts(self): + a, b, c = self._seed() + + resp, events = await self._start() + + self.assertEqual(resp.status, 200) + done = events[-1] + self.assertEqual(done["type"], "complete") + self.assertEqual(done["processed"], 2) + self.assertEqual(done["tagged"], 2) + self.assertEqual(done["skipped"], 2) + self.assertEqual(done["errors"], 0) + self.assertEqual(len(self.service.generate_calls), 2) + self.assertEqual(self.api.db.get_prompt_by_id(a)["tags"], ["1girl", "smile"]) + self.assertEqual(self.api.db.get_prompt_by_id(b)["tags"], ["1girl", "smile"]) + self.assertEqual(self.api.db.get_prompt_by_id(c)["tags"], []) + progress = [e for e in events if e["type"] == "progress"] + self.assertTrue(all(0 <= e["progress"] <= 100 for e in progress)) + self.assertTrue(any("processed" in e for e in progress)) + + async def test_already_tagged_prompts_are_skipped_unless_asked(self): + p = self._save_prompt("tagged", tags=["cat"]) + self._link_image(p, self._make_image(self.output_dir, "t.png")) + + _, events = await self._start("model_type=gguf&skip_tagged=true") + self.assertEqual(events[-1]["skipped"], 1) + self.assertEqual(self.service.generate_calls, []) + + _, events = await self._start("model_type=gguf&skip_tagged=false") + self.assertEqual(events[-1]["tagged"], 1) + self.assertEqual( + self.api.db.get_prompt_by_id(p)["tags"], ["cat", "1girl", "smile"] + ) + + async def test_prompt_already_holding_the_tags_counts_as_skipped(self): + p = self._save_prompt("same", tags=["1girl", "smile"]) + self._link_image(p, self._make_image(self.output_dir, "s.png")) + + _, events = await self._start("model_type=gguf&skip_tagged=false") + + self.assertEqual(events[-1]["skipped"], 1) + self.assertEqual(events[-1]["tagged"], 0) + + async def test_generation_errors_are_counted_and_do_not_abort(self): + self._seed() + self.service.generate_error = RuntimeError("boom") + + _, events = await self._start() + + self.assertEqual(events[-1]["type"], "complete") + # a1, b1 and b2 each fail: a failed prompt is retried on its next image + self.assertEqual(events[-1]["errors"], 3) + self.assertEqual(events[-1]["tagged"], 0) + + async def test_json_body_parameters_are_accepted(self): + self._seed() + + _, events = await self._start( + "", {"model_type": "wd14-vit", "general_threshold": 0.5, "prompt": "p"} + ) + + self.assertEqual(events[-1]["type"], "complete") + self.assertEqual(self.service.load_calls, [("wd14-vit", True)]) + self.assertEqual(self.service.generate_calls[0][2], 0.5) + self.assertEqual(self.service.custom_prompt, "p") + + async def test_keep_in_memory_false_unloads_after_batch(self): + self._seed() + + _, events = await self._start("model_type=gguf&keep_in_memory=false") + + self.assertEqual(events[-1]["model_loaded"], False) + self.assertEqual(self.service.unload_calls, 1) + + async def test_model_not_downloaded_is_an_error_event(self): + self.service.downloaded = False + + _, events = await self._start() + + self.assertEqual(events[-1]["type"], "error") + self.assertEqual(self.service.load_calls, []) + + async def test_load_failure_is_an_error_event_without_details(self): + self.service.load_error = RuntimeError("/srv/models/secret.gguf is corrupt") + + _, events = await self._start() + + self.assertEqual(events[-1]["type"], "error") + self.assertNotIn("/srv/models", events[-1]["message"]) + + async def test_non_numeric_threshold_is_400(self): + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start?general_threshold=abc" + ) + self.assertEqual(resp.status, 400) + + async def test_invalid_json_body_is_400(self): + resp = await self.client.request( + "POST", + "/prompt_manager/autotag/start", + data="{not json", + headers={"Content-Type": "application/json"}, + ) + self.assertEqual(resp.status, 400) + + async def test_client_disconnect_stops_the_batch_and_unloads(self): + self._seed() + real_write = web.StreamResponse.write + state = {"writes": 0} + + async def flaky_write(response, data): + state["writes"] += 1 + if b"Processing " in data: + raise ConnectionResetError("Cannot write to closing transport") + return await real_write(response, data) + + with patch.object(web.StreamResponse, "write", flaky_write): + try: + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start?keep_in_memory=false" + ) + await resp.read() + except Exception: + pass + + self.assertEqual(len(self.service.generate_calls), 1) + self.assertEqual(self.service.unload_calls, 1) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_autotag_engine.py b/tests/test_autotag_engine.py new file mode 100644 index 0000000..3890458 --- /dev/null +++ b/tests/test_autotag_engine.py @@ -0,0 +1,331 @@ +"""Tests for the AutoTagService tagging pipeline with an injected fake model. + +No real ML model is loaded; ``model_loader`` returns small fakes that mimic +the onnxruntime session (WD14) and the llama.cpp chat model (GGUF). +""" + +import os +import shutil +import sys +import tempfile +import unittest +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from PIL import Image + +from py.autotag import MODELS, AutoTagService + + +class FakeInput: + def __init__(self, name="input", shape=(1, 8, 8, 3)): + self.name = name + self.shape = list(shape) + + +class FakeOnnxSession: + """Mimics onnxruntime.InferenceSession for a WD14 classifier.""" + + def __init__(self, probabilities): + self.probabilities = probabilities + self.run_calls = [] + + def get_inputs(self): + return [FakeInput()] + + def get_providers(self): + return ["CPUExecutionProvider"] + + def run(self, outputs, feeds): + self.run_calls.append(feeds) + return [[list(self.probabilities)]] + + +class FakeGgufModel: + """Mimics llama_cpp.Llama.create_chat_completion.""" + + def __init__(self, text): + self.text = text + self.calls = [] + + def create_chat_completion(self, **kwargs): + self.calls.append(kwargs) + return {"choices": [{"message": {"content": self.text}}]} + + +WD14_TAGS = [ + {"name": "1girl", "category": "0"}, + {"name": "smile", "category": "0"}, + {"name": "hatsune_miku", "category": "4"}, + {"name": "rating:safe", "category": "9"}, + {"name": "", "category": "0"}, + {"name": "long_hair", "category": "0"}, +] + + +def _touch_model_files(models_dir, model_type): + config = MODELS[model_type] + subdir = Path(models_dir) / config["subdir"] + subdir.mkdir(parents=True, exist_ok=True) + names = config.get("files") or [config["filename"], config["mmproj_filename"]] + for name in names: + (subdir / name).write_bytes(b"stub") + + +class EngineTestCase(unittest.TestCase): + + def setUp(self): + self.tmp = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp, True) + self.models_dir = self.tmp / "models" + self.image_path = self.tmp / "img.png" + Image.new("RGB", (4, 6), (200, 100, 50)).save(self.image_path, "PNG") + self.loader_calls = [] + + def _service(self, loaded): + def loader(model_type, use_gpu): + self.loader_calls.append((model_type, use_gpu)) + return loaded + + return AutoTagService(models_dir=self.models_dir, model_loader=loader) + + +class TestModelInjection(EngineTestCase): + + def test_injected_loader_is_used_and_state_tracked(self): + _touch_model_files(self.models_dir, "wd14-vit") + service = self._service(("wd14", (FakeOnnxSession([0.0] * 6), WD14_TAGS))) + + self.assertTrue(service.load_model("wd14-vit", use_gpu=False)) + + self.assertEqual(self.loader_calls, [("wd14-vit", False)]) + self.assertTrue(service.is_model_loaded()) + self.assertEqual(service.get_loaded_model_type(), "wd14-vit") + + def test_default_loader_is_the_lazy_dispatcher(self): + service = AutoTagService(models_dir=self.models_dir) + self.assertEqual(service._model_loader, service._load_model) + + def test_importing_autotag_does_not_import_ml_backends(self): + for heavy in ("onnxruntime", "llama_cpp"): + self.assertNotIn(heavy, sys.modules) + + def test_default_wd14_loader_reports_missing_onnxruntime(self): + _touch_model_files(self.models_dir, "wd14-vit") + service = AutoTagService(models_dir=self.models_dir) + if "onnxruntime" in sys.modules: + self.skipTest("onnxruntime is installed in this environment") + + with self.assertRaises(RuntimeError) as ctx: + service.load_model("wd14-vit") + + self.assertIn("onnxruntime", str(ctx.exception)) + self.assertFalse(service.is_model_loaded()) + + def test_invalid_model_type_raises_value_error(self): + service = self._service(("gguf", FakeGgufModel("x"))) + with self.assertRaises(ValueError): + service.load_model("nope") + + def test_not_downloaded_raises_runtime_error_before_loading(self): + service = self._service(("gguf", FakeGgufModel("x"))) + + with self.assertRaises(RuntimeError): + service.load_model("gguf") + + self.assertEqual(self.loader_calls, []) + + def test_loader_failure_clears_state(self): + _touch_model_files(self.models_dir, "gguf") + + def failing_loader(model_type, use_gpu): + raise OSError("cuda exploded") + + service = AutoTagService( + models_dir=self.models_dir, model_loader=failing_loader + ) + + with self.assertRaises(RuntimeError): + service.load_model("gguf") + + self.assertFalse(service.is_model_loaded()) + self.assertIsNone(service.get_loaded_model_type()) + + def test_loading_again_unloads_first(self): + _touch_model_files(self.models_dir, "gguf") + _touch_model_files(self.models_dir, "wd14-swinv2") + service = self._service(("gguf", FakeGgufModel("x"))) + + service.load_model("gguf") + service.load_model("wd14-swinv2") + + self.assertEqual([c[0] for c in self.loader_calls], ["gguf", "wd14-swinv2"]) + self.assertEqual(service.get_loaded_model_type(), "wd14-swinv2") + + def test_unload_clears_state_and_is_idempotent(self): + _touch_model_files(self.models_dir, "gguf") + service = self._service(("gguf", FakeGgufModel("x"))) + service.load_model("gguf") + + service.unload_model() + service.unload_model() + + self.assertFalse(service.is_model_loaded()) + self.assertIsNone(service.get_loaded_model_type()) + + +class TestWd14Thresholding(EngineTestCase): + + def _loaded(self, probabilities): + _touch_model_files(self.models_dir, "wd14-vit") + self.session = FakeOnnxSession(probabilities) + service = self._service(("wd14", (self.session, WD14_TAGS))) + service.load_model("wd14-vit") + return service + + def test_default_thresholds_apply_per_category(self): + # 1girl 0.9 (general, keep), smile 0.3 (general, drop at 0.35), + # miku 0.8 (character, drop at 0.85), rating 0.99 (ignored), + # empty name 0.99 (ignored), long_hair 0.5 (keep) + service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5]) + + tags = service.generate_tags(str(self.image_path)) + + self.assertEqual(tags, ["1girl", "long_hair"]) + + def test_explicit_thresholds_override_defaults(self): + service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5]) + + tags = service.generate_tags( + str(self.image_path), general_threshold=0.25, character_threshold=0.75 + ) + + self.assertEqual(tags, ["1girl", "hatsune_miku", "long_hair", "smile"]) + + def test_results_sorted_by_confidence(self): + service = self._loaded([0.4, 0.95, 0.0, 0.0, 0.0, 0.6]) + + self.assertEqual( + service.generate_tags(str(self.image_path)), + ["smile", "long_hair", "1girl"], + ) + + def test_fewer_probabilities_than_tags_is_safe(self): + service = self._loaded([0.9, 0.9]) + + self.assertEqual( + service.generate_tags(str(self.image_path)), ["1girl", "smile"] + ) + + def test_image_is_fed_as_square_batch(self): + service = self._loaded([0.0] * 6) + + service.generate_tags(str(self.image_path)) + + feeds = self.session.run_calls[0] + self.assertEqual(list(feeds.keys()), ["input"]) + self.assertEqual(tuple(feeds["input"].shape), (1, 8, 8, 3)) + + def test_service_threshold_setters_clamp(self): + service = self._loaded([0.0] * 6) + service.wd14_general_threshold = 1.7 + service.wd14_character_threshold = -2 + self.assertEqual(service.wd14_general_threshold, 1.0) + self.assertEqual(service.wd14_character_threshold, 0.0) + + +class TestGgufNormalisation(EngineTestCase): + + def _loaded(self, text): + _touch_model_files(self.models_dir, "gguf") + self.model = FakeGgufModel(text) + service = self._service(("gguf", self.model)) + service.load_model("gguf") + return service + + def test_tags_are_lowercased_underscored_deduped_and_filtered(self): + service = self._loaded( + "1Girl, Smile, smile, 'quoted tag', photo_realistic, meta:x, " + "copyright:y, prepend:z, append:w, a, , Long Hair " + ) + + tags = service.generate_tags(str(self.image_path)) + + self.assertEqual(tags, ["1girl", "smile", "quoted_tag", "long_hair"]) + + def test_custom_prompt_is_sent_to_the_model(self): + service = self._loaded("tag") + service.custom_prompt = "describe tersely" + + service.generate_tags(str(self.image_path)) + + user_turn = self.model.calls[0]["messages"][1]["content"] + self.assertEqual(user_turn[0]["text"], "describe tersely") + self.assertTrue(user_turn[1]["image_url"]["url"].startswith("data:image/png")) + + def test_explicit_prompt_argument_wins(self): + service = self._loaded("tag") + + service.generate_tags(str(self.image_path), prompt="one-off") + + self.assertEqual( + self.model.calls[0]["messages"][1]["content"][0]["text"], "one-off" + ) + + +class TestGenerateTagsErrors(EngineTestCase): + + def test_no_model_loaded(self): + service = self._service(("gguf", FakeGgufModel("x"))) + with self.assertRaises(RuntimeError): + service.generate_tags(str(self.image_path)) + + def test_missing_image_error_does_not_embed_the_full_path(self): + _touch_model_files(self.models_dir, "gguf") + service = self._service(("gguf", FakeGgufModel("x"))) + service.load_model("gguf") + missing = self.tmp / "nope.png" + + with self.assertRaises(FileNotFoundError) as ctx: + service.generate_tags(str(missing)) + + self.assertNotIn(str(self.tmp), ctx.exception.strerror) + self.assertEqual(ctx.exception.filename, str(missing)) + + +class TestModelsStatus(EngineTestCase): + + def test_status_reflects_files_on_disk(self): + _touch_model_files(self.models_dir, "wd14-vit") + gguf_dir = self.models_dir / MODELS["gguf"]["subdir"] + gguf_dir.mkdir(parents=True) + (gguf_dir / MODELS["gguf"]["filename"]).write_bytes(b"stub") + service = AutoTagService(models_dir=self.models_dir) + + status = service.get_models_status() + + self.assertTrue(status["wd14-vit"]["downloaded"]) + self.assertTrue(status["wd14-vit"]["model_path"].endswith("model.onnx")) + self.assertFalse(status["wd14-swinv2"]["downloaded"]) + self.assertTrue(status["gguf"]["model_exists"]) + self.assertFalse(status["gguf"]["mmproj_exists"]) + self.assertFalse(status["gguf"]["downloaded"]) + self.assertFalse(status["hf"]["downloaded"]) + + def test_hf_local_directory_counts_as_downloaded(self): + hf_dir = self.models_dir / MODELS["hf"]["subdir"] + hf_dir.mkdir(parents=True) + (hf_dir / "config.json").write_text("{}") + service = AutoTagService(models_dir=self.models_dir) + + status = service.get_models_status() + + self.assertTrue(status["hf"]["downloaded"]) + self.assertEqual( + os.path.realpath(status["hf"]["model_path"]), os.path.realpath(str(hf_dir)) + ) + + +if __name__ == "__main__": + unittest.main()