nodes/silveroxides-ComfyUI-RR-JointTagger
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update DINOv3 node with better tag filtering options and implications mode

Browse Source
This commit is contained in:
silveroxides
2026-03-28 17:13:51 +01:00
parent 6cddec2eb6
commit 4d2e9e0984
4 changed files with 358 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.json
+3 -1
View File
@@ -51,7 +51,9 @@
"tagger-proto": {
"name": "DINOv3 Tagger Proto",
"url": "{HF_ENDPOINT}/lodestones/tagger-experiment/resolve/main/tagger_proto.safetensors",
"vocab_url": "{HF_ENDPOINT}/lodestones/tagger-experiment/resolve/main/tagger_vocab.json"
"vocab_url": "{HF_ENDPOINT}/lodestones/tagger-experiment/resolve/main/tagger_vocab.json",
"data_url": "{HF_ENDPOINT}/RedRocket/JTP-3/resolve/main/data",
"data_tag_file": "jtp-3-hydra-tags.csv"
}
}
}
nodes/rrtagger_dino.py
+11 -5
View File
@@ -34,12 +34,14 @@ class DINOv3Tagger(ComfyNodeABC):
return {"required": {
"image": (IO.IMAGE,),
"model": (models, {"default": models[0] if models else "tagger-proto"}),
"mode": (["topk", "threshold"], {"default": "topk", "tooltip": "Tag selection mode. 'topk' returns the top K tags; 'threshold' returns all tags above a score."}),
"mode": (["topk", "threshold"], {"default": "topk", "tooltip": "Tag selection mode. 'topk' returns the top K tags; 'threshold' returns all tags above a score. Threshold always acts as a minimum score floor in both modes."}),
"topk": (IO.INT, {"default": 40, "min": 1, "max": 500, "step": 1, "display": "slider", "tooltip": "Number of top tags to return (when mode is 'topk')."}),
"threshold": (IO.FLOAT, {"default": 0.35, "min": 0.01, "max": 0.99, "step": 0.01, "display": "slider", "tooltip": "Minimum score to include a tag (when mode is 'threshold')."}),
"threshold": (IO.FLOAT, {"default": 0.35, "min": 0.01, "max": 0.99, "step": 0.01, "display": "slider", "tooltip": "Minimum score to include a tag. Applied in both modes."}),
"max_size": (IO.INT, {"default": 1024, "min": 64, "max": 4096, "step": 16, "tooltip": "Long-edge pixel cap. Image is resized to fit, snapped to 16px multiples."}),
"implications_mode": (["off", "inherit", "constrain", "remove", "constrain-remove"], {"default": "off", "tooltip": "How to handle implied tags (e.g. if 'cat' is present, 'feline' is implied). Requires tag metadata CSV."}),
"max_tags": (IO.INT, {"default": 0, "min": 0, "max": 500, "step": 1, "tooltip": "Maximum number of tags in the final output (applied after implications). 0 = unlimited."}),
"exclude_tags": (IO.STRING, {"multiline": True, "tooltip": "Comma-separated tags to exclude from output."}),
"replace_underscore": (IO.BOOLEAN, {"default": True, "tooltip": "Replace underscores with spaces in tags."}),
"exclude_categories": (IO.STRING, {"multiline": True, "tooltip": "Comma-separated categories to exclude (e.g. copyright, character, species, meta, lore)."}),
"trailing_comma": (IO.BOOLEAN, {"default": False, "tooltip": "Add a trailing comma to the tag string."}),
"prefix": (IO.STRING, {"default": "", "tooltip": "Text to prepend to the tags output."}),
"seed": (IO.INT, {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Seed for deterministic execution."}),
@@ -60,9 +62,11 @@ class DINOv3Tagger(ComfyNodeABC):
topk: int,
threshold: float,
max_size: int,
implications_mode: str,
max_tags: int,
seed: int,
exclude_tags: str = "",
replace_underscore: bool = True,
exclude_categories: str = "",
trailing_comma: bool = False,
prefix: str = "",
) -> Dict[str, Any]:
@@ -87,8 +91,10 @@ class DINOv3Tagger(ComfyNodeABC):
topk=topk,
threshold=threshold,
max_size=max_size,
implications_mode=implications_mode,
exclude_tags=exclude_tags,
replace_underscore=replace_underscore,
exclude_categories=exclude_categories,
max_tags=max_tags,
trailing_comma=trailing_comma,
prefix=prefix,
seed=seed,
redrocket/classifier_dino.py
+42 -18
View File
@@ -43,8 +43,10 @@ class DINOv3Inference(metaclass=Singleton):
topk: int = 40,
threshold: float = 0.35,
max_size: int = 1024,
implications_mode: str = "off",
exclude_tags: str = "",
replace_underscore: bool = True,
exclude_categories: str = "",
max_tags: int = 0,
trailing_comma: bool = False,
prefix: str = "",
seed: int = 0,
@@ -64,13 +66,17 @@ class DINOv3Inference(metaclass=Singleton):
topk:
Number of top tags to return when *mode* is ``"topk"``.
threshold:
Minimum score to include when *mode* is ``"threshold"``.
Minimum score to include.
max_size:
Long-edge pixel cap (snapped to 16px multiples).
implications_mode:
How to handle implied tags: ``"inherit"``, ``"constrain"``,
``"remove"``, ``"constrain-remove"``, or ``"off"``.
exclude_tags:
Comma-separated tags to exclude from the output.
replace_underscore:
Replace underscores with spaces in output tags.
exclude_categories:
Comma-separated category names to exclude
(e.g. ``"copyright, character, meta"``).
trailing_comma:
Append a trailing comma to the output tag string.
prefix:
@@ -85,8 +91,8 @@ class DINOv3Inference(metaclass=Singleton):
# Build a params hash for caching
params_string = (
f"{model_name}|{mode}|{topk}|{threshold}|{max_size}|"
f"{exclude_tags}|{replace_underscore}|{trailing_comma}|"
f"{prefix}|{seed}"
f"{implications_mode}|{exclude_tags}|{exclude_categories}|"
f"{max_tags}|{trailing_comma}|{prefix}|{seed}"
)
params_key = hashlib.sha256(params_string.encode()).hexdigest()
@@ -112,6 +118,15 @@ class DINOv3Inference(metaclass=Singleton):
ComfyLogger().log("DINOv3: vocab is empty", "ERROR", True)
return "", {}
# Load tag metadata (implications/categories) if needed and available
if implications_mode != "off" and not DINOv3TagManager.has_metadata(model_name):
csv_path = os.path.join(
DINOv3TagManager().tags_basepath, f"{model_name}-tags.csv",
)
if not os.path.exists(csv_path):
DINOv3TagManager.download_metadata(model_name)
DINOv3TagManager.load_metadata(model_name)
# ----------------------------------------------------------
# 2. Ensure model is available
# ----------------------------------------------------------
@@ -157,22 +172,27 @@ class DINOv3Inference(metaclass=Singleton):
scores = torch.sigmoid(logits.float())
# ----------------------------------------------------------
# 5. Top-K or threshold selection
# 5. Top-K + threshold selection (both always applied)
# ----------------------------------------------------------
# Always apply threshold as a minimum score floor
mask = scores >= threshold
valid_indices = mask.nonzero(as_tuple=True)[0]
valid_values = scores[valid_indices]
# Sort descending
order = valid_values.argsort(descending=True)
valid_indices = valid_indices[order]
valid_values = valid_values[order]
# In topk mode, additionally cap the number of results
if mode == "topk":
k = min(topk, num_tags)
values, indices = scores.topk(k)
else:
# Threshold mode
mask = scores >= threshold
indices = mask.nonzero(as_tuple=True)[0]
values = scores[indices]
order = values.argsort(descending=True)
indices, values = indices[order], values[order]
k = min(topk, len(valid_indices))
valid_indices = valid_indices[:k]
valid_values = valid_values[:k]
raw_results: List[Tuple[str, float]] = [
(idx2tag[i], float(v))
for i, v in zip(indices.tolist(), values.tolist())
for i, v in zip(valid_indices.tolist(), valid_values.tolist())
]
# ----------------------------------------------------------
@@ -180,8 +200,12 @@ class DINOv3Inference(metaclass=Singleton):
# ----------------------------------------------------------
tag_string, scores_dict = DINOv3TagManager.process_tags(
results=raw_results,
model_name=model_name,
implications_mode=implications_mode,
threshold=threshold,
max_tags=max_tags,
exclude_tags=exclude_tags,
replace_underscore=replace_underscore,
exclude_categories=exclude_categories,
trailing_comma=trailing_comma,
prefix=prefix,
)
redrocket/tag_dino_manager.py
+302 -21
View File
@@ -1,13 +1,18 @@
"""DINOv3 Tagger vocabulary manager — download, load, process tags."""
"""DINOv3 Tagger vocabulary manager — download, load, process tags.
Supports optional implications metadata loaded from a JTP-3-style CSV.
All tag keys are normalised to spaces (no underscores) to match DINOv3's
native output format.
"""
from __future__ import annotations
import csv
import gc
import json
import os
import re
import traceback
from typing import Callable, Dict, List, Optional, Tuple
from typing import Callable, Dict, List, Optional, Set, Tuple
import requests
from tqdm import tqdm
@@ -17,11 +22,24 @@ from ..helpers.config import ComfyExtensionConfig
from ..helpers.logger import ComfyLogger
from ..helpers.metaclasses import Singleton
# Category IDs used in the JTP-3 tag CSV
TAG_CATEGORIES = {
"general": 0,
"copyright": 3,
"character": 4,
"species": 5,
"meta": 7,
"lore": 8,
}
class DINOv3TagManager(metaclass=Singleton):
"""Singleton manager for DINOv3 tagger vocabulary.
"""Singleton manager for DINOv3 tagger vocabulary and optional metadata.
The vocabulary is a JSON file with ``{"idx2tag": ["tag0", "tag1", ...]}``.
Tag metadata (implications, categories) is loaded from a JTP-3-style CSV
and normalised so all keys use spaces instead of underscores.
Cached under the ``tags_dino`` namespace.
"""
@@ -49,6 +67,10 @@ class DINOv3TagManager(metaclass=Singleton):
def is_loaded(cls, model_name: str) -> bool:
return ComfyCache.get(f"tags_dino.{model_name}") is not None
@classmethod
def has_metadata(cls, model_name: str) -> bool:
return ComfyCache.get(f"tags_dino.{model_name}.metadata") is not None
@classmethod
def get_idx2tag(cls, model_name: str) -> Optional[List[str]]:
return ComfyCache.get(f"tags_dino.{model_name}.idx2tag")
@@ -59,7 +81,7 @@ class DINOv3TagManager(metaclass=Singleton):
return len(idx2tag) if idx2tag else 0
# ------------------------------------------------------------------
# Load
# Load vocabulary (JSON)
# ------------------------------------------------------------------
@classmethod
@@ -95,7 +117,57 @@ class DINOv3TagManager(metaclass=Singleton):
return False
# ------------------------------------------------------------------
# Download
# Load metadata / implications (CSV)
# ------------------------------------------------------------------
@classmethod
def load_metadata(cls, model_name: str) -> bool:
"""Load tag metadata (categories + implications) from a JTP-3-style
CSV. All tag keys and implication references are normalised to use
spaces instead of underscores to match DINOv3 output format.
"""
csv_path = os.path.join(cls().tags_basepath, f"{model_name}-tags.csv")
if not os.path.exists(csv_path):
ComfyLogger().log(
f"DINOv3 tag metadata CSV not found: {csv_path} "
"(implications will be unavailable)",
type="WARNING", always=True,
)
return False
try:
metadata: Dict[str, Tuple[int, List[str]]] = {}
with open(csv_path, "r", encoding="utf-8", newline="") as f:
reader = csv.DictReader(f)
for row in reader:
# Normalise tag: underscores -> spaces
tag = row["tag"].replace("_", " ")
category = int(row.get("category", 0))
# Normalise implications the same way
raw_impl = row.get("implications", "")
implications = [
imp.replace("_", " ")
for imp in raw_impl.split()
if imp
]
metadata[tag] = (category, implications)
ComfyCache.set(f"tags_dino.{model_name}.metadata", metadata)
ComfyLogger().log(
f"DINOv3 tag metadata loaded: {len(metadata):,} entries for {model_name}",
type="INFO", always=True,
)
return True
except Exception as e:
ComfyLogger().log(
f"Error loading DINOv3 tag metadata: {e}\n{traceback.format_exc()}",
type="ERROR", always=True,
)
return False
# ------------------------------------------------------------------
# Download vocabulary (JSON)
# ------------------------------------------------------------------
@classmethod
@@ -162,56 +234,251 @@ class DINOv3TagManager(metaclass=Singleton):
return False
# ------------------------------------------------------------------
# Tag processing
# Download metadata (CSV)
# ------------------------------------------------------------------
@classmethod
def download_metadata(cls, model_name: str) -> bool:
"""Download tag metadata CSV from the configured ``data_url``."""
os.makedirs(cls().tags_basepath, exist_ok=True)
config = ComfyExtensionConfig().get()
hf_endpoint: str = config.get("huggingface_endpoint", "https://huggingface.co")
if not hf_endpoint.startswith("https://"):
hf_endpoint = f"https://{hf_endpoint}"
hf_endpoint = hf_endpoint.rstrip("/")
data_url = ComfyExtensionConfig().get(
property=f"models_dino.{model_name}.data_url",
)
if not data_url:
ComfyLogger().log(
f"No data_url for DINOv3 model {model_name} — "
"implications will not be available",
type="WARNING", always=True,
)
return False
data_url = data_url.replace("{HF_ENDPOINT}", hf_endpoint)
if not data_url.endswith("/"):
data_url += "/"
# The CSV filename follows the convention from the data_url source.
# The config's data_url_tag_file overrides the default if set.
tag_file = ComfyExtensionConfig().get(
property=f"models_dino.{model_name}.data_url_tag_file",
)
if not tag_file or not isinstance(tag_file, str):
# Default: derive from the JTP-3 naming convention
tag_file = ComfyExtensionConfig().get(
property=f"models_dino.{model_name}.data_tag_file",
)
if not tag_file or not isinstance(tag_file, str):
tag_file = "jtp-3-hydra-tags.csv"
url = f"{data_url}{tag_file}"
dest = os.path.join(cls().tags_basepath, f"{model_name}-tags.csv")
ComfyLogger().log(
f"Downloading DINOv3 tag metadata from {url}",
type="INFO", always=True,
)
try:
response = requests.get(url, stream=True)
if response.status_code != 200:
ComfyLogger().log(
f"Failed to download tag metadata: HTTP {response.status_code}",
type="WARNING", always=True,
)
return False
total_size = int(response.headers.get("content-length", 0))
block_size = 1024
with open(dest, "wb") as f, tqdm(
desc=f"{model_name}-tags.csv",
total=total_size,
unit="iB",
unit_scale=True,
unit_divisor=1024,
) as bar:
downloaded = 0
for chunk in response.iter_content(block_size):
f.write(chunk)
bar.update(len(chunk))
downloaded += len(chunk)
if cls().download_progress_callback and total_size > 0:
cls().download_progress_callback(
int(downloaded / total_size * 100),
f"{model_name}-tags.csv",
)
if cls().download_complete_callback:
cls().download_complete_callback(f"{model_name}-tags.csv")
return True
except Exception as e:
ComfyLogger().log(
f"Failed to download DINOv3 tag metadata: {e}\n{traceback.format_exc()}",
type="WARNING", always=True,
)
return False
# ------------------------------------------------------------------
# Tag processing (with optional implications)
# ------------------------------------------------------------------
@classmethod
def process_tags(
cls,
results: List[Tuple[str, float]],
model_name: str = "",
implications_mode: str = "off",
threshold: float = 0.0,
max_tags: int = 0,
exclude_tags: str = "",
replace_underscore: bool = True,
exclude_categories: str = "",
trailing_comma: bool = False,
prefix: str = "",
) -> Tuple[str, Dict[str, float]]:
"""Format raw (tag, score) list into a tag string and scores dict.
DINOv3 tags are already space-separated (no underscores), so no
``replace_underscore`` option is needed.
Parameters
----------
results:
List of ``(tag, score)`` tuples, already sorted by descending score.
model_name:
Config key — needed to look up metadata for implications.
implications_mode:
``"inherit"``, ``"constrain"``, ``"remove"``,
``"constrain-remove"``, or ``"off"``.
threshold:
Minimum score floor — re-applied after implications processing
so that constrained tags dropping below threshold are removed.
max_tags:
Maximum number of tags to return (applied after implications and
filtering). 0 = unlimited.
exclude_tags:
Comma-separated tags to exclude.
replace_underscore:
Replace underscores with spaces in output tags.
exclude_categories:
Comma-separated category names to exclude
(e.g. ``"copyright, character, meta"``).
trailing_comma:
Append a trailing comma to the tag string.
prefix:
Text to prepend.
"""
# Build exclusion set (normalised to spaces)
excluded = set()
metadata = ComfyCache.get(f"tags_dino.{model_name}.metadata")
# Build exclusion set (normalised to lowercase spaces)
excluded: Set[str] = set()
if exclude_tags:
for t in exclude_tags.split(","):
t = t.strip()
if t:
excluded.add(t.replace("_", " ").lower())
tag_list: List[str] = []
scores_dict: Dict[str, float] = {}
# Parse excluded categories
excluded_cats: Set[int] = set()
if exclude_categories and metadata:
for cat in exclude_categories.split(","):
cat = cat.strip().lower()
if cat in TAG_CATEGORIES:
excluded_cats.add(TAG_CATEGORIES[cat])
for tag, score in results:
# Normalised check for exclusion
# Build labels dict from results
labels: Dict[str, float] = {tag: score for tag, score in results}
# ----- Implications processing (only if metadata is available) -----
if metadata and implications_mode != "off":
def inherit_implications(lbls: Dict[str, float], antecedent: str) -> None:
if antecedent not in metadata:
return
p = lbls.get(antecedent, 0.0)
for consequent in metadata[antecedent][1]:
if lbls.get(consequent, 0.0) < p:
lbls[consequent] = p
inherit_implications(lbls, consequent)
def constrain_implications(
lbls: Dict[str, float], antecedent: str, _target: Optional[str] = None,
) -> None:
if _target is None:
_target = antecedent
if antecedent not in metadata:
return
for consequent in metadata[antecedent][1]:
p = lbls.get(consequent, 0.0)
if lbls.get(_target, 0.0) > p:
lbls[_target] = p
constrain_implications(lbls, consequent, _target=_target)
def remove_implications(lbls: Dict[str, float], antecedent: str) -> None:
if antecedent not in metadata:
return
for consequent in metadata[antecedent][1]:
lbls.pop(consequent, None)
remove_implications(lbls, consequent)
tag_keys = list(labels.keys())
if implications_mode == "inherit":
for tag in tag_keys:
inherit_implications(labels, tag)
elif implications_mode in ("constrain", "constrain-remove"):
for tag in tag_keys:
constrain_implications(labels, tag)
# ----- Filtering (threshold re-applied after implications) -----
filtered: Dict[str, float] = {}
for tag, prob in labels.items():
# Re-apply threshold (catches scores lowered by constrain)
if prob < threshold:
continue
# Check exclusion by name
tag_check = tag.replace("_", " ").lower()
if tag_check in excluded:
continue
scores_dict[tag] = round(score, 4)
# Check exclusion by category
if excluded_cats and metadata and tag in metadata:
cat_id = metadata[tag][0]
if cat_id in excluded_cats:
continue
display_tag = tag
if replace_underscore:
display_tag = display_tag.replace("_", " ")
tag_list.append(display_tag)
filtered[tag] = prob
# ----- Remove implications step (after filtering) -----
if metadata and implications_mode in ("remove", "constrain-remove"):
keys_to_check = list(filtered.keys())
for tag in keys_to_check:
if tag in filtered:
if tag not in metadata:
continue
for consequent in metadata[tag][1]:
filtered.pop(consequent, None)
# Recurse
cls._remove_implications_recursive(
filtered, consequent, metadata,
)
# ----- Format output -----
sorted_items = sorted(filtered.items(), key=lambda x: x[1], reverse=True)
# Apply max_tags cap (post-implications limit)
if max_tags > 0 and len(sorted_items) > max_tags:
sorted_items = sorted_items[:max_tags]
tag_list: List[str] = []
scores_dict: Dict[str, float] = {}
for tag, score in sorted_items:
scores_dict[tag] = round(score, 4)
tag_list.append(tag)
# Prefix handling
if prefix:
@@ -226,3 +493,17 @@ class DINOv3TagManager(metaclass=Singleton):
tag_string += ","
return tag_string, scores_dict
@staticmethod
def _remove_implications_recursive(
labels: Dict[str, float],
antecedent: str,
metadata: Dict[str, Tuple[int, List[str]]],
) -> None:
if antecedent not in metadata:
return
for consequent in metadata[antecedent][1]:
labels.pop(consequent, None)
DINOv3TagManager._remove_implications_recursive(
labels, consequent, metadata,
)