Vito
ae755d2ce0
* feat(system): configure auto-tagging natively * test(system): compare Auto-tagging with the concept against the real backend
557 lines
22 KiB
Python
557 lines
22 KiB
Python
"""Scored top-N tag suggestions for a captured image (#249 B2).
|
||
|
||
``POST /prompt_manager/image-annotations/suggest`` runs the selected legacy
|
||
engine on the private capture (never the current file), behind the one
|
||
process-wide model gate shared with the legacy Auto-tag routes. It never
|
||
downloads a model and never writes. No real model is loaded: the engine is a
|
||
fake service.
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import math
|
||
import os
|
||
import threading
|
||
import unittest
|
||
import uuid
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
try:
|
||
from test_api_annotations import AnnotationAPITestCase
|
||
from test_api_images import make_png
|
||
except ImportError:
|
||
from tests.test_api_annotations import AnnotationAPITestCase
|
||
from tests.test_api_images import make_png
|
||
|
||
try: # guarded so a missing module fails assertions, not test loading
|
||
from py import model_gate
|
||
except ImportError:
|
||
model_gate = None
|
||
|
||
from py import config as config_module
|
||
|
||
AutotagConfig = getattr(config_module, "AutotagConfig", None)
|
||
|
||
SUGGEST = "/prompt_manager/image-annotations/suggest"
|
||
WAIT = 30.0
|
||
WD14 = {"general_threshold": 0.35, "character_threshold": 0.85}
|
||
LLM = {"instructions": "Generate tags."}
|
||
UNAVAILABLE = "Download it in System › Auto-tagging."
|
||
|
||
SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [
|
||
(f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150)
|
||
]
|
||
|
||
|
||
class FakeTagger:
|
||
"""Stand-in for ``AutoTagService`` with the suggestion entry point."""
|
||
|
||
default_prompt = "server default prompt"
|
||
custom_prompt = "legacy prompt"
|
||
wd14_general_threshold = 0.35
|
||
wd14_character_threshold = 0.85
|
||
|
||
def __init__(self):
|
||
self.models_config = {
|
||
"wd14-swinv2": {},
|
||
"wd14-vit": {},
|
||
"gguf": {},
|
||
"hf": {},
|
||
}
|
||
self.downloaded = {name: True for name in self.models_config}
|
||
self.loaded = None
|
||
self.load_calls = []
|
||
self.unload_calls = 0
|
||
self.calls = []
|
||
self.read_paths = []
|
||
self.read_bytes = []
|
||
self.download_model = MagicMock(side_effect=AssertionError("downloaded"))
|
||
self.load_error = None
|
||
self.block = None # (entered, proceed) events
|
||
|
||
def get_models_status(self):
|
||
return {
|
||
name: {"downloaded": self.downloaded[name], "model_path": None}
|
||
for name in self.models_config
|
||
}
|
||
|
||
def is_model_loaded(self):
|
||
return self.loaded is not None
|
||
|
||
def get_loaded_model_type(self):
|
||
return self.loaded
|
||
|
||
def load_model(self, model_type, use_gpu=True):
|
||
self.load_calls.append((model_type, use_gpu))
|
||
if self.load_error is not None:
|
||
raise self.load_error
|
||
self.loaded = model_type
|
||
return True
|
||
|
||
def unload_model(self):
|
||
self.unload_calls += 1
|
||
self.loaded = None
|
||
|
||
def generate_tags(self, image_path, **kwargs):
|
||
return ["legacy"]
|
||
|
||
def generate_tag_suggestions(
|
||
self,
|
||
image_path,
|
||
*,
|
||
limit,
|
||
prompt=None,
|
||
general_threshold=None,
|
||
character_threshold=None,
|
||
):
|
||
self.read_paths.append(image_path)
|
||
with open(image_path, "rb") as handle:
|
||
self.read_bytes.append(handle.read())
|
||
self.calls.append(
|
||
{
|
||
"limit": limit,
|
||
"prompt": prompt,
|
||
"general_threshold": general_threshold,
|
||
"character_threshold": character_threshold,
|
||
"model": self.loaded,
|
||
}
|
||
)
|
||
if self.block is not None:
|
||
entered, proceed = self.block
|
||
entered.set()
|
||
proceed.wait(WAIT)
|
||
if self.loaded.startswith("wd14"):
|
||
return list(SCORED[:limit])
|
||
return [(name, None) for name, _ in SCORED[:limit]]
|
||
|
||
|
||
class SuggestTestCase(AnnotationAPITestCase):
|
||
async def get_application(self):
|
||
app = await super().get_application()
|
||
self.assertIsNotNone(model_gate, "py.model_gate is not implemented")
|
||
self.assertIsNotNone(AutotagConfig, "AutotagConfig is not implemented")
|
||
self.tagger = FakeTagger()
|
||
patcher = patch("py.autotag.get_autotag_service", return_value=self.tagger)
|
||
patcher.start()
|
||
self.addCleanup(patcher.stop)
|
||
AutotagConfig.SUGGESTION_LIMIT = 10
|
||
self.addCleanup(setattr, AutotagConfig, "SUGGESTION_LIMIT", 10)
|
||
return app
|
||
|
||
async def tearDownAsync(self):
|
||
leaked = model_gate.holder() if model_gate is not None else None
|
||
if leaked is not None:
|
||
model_gate.release(leaked)
|
||
await super().tearDownAsync()
|
||
self.assertIsNone(leaked, "the suggest route left the model gate held")
|
||
|
||
async def observed(self, path="sub/a.png", size=(8, 6)):
|
||
make_png(self.output_dir / path, size=size)
|
||
return await self.observe_ok(path)
|
||
|
||
def suggest_body(self, observed, model_type="wd14-vit", options=None, **changes):
|
||
if options is None:
|
||
options = dict(WD14 if model_type.startswith("wd14") else LLM)
|
||
return {
|
||
"target": dict(observed["target"]),
|
||
"capture_token": observed["capture"]["token"],
|
||
"expected_epoch": observed["library_epoch"],
|
||
"model_type": model_type,
|
||
"options": options,
|
||
**changes,
|
||
}
|
||
|
||
async def suggest(self, body):
|
||
response = await self.client.post(SUGGEST, json=body)
|
||
return response.status, await self.json_of(response)
|
||
|
||
async def suggest_ok(self, body):
|
||
status, data = await self.suggest(body)
|
||
self.assertEqual(status, 200, data)
|
||
return data
|
||
|
||
def capture_path(self, observed):
|
||
return self.api._captures.get(observed["capture"]["token"]).path
|
||
|
||
async def wait_until_free(self):
|
||
for _ in range(3000):
|
||
if model_gate.holder() is None:
|
||
return
|
||
await asyncio.sleep(0.01)
|
||
self.fail(f"model gate still held by {model_gate.holder()}")
|
||
|
||
|
||
class TestSuggestions(SuggestTestCase):
|
||
async def test_wd14_suggestions_are_scored_top_n_from_the_capture(self):
|
||
observed = await self.observed()
|
||
|
||
body = await self.suggest_ok(self.suggest_body(observed))
|
||
|
||
self.assertEqual(
|
||
body["suggestions"][:2],
|
||
[
|
||
{"name": "sandstone", "confidence": 0.98},
|
||
{"name": "desert", "confidence": 0.96},
|
||
],
|
||
)
|
||
self.assertEqual(len(body["suggestions"]), 10)
|
||
self.assertEqual(body["model_type"], "wd14-vit")
|
||
self.assertEqual(
|
||
body["options"],
|
||
{
|
||
"general_threshold": 0.35,
|
||
"character_threshold": 0.85,
|
||
"use_gpu": True,
|
||
"suggestion_limit": 10,
|
||
},
|
||
)
|
||
self.assertEqual(self.tagger.read_paths, [self.capture_path(observed)])
|
||
self.assertEqual(
|
||
self.tagger.calls[0],
|
||
{
|
||
"limit": 10,
|
||
"prompt": None,
|
||
"general_threshold": 0.35,
|
||
"character_threshold": 0.85,
|
||
"model": "wd14-vit",
|
||
},
|
||
)
|
||
self.assertEqual(self.tagger.load_calls, [("wd14-vit", True)])
|
||
self.tagger.download_model.assert_not_called()
|
||
self.assertIsNone(model_gate.holder())
|
||
self.assert_no_paths(body)
|
||
|
||
async def test_suggestions_keep_the_requested_limit(self):
|
||
observed = await self.observed()
|
||
for limit in (1, 10, 100):
|
||
with self.subTest(limit=limit):
|
||
options = {**WD14, "suggestion_limit": limit}
|
||
body = await self.suggest_ok(
|
||
self.suggest_body(observed, options=options)
|
||
)
|
||
self.assertEqual(len(body["suggestions"]), limit)
|
||
self.assertEqual(body["options"]["suggestion_limit"], limit)
|
||
# Omitted: the persisted autotag.defaults.suggestion_limit.
|
||
AutotagConfig.SUGGESTION_LIMIT = 3
|
||
body = await self.suggest_ok(self.suggest_body(observed))
|
||
self.assertEqual(len(body["suggestions"]), 3)
|
||
self.assertEqual(body["options"]["suggestion_limit"], 3)
|
||
|
||
async def test_llm_suggestions_send_instructions_and_null_confidence(self):
|
||
observed = await self.observed()
|
||
self.tagger.custom_prompt = "legacy prompt"
|
||
for model_type in ("gguf", "hf"):
|
||
with self.subTest(model_type=model_type):
|
||
options = {"instructions": f"tags for {model_type}", "use_gpu": False}
|
||
body = await self.suggest_ok(
|
||
self.suggest_body(observed, model_type, options)
|
||
)
|
||
self.assertEqual(body["model_type"], model_type)
|
||
self.assertEqual(body["options"], {**options, "suggestion_limit": 10})
|
||
self.assertEqual(
|
||
body["suggestions"][:2],
|
||
[
|
||
{"name": "sandstone", "confidence": None},
|
||
{"name": "desert", "confidence": None},
|
||
],
|
||
)
|
||
call = self.tagger.calls[-1]
|
||
self.assertEqual(call["prompt"], f"tags for {model_type}")
|
||
self.assertIsNone(call["general_threshold"])
|
||
self.assertEqual(self.tagger.load_calls[-1], (model_type, False))
|
||
self.assertEqual(self.tagger.custom_prompt, "legacy prompt")
|
||
|
||
async def test_loaded_model_is_reused_and_another_is_swapped_in(self):
|
||
observed = await self.observed()
|
||
self.tagger.loaded = "wd14-vit"
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
self.assertEqual(self.tagger.load_calls, [])
|
||
await self.suggest_ok(self.suggest_body(observed, "wd14-swinv2"))
|
||
self.assertEqual(self.tagger.load_calls, [("wd14-swinv2", True)])
|
||
self.assertEqual(self.tagger.calls[-1]["model"], "wd14-swinv2")
|
||
|
||
async def test_suggest_uses_capture_bytes_not_current_file(self):
|
||
observed = await self.observed()
|
||
original = (self.output_dir / "sub" / "a.png").read_bytes()
|
||
make_png(self.output_dir / "sub" / "a.png", size=(9, 9), color=(0, 0, 255))
|
||
self.assertNotEqual((self.output_dir / "sub" / "a.png").read_bytes(), original)
|
||
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
|
||
capture_path = self.capture_path(observed)
|
||
self.assertEqual(self.tagger.read_paths, [capture_path])
|
||
self.assertEqual(self.tagger.read_bytes, [original])
|
||
self.assertNotEqual(
|
||
os.path.normcase(capture_path),
|
||
os.path.normcase(str(self.output_dir / "sub" / "a.png")),
|
||
)
|
||
|
||
async def test_suggest_writes_nothing(self):
|
||
observed = await self.observed()
|
||
before = self.db_snapshot()
|
||
rows = self.bytes_rows()
|
||
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
await self.suggest_ok(self.suggest_body(observed, "gguf"))
|
||
|
||
self.assertEqual(self.db_snapshot(), before)
|
||
self.assertEqual(self.bytes_rows(), rows)
|
||
# The capture stays usable for Save after Generate.
|
||
self.assertIsNotNone(self.api._captures.get(observed["capture"]["token"]))
|
||
|
||
|
||
class TestModelAvailability(SuggestTestCase):
|
||
async def test_suggest_never_downloads_and_reports_model_unavailable(self):
|
||
observed = await self.observed()
|
||
self.tagger.downloaded["wd14-vit"] = False
|
||
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
|
||
self.assertEqual((status, body["code"]), (409, "model-unavailable"))
|
||
self.assertIn(UNAVAILABLE, body["error"])
|
||
self.assertIn("WD14 ViT", body["error"])
|
||
self.tagger.download_model.assert_not_called()
|
||
self.assertEqual(self.tagger.load_calls, [])
|
||
self.assertEqual(self.tagger.calls, [])
|
||
self.assertIsNone(model_gate.holder())
|
||
|
||
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}")
|
||
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
|
||
self.assertEqual((status, body["code"]), (500, "suggest-failed"))
|
||
self.assert_no_paths(body)
|
||
self.assertIsNone(model_gate.holder())
|
||
entry = self.api._captures._entries[observed["capture"]["token"]]
|
||
self.assertEqual(entry.uses, 0)
|
||
|
||
async def test_gate_work_that_never_starts_releases_the_capture_pin(self):
|
||
observed = await self.observed()
|
||
loop = asyncio.get_running_loop()
|
||
real = loop.run_in_executor
|
||
|
||
def refuse_gate_work(executor, func, *args):
|
||
if getattr(func, "__qualname__", "").startswith("ModelGate.run"):
|
||
raise RuntimeError("cannot schedule new futures after shutdown")
|
||
return real(executor, func, *args)
|
||
|
||
with patch.object(loop, "run_in_executor", refuse_gate_work):
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
|
||
self.assertEqual((status, body["code"]), (500, "suggest-failed"))
|
||
self.assertIsNone(model_gate.holder())
|
||
entry = self.api._captures._entries[observed["capture"]["token"]]
|
||
self.assertEqual(entry.uses, 0)
|
||
self.assertEqual(self.tagger.calls, [])
|
||
|
||
|
||
class TestModelGate(SuggestTestCase):
|
||
async def _blocked_suggest(self, observed, model_type="wd14-vit"):
|
||
entered, proceed = threading.Event(), threading.Event()
|
||
self.tagger.block = (entered, proceed)
|
||
task = asyncio.ensure_future(
|
||
self.client.post(SUGGEST, json=self.suggest_body(observed, model_type))
|
||
)
|
||
await asyncio.get_running_loop().run_in_executor(None, entered.wait, WAIT)
|
||
self.assertTrue(entered.is_set())
|
||
return task, proceed
|
||
|
||
async def test_suggest_reports_busy_while_legacy_batch_runs(self):
|
||
observed = await self.observed()
|
||
prompt_id = self._save_prompt("batch prompt")
|
||
self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png"))
|
||
started, proceed = threading.Event(), threading.Event()
|
||
|
||
def slow_generate(image_path, **kwargs):
|
||
started.set()
|
||
proceed.wait(WAIT)
|
||
return ["batch"]
|
||
|
||
self.tagger.generate_tags = slow_generate
|
||
batch = asyncio.ensure_future(
|
||
self.client.post("/prompt_manager/autotag/start?model_type=wd14-vit")
|
||
)
|
||
try:
|
||
await asyncio.get_running_loop().run_in_executor(None, started.wait, WAIT)
|
||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||
self.assertEqual((status, body["code"]), (409, "model-busy"))
|
||
self.assertEqual(self.tagger.calls, [])
|
||
self.assertEqual(self.tagger.loaded, "wd14-vit") # not swapped
|
||
finally:
|
||
proceed.set()
|
||
response = await batch
|
||
await response.read()
|
||
await self.wait_until_free()
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
|
||
async def test_legacy_unload_returns_model_busy_while_native_generate_runs(self):
|
||
observed = await self.observed()
|
||
task, proceed = await self._blocked_suggest(observed)
|
||
try:
|
||
response = await self.client.post("/prompt_manager/autotag/unload")
|
||
self.assertEqual(response.status, 409)
|
||
self.assertEqual((await response.json())["code"], "model-busy")
|
||
response = await self.client.post(
|
||
"/prompt_manager/autotag/single",
|
||
json={"root_index": 0, "path": "sub/a.png", "model_type": "gguf"},
|
||
)
|
||
self.assertEqual(response.status, 409)
|
||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||
self.assertEqual((status, body["code"]), (409, "model-busy"))
|
||
self.assertEqual(self.tagger.unload_calls, 0)
|
||
self.assertEqual(self.tagger.loaded, "wd14-vit")
|
||
finally:
|
||
proceed.set()
|
||
response = await task
|
||
self.assertEqual(response.status, 200)
|
||
await self.wait_until_free()
|
||
self.assertEqual([call["model"] for call in self.tagger.calls], ["wd14-vit"])
|
||
|
||
async def test_gate_held_after_client_disconnect_until_executor_finishes(self):
|
||
observed = await self.observed()
|
||
token = observed["capture"]["token"]
|
||
task, proceed = await self._blocked_suggest(observed)
|
||
task.cancel()
|
||
with self.assertRaises(asyncio.CancelledError):
|
||
await task
|
||
try:
|
||
# The client is gone; the inference still runs, so the gate and
|
||
# the capture pin stay until the worker thread exits.
|
||
self.assertIsNotNone(model_gate.holder())
|
||
response = await self.client.post("/prompt_manager/autotag/unload")
|
||
self.assertEqual(response.status, 409)
|
||
self.assertEqual(self.api._captures._entries[token].uses, 1)
|
||
finally:
|
||
proceed.set()
|
||
await self.wait_until_free()
|
||
for _ in range(3000):
|
||
if self.api._captures._entries[token].uses == 0:
|
||
break
|
||
await asyncio.sleep(0.01)
|
||
self.assertEqual(self.api._captures._entries[token].uses, 0)
|
||
response = await self.client.post("/prompt_manager/autotag/unload")
|
||
self.assertEqual(response.status, 200)
|
||
|
||
|
||
class TestRequestValidation(SuggestTestCase):
|
||
async def test_strict_option_validation(self):
|
||
observed = await self.observed()
|
||
bad_options = {
|
||
"wd14-vit": [
|
||
{**WD14, "general_threshold": True},
|
||
{**WD14, "general_threshold": math.nan},
|
||
{**WD14, "general_threshold": math.inf},
|
||
{**WD14, "general_threshold": "0.3"},
|
||
{**WD14, "general_threshold": 1.5},
|
||
{**WD14, "character_threshold": -0.1},
|
||
{"general_threshold": 0.35},
|
||
{**WD14, "instructions": "x"},
|
||
{**WD14, "unknown": 1},
|
||
{**WD14, "use_gpu": "yes"},
|
||
{**WD14, "suggestion_limit": 0},
|
||
{**WD14, "suggestion_limit": 101},
|
||
{**WD14, "suggestion_limit": 2.5},
|
||
{**WD14, "suggestion_limit": "10"},
|
||
{**WD14, "suggestion_limit": True},
|
||
],
|
||
"gguf": [
|
||
{**LLM, "general_threshold": 0.35},
|
||
{**LLM, "character_threshold": 0.85},
|
||
{"instructions": "x" * 4001},
|
||
{"instructions": "\ud800"}, # not encodable as UTF-8
|
||
{"instructions": " "},
|
||
{"instructions": 5},
|
||
{},
|
||
{**LLM, "suggestion_limit": 0},
|
||
],
|
||
}
|
||
for model_type, cases in bad_options.items():
|
||
for options in cases:
|
||
with self.subTest(model_type=model_type, options=repr(options)[:80]):
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, model_type, options)
|
||
)
|
||
self.assertEqual(status, 400, body)
|
||
for body in (
|
||
{**self.suggest_body(observed), "extra": 1},
|
||
{k: v for k, v in self.suggest_body(observed).items() if k != "options"},
|
||
self.suggest_body(observed, options=[]),
|
||
self.suggest_body(observed, model_type="ollama"),
|
||
self.suggest_body(observed, model_type=7, options=dict(WD14)),
|
||
self.suggest_body(observed, capture_token="../x"),
|
||
self.suggest_body(observed, expected_epoch="nope"),
|
||
self.suggest_body(observed, target={"root_index": 0}),
|
||
):
|
||
with self.subTest(body=json.dumps(body)[:120]):
|
||
status, _ = await self.suggest(body)
|
||
self.assertEqual(status, 400)
|
||
# The edges that are allowed.
|
||
await self.suggest_ok(
|
||
self.suggest_body(
|
||
observed,
|
||
options={
|
||
"general_threshold": 0,
|
||
"character_threshold": 1,
|
||
"suggestion_limit": 100,
|
||
},
|
||
)
|
||
)
|
||
await self.suggest_ok(
|
||
self.suggest_body(observed, "hf", options={"instructions": "x" * 4000})
|
||
)
|
||
self.assertEqual(len(self.tagger.calls), 2)
|
||
self.tagger.download_model.assert_not_called()
|
||
|
||
async def test_capture_token_must_match_target_and_epoch(self):
|
||
observed = await self.observed()
|
||
make_png(self.output_dir / "sub" / "b.png")
|
||
other = await self.observe_ok("sub/b.png")
|
||
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, capture_token="x" * 43)
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "capture-expired"))
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, capture_token=other["capture"]["token"])
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "capture-expired"))
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, expected_epoch=str(uuid.uuid4()))
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "storage-changed"))
|
||
self.assertEqual(body["library_epoch"], self.epoch)
|
||
status, body = await self.suggest(
|
||
self.suggest_body(
|
||
observed, target={**observed["target"], "root_scope": "0" * 64}
|
||
)
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "source-changed"))
|
||
self.api._captures.release(observed["capture"]["token"])
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
self.assertEqual((status, body["code"]), (409, "capture-expired"))
|
||
self.assertEqual(self.tagger.calls, [])
|
||
for token in (observed, other):
|
||
entry = self.api._captures._entries.get(token["capture"]["token"])
|
||
self.assertTrue(entry is None or entry.uses == 0)
|
||
|
||
async def test_video_and_oversized_dimensions_unsupported_media(self):
|
||
(self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 4)
|
||
video = await self.observe_ok("clip.mp4")
|
||
self.assertFalse(video["capture"]["generate"]["supported"])
|
||
big = await self.observed("big.png", size=(4097, 1))
|
||
self.assertFalse(big["capture"]["generate"]["supported"])
|
||
|
||
for observed in (video, big):
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
self.assertEqual((status, body["code"]), (415, "unsupported-media"))
|
||
self.assertEqual(body["error"], observed["capture"]["generate"]["reason"])
|
||
self.assertEqual(self.tagger.calls, [])
|
||
self.assertEqual(self.tagger.load_calls, [])
|
||
self.assertIsNone(model_gate.holder())
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|