Vito Sansevero
489d449ec8
- py/model_gate.py: one process-wide, non-blocking gate for loading,
swapping, running, downloading and unloading a tagging model. Worker
calls run shielded; a release requested while one still runs takes
effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
gate and answer 409 {code: "model-busy"} while another operation holds
it; their bodies are otherwise unchanged. /autotag/models gains
suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
first, GGUF/HF (name, None) in output order, first `limit` distinct
case-insensitively; generate_tags keeps its return; the image file is
closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
pinned capture, with strict per-engine options, never downloads
(409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
metadata over the inspection budgets no longer blocks it.
1007 lines
35 KiB
Python
1007 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]]
|
|
|
|
# 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 = (
|
|
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
|
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
|
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
|
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
|
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
|
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
|
"Prioritize visual accuracy.\n"
|
|
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
|
"descriptive."
|
|
)
|
|
|
|
# Upper bound of one suggestion request: the image tag set cap (100 names).
|
|
MAX_SUGGESTION_LIMIT = 100
|
|
|
|
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
|
|
|
|
# WD14 threshold defaults
|
|
self._wd14_general_threshold: float = 0.35
|
|
self._wd14_character_threshold: float = 0.85
|
|
|
|
@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:
|
|
return self._wd14_general_threshold
|
|
|
|
@wd14_general_threshold.setter
|
|
def wd14_general_threshold(self, value: float):
|
|
self._wd14_general_threshold = max(0.0, min(1.0, value))
|
|
|
|
@property
|
|
def wd14_character_threshold(self) -> float:
|
|
return self._wd14_character_threshold
|
|
|
|
@wd14_character_threshold.setter
|
|
def wd14_character_threshold(self, value: float):
|
|
self._wd14_character_threshold = max(0.0, min(1.0, value))
|
|
|
|
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).
|
|
``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 name, score in scored:
|
|
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 self._wd14_general_threshold
|
|
)
|
|
char_thresh = (
|
|
character_threshold
|
|
if character_threshold is not None
|
|
else self._wd14_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
|