feat(caption): add WD14 tagger with Booru Tags tab

Add SmilingWolf's WD14/WaifuDiffusion tagger models for anime/illustration
tagging as a new "Booru Tags" tab in the Caption panel.

- Support 9 models (v2 and v3 variants) via HuggingFace
- ONNX backend chosen due to safetensors v3 variants exhibiting
  unacceptable accuracy loss
- Separate thresholds for general/character tags
- Batch processing with progress bar
- Consolidate debug env var to SD_INTERROGATE_DEBUG

modules/interrogate/moondream3.py

@@ -11,7 +11,7 @@ from modules.interrogate import vqa_detection
 
 
 # Debug logging - function-based to avoid circular import
-debug_enabled = os.environ.get('SD_VQA_DEBUG', None) is not None
+debug_enabled = os.environ.get('SD_INTERROGATE_DEBUG', None) is not None
 
 def debug(*args, **kwargs):
     if debug_enabled:

modules/interrogate/openclip.py

@@ -1,4 +1,5 @@
 import os
+import time
 from collections import namedtuple
 import threading
 import re
@@ -7,6 +8,23 @@ from PIL import Image
 from modules import devices, paths, shared, errors, sd_models
 
 
+debug_enabled = os.environ.get('SD_INTERROGATE_DEBUG', None) is not None
+debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
+
+
+def _apply_blip2_fix(model, processor):
+    """Apply compatibility fix for BLIP2 models with newer transformers versions."""
+    from transformers import AddedToken
+    if not hasattr(model.config, 'num_query_tokens'):
+        return
+    processor.num_query_tokens = model.config.num_query_tokens
+    image_token = AddedToken("<image>", normalized=False, special=True)
+    processor.tokenizer.add_tokens([image_token], special_tokens=True)
+    model.resize_token_embeddings(len(processor.tokenizer), pad_to_multiple_of=64)
+    model.config.image_token_index = len(processor.tokenizer) - 1
+    debug_log(f'CLIP load: applied BLIP2 tokenizer fix num_query_tokens={model.config.num_query_tokens}')
+
+
 caption_models = {
     'blip-base': 'Salesforce/blip-image-captioning-base',
     'blip-large': 'Salesforce/blip-image-captioning-large',
@@ -79,10 +97,14 @@ def load_interrogator(clip_model, blip_model):
     clip_interrogator.clip_interrogator.CAPTION_MODELS = caption_models
     global ci # pylint: disable=global-statement
     if ci is None:
-        shared.log.debug(f'Interrogate load: clip="{clip_model}" blip="{blip_model}"')
+        t0 = time.time()
+        device = devices.get_optimal_device()
+        cache_path = os.path.join(paths.models_path, 'Interrogator')
+        shared.log.info(f'CLIP load: clip="{clip_model}" blip="{blip_model}" device={device}')
+        debug_log(f'CLIP load: cache_path="{cache_path}" max_length={shared.opts.interrogate_clip_max_length} chunk_size={shared.opts.interrogate_clip_chunk_size} flavor_count={shared.opts.interrogate_clip_flavor_count} offload={shared.opts.interrogate_offload}')
         interrogator_config = clip_interrogator.Config(
-            device=devices.get_optimal_device(),
-            cache_path=os.path.join(paths.models_path, 'Interrogator'),
+            device=device,
+            cache_path=cache_path,
             clip_model_name=clip_model,
             caption_model_name=blip_model,
             quiet=True,
@@ -93,22 +115,39 @@ def load_interrogator(clip_model, blip_model):
             caption_offload=shared.opts.interrogate_offload,
         )
         ci = clip_interrogator.Interrogator(interrogator_config)
+        if blip_model.startswith('blip2-'):
+            _apply_blip2_fix(ci.caption_model, ci.caption_processor)
+        shared.log.debug(f'CLIP load: time={time.time()-t0:.2f}s')
     elif clip_model != ci.config.clip_model_name or blip_model != ci.config.caption_model_name:
-        ci.config.clip_model_name = clip_model
-        ci.config.clip_model = None
-        ci.load_clip_model()
-        ci.config.caption_model_name = blip_model
-        ci.config.caption_model = None
-        ci.load_caption_model()
+        t0 = time.time()
+        if clip_model != ci.config.clip_model_name:
+            shared.log.info(f'CLIP load: clip="{clip_model}" reloading')
+            debug_log(f'CLIP load: previous clip="{ci.config.clip_model_name}"')
+            ci.config.clip_model_name = clip_model
+            ci.config.clip_model = None
+            ci.load_clip_model()
+        if blip_model != ci.config.caption_model_name:
+            shared.log.info(f'CLIP load: blip="{blip_model}" reloading')
+            debug_log(f'CLIP load: previous blip="{ci.config.caption_model_name}"')
+            ci.config.caption_model_name = blip_model
+            ci.config.caption_model = None
+            ci.load_caption_model()
+            if blip_model.startswith('blip2-'):
+                _apply_blip2_fix(ci.caption_model, ci.caption_processor)
+        shared.log.debug(f'CLIP load: time={time.time()-t0:.2f}s')
+    else:
+        debug_log(f'CLIP: models already loaded clip="{clip_model}" blip="{blip_model}"')
 
 
 def unload_clip_model():
     if ci is not None and shared.opts.interrogate_offload:
+        shared.log.debug('CLIP unload: offloading models to CPU')
         sd_models.move_model(ci.caption_model, devices.cpu)
         sd_models.move_model(ci.clip_model, devices.cpu)
         ci.caption_offloaded = True
         ci.clip_offloaded = True
         devices.torch_gc()
+        debug_log('CLIP unload: complete')
 
 
 def interrogate(image, mode, caption=None):
@@ -119,6 +158,8 @@ def interrogate(image, mode, caption=None):
     if image is None:
         return ''
     image = image.convert("RGB")
+    t0 = time.time()
+    debug_log(f'CLIP: mode="{mode}" image_size={image.size} caption={caption is not None} min_flavors={shared.opts.interrogate_clip_min_flavors} max_flavors={shared.opts.interrogate_clip_max_flavors}')
     if mode == 'best':
         prompt = ci.interrogate(image, caption=caption, min_flavors=shared.opts.interrogate_clip_min_flavors, max_flavors=shared.opts.interrogate_clip_max_flavors, )
     elif mode == 'caption':
@@ -131,22 +172,27 @@ def interrogate(image, mode, caption=None):
         prompt = ci.interrogate_negative(image, max_flavors=shared.opts.interrogate_clip_max_flavors)
     else:
         raise RuntimeError(f"Unknown mode {mode}")
+    debug_log(f'CLIP: mode="{mode}" time={time.time()-t0:.2f}s result="{prompt[:100]}..."' if len(prompt) > 100 else f'CLIP: mode="{mode}" time={time.time()-t0:.2f}s result="{prompt}"')
     return prompt
 
 
 def interrogate_image(image, clip_model, blip_model, mode):
     jobid = shared.state.begin('Interrogate CLiP')
+    t0 = time.time()
+    shared.log.info(f'CLIP: mode="{mode}" clip="{clip_model}" blip="{blip_model}" image_size={image.size if image else None}')
     try:
         if shared.sd_loaded:
             from modules.sd_models import apply_balanced_offload # prevent circular import
             apply_balanced_offload(shared.sd_model)
+            debug_log('CLIP: applied balanced offload to sd_model')
         load_interrogator(clip_model, blip_model)
         image = image.convert('RGB')
         prompt = interrogate(image, mode)
         devices.torch_gc()
+        shared.log.debug(f'CLIP: complete time={time.time()-t0:.2f}s')
     except Exception as e:
         prompt = f"Exception {type(e)}"
-        shared.log.error(f'Interrogate: {e}')
+        shared.log.error(f'CLIP: {e}')
         errors.display(e, 'Interrogate')
     shared.state.end(jobid)
     return prompt
@@ -162,8 +208,11 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod
         from modules.files_cache import list_files
         files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp', '.jxl'], recursive=recursive))
     if len(files) == 0:
-        shared.log.warning('Interrogate batch: type=clip no images')
+        shared.log.warning('CLIP batch: no images found')
         return ''
+    t0 = time.time()
+    shared.log.info(f'CLIP batch: mode="{mode}" images={len(files)} clip="{clip_model}" blip="{blip_model}" write={write} append={append}')
+    debug_log(f'CLIP batch: recursive={recursive} files={files[:5]}{"..." if len(files) > 5 else ""}')
     jobid = shared.state.begin('Interrogate batch')
     prompts = []
 
@@ -171,6 +220,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod
     if write:
         file_mode = 'w' if not append else 'a'
         writer = BatchWriter(os.path.dirname(files[0]), mode=file_mode)
+        debug_log(f'CLIP batch: writing to "{os.path.dirname(files[0])}" mode="{file_mode}"')
     import rich.progress as rp
     pbar = rp.Progress(rp.TextColumn('[cyan]Caption:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
     with pbar:
@@ -179,6 +229,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod
             pbar.update(task, advance=1, description=file)
             try:
                 if shared.state.interrupted:
+                    shared.log.info('CLIP batch: interrupted')
                     break
                 image = Image.open(file).convert('RGB')
                 prompt = interrogate(image, mode)
@@ -186,19 +237,23 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod
                 if write:
                     writer.add(file, prompt)
             except OSError as e:
-                shared.log.error(f'Interrogate batch: {e}')
+                shared.log.error(f'CLIP batch: file="{file}" error={e}')
     if write:
         writer.close()
     ci.config.quiet = False
     unload_clip_model()
     shared.state.end(jobid)
+    shared.log.info(f'CLIP batch: complete images={len(prompts)} time={time.time()-t0:.2f}s')
     return '\n\n'.join(prompts)
 
 
 def analyze_image(image, clip_model, blip_model):
+    t0 = time.time()
+    shared.log.info(f'CLIP analyze: clip="{clip_model}" blip="{blip_model}" image_size={image.size if image else None}')
     load_interrogator(clip_model, blip_model)
     image = image.convert('RGB')
     image_features = ci.image_to_features(image)
+    debug_log(f'CLIP analyze: features shape={image_features.shape if hasattr(image_features, "shape") else "unknown"}')
     top_mediums = ci.mediums.rank(image_features, 5)
     top_artists = ci.artists.rank(image_features, 5)
     top_movements = ci.movements.rank(image_features, 5)
@@ -209,6 +264,7 @@ def analyze_image(image, clip_model, blip_model):
     movement_ranks = dict(sorted(zip(top_movements, ci.similarities(image_features, top_movements)), key=lambda x: x[1], reverse=True))
     trending_ranks = dict(sorted(zip(top_trendings, ci.similarities(image_features, top_trendings)), key=lambda x: x[1], reverse=True))
     flavor_ranks = dict(sorted(zip(top_flavors, ci.similarities(image_features, top_flavors)), key=lambda x: x[1], reverse=True))
+    shared.log.debug(f'CLIP analyze: complete time={time.time()-t0:.2f}s')
 
     # Format labels as text
     def format_category(name, ranks):

modules/interrogate/vqa.py

@@ -13,7 +13,7 @@ from modules.interrogate import vqa_detection
 
 
 # Debug logging - function-based to avoid circular import
-debug_enabled = os.environ.get('SD_VQA_DEBUG', None) is not None
+debug_enabled = os.environ.get('SD_INTERROGATE_DEBUG', None) is not None
 
 def debug(*args, **kwargs):
     if debug_enabled:

modules/interrogate/wd14.py

@@ -0,0 +1,535 @@
+# WD14/WaifuDiffusion Tagger - ONNX-based anime/illustration tagging
+# Based on SmilingWolf's tagger models: https://huggingface.co/SmilingWolf
+
+import os
+import re
+import time
+import threading
+import numpy as np
+from PIL import Image
+from modules import shared, devices, errors
+
+
+# Debug logging - enable with SD_INTERROGATE_DEBUG environment variable
+debug_enabled = os.environ.get('SD_INTERROGATE_DEBUG', None) is not None
+debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
+
+re_special = re.compile(r'([\\()])')
+load_lock = threading.Lock()
+
+# WD14 model repository mappings
+WD14_MODELS = {
+    # v3 models (latest, recommended)
+    "wd-eva02-large-tagger-v3": "SmilingWolf/wd-eva02-large-tagger-v3",
+    "wd-vit-tagger-v3": "SmilingWolf/wd-vit-tagger-v3",
+    "wd-convnext-tagger-v3": "SmilingWolf/wd-convnext-tagger-v3",
+    "wd-swinv2-tagger-v3": "SmilingWolf/wd-swinv2-tagger-v3",
+    # v2 models
+    "wd-v1-4-moat-tagger-v2": "SmilingWolf/wd-v1-4-moat-tagger-v2",
+    "wd-v1-4-swinv2-tagger-v2": "SmilingWolf/wd-v1-4-swinv2-tagger-v2",
+    "wd-v1-4-convnext-tagger-v2": "SmilingWolf/wd-v1-4-convnext-tagger-v2",
+    "wd-v1-4-convnextv2-tagger-v2": "SmilingWolf/wd-v1-4-convnextv2-tagger-v2",
+    "wd-v1-4-vit-tagger-v2": "SmilingWolf/wd-v1-4-vit-tagger-v2",
+}
+
+# Tag categories from selected_tags.csv
+CATEGORY_GENERAL = 0
+CATEGORY_CHARACTER = 4
+CATEGORY_RATING = 9
+
+
+class WD14Tagger:
+    """WD14/WaifuDiffusion Tagger using ONNX inference."""
+
+    def __init__(self):
+        self.session = None
+        self.tags = None
+        self.tag_categories = None
+        self.model_name = None
+        self.model_path = None
+        self.image_size = 448  # Standard for WD models
+
+    def load(self, model_name: str = None):
+        """Load the ONNX model and tags from HuggingFace."""
+        import huggingface_hub
+
+        if model_name is None:
+            model_name = shared.opts.wd14_model
+        if model_name not in WD14_MODELS:
+            shared.log.error(f'WD14: unknown model "{model_name}"')
+            return False
+
+        with load_lock:
+            if self.session is not None and self.model_name == model_name:
+                debug_log(f'WD14: model already loaded model="{model_name}"')
+                return True  # Already loaded
+
+            # Unload previous model if different
+            if self.model_name != model_name and self.session is not None:
+                debug_log(f'WD14: switching model from "{self.model_name}" to "{model_name}"')
+                self.unload()
+
+            repo_id = WD14_MODELS[model_name]
+            t0 = time.time()
+            shared.log.info(f'WD14 load: model="{model_name}" repo="{repo_id}"')
+
+            try:
+                # Download only ONNX model and tags CSV (skip safetensors/msgpack variants)
+                debug_log(f'WD14 load: downloading from HuggingFace cache_dir="{shared.opts.hfcache_dir}"')
+                self.model_path = huggingface_hub.snapshot_download(
+                    repo_id,
+                    cache_dir=shared.opts.hfcache_dir,
+                    allow_patterns=["model.onnx", "selected_tags.csv"],
+                )
+                debug_log(f'WD14 load: model_path="{self.model_path}"')
+
+                # Load ONNX model
+                model_file = os.path.join(self.model_path, "model.onnx")
+                if not os.path.exists(model_file):
+                    shared.log.error(f'WD14 load: model file not found: {model_file}')
+                    return False
+
+                import onnxruntime as ort
+                providers = []
+                if devices.backend == 'cuda':
+                    providers.append('CUDAExecutionProvider')
+                providers.append('CPUExecutionProvider')
+                debug_log(f'WD14 load: onnxruntime version={ort.__version__} providers={providers}')
+
+                self.session = ort.InferenceSession(model_file, providers=providers)
+                self.model_name = model_name
+
+                # Get actual providers used
+                actual_providers = self.session.get_providers()
+                debug_log(f'WD14 load: active providers={actual_providers}')
+
+                # Load tags from CSV
+                self._load_tags()
+
+                load_time = time.time() - t0
+                shared.log.debug(f'WD14 load: time={load_time:.2f}s tags={len(self.tags)}')
+                debug_log(f'WD14 load: input_name={self.session.get_inputs()[0].name} output_name={self.session.get_outputs()[0].name}')
+                return True
+
+            except Exception as e:
+                shared.log.error(f'WD14 load: failed error={e}')
+                errors.display(e, 'WD14 load')
+                self.unload()
+                return False
+
+    def _load_tags(self):
+        """Load tags and categories from selected_tags.csv."""
+        import csv
+
+        csv_path = os.path.join(self.model_path, "selected_tags.csv")
+        if not os.path.exists(csv_path):
+            shared.log.error(f'WD14 load: tags file not found: {csv_path}')
+            return
+
+        self.tags = []
+        self.tag_categories = []
+
+        with open(csv_path, 'r', encoding='utf-8') as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                self.tags.append(row['name'])
+                self.tag_categories.append(int(row['category']))
+
+        # Count tags by category
+        category_counts = {}
+        for cat in self.tag_categories:
+            category_counts[cat] = category_counts.get(cat, 0) + 1
+        debug_log(f'WD14 load: tag categories={category_counts}')
+
+    def unload(self):
+        """Unload the model and free resources."""
+        if self.session is not None:
+            shared.log.debug(f'WD14 unload: model="{self.model_name}"')
+            self.session = None
+            self.tags = None
+            self.tag_categories = None
+            self.model_name = None
+            self.model_path = None
+            devices.torch_gc(force=True)
+            debug_log('WD14 unload: complete')
+        else:
+            debug_log('WD14 unload: no model loaded')
+
+    def preprocess_image(self, image: Image.Image) -> np.ndarray:
+        """Preprocess image for WD14 model input.
+
+        - Resize to 448x448 (standard for WD models)
+        - Pad to square with white background
+        - Normalize to [0, 1] range
+        - BGR channel order (as used by these models)
+        """
+        original_size = image.size
+        original_mode = image.mode
+
+        # Convert to RGB if needed
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        # Pad to square with white background
+        w, h = image.size
+        max_dim = max(w, h)
+        pad_left = (max_dim - w) // 2
+        pad_top = (max_dim - h) // 2
+
+        padded = Image.new('RGB', (max_dim, max_dim), (255, 255, 255))
+        padded.paste(image, (pad_left, pad_top))
+
+        # Resize to model input size
+        if max_dim != self.image_size:
+            padded = padded.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+
+        # Convert to numpy array and normalize
+        img_array = np.array(padded, dtype=np.float32)
+
+        # Convert RGB to BGR (model expects BGR)
+        img_array = img_array[:, :, ::-1]
+
+        # Add batch dimension
+        img_array = np.expand_dims(img_array, axis=0)
+
+        debug_log(f'WD14 preprocess: original_size={original_size} mode={original_mode} padded_size={max_dim} output_shape={img_array.shape}')
+        return img_array
+
+    def predict(
+        self,
+        image: Image.Image,
+        general_threshold: float = None,
+        character_threshold: float = None,
+        include_rating: bool = None,
+        exclude_tags: str = None,
+        max_tags: int = None,
+        sort_alpha: bool = None,
+        use_spaces: bool = None,
+        escape_brackets: bool = None,
+    ) -> str:
+        """Run inference and return formatted tag string.
+
+        Args:
+            image: PIL Image to tag
+            general_threshold: Threshold for general tags (0-1)
+            character_threshold: Threshold for character tags (0-1)
+            include_rating: Whether to include rating tags
+            exclude_tags: Comma-separated tags to exclude
+            max_tags: Maximum number of tags to return
+            sort_alpha: Sort tags alphabetically vs by confidence
+            use_spaces: Use spaces instead of underscores
+            escape_brackets: Escape parentheses/brackets in tags
+
+        Returns:
+            Formatted tag string
+        """
+        t0 = time.time()
+
+        # Use settings defaults if not specified
+        if general_threshold is None:
+            general_threshold = shared.opts.wd14_general_threshold
+        if character_threshold is None:
+            character_threshold = shared.opts.wd14_character_threshold
+        if include_rating is None:
+            include_rating = shared.opts.wd14_include_rating
+        if exclude_tags is None:
+            exclude_tags = shared.opts.wd14_exclude_tags
+        if max_tags is None:
+            max_tags = shared.opts.wd14_max_tags
+        if sort_alpha is None:
+            sort_alpha = shared.opts.wd14_sort_alpha
+        if use_spaces is None:
+            use_spaces = shared.opts.wd14_use_spaces
+        if escape_brackets is None:
+            escape_brackets = shared.opts.wd14_escape
+
+        debug_log(f'WD14 predict: general_threshold={general_threshold} character_threshold={character_threshold} max_tags={max_tags} include_rating={include_rating} sort_alpha={sort_alpha}')
+
+        # Handle input variations
+        if isinstance(image, list):
+            image = image[0] if len(image) > 0 else None
+        if isinstance(image, dict) and 'name' in image:
+            image = Image.open(image['name'])
+        if image is None:
+            shared.log.error('WD14 predict: no image provided')
+            return ''
+
+        # Load model if needed
+        if self.session is None:
+            if not self.load():
+                return ''
+
+        # Preprocess image
+        img_input = self.preprocess_image(image)
+
+        # Run inference
+        t_infer = time.time()
+        input_name = self.session.get_inputs()[0].name
+        output_name = self.session.get_outputs()[0].name
+        probs = self.session.run([output_name], {input_name: img_input})[0][0]
+        infer_time = time.time() - t_infer
+        debug_log(f'WD14 predict: inference time={infer_time:.3f}s output_shape={probs.shape}')
+
+        # Build tag list with probabilities
+        tag_probs = {}
+        exclude_set = {x.strip().replace(' ', '_').lower() for x in exclude_tags.split(',') if x.strip()}
+        if exclude_set:
+            debug_log(f'WD14 predict: exclude_tags={exclude_set}')
+
+        general_count = 0
+        character_count = 0
+        rating_count = 0
+
+        for i, (tag_name, prob) in enumerate(zip(self.tags, probs)):
+            category = self.tag_categories[i]
+            tag_lower = tag_name.lower()
+
+            # Skip excluded tags
+            if tag_lower in exclude_set:
+                continue
+
+            # Apply category-specific thresholds
+            if category == CATEGORY_RATING:
+                if not include_rating:
+                    continue
+                # Always include rating if enabled
+                tag_probs[tag_name] = float(prob)
+                rating_count += 1
+            elif category == CATEGORY_CHARACTER:
+                if prob >= character_threshold:
+                    tag_probs[tag_name] = float(prob)
+                    character_count += 1
+            elif category == CATEGORY_GENERAL:
+                if prob >= general_threshold:
+                    tag_probs[tag_name] = float(prob)
+                    general_count += 1
+            else:
+                # Other categories use general threshold
+                if prob >= general_threshold:
+                    tag_probs[tag_name] = float(prob)
+
+        debug_log(f'WD14 predict: matched tags general={general_count} character={character_count} rating={rating_count} total={len(tag_probs)}')
+
+        # Sort tags
+        if sort_alpha:
+            sorted_tags = sorted(tag_probs.keys())
+        else:
+            sorted_tags = [t for t, _ in sorted(tag_probs.items(), key=lambda x: -x[1])]
+
+        # Limit number of tags
+        if max_tags > 0 and len(sorted_tags) > max_tags:
+            sorted_tags = sorted_tags[:max_tags]
+            debug_log(f'WD14 predict: limited to max_tags={max_tags}')
+
+        # Format output
+        result = []
+        for tag_name in sorted_tags:
+            formatted_tag = tag_name
+            if use_spaces:
+                formatted_tag = formatted_tag.replace('_', ' ')
+            if escape_brackets:
+                formatted_tag = re.sub(re_special, r'\\\1', formatted_tag)
+            if shared.opts.interrogate_score:
+                formatted_tag = f"({formatted_tag}:{tag_probs[tag_name]:.2f})"
+            result.append(formatted_tag)
+
+        output = ', '.join(result)
+        total_time = time.time() - t0
+        debug_log(f'WD14 predict: complete tags={len(result)} time={total_time:.2f}s result="{output[:100]}..."' if len(output) > 100 else f'WD14 predict: complete tags={len(result)} time={total_time:.2f}s result="{output}"')
+
+        return output
+
+    def tag(self, image: Image.Image, **kwargs) -> str:
+        """Alias for predict() to match deepbooru interface."""
+        return self.predict(image, **kwargs)
+
+
+# Global tagger instance
+tagger = WD14Tagger()
+
+
+def get_models() -> list:
+    """Return list of available WD14 model names."""
+    return list(WD14_MODELS.keys())
+
+
+def refresh_models() -> list:
+    """Refresh and return list of available models."""
+    # For now, just return the static list
+    # Could be extended to check for locally cached models
+    return get_models()
+
+
+def load_model(model_name: str = None) -> bool:
+    """Load the specified WD14 model."""
+    return tagger.load(model_name)
+
+
+def unload_model():
+    """Unload the current WD14 model."""
+    tagger.unload()
+
+
+def tag(image: Image.Image, model_name: str = None, **kwargs) -> str:
+    """Tag an image using WD14 tagger.
+
+    Args:
+        image: PIL Image to tag
+        model_name: Model to use (loads if needed)
+        **kwargs: Additional arguments passed to predict()
+
+    Returns:
+        Formatted tag string
+    """
+    t0 = time.time()
+    jobid = shared.state.begin('WD14 Tag')
+    shared.log.info(f'WD14: model="{model_name or tagger.model_name or shared.opts.wd14_model}" image_size={image.size if image else None}')
+
+    try:
+        if model_name and model_name != tagger.model_name:
+            tagger.load(model_name)
+        result = tagger.predict(image, **kwargs)
+        shared.log.debug(f'WD14: complete time={time.time()-t0:.2f}s tags={len(result.split(", ")) if result else 0}')
+    except Exception as e:
+        result = f"Exception {type(e)}"
+        shared.log.error(f'WD14: {e}')
+        errors.display(e, 'WD14 Tag')
+
+    shared.state.end(jobid)
+    return result
+
+
+def batch(
+    model_name: str,
+    batch_files: list,
+    batch_folder: str,
+    batch_str: str,
+    save_output: bool = True,
+    save_append: bool = False,
+    recursive: bool = False,
+    **kwargs
+) -> str:
+    """Process multiple images in batch mode.
+
+    Args:
+        model_name: Model to use
+        batch_files: List of file paths
+        batch_folder: Folder path from file picker
+        batch_str: Folder path as string
+        save_output: Save caption to .txt files
+        save_append: Append to existing caption files
+        recursive: Recursively process subfolders
+        **kwargs: Additional arguments passed to predict()
+
+    Returns:
+        Combined tag results
+    """
+    from pathlib import Path
+
+    # Load model
+    if model_name:
+        tagger.load(model_name)
+    elif tagger.session is None:
+        tagger.load()
+
+    # Collect image files
+    image_files = []
+    image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
+
+    # From file picker
+    if batch_files:
+        for f in batch_files:
+            if isinstance(f, dict):
+                image_files.append(Path(f['name']))
+            elif hasattr(f, 'name'):
+                image_files.append(Path(f.name))
+            else:
+                image_files.append(Path(f))
+
+    # From folder picker
+    if batch_folder:
+        folder_path = None
+        if isinstance(batch_folder, list) and len(batch_folder) > 0:
+            f = batch_folder[0]
+            if isinstance(f, dict):
+                folder_path = Path(f['name']).parent
+            elif hasattr(f, 'name'):
+                folder_path = Path(f.name).parent
+        if folder_path and folder_path.is_dir():
+            if recursive:
+                for ext in image_extensions:
+                    image_files.extend(folder_path.rglob(f'*{ext}'))
+            else:
+                for ext in image_extensions:
+                    image_files.extend(folder_path.glob(f'*{ext}'))
+
+    # From string path
+    if batch_str and batch_str.strip():
+        folder_path = Path(batch_str.strip())
+        if folder_path.is_dir():
+            if recursive:
+                for ext in image_extensions:
+                    image_files.extend(folder_path.rglob(f'*{ext}'))
+            else:
+                for ext in image_extensions:
+                    image_files.extend(folder_path.glob(f'*{ext}'))
+
+    # Remove duplicates while preserving order
+    seen = set()
+    unique_files = []
+    for f in image_files:
+        f_resolved = f.resolve()
+        if f_resolved not in seen:
+            seen.add(f_resolved)
+            unique_files.append(f)
+    image_files = unique_files
+
+    if not image_files:
+        shared.log.warning('WD14 batch: no images found')
+        return ''
+
+    t0 = time.time()
+    jobid = shared.state.begin('WD14 Batch')
+    shared.log.info(f'WD14 batch: model="{tagger.model_name}" images={len(image_files)} write={save_output} append={save_append} recursive={recursive}')
+    debug_log(f'WD14 batch: files={[str(f) for f in image_files[:5]]}{"..." if len(image_files) > 5 else ""}')
+
+    results = []
+
+    # Progress bar
+    import rich.progress as rp
+    pbar = rp.Progress(rp.TextColumn('[cyan]WD14:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
+
+    with pbar:
+        task = pbar.add_task(total=len(image_files), description='starting...')
+        for img_path in image_files:
+            pbar.update(task, advance=1, description=str(img_path.name))
+            try:
+                if shared.state.interrupted:
+                    shared.log.info('WD14 batch: interrupted')
+                    break
+
+                image = Image.open(img_path)
+                tags_str = tagger.predict(image, **kwargs)
+
+                if save_output:
+                    txt_path = img_path.with_suffix('.txt')
+                    if save_append and txt_path.exists():
+                        with open(txt_path, 'a', encoding='utf-8') as f:
+                            f.write(f', {tags_str}')
+                        debug_log(f'WD14 batch: appended to "{txt_path}"')
+                    else:
+                        with open(txt_path, 'w', encoding='utf-8') as f:
+                            f.write(tags_str)
+                        debug_log(f'WD14 batch: wrote to "{txt_path}"')
+
+                results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}')
+
+            except Exception as e:
+                shared.log.error(f'WD14 batch: file="{img_path}" error={e}')
+                results.append(f'{img_path.name}: ERROR - {e}')
+
+    elapsed = time.time() - t0
+    shared.log.info(f'WD14 batch: complete images={len(results)} time={elapsed:.1f}s')
+    shared.state.end(jobid)
+
+    return '\n'.join(results)

modules/shared.py

@@ -711,6 +711,17 @@ options_templates.update(options_section(('interrogate', "Interrogate"), {
     "deepbooru_use_spaces": OptionInfo(False, "DeepBooru: use spaces for tags"),
     "deepbooru_escape": OptionInfo(True, "DeepBooru: escape brackets"),
     "deepbooru_filter_tags": OptionInfo("", "DeepBooru: exclude tags"),
+
+    "wd14_sep": OptionInfo("<h2>WD14 Tagger</h2>", "", gr.HTML),
+    "wd14_model": OptionInfo("wd-eva02-large-tagger-v3", "WD14: default model", gr.Dropdown, {"choices": []}),
+    "wd14_general_threshold": OptionInfo(0.35, "WD14: general tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
+    "wd14_character_threshold": OptionInfo(0.85, "WD14: character tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
+    "wd14_max_tags": OptionInfo(74, "WD14: max tags", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1}),
+    "wd14_include_rating": OptionInfo(False, "WD14: include rating tags"),
+    "wd14_sort_alpha": OptionInfo(False, "WD14: sort alphabetically"),
+    "wd14_use_spaces": OptionInfo(False, "WD14: use spaces for tags"),
+    "wd14_escape": OptionInfo(True, "WD14: escape brackets"),
+    "wd14_exclude_tags": OptionInfo("", "WD14: exclude tags"),
 }))
 
 options_templates.update(options_section(('huggingface', "Huggingface"), {

modules/ui_caption.py

@@ -43,6 +43,45 @@ def update_vlm_params(*args):
     shared.opts.save()
 
 
+def wd14_tag_wrapper(image, model_name, general_threshold, character_threshold, include_rating, exclude_tags, max_tags, sort_alpha, use_spaces, escape_brackets):
+    """Wrapper for wd14.tag that maps UI inputs to function parameters."""
+    from modules.interrogate import wd14
+    return wd14.tag(
+        image=image,
+        model_name=model_name,
+        general_threshold=general_threshold,
+        character_threshold=character_threshold,
+        include_rating=include_rating,
+        exclude_tags=exclude_tags,
+        max_tags=int(max_tags),
+        sort_alpha=sort_alpha,
+        use_spaces=use_spaces,
+        escape_brackets=escape_brackets,
+    )
+
+
+def wd14_batch_wrapper(model_name, batch_files, batch_folder, batch_str, save_output, save_append, recursive, general_threshold, character_threshold, include_rating, exclude_tags, max_tags, sort_alpha, use_spaces, escape_brackets):
+    """Wrapper for wd14.batch that maps UI inputs to function parameters."""
+    from modules.interrogate import wd14
+    return wd14.batch(
+        model_name=model_name,
+        batch_files=batch_files,
+        batch_folder=batch_folder,
+        batch_str=batch_str,
+        save_output=save_output,
+        save_append=save_append,
+        recursive=recursive,
+        general_threshold=general_threshold,
+        character_threshold=character_threshold,
+        include_rating=include_rating,
+        exclude_tags=exclude_tags,
+        max_tags=int(max_tags),
+        sort_alpha=sort_alpha,
+        use_spaces=use_spaces,
+        escape_brackets=escape_brackets,
+    )
+
+
 def update_clip_params(*args):
     clip_min_length, clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams = args
     shared.opts.interrogate_clip_min_length = int(clip_min_length)
@@ -158,6 +197,42 @@ def create_ui():
                     with gr.Row():
                         btn_clip_interrogate_img = gr.Button("Interrogate", variant='primary', elem_id="btn_clip_interrogate_img")
                         btn_clip_analyze_img = gr.Button("Analyze", variant='primary', elem_id="btn_clip_analyze_img")
+                with gr.Tab("Booru Tags", elem_id='tab_booru_tags'):
+                    from modules.interrogate import wd14
+                    with gr.Row():
+                        wd_model = gr.Dropdown(wd14.get_models(), value=shared.opts.wd14_model, label='Tagger Model', elem_id='wd_model')
+                        ui_common.create_refresh_button(wd_model, wd14.refresh_models, lambda: {"choices": wd14.get_models()}, 'wd_models_refresh')
+                    with gr.Row():
+                        wd_load_btn = gr.Button(value='Load', elem_id='wd_load', variant='secondary')
+                        wd_unload_btn = gr.Button(value='Unload', elem_id='wd_unload', variant='secondary')
+                    with gr.Accordion(label='Tagger: Advanced Options', open=True, visible=True):
+                        with gr.Row():
+                            wd_general_threshold = gr.Slider(label='General threshold', value=shared.opts.wd14_general_threshold, minimum=0.0, maximum=1.0, step=0.01, elem_id='wd_general_threshold')
+                            wd_character_threshold = gr.Slider(label='Character threshold', value=shared.opts.wd14_character_threshold, minimum=0.0, maximum=1.0, step=0.01, elem_id='wd_character_threshold')
+                        with gr.Row():
+                            wd_max_tags = gr.Slider(label='Max tags', value=shared.opts.wd14_max_tags, minimum=1, maximum=512, step=1, elem_id='wd_max_tags')
+                            wd_include_rating = gr.Checkbox(label='Include rating', value=shared.opts.wd14_include_rating, elem_id='wd_include_rating')
+                        with gr.Row():
+                            wd_sort_alpha = gr.Checkbox(label='Sort alphabetically', value=shared.opts.wd14_sort_alpha, elem_id='wd_sort_alpha')
+                            wd_use_spaces = gr.Checkbox(label='Use spaces', value=shared.opts.wd14_use_spaces, elem_id='wd_use_spaces')
+                            wd_escape = gr.Checkbox(label='Escape brackets', value=shared.opts.wd14_escape, elem_id='wd_escape')
+                        with gr.Row():
+                            wd_exclude_tags = gr.Textbox(label='Exclude tags', value=shared.opts.wd14_exclude_tags, placeholder='Comma-separated tags to exclude', elem_id='wd_exclude_tags')
+                    with gr.Accordion(label='Tagger: Batch', open=False, visible=True):
+                        with gr.Row():
+                            wd_batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], interactive=True, height=100, elem_id='wd_batch_files')
+                        with gr.Row():
+                            wd_batch_folder = gr.File(label="Folder", show_label=True, file_count='directory', file_types=['image'], interactive=True, height=100, elem_id='wd_batch_folder')
+                        with gr.Row():
+                            wd_batch_str = gr.Textbox(label="Folder", value="", interactive=True, elem_id='wd_batch_str')
+                        with gr.Row():
+                            wd_save_output = gr.Checkbox(label='Save Caption Files', value=True, elem_id="wd_save_output")
+                            wd_save_append = gr.Checkbox(label='Append Caption Files', value=False, elem_id="wd_save_append")
+                            wd_folder_recursive = gr.Checkbox(label='Recursive', value=False, elem_id="wd_folder_recursive")
+                        with gr.Row():
+                            btn_wd_tag_batch = gr.Button("Batch Tag", variant='primary', elem_id="btn_wd_tag_batch")
+                    with gr.Row():
+                        btn_wd_tag = gr.Button("Tag", variant='primary', elem_id="btn_wd_tag")
         with gr.Column(variant='compact', elem_id='interrogate_output'):
             with gr.Row(elem_id='interrogate_output_prompt'):
                 prompt = gr.Textbox(label="Answer", lines=12, placeholder="ai generated image description")
@@ -178,6 +253,8 @@ def create_ui():
     btn_clip_interrogate_batch.click(fn=openclip.interrogate_batch, inputs=[clip_batch_files, clip_batch_folder, clip_batch_str, clip_model, blip_model, clip_mode, clip_save_output, clip_save_append, clip_folder_recursive], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])
     btn_vlm_caption.click(fn=vlm_caption_wrapper, inputs=[vlm_question, vlm_system, vlm_prompt, image, vlm_model, vlm_prefill, vlm_thinking_mode], outputs=[prompt, output_image])
     btn_vlm_caption_batch.click(fn=vqa.batch, inputs=[vlm_model, vlm_system, vlm_batch_files, vlm_batch_folder, vlm_batch_str, vlm_question, vlm_prompt, vlm_save_output, vlm_save_append, vlm_folder_recursive, vlm_prefill, vlm_thinking_mode], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])
+    btn_wd_tag.click(fn=wd14_tag_wrapper, inputs=[image, wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])
+    btn_wd_tag_batch.click(fn=wd14_batch_wrapper, inputs=[wd_model, wd_batch_files, wd_batch_folder, wd_batch_str, wd_save_output, wd_save_append, wd_folder_recursive, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image])
 
     # Dynamic UI updates based on selected model and task
     vlm_model.change(fn=update_vlm_prompts_for_model, inputs=[vlm_model], outputs=[vlm_question])
@@ -186,6 +263,14 @@ def create_ui():
     # Load/Unload model buttons
     vlm_load_btn.click(fn=vqa.load_model, inputs=[vlm_model], outputs=[])
     vlm_unload_btn.click(fn=vqa.unload_model, inputs=[], outputs=[])
+    def wd14_load_wrapper(model_name):
+        from modules.interrogate import wd14
+        return wd14.load_model(model_name)
+    def wd14_unload_wrapper():
+        from modules.interrogate import wd14
+        return wd14.unload_model()
+    wd_load_btn.click(fn=wd14_load_wrapper, inputs=[wd_model], outputs=[])
+    wd_unload_btn.click(fn=wd14_unload_wrapper, inputs=[], outputs=[])
 
     for tabname, button in copy_interrogate_buttons.items():
         generation_parameters_copypaste.register_paste_params_button(generation_parameters_copypaste.ParamBinding(paste_button=button, tabname=tabname, source_text_component=prompt, source_image_component=image,))