Vito Sansevero
489d449ec8
- py/model_gate.py: one process-wide, non-blocking gate for loading,
swapping, running, downloading and unloading a tagging model. Worker
calls run shielded; a release requested while one still runs takes
effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
gate and answer 409 {code: "model-busy"} while another operation holds
it; their bodies are otherwise unchanged. /autotag/models gains
suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
first, GGUF/HF (name, None) in output order, first `limit` distinct
case-insensitively; generate_tags keeps its return; the image file is
closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
pinned capture, with strict per-engine options, never downloads
(409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
metadata over the inspection budgets no longer blocks it.
538 lines
22 KiB
Python
538 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 from the legacy Auto-tag page."
|
|
|
|
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)
|
|
|
|
|
|
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()
|