Joe Gores
8dbdf51d97
Add two new nodes that can be wired together to replace the original NeuralBooru node, giving users control over LLM and validation steps independently. The original NeuralBooru class is unchanged. - llm_node.py: NeuralBooruLLM handles LLM communication (Ollama native API + OpenAI-compatible), outputs (raw_tags, is_fallback) - validator_node.py: NeuralBooruValidator handles tag validation + template application, outputs (prompt, dropped_tags, tags) - __init__.py: register new nodes via try/except ImportError block - tests: 32 new tests (19 LLM + 13 validator), all 91 pass - workflows/AI_Anime_split.json: example workflow using split nodes
137 lines
5.8 KiB
Python
137 lines
5.8 KiB
Python
"""
|
|
NeuralBooruValidator - Danbooru tag validation node.
|
|
|
|
Validates comma-separated tags against the real Danbooru vocabulary,
|
|
applies aliasing/sorting/fuzzy matching, and wraps the result in a
|
|
prompt template. Has no LLM awareness; works standalone with hand-typed
|
|
tags or wired after NeuralBooruLLM for the full original pipeline.
|
|
|
|
When wired in sequence: NeuralBooruLLM.raw_tags -> NeuralBooruValidator.tags
|
|
the two nodes reproduce the original NeuralBooru node exactly.
|
|
"""
|
|
|
|
try:
|
|
from . import tags as tagdb
|
|
from . import (TEMPLATE_PRESETS, NOVA_ANIME_XL_TEMPLATE,
|
|
_apply_template, _parse_categories)
|
|
except ImportError:
|
|
import tags as tagdb
|
|
from neuralbooru import (TEMPLATE_PRESETS, NOVA_ANIME_XL_TEMPLATE,
|
|
_apply_template, _parse_categories)
|
|
|
|
|
|
class NeuralBooruValidator:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"tags": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Comma-separated tags to validate. From NeuralBooruLLM "
|
|
"or typed manually."
|
|
}),
|
|
"prompt_template": ("STRING", {
|
|
"multiline": True,
|
|
"default": NOVA_ANIME_XL_TEMPLATE,
|
|
"tooltip": "Final prompt wrapper; {prompt} is replaced with the "
|
|
"validated tags. If {prompt} is missing, tags are appended."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"is_fallback": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "When True (set by NeuralBooruLLM on error), skip "
|
|
"validation and pass tags through into the template. "
|
|
"Reproduces the original node's use_input_text fallback."
|
|
}),
|
|
"validate_tags": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Check every tag against the real Danbooru vocabulary."
|
|
}),
|
|
"strict_tags": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Drop candidates that match no real tag. Off keeps them raw."
|
|
}),
|
|
"fuzzy_cutoff": ("FLOAT", {
|
|
"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"tooltip": "Remap near-miss tags by similarity (0.8+ recommended). "
|
|
"0 disables fuzzy matching."
|
|
}),
|
|
"min_post_count": ("INT", {
|
|
"default": 0, "min": 0, "max": 1000000, "step": 100,
|
|
"tooltip": "Drop tags with fewer Danbooru posts than this. 0 keeps all."
|
|
}),
|
|
"max_tags": ("INT", {
|
|
"default": 0, "min": 0, "max": 200, "step": 1,
|
|
"tooltip": "Keep at most this many tags. 0 means no limit."
|
|
}),
|
|
"exclude_categories": ("STRING", {
|
|
"multiline": False,
|
|
"default": "",
|
|
"tooltip": "Tag categories to drop, comma-separated: "
|
|
"artist, character, copyright, meta, general."
|
|
}),
|
|
"template_preset": (["custom"] + sorted(TEMPLATE_PRESETS), {
|
|
"default": "custom",
|
|
"tooltip": "Quality-tag wrapper for common model families. "
|
|
"custom uses the prompt_template field above."
|
|
}),
|
|
"sort_tags": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Reorder tags Danbooru-style: people counts, "
|
|
"character, copyright, artist, general, meta. "
|
|
"Booru-trained models respond to this order."
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("prompt", "dropped_tags", "tags")
|
|
OUTPUT_TOOLTIPS = (
|
|
"Validated tags wrapped in the prompt template, for your CLIP encoder.",
|
|
"Candidates that failed validation, comma-separated.",
|
|
"The validated tags alone, without the template.",
|
|
)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "NeuralBooru"
|
|
OUTPUT_NODE = False
|
|
|
|
def generate(self, tags, prompt_template,
|
|
is_fallback=False, validate_tags=True, strict_tags=True,
|
|
fuzzy_cutoff=0.0, min_post_count=0, max_tags=0,
|
|
exclude_categories="", template_preset="custom",
|
|
sort_tags=True):
|
|
|
|
template = TEMPLATE_PRESETS.get(template_preset, prompt_template)
|
|
|
|
# Error fallback or validation disabled: pass through into template
|
|
if is_fallback or not validate_tags:
|
|
final = _apply_template(template, tags)
|
|
print(f"[NeuralBooruValidator] {final[:160]}...")
|
|
return (final, "", tags)
|
|
|
|
db = tagdb.get_db()
|
|
if db is None:
|
|
# Tag database unavailable; pass through
|
|
final = _apply_template(template, tags)
|
|
print(f"[NeuralBooruValidator] {final[:160]}...")
|
|
return (final, "", tags)
|
|
|
|
validated, _kept, dropped = db.validate(
|
|
tags,
|
|
strict=strict_tags,
|
|
fuzzy_cutoff=fuzzy_cutoff,
|
|
min_post_count=min_post_count,
|
|
max_tags=max_tags,
|
|
exclude_categories=_parse_categories(exclude_categories),
|
|
sort_tags=sort_tags,
|
|
)
|
|
|
|
final = _apply_template(template, validated)
|
|
print(f"[NeuralBooruValidator] {final[:160]}...")
|
|
if dropped:
|
|
print(f"[NeuralBooruValidator] dropped {len(dropped)} non-tags: {dropped}")
|
|
return (final, ", ".join(dropped), validated)
|