mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user