Vito
1b8fbdbb75
* feat(annotations): add scoped observe, capture preview and tag save routes
POST /prompt_manager/image-annotations/observe captures the file named by
the listing identity {root_index, root_scope, path} (scope checked against
the current roots before any file access) and answers the saved tag set of
exactly those bytes, refreshing a matching record's verified fingerprint.
GET/DELETE .../capture/{token} serve the private copy (no-store, nosniff,
pinned with acquire/release_use while it streams) and release it.
PUT .../tags re-hashes the current file and writes only when it still
holds the captured bytes, under the storage epoch and tag revision
(file-changed, storage-changed, tag-conflict, tag-limit, replayed).
Bodies are capped at 64 KiB with unknown fields refused; envelopes never
carry server paths or locators. Restores invalidate every capture and the
host's cleanup removes the capture directory.
* feat(autotag): scored suggestions with a configurable count behind a shared model gate
- 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.
* feat(images): match image tags in the Tags filter, search and inspector
The workspace Images listing now admits a file under a Tags filter when
every requested tag is carried by its linked positive or its current image
tag set, and q also matches image tag names (case-folded substring). Image
tags count only while a byte record is verified at the file's present
(size, mtime_ns) (spec amendment S1): the tagged records are streamed
once, only listed files carrying a filter tag or q hit are stat'ed, and
their linked positives are read through the file_path index so tags,
category, rating and the positive count combine.
The rooted inspector envelope gains a validated image_tags facet
(current / none / unavailable) read at the inspection's own stable stat.
A failed image-tag read fails a Tags filter like a failed link read and
degrades a q search to no image tags; plain listings read none. Legacy
gallery and search responses are unchanged.
* fix(annotations): answer a fixed coded message when the filtered image list fails
A failed prompt-link or image-tag read under a Tags, category or rating
filter returned the raw exception text (SQL or database messages) in the
500 body. The workspace listing now answers code listing-failed with a
fixed sentence and logs the traceback. The legacy output listing and the
inspector image-tag read now log with exc_info as well.
* fix(annotations): drop suggestion names that Save would reject
A JoyCaption part over 50 characters or holding a newline came back as a
preselected chip, so the first Save of the dialog failed with a 400 and
nothing was stored. generate_tag_suggestions now applies the vocabulary
tag-name rules before the case-insensitive dedupe and the limit count, so
the requested number of valid chips still comes back.
* fix(annotations): stop changing the process warnings filters during capture
header_dimensions wrapped Image.open in warnings.catch_warnings with an
error filter for DecompressionBombWarning. catch_warnings is not
thread-safe, so two overlapping captures could leave that error filter
installed for every later Pillow open in the ComfyUI process. Only
Pillow's DecompressionBombError is now mapped to HeaderTooLarge; a size
that only warns is refused by the existing inference pixel cap.
* fix(annotations): free pins on cancellation, serve preview ranges, code every storage error
- A preview cancelled while its file opens closes the handle and only
then ends the pin, so a doomed capture is never unlinked while open.
- An observe cancelled while capturing or reading releases the capture
instead of leaving it until the idle expiry.
- The model gate takes an on_skipped callback that runs once when the
work never starts; Generate passes the capture's release_use, so an
executor refusal no longer leaks a pin against the capture budget.
- The preview answers a single byte range with 206 and Content-Range
(416 for an unsatisfiable or multi-range header), still streamed inside
the pin, so Safari and WebKit can play video captures.
- _coded maps tag-limit (400), metadata-conflict and vocabulary-conflict,
and falls back to 409 with a fixed sentence for any other code instead
of raising KeyError.
* fix(annotations): deny case and Unicode spellings of a listed file
On macOS (case- and normalization-insensitive APFS) a crafted target such
as A.PNG or an NFD name opened the listed a.png or NFC file while
realpath kept the requested spelling, so the tags were stored under a
second locator the Images filter never shows. Capture and verify now
require each path part to be spelled as its directory lists it and deny
any other case or Unicode spelling as an alias (409 source-changed, like
a symlink alias). A missing file with no other spelling keeps its usual
outcome.
* test(annotations): one shared folder_paths stub for the API test modules
test_api_images and test_api_prompts installed a MagicMock folder_paths,
so whenever one of them was imported before test_api_autotag the later
patch("folder_paths.get_output_directory") in test_api_history_outputs
left a mock attribute behind and two test_diagnostics_config tests
failed. tests/host_stubs.py now installs one real-module stub, used by
test_api_images, test_api_prompts and test_api_autotag (and through
test_api_images by the annotation modules, whose copy is gone), so the
modules pass in either import order.
* fix(annotations): derive the suggestion cap from the image tag set cap
py/autotag.py now takes MAX_SUGGESTION_LIMIT from database.annotations
MAX_IMAGE_TAGS. The two copies that cannot import it (py/config.py loads
before the database, utils/inspection_response.py stays storage-free) are
pinned equal by a test, so the 100 limits cannot drift apart.
* docs(testing): record capture aliases, preview ranges and the PR B follow-ups
TESTING.md now states that case and Unicode spellings are aliases, that
the preview serves byte ranges inside the capture pin, where the
cancellation tests live, and two tracked follow-ups: the legacy batch's
ungated cleanup unload under its own last inference, and the three files
over the 800-line ceiling (split py/autotag.py first).
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 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)
|
|
|
|
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()
|