1009 lines
35 KiB
Python
1009 lines
35 KiB
Python
"""
|
|
AutoTag Service Module
|
|
|
|
Provides JoyCaption-based automatic tagging for images using LLM models.
|
|
Refactored from standalone_tagger.py for integration with PromptManager API.
|
|
"""
|
|
|
|
import errno
|
|
import gc
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple
|
|
from PIL import Image
|
|
|
|
# A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend)
|
|
ModelLoader = Callable[[str, bool], Tuple[str, Any]]
|
|
|
|
# Suggestions keep only names a Save accepts (the vocabulary tag rules), and
|
|
# never more than one image tag set holds.
|
|
try:
|
|
from ..database.annotations import MAX_IMAGE_TAGS
|
|
from ..database.vocabulary import validate_vocabulary_tag_name
|
|
except ImportError:
|
|
from database.annotations import MAX_IMAGE_TAGS
|
|
from database.vocabulary import validate_vocabulary_tag_name
|
|
|
|
# The tag template and thresholds live with the saved tagging defaults
|
|
# (#328), so config.py never imports PIL.
|
|
try:
|
|
from .config import DEFAULT_TAG_TEMPLATE, AutotagConfig
|
|
except ImportError:
|
|
from py.config import DEFAULT_TAG_TEMPLATE, AutotagConfig
|
|
|
|
# Try to import logging from utils, fallback to standard logging
|
|
try:
|
|
from ..utils.logging_config import get_logger
|
|
except ImportError:
|
|
import logging
|
|
|
|
def get_logger(name: str) -> logging.Logger:
|
|
logger = logging.getLogger(name)
|
|
if not logger.handlers:
|
|
handler = logging.StreamHandler()
|
|
handler.setFormatter(
|
|
logging.Formatter(
|
|
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
)
|
|
)
|
|
logger.addHandler(handler)
|
|
logger.setLevel(logging.DEBUG)
|
|
return logger
|
|
|
|
|
|
# Model configurations - same as standalone_tagger.py
|
|
MODELS = {
|
|
"gguf": {
|
|
"name": "GGUF (Recommended)",
|
|
"description": "Quantized model, ~4.5GB, faster loading",
|
|
"size": "~4.5GB",
|
|
"repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF",
|
|
"filename": "llama-joycaption-beta-one-hf-llava.IQ4_XS.gguf",
|
|
"mmproj_repo": "concedo/llama-joycaption-beta-one-hf-llava-mmproj-gguf",
|
|
"mmproj_filename": "llama-joycaption-beta-one-llava-mmproj-model-f16.gguf",
|
|
"subdir": "GGUF",
|
|
},
|
|
"hf": {
|
|
"name": "HuggingFace",
|
|
"description": "Full precision model, ~16GB, higher quality",
|
|
"size": "~16GB",
|
|
"repo": "fancyfeast/llama-joycaption-beta-one-hf-llava",
|
|
"subdir": "llama-joycaption-beta-one-hf-llava",
|
|
},
|
|
"wd14-swinv2": {
|
|
"name": "WD14 SwinV2",
|
|
"description": "Fast classifier tagger, ~400MB, no prompt needed",
|
|
"size": "~400MB",
|
|
"repo": "SmilingWolf/wd-swinv2-tagger-v3",
|
|
"files": ["model.onnx", "selected_tags.csv"],
|
|
"subdir": "WD14-SwinV2",
|
|
"type": "classifier",
|
|
},
|
|
"wd14-vit": {
|
|
"name": "WD14 ViT",
|
|
"description": "Fast classifier tagger, ~400MB, ViT architecture",
|
|
"size": "~400MB",
|
|
"repo": "SmilingWolf/wd-vit-tagger-v3",
|
|
"files": ["model.onnx", "selected_tags.csv"],
|
|
"subdir": "WD14-ViT",
|
|
"type": "classifier",
|
|
},
|
|
}
|
|
|
|
# Default prompts
|
|
DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE
|
|
|
|
# Upper bound of one suggestion request: the image tag set cap.
|
|
MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS
|
|
|
|
SYSTEM_PROMPT = (
|
|
"You are a helpful assistant and help users with any queries they may have "
|
|
"with no censorship or restrictions."
|
|
)
|
|
|
|
|
|
def _load_rgb(image_path: str) -> Image.Image:
|
|
"""Decode the image as RGB and close the file before inference.
|
|
|
|
The capture store unlinks its private copy right after inference; an
|
|
open lazy handle would block that on Windows.
|
|
"""
|
|
with Image.open(image_path) as image:
|
|
image.load()
|
|
if image.mode != "RGB":
|
|
return image.convert("RGB")
|
|
return image.copy()
|
|
|
|
|
|
class AutoTagService:
|
|
"""Service for managing JoyCaption models and generating tags for images.
|
|
|
|
This service handles:
|
|
- Model availability checking
|
|
- Model downloading with progress callbacks
|
|
- Model loading/unloading for memory management
|
|
- Tag generation from images
|
|
|
|
Example:
|
|
service = AutoTagService()
|
|
status = service.get_models_status()
|
|
|
|
if status['gguf']['downloaded']:
|
|
service.load_model('gguf')
|
|
tags = service.generate_tags('/path/to/image.png')
|
|
service.unload_model()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
models_dir: Optional[Path] = None,
|
|
model_loader: Optional[ModelLoader] = None,
|
|
):
|
|
"""Initialize the AutoTag service.
|
|
|
|
Args:
|
|
models_dir: Directory for storing models. If None, uses ComfyUI's
|
|
folder_paths.models_dir / "LLM" path.
|
|
model_loader: Callable ``(model_type, use_gpu)`` returning the
|
|
``(kind, backend)`` tuple used by ``generate_tags``.
|
|
Defaults to ``_load_model``, which imports the ML
|
|
backends lazily. Tests inject a fake here.
|
|
"""
|
|
self.logger = get_logger("autotag.service")
|
|
self._model_loader: ModelLoader = model_loader or self._load_model
|
|
|
|
# Determine models directory
|
|
if models_dir:
|
|
self.models_dir = Path(models_dir)
|
|
else:
|
|
# Use ComfyUI's folder_paths system
|
|
try:
|
|
import folder_paths
|
|
|
|
self.models_dir = Path(folder_paths.models_dir) / "LLM"
|
|
except ImportError:
|
|
# Fallback for standalone usage (not running in ComfyUI)
|
|
self.logger.warning("folder_paths not available, using fallback path")
|
|
self.models_dir = (
|
|
Path(__file__).parent.parent.parent.parent / "models" / "LLM"
|
|
)
|
|
|
|
self.logger.info(f"AutoTag service initialized. Models dir: {self.models_dir}")
|
|
|
|
# Current loaded tagger instance
|
|
self._tagger = None
|
|
self._current_model_type: Optional[str] = None
|
|
self._custom_prompt: str = DEFAULT_PROMPT
|
|
|
|
@property
|
|
def models_config(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Get the models configuration dictionary."""
|
|
return MODELS
|
|
|
|
@property
|
|
def default_prompt(self) -> str:
|
|
"""Get the default tag generation prompt."""
|
|
return DEFAULT_PROMPT
|
|
|
|
@property
|
|
def custom_prompt(self) -> str:
|
|
"""Get the current custom prompt."""
|
|
return self._custom_prompt
|
|
|
|
@custom_prompt.setter
|
|
def custom_prompt(self, value: str):
|
|
"""Set a custom prompt for tag generation."""
|
|
self._custom_prompt = value
|
|
|
|
@property
|
|
def wd14_general_threshold(self) -> float:
|
|
"""The saved General-tag confidence (``autotag.defaults``)."""
|
|
return AutotagConfig.GENERAL_THRESHOLD
|
|
|
|
@property
|
|
def wd14_character_threshold(self) -> float:
|
|
"""The saved Character-tag confidence (``autotag.defaults``)."""
|
|
return AutotagConfig.CHARACTER_THRESHOLD
|
|
|
|
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Get availability status for all model types.
|
|
|
|
Returns:
|
|
Dictionary with status for each model type:
|
|
{
|
|
'gguf': {
|
|
'name': 'GGUF (Recommended)',
|
|
'description': '...',
|
|
'size': '~4.5GB',
|
|
'downloaded': True,
|
|
'model_exists': True,
|
|
'mmproj_exists': True, # GGUF only
|
|
'model_path': '/path/to/model'
|
|
},
|
|
'hf': {...}
|
|
}
|
|
"""
|
|
status = {}
|
|
|
|
for model_type, config in MODELS.items():
|
|
model_status = {
|
|
"name": config["name"],
|
|
"description": config["description"],
|
|
"size": config["size"],
|
|
"downloaded": False,
|
|
"model_path": None,
|
|
}
|
|
|
|
if model_type == "gguf":
|
|
model_exists, mmproj_exists = self._check_gguf_models()
|
|
model_status["model_exists"] = model_exists
|
|
model_status["mmproj_exists"] = mmproj_exists
|
|
model_status["downloaded"] = model_exists and mmproj_exists
|
|
if model_status["downloaded"]:
|
|
model_status["model_path"] = str(
|
|
self.models_dir / config["subdir"] / config["filename"]
|
|
)
|
|
elif model_type.startswith("wd14"):
|
|
model_status["downloaded"] = self._check_wd14_model(model_type)
|
|
if model_status["downloaded"]:
|
|
model_status["model_path"] = str(
|
|
self.models_dir / config["subdir"] / "model.onnx"
|
|
)
|
|
else: # hf
|
|
model_status["downloaded"] = self._check_hf_model()
|
|
if model_status["downloaded"]:
|
|
# Get the actual path (local or cache)
|
|
model_status["model_path"] = str(self._get_hf_model_path())
|
|
|
|
status[model_type] = model_status
|
|
|
|
return status
|
|
|
|
def _check_gguf_models(self) -> Tuple[bool, bool]:
|
|
"""Check if GGUF model and mmproj files exist.
|
|
|
|
Returns:
|
|
Tuple of (model_exists, mmproj_exists)
|
|
"""
|
|
config = MODELS["gguf"]
|
|
gguf_dir = self.models_dir / config["subdir"]
|
|
model_path = gguf_dir / config["filename"]
|
|
mmproj_path = gguf_dir / config["mmproj_filename"]
|
|
return model_path.exists(), mmproj_path.exists()
|
|
|
|
def _check_hf_model(self) -> bool:
|
|
"""Check if HuggingFace model exists.
|
|
|
|
Checks both the local models directory and the HuggingFace cache.
|
|
|
|
Returns:
|
|
True if model directory contains config.json
|
|
"""
|
|
config = MODELS["hf"]
|
|
|
|
# Check local directory first
|
|
model_dir = self.models_dir / config["subdir"]
|
|
self.logger.debug(f"Checking local HF model path: {model_dir}")
|
|
if (model_dir / "config.json").exists():
|
|
self.logger.debug("Found HF model in local directory")
|
|
return True
|
|
|
|
# Check HuggingFace cache as fallback
|
|
hf_cache_path = self._get_hf_cache_path(config["repo"])
|
|
if hf_cache_path:
|
|
self.logger.debug(f"Found HF model in cache: {hf_cache_path}")
|
|
return True
|
|
|
|
self.logger.debug("HF model not found in local dir or cache")
|
|
return False
|
|
|
|
def _get_hf_cache_path(self, repo_id: str) -> Optional[Path]:
|
|
"""Get the path to a model in the HuggingFace cache.
|
|
|
|
Args:
|
|
repo_id: The HuggingFace repo ID
|
|
(e.g., 'fancyfeast/llama-joycaption-beta-one-hf-llava')
|
|
|
|
Returns:
|
|
Path to the cached model directory, or None if not found
|
|
"""
|
|
try:
|
|
from huggingface_hub import scan_cache_dir
|
|
except ImportError:
|
|
self.logger.debug("huggingface_hub not available for cache check")
|
|
return None
|
|
|
|
try:
|
|
cache_info = scan_cache_dir()
|
|
for repo in cache_info.repos:
|
|
if repo.repo_id == repo_id and repo.repo_type == "model":
|
|
# Get the latest revision's snapshot path
|
|
for revision in repo.revisions:
|
|
snapshot_path = revision.snapshot_path
|
|
if (Path(snapshot_path) / "config.json").exists():
|
|
return Path(snapshot_path)
|
|
except Exception as e:
|
|
self.logger.debug(f"Error scanning HF cache: {e}")
|
|
|
|
return None
|
|
|
|
def _get_hf_model_path(self) -> Optional[Path]:
|
|
"""Get the actual path to the HuggingFace model.
|
|
|
|
Checks local directory first, then HuggingFace cache.
|
|
|
|
Returns:
|
|
Path to the model directory, or None if not found
|
|
"""
|
|
config = MODELS["hf"]
|
|
|
|
# Check local directory first
|
|
model_dir = self.models_dir / config["subdir"]
|
|
if (model_dir / "config.json").exists():
|
|
return model_dir
|
|
|
|
# Check HuggingFace cache as fallback
|
|
cache_path = self._get_hf_cache_path(config["repo"])
|
|
if cache_path:
|
|
return cache_path
|
|
|
|
return None
|
|
|
|
def _check_wd14_model(self, model_type: str) -> bool:
|
|
"""Check if WD14 model files exist."""
|
|
config = MODELS[model_type]
|
|
model_dir = self.models_dir / config["subdir"]
|
|
return all((model_dir / f).exists() for f in config["files"])
|
|
|
|
def _download_wd14_model(
|
|
self,
|
|
model_type: str,
|
|
progress_callback: Optional[Callable[[str, float], None]] = None,
|
|
) -> bool:
|
|
"""Download WD14 model files from HuggingFace."""
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
config = MODELS[model_type]
|
|
model_dir = self.models_dir / config["subdir"]
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
files = config["files"]
|
|
for i, filename in enumerate(files):
|
|
file_path = model_dir / filename
|
|
if not file_path.exists():
|
|
progress = int((i / len(files)) * 80) + 10
|
|
if progress_callback:
|
|
progress_callback(f"Downloading {filename}...", progress)
|
|
|
|
self.logger.info(f"Downloading WD14 file: {filename}")
|
|
hf_hub_download(
|
|
repo_id=config["repo"],
|
|
filename=filename,
|
|
local_dir=str(model_dir),
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
|
|
if progress_callback:
|
|
progress_callback("Download complete", 100)
|
|
return True
|
|
|
|
def _load_wd14_tagger(self, model_type: str, use_gpu: bool = True):
|
|
"""Load WD14 ONNX model and tag list."""
|
|
try:
|
|
import onnxruntime as ort
|
|
except ImportError:
|
|
raise RuntimeError(
|
|
"onnxruntime is required for WD14 models. "
|
|
"Install with: pip install onnxruntime (or onnxruntime-gpu for CUDA)"
|
|
)
|
|
import csv
|
|
|
|
config = MODELS[model_type]
|
|
model_dir = self.models_dir / config["subdir"]
|
|
model_path = model_dir / "model.onnx"
|
|
tags_path = model_dir / "selected_tags.csv"
|
|
|
|
self.logger.info(f"Loading WD14 model from {model_dir}...")
|
|
|
|
# Set up ONNX session
|
|
providers = []
|
|
if use_gpu:
|
|
providers.append("CUDAExecutionProvider")
|
|
providers.append("CPUExecutionProvider")
|
|
|
|
session = ort.InferenceSession(str(model_path), providers=providers)
|
|
|
|
# Read tag list from CSV
|
|
tags_list = []
|
|
with open(tags_path, "r", encoding="utf-8") as f:
|
|
reader = csv.DictReader(f)
|
|
for row in reader:
|
|
tags_list.append(row)
|
|
|
|
self.logger.info(
|
|
f"WD14 model loaded ({len(tags_list)} tags, "
|
|
f"provider: {session.get_providers()[0]})"
|
|
)
|
|
return ("wd14", (session, tags_list))
|
|
|
|
def _generate_wd14(
|
|
self,
|
|
session,
|
|
tags_list: List[Dict[str, str]],
|
|
image: Image.Image,
|
|
general_threshold: float,
|
|
character_threshold: float,
|
|
) -> List[Tuple[str, float]]:
|
|
"""WD14 tags over their thresholds with probabilities, best first."""
|
|
import numpy as np
|
|
|
|
# Get model input dimensions
|
|
input_shape = session.get_inputs()[0].shape
|
|
target_size = input_shape[2] # e.g. 448
|
|
|
|
# Pad to square with white background
|
|
max_dim = max(image.size)
|
|
padded = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
|
|
padded.paste(
|
|
image, ((max_dim - image.width) // 2, (max_dim - image.height) // 2)
|
|
)
|
|
|
|
# Resize to model input size
|
|
padded = padded.resize((target_size, target_size), Image.Resampling.BILINEAR)
|
|
|
|
# Convert to numpy: RGB -> BGR, float32, add batch dim
|
|
img_array = np.array(padded, dtype=np.float32)
|
|
img_array = img_array[:, :, ::-1] # RGB to BGR
|
|
img_array = np.expand_dims(img_array, axis=0)
|
|
|
|
# Run inference
|
|
input_name = session.get_inputs()[0].name
|
|
output = session.run(None, {input_name: img_array})
|
|
probabilities = output[0][0] # First output, first batch
|
|
|
|
# Filter tags by category and threshold
|
|
result_tags = []
|
|
for i, tag_info in enumerate(tags_list):
|
|
if i >= len(probabilities):
|
|
break
|
|
prob = float(probabilities[i])
|
|
category = tag_info.get("category", "0")
|
|
tag_name = tag_info.get("name", "").strip()
|
|
|
|
if not tag_name:
|
|
continue
|
|
|
|
# Category 0 = general, 4 = character, 9 = rating
|
|
if category == "0" and prob >= general_threshold:
|
|
result_tags.append((tag_name, prob))
|
|
elif category == "4" and prob >= character_threshold:
|
|
result_tags.append((tag_name, prob))
|
|
|
|
# Sort by confidence descending
|
|
result_tags.sort(key=lambda x: x[1], reverse=True)
|
|
return result_tags
|
|
|
|
def download_model(
|
|
self,
|
|
model_type: str,
|
|
progress_callback: Optional[Callable[[str, float], None]] = None,
|
|
) -> bool:
|
|
"""Download a model with optional progress updates.
|
|
|
|
Args:
|
|
model_type: Either 'gguf' or 'hf'
|
|
progress_callback: Optional callback(status_message, progress_percent)
|
|
|
|
Returns:
|
|
True if download successful, False otherwise
|
|
|
|
Raises:
|
|
ValueError: If model_type is not valid
|
|
"""
|
|
if model_type not in MODELS:
|
|
raise ValueError(
|
|
f"Invalid model type: {model_type}. "
|
|
f"Must be one of: {', '.join(MODELS.keys())}"
|
|
)
|
|
|
|
try:
|
|
from huggingface_hub import ( # noqa: F401 - availability probe
|
|
hf_hub_download,
|
|
snapshot_download,
|
|
)
|
|
except ImportError:
|
|
self.logger.error("huggingface_hub not installed")
|
|
if progress_callback:
|
|
progress_callback("Error: huggingface_hub not installed", 0)
|
|
return False
|
|
|
|
try:
|
|
if model_type == "gguf":
|
|
return self._download_gguf_models(progress_callback)
|
|
elif model_type.startswith("wd14"):
|
|
return self._download_wd14_model(model_type, progress_callback)
|
|
else:
|
|
return self._download_hf_model(progress_callback)
|
|
except Exception as e:
|
|
self.logger.error(f"Download failed: {e}")
|
|
if progress_callback:
|
|
progress_callback(f"Error: {str(e)}", 0)
|
|
return False
|
|
|
|
def _download_gguf_models(
|
|
self, progress_callback: Optional[Callable[[str, float], None]] = None
|
|
) -> bool:
|
|
"""Download GGUF model and mmproj files."""
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
config = MODELS["gguf"]
|
|
gguf_dir = self.models_dir / config["subdir"]
|
|
gguf_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
model_path = gguf_dir / config["filename"]
|
|
mmproj_path = gguf_dir / config["mmproj_filename"]
|
|
|
|
# Download main model
|
|
if not model_path.exists():
|
|
if progress_callback:
|
|
progress_callback(f"Downloading {config['filename']}...", 10)
|
|
|
|
self.logger.info(f"Downloading GGUF model: {config['filename']}")
|
|
hf_hub_download(
|
|
repo_id=config["repo"],
|
|
filename=config["filename"],
|
|
local_dir=str(gguf_dir),
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
self.logger.info("GGUF model downloaded")
|
|
|
|
if progress_callback:
|
|
progress_callback("Main model ready", 50)
|
|
|
|
# Download mmproj
|
|
if not mmproj_path.exists():
|
|
if progress_callback:
|
|
progress_callback(f"Downloading {config['mmproj_filename']}...", 60)
|
|
|
|
self.logger.info(f"Downloading mmproj: {config['mmproj_filename']}")
|
|
hf_hub_download(
|
|
repo_id=config["mmproj_repo"],
|
|
filename=config["mmproj_filename"],
|
|
local_dir=str(gguf_dir),
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
self.logger.info("mmproj downloaded")
|
|
|
|
if progress_callback:
|
|
progress_callback("Download complete", 100)
|
|
|
|
return True
|
|
|
|
def _download_hf_model(
|
|
self, progress_callback: Optional[Callable[[str, float], None]] = None
|
|
) -> bool:
|
|
"""Download HuggingFace model."""
|
|
from huggingface_hub import snapshot_download
|
|
|
|
config = MODELS["hf"]
|
|
model_dir = self.models_dir / config["subdir"]
|
|
|
|
if not self._check_hf_model():
|
|
if progress_callback:
|
|
progress_callback(f"Downloading {config['repo']}...", 10)
|
|
|
|
self.logger.info(f"Downloading HF model: {config['repo']}")
|
|
snapshot_download(
|
|
repo_id=config["repo"],
|
|
local_dir=str(model_dir),
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
self.logger.info("HF model downloaded")
|
|
|
|
if progress_callback:
|
|
progress_callback("Download complete", 100)
|
|
|
|
return True
|
|
|
|
def load_model(self, model_type: str, use_gpu: bool = True) -> bool:
|
|
"""Load a model into memory for tag generation.
|
|
|
|
Args:
|
|
model_type: Either 'gguf' or 'hf'
|
|
use_gpu: Whether to use GPU acceleration (default True)
|
|
|
|
Returns:
|
|
True if model loaded successfully
|
|
|
|
Raises:
|
|
ValueError: If model_type is invalid
|
|
RuntimeError: If model not downloaded or loading fails
|
|
"""
|
|
if model_type not in MODELS:
|
|
raise ValueError(f"Invalid model type: {model_type}")
|
|
|
|
# Unload existing model first
|
|
if self._tagger is not None:
|
|
self.unload_model()
|
|
|
|
status = self.get_models_status()
|
|
if not status[model_type]["downloaded"]:
|
|
raise RuntimeError(f"Model {model_type} not downloaded")
|
|
|
|
try:
|
|
self._tagger = self._model_loader(model_type, use_gpu)
|
|
self._current_model_type = model_type
|
|
self.logger.info(f"Model {model_type} loaded successfully")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to load model {model_type}: {e}")
|
|
self._tagger = None
|
|
self._current_model_type = None
|
|
raise RuntimeError(f"Failed to load model: {e}")
|
|
|
|
def _load_model(self, model_type: str, use_gpu: bool = True) -> Tuple[str, Any]:
|
|
"""Default loader: dispatch to a backend, importing it only now."""
|
|
if model_type == "gguf":
|
|
return self._load_gguf_tagger(use_gpu)
|
|
if model_type.startswith("wd14"):
|
|
return self._load_wd14_tagger(model_type, use_gpu)
|
|
return self._load_hf_tagger()
|
|
|
|
def _load_gguf_tagger(self, use_gpu: bool = True):
|
|
"""Load GGUF-based tagger."""
|
|
from llama_cpp import Llama
|
|
from llama_cpp.llama_chat_format import Llava15ChatHandler
|
|
|
|
config = MODELS["gguf"]
|
|
gguf_dir = self.models_dir / config["subdir"]
|
|
model_path = gguf_dir / config["filename"]
|
|
mmproj_path = gguf_dir / config["mmproj_filename"]
|
|
|
|
self.logger.info("Loading GGUF model...")
|
|
n_gpu_layers = -1 if use_gpu else 0
|
|
|
|
tagger = Llama(
|
|
model_path=str(model_path),
|
|
n_ctx=4096,
|
|
n_batch=2048,
|
|
n_threads=4,
|
|
n_gpu_layers=n_gpu_layers,
|
|
verbose=False,
|
|
chat_handler=Llava15ChatHandler(clip_model_path=str(mmproj_path)),
|
|
offload_kqv=True,
|
|
)
|
|
|
|
self.logger.info("GGUF model loaded")
|
|
return ("gguf", tagger)
|
|
|
|
def _load_hf_tagger(self, quantization: str = "8bit"):
|
|
"""Load HuggingFace-based tagger."""
|
|
import torch
|
|
from transformers import (
|
|
AutoProcessor,
|
|
LlavaForConditionalGeneration,
|
|
BitsAndBytesConfig,
|
|
)
|
|
|
|
# Get the actual model path (local or cache)
|
|
model_path = self._get_hf_model_path()
|
|
if model_path is None:
|
|
raise RuntimeError(
|
|
"HuggingFace model not found in local directory or cache"
|
|
)
|
|
|
|
self.logger.info(f"Loading HuggingFace model from {model_path}...")
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
processor = AutoProcessor.from_pretrained(str(model_path))
|
|
model_kwargs = {"device_map": "cuda" if device == "cuda" else "cpu"}
|
|
|
|
if quantization == "8bit":
|
|
qnt_config = BitsAndBytesConfig(
|
|
load_in_8bit=True,
|
|
bnb_8bit_compute_dtype=torch.float16,
|
|
bnb_8bit_use_double_quant=True,
|
|
llm_int8_skip_modules=["vision_tower", "multi_modal_projector"],
|
|
)
|
|
model = LlavaForConditionalGeneration.from_pretrained(
|
|
str(model_path),
|
|
torch_dtype=torch.float16,
|
|
quantization_config=qnt_config,
|
|
**model_kwargs,
|
|
)
|
|
else:
|
|
model = LlavaForConditionalGeneration.from_pretrained(
|
|
str(model_path), torch_dtype=torch.bfloat16, **model_kwargs
|
|
)
|
|
|
|
model.eval()
|
|
self.logger.info("HuggingFace model loaded")
|
|
# Track the compute dtype for pixel_values conversion
|
|
compute_dtype = torch.float16 if quantization == "8bit" else torch.bfloat16
|
|
return ("hf", (model, processor, device, compute_dtype))
|
|
|
|
def unload_model(self):
|
|
"""Unload the current model and free memory."""
|
|
if self._tagger is not None:
|
|
self.logger.info(f"Unloading model: {self._current_model_type}")
|
|
del self._tagger
|
|
self._tagger = None
|
|
self._current_model_type = None
|
|
gc.collect()
|
|
|
|
# Clear the CUDA cache only if torch is already loaded; importing
|
|
# it here would pull in a heavy dependency the backend never used.
|
|
torch = sys.modules.get("torch")
|
|
if torch is not None and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
self.logger.info("Model unloaded, memory freed")
|
|
|
|
def is_model_loaded(self) -> bool:
|
|
"""Check if a model is currently loaded."""
|
|
return self._tagger is not None
|
|
|
|
def get_loaded_model_type(self) -> Optional[str]:
|
|
"""Get the type of currently loaded model."""
|
|
return self._current_model_type
|
|
|
|
def generate_tags(
|
|
self,
|
|
image_path: str,
|
|
prompt: Optional[str] = None,
|
|
general_threshold: Optional[float] = None,
|
|
character_threshold: Optional[float] = None,
|
|
) -> List[str]:
|
|
"""Generate tags for an image.
|
|
|
|
Args:
|
|
image_path: Path to the image file
|
|
prompt: Custom prompt for tag generation (LLM models only).
|
|
Uses default if None.
|
|
general_threshold: Confidence threshold for general tags (WD14 only).
|
|
character_threshold: Confidence threshold for character tags (WD14 only).
|
|
|
|
Returns:
|
|
List of generated tags
|
|
|
|
Raises:
|
|
RuntimeError: If no model is loaded
|
|
FileNotFoundError: If image doesn't exist
|
|
"""
|
|
scored = self._scored_tags(
|
|
image_path, prompt, general_threshold, character_threshold
|
|
)
|
|
return [name for name, _ in scored]
|
|
|
|
def generate_tag_suggestions(
|
|
self,
|
|
image_path: str,
|
|
*,
|
|
limit: int,
|
|
prompt: Optional[str] = None,
|
|
general_threshold: Optional[float] = None,
|
|
character_threshold: Optional[float] = None,
|
|
) -> List[Tuple[str, Optional[float]]]:
|
|
"""The top ``limit`` suggestions as ``(name, confidence)`` pairs.
|
|
|
|
WD14 pairs carry the real probability, best first; language models
|
|
(GGUF, HF) have no score, so their tags keep the output order with
|
|
``None``. Names are distinct case-insensitively (first one wins).
|
|
A name Save would reject (over 50 characters, a control character
|
|
such as a newline, a comma) is skipped before the dedupe and the
|
|
``limit`` count. ``prompt`` is used for this call only;
|
|
``custom_prompt`` is never changed.
|
|
"""
|
|
if type(limit) is not int or not 1 <= limit <= MAX_SUGGESTION_LIMIT:
|
|
raise ValueError(
|
|
f"limit must be an integer from 1 to {MAX_SUGGESTION_LIMIT}"
|
|
)
|
|
scored = self._scored_tags(
|
|
image_path, prompt, general_threshold, character_threshold
|
|
)
|
|
suggestions: List[Tuple[str, Optional[float]]] = []
|
|
seen = set()
|
|
for raw, score in scored:
|
|
try:
|
|
name = validate_vocabulary_tag_name(raw)
|
|
except ValueError:
|
|
continue
|
|
key = name.casefold()
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
suggestions.append((name, score))
|
|
if len(suggestions) == limit:
|
|
break
|
|
return suggestions
|
|
|
|
def _scored_tags(
|
|
self,
|
|
image_path: str,
|
|
prompt: Optional[str],
|
|
general_threshold: Optional[float],
|
|
character_threshold: Optional[float],
|
|
) -> List[Tuple[str, Optional[float]]]:
|
|
"""Run the loaded model: WD14 ``(name, probability)``, LLM ``(name, None)``."""
|
|
if self._tagger is None:
|
|
raise RuntimeError("No model loaded. Call load_model() first.")
|
|
|
|
if not os.path.exists(image_path):
|
|
raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path))
|
|
|
|
use_prompt = prompt or self._custom_prompt
|
|
image = _load_rgb(image_path)
|
|
|
|
# Generate based on model type
|
|
model_type, tagger_obj = self._tagger
|
|
|
|
if model_type == "gguf":
|
|
raw_tags = self._generate_gguf(tagger_obj, image, use_prompt)
|
|
return [(tag, None) for tag in self._parse_tags(raw_tags)]
|
|
elif model_type == "wd14":
|
|
session, tags_list = tagger_obj
|
|
gen_thresh = (
|
|
general_threshold
|
|
if general_threshold is not None
|
|
else AutotagConfig.GENERAL_THRESHOLD
|
|
)
|
|
char_thresh = (
|
|
character_threshold
|
|
if character_threshold is not None
|
|
else AutotagConfig.CHARACTER_THRESHOLD
|
|
)
|
|
return self._generate_wd14(
|
|
session, tags_list, image, gen_thresh, char_thresh
|
|
)
|
|
else:
|
|
model, processor, device, compute_dtype = tagger_obj
|
|
raw_tags = self._generate_hf(
|
|
model, processor, device, compute_dtype, image, use_prompt
|
|
)
|
|
return [(tag, None) for tag in self._parse_tags(raw_tags)]
|
|
|
|
def _generate_gguf(self, model, image: Image.Image, prompt: str) -> str:
|
|
"""Generate tags using GGUF model."""
|
|
import base64
|
|
import io
|
|
|
|
# Resize image
|
|
image = image.resize((336, 336), Image.Resampling.BILINEAR)
|
|
|
|
# Encode to base64
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="PNG")
|
|
buffer.seek(0)
|
|
img_base64 = base64.b64encode(buffer.read()).decode("utf-8")
|
|
data_uri = f"data:image/png;base64,{img_base64}"
|
|
|
|
# Create message
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": prompt},
|
|
{"type": "image_url", "image_url": {"url": data_uri}},
|
|
],
|
|
},
|
|
]
|
|
|
|
# Generate
|
|
response = model.create_chat_completion(
|
|
messages=messages,
|
|
max_tokens=512,
|
|
temperature=0.6,
|
|
top_p=0.9,
|
|
stop=["</s>", "User:", "Assistant:"],
|
|
stream=False,
|
|
)
|
|
|
|
return response["choices"][0]["message"]["content"].strip()
|
|
|
|
def _generate_hf(
|
|
self,
|
|
model,
|
|
processor,
|
|
device: str,
|
|
compute_dtype,
|
|
image: Image.Image,
|
|
prompt: str,
|
|
) -> str:
|
|
"""Generate tags using HuggingFace model."""
|
|
import torch
|
|
|
|
# Resize image
|
|
image = image.resize((336, 336), Image.Resampling.LANCZOS)
|
|
|
|
convo = [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": prompt},
|
|
]
|
|
convo_string = processor.apply_chat_template(
|
|
convo, tokenize=False, add_generation_prompt=True
|
|
)
|
|
|
|
inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(
|
|
device
|
|
)
|
|
|
|
# Convert pixel_values to the model's compute dtype
|
|
# (float16 for 8-bit quantized, bfloat16 for non-quantized)
|
|
if "pixel_values" in inputs and inputs["pixel_values"] is not None:
|
|
inputs["pixel_values"] = inputs["pixel_values"].to(compute_dtype)
|
|
|
|
with torch.inference_mode(), torch.cuda.amp.autocast(enabled=True):
|
|
generate_ids = model.generate(
|
|
**inputs,
|
|
max_new_tokens=512,
|
|
do_sample=True,
|
|
temperature=0.6,
|
|
top_p=0.9,
|
|
use_cache=True,
|
|
)[0]
|
|
|
|
generate_ids = generate_ids[inputs["input_ids"].shape[1] :]
|
|
return processor.tokenizer.decode(
|
|
generate_ids, skip_special_tokens=True
|
|
).strip()
|
|
|
|
def _parse_tags(self, raw_output: str) -> List[str]:
|
|
"""Parse raw model output into a list of clean tags.
|
|
|
|
Args:
|
|
raw_output: Raw text output from the model
|
|
|
|
Returns:
|
|
List of cleaned, deduplicated tags
|
|
"""
|
|
# Patterns to exclude
|
|
exclude_prefixes = (
|
|
"copyright:",
|
|
"meta:",
|
|
"photo_",
|
|
"photo:",
|
|
"prepend:",
|
|
"append:",
|
|
)
|
|
|
|
# Split by common delimiters
|
|
tags = []
|
|
|
|
# Handle comma-separated tags
|
|
for part in raw_output.split(","):
|
|
tag = part.strip().lower()
|
|
# Remove any quotes or extra characters
|
|
tag = tag.strip("\"'")
|
|
# Replace spaces with underscores (Danbooru style)
|
|
tag = tag.replace(" ", "_")
|
|
# Remove empty tags
|
|
if tag and len(tag) > 1:
|
|
# Filter out unwanted tag patterns
|
|
if not tag.startswith(exclude_prefixes):
|
|
tags.append(tag)
|
|
|
|
# Deduplicate while preserving order
|
|
seen = set()
|
|
unique_tags = []
|
|
for tag in tags:
|
|
if tag not in seen:
|
|
seen.add(tag)
|
|
unique_tags.append(tag)
|
|
|
|
return unique_tags
|
|
|
|
|
|
# Singleton instance for API use
|
|
_service_instance: Optional[AutoTagService] = None
|
|
|
|
|
|
def get_autotag_service() -> AutoTagService:
|
|
"""Get or create the singleton AutoTagService instance."""
|
|
global _service_instance
|
|
if _service_instance is None:
|
|
_service_instance = AutoTagService()
|
|
return _service_instance
|