diff --git a/cli/test-tagger.py b/cli/test-tagger.py index 2694a0e1f..9c5f858bf 100644 --- a/cli/test-tagger.py +++ b/cli/test-tagger.py @@ -2,7 +2,7 @@ """ Tagger Settings Test Suite -Tests all WD14 and DeepBooru tagger settings to verify they're properly +Tests all WaifuDiffusion and DeepBooru tagger settings to verify they're properly mapped and affect output correctly. Usage: @@ -71,7 +71,7 @@ class TaggerTest: def __init__(self): self.results = {'passed': [], 'failed': [], 'skipped': []} self.test_image = None - self.wd14_loaded = False + self.waifudiffusion_loaded = False self.deepbooru_loaded = False def log_pass(self, msg): @@ -116,11 +116,11 @@ class TaggerTest: # Load models print("\nLoading models...") - from modules.interrogate import wd14, deepbooru + from modules.interrogate import waifudiffusion, deepbooru t0 = time.time() - self.wd14_loaded = wd14.load_model() - print(f" WD14: {'loaded' if self.wd14_loaded else 'FAILED'} ({time.time()-t0:.1f}s)") + self.waifudiffusion_loaded = waifudiffusion.load_model() + print(f" WaifuDiffusion: {'loaded' if self.waifudiffusion_loaded else 'FAILED'} ({time.time()-t0:.1f}s)") t0 = time.time() self.deepbooru_loaded = deepbooru.load_model() @@ -132,10 +132,10 @@ class TaggerTest: print("CLEANUP") print("=" * 70) - from modules.interrogate import wd14, deepbooru + from modules.interrogate import waifudiffusion, deepbooru from modules import devices - wd14.unload_model() + waifudiffusion.unload_model() deepbooru.unload_model() devices.torch_gc(force=True) print(" Models unloaded") @@ -205,14 +205,14 @@ class TaggerTest: else: self.log_fail(f"Provider '{provider}' configured but not available") - # Test 5: If WD14 loaded, check session providers - if self.wd14_loaded: - from modules.interrogate import wd14 - if wd14.tagger.session is not None: - session_providers = wd14.tagger.session.get_providers() - self.log_pass(f"WD14 session providers: {session_providers}") + # Test 5: If WaifuDiffusion loaded, check session providers + if self.waifudiffusion_loaded: + from modules.interrogate import waifudiffusion + if waifudiffusion.tagger.session is not None: + session_providers = waifudiffusion.tagger.session.get_providers() + self.log_pass(f"WaifuDiffusion session providers: {session_providers}") else: - self.log_skip("WD14 session not initialized") + self.log_skip("WaifuDiffusion session not initialized") # ========================================================================= # TEST: Memory Management (Offload/Reload/Unload) @@ -251,7 +251,7 @@ class TaggerTest: import torch import gc from modules import devices - from modules.interrogate import wd14, deepbooru + from modules.interrogate import waifudiffusion, deepbooru # Memory leak tolerance (MB) - some variance is expected GPU_LEAK_TOLERANCE_MB = 50 @@ -362,10 +362,10 @@ class TaggerTest: deepbooru.load_model() # ===================================================================== - # WD14: Test session lifecycle with memory monitoring + # WaifuDiffusion: Test session lifecycle with memory monitoring # ===================================================================== - if self.wd14_loaded: - print("\n WD14 Memory Management:") + if self.waifudiffusion_loaded: + print("\n WaifuDiffusion Memory Management:") # Baseline memory gc.collect() @@ -375,74 +375,74 @@ class TaggerTest: print(f" Baseline: GPU={baseline['gpu_allocated']:.1f}MB, RAM={baseline['ram_used']:.1f}MB") # Test 1: Session exists - if wd14.tagger.session is not None: - self.log_pass("WD14: session loaded") + if waifudiffusion.tagger.session is not None: + self.log_pass("WaifuDiffusion: session loaded") else: - self.log_fail("WD14: session not loaded") + self.log_fail("WaifuDiffusion: session not loaded") return # Test 2: Get current providers - providers = wd14.tagger.session.get_providers() + providers = waifudiffusion.tagger.session.get_providers() print(f" Active providers: {providers}") - self.log_pass(f"WD14: using providers {providers}") + self.log_pass(f"WaifuDiffusion: using providers {providers}") # Test 3: Run inference try: - tags = wd14.tagger.predict(self.test_image, max_tags=3) + tags = waifudiffusion.tagger.predict(self.test_image, max_tags=3) after_infer = self.get_memory_stats() print(f" After inference: GPU={after_infer['gpu_allocated']:.1f}MB, RAM={after_infer['ram_used']:.1f}MB") if tags: - self.log_pass(f"WD14: inference works ({tags[:30]}...)") + self.log_pass(f"WaifuDiffusion: inference works ({tags[:30]}...)") else: - self.log_fail("WD14: inference returned empty") + self.log_fail("WaifuDiffusion: inference returned empty") except Exception as e: - self.log_fail(f"WD14: inference failed: {e}") + self.log_fail(f"WaifuDiffusion: inference failed: {e}") # Test 4: Unload session with memory check - model_name = wd14.tagger.model_name - wd14.unload_model() + model_name = waifudiffusion.tagger.model_name + waifudiffusion.unload_model() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() after_unload = self.get_memory_stats() print(f" After unload: GPU={after_unload['gpu_allocated']:.1f}MB, RAM={after_unload['ram_used']:.1f}MB") - if wd14.tagger.session is None: - self.log_pass("WD14: unload successful") + if waifudiffusion.tagger.session is None: + self.log_pass("WaifuDiffusion: unload successful") else: - self.log_fail("WD14: unload failed, session still exists") + self.log_fail("WaifuDiffusion: unload failed, session still exists") # Check for memory leaks after unload gpu_leak = after_unload['gpu_allocated'] - baseline['gpu_allocated'] ram_leak = after_unload['ram_used'] - baseline['ram_used'] if gpu_leak <= GPU_LEAK_TOLERANCE_MB: - self.log_pass(f"WD14: no GPU memory leak after unload (diff={gpu_leak:.1f}MB)") + self.log_pass(f"WaifuDiffusion: no GPU memory leak after unload (diff={gpu_leak:.1f}MB)") else: - self.log_fail(f"WD14: GPU memory leak detected (diff={gpu_leak:.1f}MB > {GPU_LEAK_TOLERANCE_MB}MB)") + self.log_fail(f"WaifuDiffusion: GPU memory leak detected (diff={gpu_leak:.1f}MB > {GPU_LEAK_TOLERANCE_MB}MB)") if ram_leak <= RAM_LEAK_TOLERANCE_MB: - self.log_pass(f"WD14: no RAM leak after unload (diff={ram_leak:.1f}MB)") + self.log_pass(f"WaifuDiffusion: no RAM leak after unload (diff={ram_leak:.1f}MB)") else: - self.log_warn(f"WD14: RAM increased after unload (diff={ram_leak:.1f}MB) - may be caching") + self.log_warn(f"WaifuDiffusion: RAM increased after unload (diff={ram_leak:.1f}MB) - may be caching") # Test 5: Reload session - wd14.load_model(model_name) + waifudiffusion.load_model(model_name) after_reload = self.get_memory_stats() print(f" After reload: GPU={after_reload['gpu_allocated']:.1f}MB, RAM={after_reload['ram_used']:.1f}MB") - if wd14.tagger.session is not None: - self.log_pass("WD14: reload successful") + if waifudiffusion.tagger.session is not None: + self.log_pass("WaifuDiffusion: reload successful") else: - self.log_fail("WD14: reload failed") + self.log_fail("WaifuDiffusion: reload failed") # Test 6: Inference after reload try: - tags = wd14.tagger.predict(self.test_image, max_tags=3) + tags = waifudiffusion.tagger.predict(self.test_image, max_tags=3) if tags: - self.log_pass("WD14: inference after reload works") + self.log_pass("WaifuDiffusion: inference after reload works") else: - self.log_fail("WD14: inference after reload returned empty") + self.log_fail("WaifuDiffusion: inference after reload returned empty") except Exception as e: - self.log_fail(f"WD14: inference after reload failed: {e}") + self.log_fail(f"WaifuDiffusion: inference after reload failed: {e}") # Final memory check after full cycle gc.collect() @@ -471,8 +471,8 @@ class TaggerTest: ('tagger_escape_brackets', bool), ('tagger_exclude_tags', str), ('tagger_show_scores', bool), - ('wd14_model', str), - ('wd14_character_threshold', float), + ('waifudiffusion_model', str), + ('waifudiffusion_character_threshold', float), ('interrogate_offload', bool), ] @@ -486,23 +486,23 @@ class TaggerTest: # ========================================================================= # TEST: Parameter Effect - Tests a single parameter on both taggers # ========================================================================= - def test_parameter(self, param_name, test_func, wd14_supported=True, deepbooru_supported=True): - """Test a parameter on both WD14 and DeepBooru.""" + def test_parameter(self, param_name, test_func, waifudiffusion_supported=True, deepbooru_supported=True): + """Test a parameter on both WaifuDiffusion and DeepBooru.""" print(f"\n Testing: {param_name}") - if wd14_supported and self.wd14_loaded: + if waifudiffusion_supported and self.waifudiffusion_loaded: try: - result = test_func('wd14') + result = test_func('waifudiffusion') if result is True: - self.log_pass(f"WD14: {param_name}") + self.log_pass(f"WaifuDiffusion: {param_name}") elif result is False: - self.log_fail(f"WD14: {param_name}") + self.log_fail(f"WaifuDiffusion: {param_name}") else: - self.log_skip(f"WD14: {param_name} - {result}") + self.log_skip(f"WaifuDiffusion: {param_name} - {result}") except Exception as e: - self.log_fail(f"WD14: {param_name} - {e}") - elif wd14_supported: - self.log_skip(f"WD14: {param_name} - model not loaded") + self.log_fail(f"WaifuDiffusion: {param_name} - {e}") + elif waifudiffusion_supported: + self.log_skip(f"WaifuDiffusion: {param_name} - model not loaded") if deepbooru_supported and self.deepbooru_loaded: try: @@ -520,9 +520,9 @@ class TaggerTest: def tag(self, tagger, **kwargs): """Helper to call the appropriate tagger.""" - if tagger == 'wd14': - from modules.interrogate import wd14 - return wd14.tagger.predict(self.test_image, **kwargs) + if tagger == 'waifudiffusion': + from modules.interrogate import waifudiffusion + return waifudiffusion.tagger.predict(self.test_image, **kwargs) else: from modules.interrogate import deepbooru return deepbooru.model.tag(self.test_image, **kwargs) @@ -759,16 +759,16 @@ class TaggerTest: self.test_parameter('include_rating', check_include_rating) # ========================================================================= - # TEST: character_threshold (WD14 only) + # TEST: character_threshold (WaifuDiffusion only) # ========================================================================= def test_character_threshold(self): - """Test that character_threshold affects character tag count (WD14 only).""" + """Test that character_threshold affects character tag count (WaifuDiffusion only).""" print("\n" + "=" * 70) - print("TEST: character_threshold effect (WD14 only)") + print("TEST: character_threshold effect (WaifuDiffusion only)") print("=" * 70) def check_character_threshold(tagger): - if tagger != 'wd14': + if tagger != 'waifudiffusion': return "not supported" # Character threshold only affects character tags @@ -796,17 +796,17 @@ class TaggerTest: from modules.interrogate import tagger - # Test WD14 through unified interface - if self.wd14_loaded: + # Test WaifuDiffusion through unified interface + if self.waifudiffusion_loaded: try: models = tagger.get_models() - wd14_model = next((m for m in models if m != 'DeepBooru'), None) - if wd14_model: - tags = tagger.tag(self.test_image, model_name=wd14_model, max_tags=5) - print(f" WD14 ({wd14_model}): {tags[:50]}...") - self.log_pass("Unified interface: WD14") + waifudiffusion_model = next((m for m in models if m != 'DeepBooru'), None) + if waifudiffusion_model: + tags = tagger.tag(self.test_image, model_name=waifudiffusion_model, max_tags=5) + print(f" WaifuDiffusion ({waifudiffusion_model}): {tags[:50]}...") + self.log_pass("Unified interface: WaifuDiffusion") except Exception as e: - self.log_fail(f"Unified interface: WD14 - {e}") + self.log_fail(f"Unified interface: WaifuDiffusion - {e}") # Test DeepBooru through unified interface if self.deepbooru_loaded: diff --git a/modules/interrogate/deepbooru.py b/modules/interrogate/deepbooru.py index 5ee2df751..88f0ab44f 100644 --- a/modules/interrogate/deepbooru.py +++ b/modules/interrogate/deepbooru.py @@ -80,20 +80,13 @@ class DeepDanbooru: Formatted tag string """ # Use settings defaults if not specified - if general_threshold is None: - general_threshold = shared.opts.tagger_threshold - if include_rating is None: - include_rating = shared.opts.tagger_include_rating - if exclude_tags is None: - exclude_tags = shared.opts.tagger_exclude_tags - if max_tags is None: - max_tags = shared.opts.tagger_max_tags - if sort_alpha is None: - sort_alpha = shared.opts.tagger_sort_alpha - if use_spaces is None: - use_spaces = shared.opts.tagger_use_spaces - if escape_brackets is None: - escape_brackets = shared.opts.tagger_escape_brackets + general_threshold = general_threshold or shared.opts.tagger_threshold + include_rating = include_rating if include_rating is not None else shared.opts.tagger_include_rating + exclude_tags = exclude_tags or shared.opts.tagger_exclude_tags + max_tags = max_tags or shared.opts.tagger_max_tags + sort_alpha = sort_alpha if sort_alpha is not None else shared.opts.tagger_sort_alpha + use_spaces = use_spaces if use_spaces is not None else shared.opts.tagger_use_spaces + escape_brackets = escape_brackets if escape_brackets is not None else shared.opts.tagger_escape_brackets if isinstance(pil_image, list): pil_image = pil_image[0] if len(pil_image) > 0 else None @@ -137,6 +130,31 @@ class DeepDanbooru: model = DeepDanbooru() +def _save_tags_to_file(img_path, tags_str: str, save_append: bool) -> bool: + """Save tags to a text file with error handling. + + Args: + img_path: Path to the image file + tags_str: Tags string to save + save_append: If True, append to existing file; otherwise overwrite + + Returns: + True if save succeeded, False otherwise + """ + try: + 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}') + else: + with open(txt_path, 'w', encoding='utf-8') as f: + f.write(tags_str) + return True + except Exception as e: + shared.log.error(f'DeepBooru batch: failed to save file="{img_path}" error={e}') + return False + + def get_models() -> list: """Return list of available DeepBooru models (just one).""" return ["DeepBooru"] @@ -179,7 +197,7 @@ def tag(image, **kwargs) -> str: try: result = model.tag(image, **kwargs) - shared.log.debug(f'DeepBooru: complete time={time.time()-t0:.2f}s tags={len(result.split(", ")) if result else 0}') + shared.log.debug(f'DeepBooru: complete time={time.time()-t0:.2f} tags={len(result.split(", ")) if result else 0}') except Exception as e: result = f"Exception {type(e)}" shared.log.error(f'DeepBooru: {e}') @@ -299,13 +317,7 @@ def batch( tags_str = model.tag_multi(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}') - else: - with open(txt_path, 'w', encoding='utf-8') as f: - f.write(tags_str) + _save_tags_to_file(img_path, tags_str, save_append) results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}') diff --git a/modules/interrogate/interrogate.py b/modules/interrogate/interrogate.py index 4efc32732..4e06fb36f 100644 --- a/modules/interrogate/interrogate.py +++ b/modules/interrogate/interrogate.py @@ -21,13 +21,13 @@ def interrogate(image): shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"') return prompt elif shared.opts.interrogate_default_type == 'Tagger': - shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} model="{shared.opts.wd14_model}"') + shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} model="{shared.opts.waifudiffusion_model}"') from modules.interrogate import tagger prompt = tagger.tag( image=image, - model_name=shared.opts.wd14_model, + model_name=shared.opts.waifudiffusion_model, general_threshold=shared.opts.tagger_threshold, - character_threshold=shared.opts.wd14_character_threshold, + character_threshold=shared.opts.waifudiffusion_character_threshold, include_rating=shared.opts.tagger_include_rating, exclude_tags=shared.opts.tagger_exclude_tags, max_tags=shared.opts.tagger_max_tags, diff --git a/modules/interrogate/openclip.py b/modules/interrogate/openclip.py index 68de085b0..ca69ad8dd 100644 --- a/modules/interrogate/openclip.py +++ b/modules/interrogate/openclip.py @@ -117,7 +117,7 @@ def load_interrogator(clip_model, blip_model): 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') + shared.log.debug(f'CLIP load: time={time.time()-t0:.2f}') elif clip_model != ci.config.clip_model_name or blip_model != ci.config.caption_model_name: t0 = time.time() if clip_model != ci.config.clip_model_name: @@ -134,7 +134,7 @@ def load_interrogator(clip_model, blip_model): 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') + shared.log.debug(f'CLIP load: time={time.time()-t0:.2f}') else: debug_log(f'CLIP: models already loaded clip="{clip_model}" blip="{blip_model}"') @@ -172,7 +172,7 @@ 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}"') + debug_log(f'CLIP: mode="{mode}" time={time.time()-t0:.2f} result="{prompt[:100]}..."' if len(prompt) > 100 else f'CLIP: mode="{mode}" time={time.time()-t0:.2f} result="{prompt}"') return prompt @@ -189,7 +189,7 @@ def interrogate_image(image, clip_model, blip_model, mode): image = image.convert('RGB') prompt = interrogate(image, mode) devices.torch_gc() - shared.log.debug(f'CLIP: complete time={time.time()-t0:.2f}s') + shared.log.debug(f'CLIP: complete time={time.time()-t0:.2f}') except Exception as e: prompt = f"Exception {type(e)}" shared.log.error(f'CLIP: {e}') @@ -243,7 +243,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod 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') + shared.log.info(f'CLIP batch: complete images={len(prompts)} time={time.time()-t0:.2f}') return '\n\n'.join(prompts) @@ -264,7 +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') + shared.log.debug(f'CLIP analyze: complete time={time.time()-t0:.2f}') # Format labels as text def format_category(name, ranks): diff --git a/modules/interrogate/tagger.py b/modules/interrogate/tagger.py index cd2374d04..51516adaa 100644 --- a/modules/interrogate/tagger.py +++ b/modules/interrogate/tagger.py @@ -1,4 +1,4 @@ -# Unified Tagger Interface - Dispatches to WD14 or DeepBooru based on model selection +# Unified Tagger Interface - Dispatches to WaifuDiffusion or DeepBooru based on model selection # Provides a common interface for the Booru Tags tab from modules import shared @@ -7,9 +7,9 @@ DEEPBOORU_MODEL = "DeepBooru" def get_models() -> list: - """Return combined list: DeepBooru + WD14 models.""" - from modules.interrogate import wd14 - return [DEEPBOORU_MODEL] + wd14.get_models() + """Return combined list: DeepBooru + WaifuDiffusion models.""" + from modules.interrogate import waifudiffusion + return [DEEPBOORU_MODEL] + waifudiffusion.get_models() def refresh_models() -> list: @@ -28,15 +28,15 @@ def load_model(model_name: str) -> bool: from modules.interrogate import deepbooru return deepbooru.load_model() else: - from modules.interrogate import wd14 - return wd14.load_model(model_name) + from modules.interrogate import waifudiffusion + return waifudiffusion.load_model(model_name) def unload_model(): """Unload both backends to ensure memory is freed.""" - from modules.interrogate import deepbooru, wd14 + from modules.interrogate import deepbooru, waifudiffusion deepbooru.unload_model() - wd14.unload_model() + waifudiffusion.unload_model() def tag(image, model_name: str = None, **kwargs) -> str: @@ -44,28 +44,28 @@ def tag(image, model_name: str = None, **kwargs) -> str: Args: image: PIL Image to tag - model_name: Model to use (DeepBooru or WD14 model name) + model_name: Model to use (DeepBooru or WaifuDiffusion model name) **kwargs: Additional arguments passed to the backend Returns: Formatted tag string """ if model_name is None: - model_name = shared.opts.wd14_model + model_name = shared.opts.waifudiffusion_model if is_deepbooru(model_name): from modules.interrogate import deepbooru return deepbooru.tag(image, **kwargs) else: - from modules.interrogate import wd14 - return wd14.tag(image, model_name=model_name, **kwargs) + from modules.interrogate import waifudiffusion + return waifudiffusion.tag(image, model_name=model_name, **kwargs) def batch(model_name: str, **kwargs) -> str: """Unified batch processing. Args: - model_name: Model to use (DeepBooru or WD14 model name) + model_name: Model to use (DeepBooru or WaifuDiffusion model name) **kwargs: Additional arguments passed to the backend Returns: @@ -75,5 +75,5 @@ def batch(model_name: str, **kwargs) -> str: from modules.interrogate import deepbooru return deepbooru.batch(model_name=model_name, **kwargs) else: - from modules.interrogate import wd14 - return wd14.batch(model_name=model_name, **kwargs) + from modules.interrogate import waifudiffusion + return waifudiffusion.batch(model_name=model_name, **kwargs) diff --git a/modules/interrogate/wd14.py b/modules/interrogate/waifudiffusion.py similarity index 69% rename from modules/interrogate/wd14.py rename to modules/interrogate/waifudiffusion.py index fcdb360b0..71951a47f 100644 --- a/modules/interrogate/wd14.py +++ b/modules/interrogate/waifudiffusion.py @@ -1,4 +1,4 @@ -# WD14/WaifuDiffusion Tagger - ONNX-based anime/illustration tagging +# WaifuDiffusion Tagger - ONNX-based anime/illustration tagging # Based on SmilingWolf's tagger models: https://huggingface.co/SmilingWolf import os @@ -17,8 +17,8 @@ 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 = { +# WaifuDiffusion model repository mappings +WAIFUDIFFUSION_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", @@ -38,8 +38,8 @@ CATEGORY_CHARACTER = 4 CATEGORY_RATING = 9 -class WD14Tagger: - """WD14/WaifuDiffusion Tagger using ONNX inference.""" +class WaifuDiffusionTagger: + """WaifuDiffusion Tagger using ONNX inference.""" def __init__(self): self.session = None @@ -54,63 +54,63 @@ class WD14Tagger: 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}"') + model_name = shared.opts.waifudiffusion_model + if model_name not in WAIFUDIFFUSION_MODELS: + shared.log.error(f'WaifuDiffusion: 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}"') + debug_log(f'WaifuDiffusion: 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}"') + debug_log(f'WaifuDiffusion: switching model from "{self.model_name}" to "{model_name}"') self.unload() - repo_id = WD14_MODELS[model_name] + repo_id = WAIFUDIFFUSION_MODELS[model_name] t0 = time.time() - shared.log.info(f'WD14 load: model="{model_name}" repo="{repo_id}"') + shared.log.info(f'WaifuDiffusion 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}"') + debug_log(f'WaifuDiffusion 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}"') + debug_log(f'WaifuDiffusion 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}') + shared.log.error(f'WaifuDiffusion load: model file not found: {model_file}') return False import onnxruntime as ort - debug_log(f'WD14 load: onnxruntime version={ort.__version__}') + debug_log(f'WaifuDiffusion load: onnxruntime version={ort.__version__}') self.session = ort.InferenceSession(model_file, providers=devices.onnx) self.model_name = model_name # Get actual providers used actual_providers = self.session.get_providers() - debug_log(f'WD14 load: active providers={actual_providers}') + debug_log(f'WaifuDiffusion 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}') + shared.log.debug(f'WaifuDiffusion load: time={load_time:.2f} tags={len(self.tags)}') + debug_log(f'WaifuDiffusion 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') + shared.log.error(f'WaifuDiffusion load: failed error={e}') + errors.display(e, 'WaifuDiffusion load') self.unload() return False @@ -120,7 +120,7 @@ class WD14Tagger: 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}') + shared.log.error(f'WaifuDiffusion load: tags file not found: {csv_path}') return self.tags = [] @@ -136,24 +136,24 @@ class WD14Tagger: 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}') + debug_log(f'WaifuDiffusion 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}"') + shared.log.debug(f'WaifuDiffusion 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') + debug_log('WaifuDiffusion unload: complete') else: - debug_log('WD14 unload: no model loaded') + debug_log('WaifuDiffusion unload: no model loaded') def preprocess_image(self, image: Image.Image) -> np.ndarray: - """Preprocess image for WD14 model input. + """Preprocess image for WaifuDiffusion model input. - Resize to 448x448 (standard for WD models) - Pad to square with white background @@ -189,7 +189,7 @@ class WD14Tagger: # 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}') + debug_log(f'WaifuDiffusion preprocess: original_size={original_size} mode={original_mode} padded_size={max_dim} output_shape={img_array.shape}') return img_array def predict( @@ -223,24 +223,16 @@ class WD14Tagger: t0 = time.time() # Use settings defaults if not specified - if general_threshold is None: - general_threshold = shared.opts.tagger_threshold - if character_threshold is None: - character_threshold = shared.opts.wd14_character_threshold - if include_rating is None: - include_rating = shared.opts.tagger_include_rating - if exclude_tags is None: - exclude_tags = shared.opts.tagger_exclude_tags - if max_tags is None: - max_tags = shared.opts.tagger_max_tags - if sort_alpha is None: - sort_alpha = shared.opts.tagger_sort_alpha - if use_spaces is None: - use_spaces = shared.opts.tagger_use_spaces - if escape_brackets is None: - escape_brackets = shared.opts.tagger_escape_brackets + general_threshold = general_threshold or shared.opts.tagger_threshold + character_threshold = character_threshold or shared.opts.waifudiffusion_character_threshold + include_rating = include_rating if include_rating is not None else shared.opts.tagger_include_rating + exclude_tags = exclude_tags or shared.opts.tagger_exclude_tags + max_tags = max_tags or shared.opts.tagger_max_tags + sort_alpha = sort_alpha if sort_alpha is not None else shared.opts.tagger_sort_alpha + use_spaces = use_spaces if use_spaces is not None else shared.opts.tagger_use_spaces + escape_brackets = escape_brackets if escape_brackets is not None else shared.opts.tagger_escape_brackets - 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}') + debug_log(f'WaifuDiffusion 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): @@ -248,7 +240,7 @@ class WD14Tagger: 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') + shared.log.error('WaifuDiffusion predict: no image provided') return '' # Load model if needed @@ -265,13 +257,13 @@ class WD14Tagger: 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}') + debug_log(f'WaifuDiffusion 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}') + debug_log(f'WaifuDiffusion predict: exclude_tags={exclude_set}') general_count = 0 character_count = 0 @@ -305,7 +297,7 @@ class WD14Tagger: 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)}') + debug_log(f'WaifuDiffusion predict: matched tags general={general_count} character={character_count} rating={rating_count} total={len(tag_probs)}') # Sort tags if sort_alpha: @@ -316,7 +308,7 @@ class WD14Tagger: # 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}') + debug_log(f'WaifuDiffusion predict: limited to max_tags={max_tags}') # Format output result = [] @@ -332,7 +324,7 @@ class WD14Tagger: 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}"') + debug_log(f'WaifuDiffusion predict: complete tags={len(result)} time={total_time:.2f} result="{output[:100]}..."' if len(output) > 100 else f'WaifuDiffusion predict: complete tags={len(result)} time={total_time:.2f} result="{output}"') return output @@ -342,12 +334,37 @@ class WD14Tagger: # Global tagger instance -tagger = WD14Tagger() +tagger = WaifuDiffusionTagger() + + +def _save_tags_to_file(img_path, tags_str: str, save_append: bool) -> bool: + """Save tags to a text file with error handling. + + Args: + img_path: Path to the image file + tags_str: Tags string to save + save_append: If True, append to existing file; otherwise overwrite + + Returns: + True if save succeeded, False otherwise + """ + try: + 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}') + else: + with open(txt_path, 'w', encoding='utf-8') as f: + f.write(tags_str) + return True + except Exception as e: + shared.log.error(f'WaifuDiffusion batch: failed to save file="{img_path}" error={e}') + return False def get_models() -> list: - """Return list of available WD14 model names.""" - return list(WD14_MODELS.keys()) + """Return list of available WaifuDiffusion model names.""" + return list(WAIFUDIFFUSION_MODELS.keys()) def refresh_models() -> list: @@ -358,17 +375,17 @@ def refresh_models() -> list: def load_model(model_name: str = None) -> bool: - """Load the specified WD14 model.""" + """Load the specified WaifuDiffusion model.""" return tagger.load(model_name) def unload_model(): - """Unload the current WD14 model.""" + """Unload the current WaifuDiffusion model.""" tagger.unload() def tag(image: Image.Image, model_name: str = None, **kwargs) -> str: - """Tag an image using WD14 tagger. + """Tag an image using WaifuDiffusion tagger. Args: image: PIL Image to tag @@ -379,21 +396,21 @@ def tag(image: Image.Image, model_name: str = None, **kwargs) -> str: 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}') + jobid = shared.state.begin('WaifuDiffusion Tag') + shared.log.info(f'WaifuDiffusion: model="{model_name or tagger.model_name or shared.opts.waifudiffusion_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}') + shared.log.debug(f'WaifuDiffusion: complete time={time.time()-t0:.2f} tags={len(result.split(", ")) if result else 0}') # Offload model if setting enabled if shared.opts.interrogate_offload: tagger.unload() except Exception as e: result = f"Exception {type(e)}" - shared.log.error(f'WD14: {e}') - errors.display(e, 'WD14 Tag') + shared.log.error(f'WaifuDiffusion: {e}') + errors.display(e, 'WaifuDiffusion Tag') shared.state.end(jobid) return result @@ -485,19 +502,19 @@ def batch( image_files = unique_files if not image_files: - shared.log.warning('WD14 batch: no images found') + shared.log.warning('WaifuDiffusion 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 ""}') + jobid = shared.state.begin('WaifuDiffusion Batch') + shared.log.info(f'WaifuDiffusion batch: model="{tagger.model_name}" images={len(image_files)} write={save_output} append={save_append} recursive={recursive}') + debug_log(f'WaifuDiffusion 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) + pbar = rp.Progress(rp.TextColumn('[cyan]WaifuDiffusion:'), 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...') @@ -505,31 +522,23 @@ def batch( pbar.update(task, advance=1, description=str(img_path.name)) try: if shared.state.interrupted: - shared.log.info('WD14 batch: interrupted') + shared.log.info('WaifuDiffusion 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}"') + _save_tags_to_file(img_path, tags_str, save_append) 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}') + shared.log.error(f'WaifuDiffusion 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.log.info(f'WaifuDiffusion 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 b143d65aa..2a6cd5ac0 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -774,8 +774,8 @@ options_templates.update(options_section(('hidden_options', "Hidden options"), { "tagger_use_spaces": OptionInfo(False, "Tagger: use spaces for tags", gr.Checkbox, {"visible": False}), "tagger_escape_brackets": OptionInfo(True, "Tagger: escape brackets", gr.Checkbox, {"visible": False}), "tagger_exclude_tags": OptionInfo("", "Tagger: exclude tags", gr.Textbox, {"visible": False}), - "wd14_model": OptionInfo("wd-eva02-large-tagger-v3", "WD14: default model", gr.Dropdown, {"choices": [], "visible": False}), - "wd14_character_threshold": OptionInfo(0.85, "WD14: character tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False}), + "waifudiffusion_model": OptionInfo("wd-eva02-large-tagger-v3", "WaifuDiffusion: default model", gr.Dropdown, {"choices": [], "visible": False}), + "waifudiffusion_character_threshold": OptionInfo(0.85, "WaifuDiffusion: character tag threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False}), # control settings are handled separately "control_hires": OptionInfo(False, "Hires use Control", gr.Checkbox, {"visible": False}), diff --git a/modules/ui_caption.py b/modules/ui_caption.py index 4821be690..5ab4d74b7 100644 --- a/modules/ui_caption.py +++ b/modules/ui_caption.py @@ -98,9 +98,9 @@ def update_tagger_ui(model_name): def update_tagger_params(model_name, general_threshold, character_threshold, include_rating, max_tags, sort_alpha, use_spaces, escape_brackets, exclude_tags, show_scores): """Save all tagger parameters to shared.opts when UI controls change.""" - shared.opts.wd14_model = model_name + shared.opts.waifudiffusion_model = model_name shared.opts.tagger_threshold = float(general_threshold) - shared.opts.wd14_character_threshold = float(character_threshold) + shared.opts.waifudiffusion_character_threshold = float(character_threshold) shared.opts.tagger_include_rating = bool(include_rating) shared.opts.tagger_max_tags = int(max_tags) shared.opts.tagger_sort_alpha = bool(sort_alpha) @@ -250,7 +250,7 @@ def create_ui(): with gr.Tab("Tagger", elem_id='tab_tagger'): from modules.interrogate import tagger with gr.Row(): - wd_model = gr.Dropdown(tagger.get_models(), value=shared.opts.wd14_model, label='Tagger Model', elem_id='wd_model') + wd_model = gr.Dropdown(tagger.get_models(), value=shared.opts.waifudiffusion_model, label='Tagger Model', elem_id='wd_model') ui_common.create_refresh_button(wd_model, tagger.refresh_models, lambda: {"choices": tagger.get_models()}, 'wd_models_refresh') with gr.Row(): wd_load_btn = gr.Button(value='Load', elem_id='wd_load', variant='secondary') @@ -258,7 +258,7 @@ def create_ui(): 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.tagger_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') + wd_character_threshold = gr.Slider(label='Character threshold', value=shared.opts.waifudiffusion_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.tagger_max_tags, minimum=1, maximum=512, step=1, elem_id='wd_max_tags') wd_include_rating = gr.Checkbox(label='Include rating', value=shared.opts.tagger_include_rating, elem_id='wd_include_rating')