nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: autotag model loading is injectable

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-08 18:01:34 -07:00
parent 2d3fc1cd08
commit a6d4facfbf
4 changed files with 756 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/autotag_routes.py
+249 -240
View File
@@ -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."""
py/autotag.py
+24 -9
View File
@@ -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
tests/test_api_autotag.py
+152
View File
@@ -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()
tests/test_autotag_engine.py
+331
View File
@@ -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()