From db97c42320512fb811e78ab667eafe4fe9ecfcbc Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 19 Jan 2026 03:16:11 +0000 Subject: [PATCH] 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 | 2 +- modules/interrogate/openclip.py | 80 ++++- modules/interrogate/vqa.py | 2 +- modules/interrogate/wd14.py | 535 ++++++++++++++++++++++++++++++ modules/shared.py | 11 + modules/ui_caption.py | 85 +++++ 6 files changed, 701 insertions(+), 14 deletions(-) create mode 100644 modules/interrogate/wd14.py diff --git a/modules/interrogate/moondream3.py b/modules/interrogate/moondream3.py index 0ba0ecd04..f760b3233 100644 --- a/modules/interrogate/moondream3.py +++ b/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: diff --git a/modules/interrogate/openclip.py b/modules/interrogate/openclip.py index 2350dc440..68de085b0 100644 --- a/modules/interrogate/openclip.py +++ b/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("", 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): diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index 036fd7ced..e71cda612 100644 --- a/modules/interrogate/vqa.py +++ b/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: diff --git a/modules/interrogate/wd14.py b/modules/interrogate/wd14.py new file mode 100644 index 000000000..af6f54729 --- /dev/null +++ b/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) diff --git a/modules/shared.py b/modules/shared.py index 0ac6511a1..da08559e4 100644 --- a/modules/shared.py +++ b/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("

WD14 Tagger

", "", 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"), { diff --git a/modules/ui_caption.py b/modules/ui_caption.py index d27b76ce6..1c61cfe88 100644 --- a/modules/ui_caption.py +++ b/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,))