From 4d2e9e098458352fe0abe2fea93e8cefef9bf68b Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sat, 28 Mar 2026 17:13:51 +0100 Subject: [PATCH] Update DINOv3 node with better tag filtering options and implications mode --- config.json | 4 +- nodes/rrtagger_dino.py | 16 +- redrocket/classifier_dino.py | 60 +++++-- redrocket/tag_dino_manager.py | 323 +++++++++++++++++++++++++++++++--- 4 files changed, 358 insertions(+), 45 deletions(-) diff --git a/config.json b/config.json index e1f45b5..b706660 100644 --- a/config.json +++ b/config.json @@ -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" } } } diff --git a/nodes/rrtagger_dino.py b/nodes/rrtagger_dino.py index 62aad6f..259dc91 100644 --- a/nodes/rrtagger_dino.py +++ b/nodes/rrtagger_dino.py @@ -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, diff --git a/redrocket/classifier_dino.py b/redrocket/classifier_dino.py index fc2a000..8d3a1ad 100644 --- a/redrocket/classifier_dino.py +++ b/redrocket/classifier_dino.py @@ -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, ) diff --git a/redrocket/tag_dino_manager.py b/redrocket/tag_dino_manager.py index 15b3dc2..a119d24 100644 --- a/redrocket/tag_dino_manager.py +++ b/redrocket/tag_dino_manager.py @@ -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, + )