diff --git a/py/autotag.py b/py/autotag.py index 33170526..0c807f06 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -16,6 +16,12 @@ from PIL import Image # A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend) ModelLoader = Callable[[str, bool], Tuple[str, Any]] +# Suggestions keep only names a Save accepts (the vocabulary tag rules). +try: + from ..database.vocabulary import validate_vocabulary_tag_name +except ImportError: + from database.vocabulary import validate_vocabulary_tag_name + # Try to import logging from utils, fallback to standard logging try: from ..utils.logging_config import get_logger @@ -795,8 +801,10 @@ class AutoTagService: WD14 pairs carry the real probability, best first; language models (GGUF, HF) have no score, so their tags keep the output order with ``None``. Names are distinct case-insensitively (first one wins). - ``prompt`` is used for this call only; ``custom_prompt`` is never - changed. + A name Save would reject (over 50 characters, a control character + such as a newline, a comma) is skipped before the dedupe and the + ``limit`` count. ``prompt`` is used for this call only; + ``custom_prompt`` is never changed. """ if type(limit) is not int or not 1 <= limit <= MAX_SUGGESTION_LIMIT: raise ValueError( @@ -807,7 +815,11 @@ class AutoTagService: ) suggestions: List[Tuple[str, Optional[float]]] = [] seen = set() - for name, score in scored: + for raw, score in scored: + try: + name = validate_vocabulary_tag_name(raw) + except ValueError: + continue key = name.casefold() if key in seen: continue diff --git a/tests/test_autotag_engine.py b/tests/test_autotag_engine.py index a16f8ba2..c76217ab 100644 --- a/tests/test_autotag_engine.py +++ b/tests/test_autotag_engine.py @@ -387,6 +387,22 @@ class TestTagSuggestions(EngineTestCase): suggestions, [("zebra", None), ("apple", None), ("mango_tree", None)] ) + def test_suggestions_skip_names_save_would_reject_before_the_limit(self): + from database.annotations import normalize_image_tags + + too_long = "x" * 60 + service = self._gguf(f"{too_long}, sky\ncloud, zebra, apple, mango, kiwi, plum") + + suggestions = service.generate_tag_suggestions( + str(self.image_path), limit=3, prompt="tags please" + ) + + # The 60-character part and the part holding a newline are dropped + # before the count, so three valid chips still come back. + names = [name for name, _ in suggestions] + self.assertEqual(names, ["zebra", "apple", "mango"]) + self.assertEqual(normalize_image_tags(names), names) + def test_generate_tags_return_shape_unchanged(self): service = self._wd14() self.assertEqual(