mirror of
https://github.com/vladmandic/automatic
synced 2026-09-10 23:08:43 +02:00
refactor(caption): address PR review feedback
Rename WD14 module and settings to WaifuDiffusion: - Rename wd14.py to waifudiffusion.py - Rename WD14Tagger class to WaifuDiffusionTagger - Rename WD14_MODELS constant to WAIFUDIFFUSION_MODELS - Rename settings: wd14_model -> waifudiffusion_model, wd14_character_threshold -> waifudiffusion_character_threshold - Update all log messages from "WD14" to "WaifuDiffusion" Code quality improvements: - Simplify threshold parameter defaulting using `or` operator - Extract save_output logic into _save_tags_to_file() helper with isolated error handling to prevent single file failures from impacting entire batch - Fix timing log format consistency (remove 's' suffix)
This commit is contained in:
CalamitousFelicitousness
+71
-71
@@ -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:
|
||||
|
||||
@@ -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}')
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
+2
-2
@@ -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}),
|
||||
|
||||
@@ -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')
|
||||
|
||||
Reference in New Issue
Block a user