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
2026-01-21 09:46:08 +00:00
parent 5abb794462
commit 6b10f0df4f
8 changed files with 223 additions and 202 deletions
+15 -15
View File
@@ -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)