Vito
9a90c7677f
* fix(exit): stop running jobs before Python joins the executor threads (#376) ComfyUI returns from main.py after Ctrl+C without closing its event loop, and Python joins every executor thread before any atexit callback runs, so a scan batch, thumbnail batch or model download in flight kept ComfyUI alive until it ended (on Windows each ffmpeg up to its 30 s timeout). The package now registers a threading._register_atexit hook (after concurrent.futures', so it runs before that join) that sets an exit signal and kills the running encoder processes. Job.cancelled follows the signal, so every check a job already makes between files stops it at exit; map_parallel starts no further item, the scan walk stops between folders and a tagging model download stops before its next chunk. Without the internal hook exit behaves as before. * fix(exit): abort Xet downloads and stop LoRA staging at exit (#376) Review follow-ups: Xet transfers report progress from native code and ignore the chunk checkpoint, so the exit hook calls huggingface_hub's own abort_xet_session when it is loaded; an aborted download logs as stopped. LoRA import staging reads the exit signal itself, so the legacy LoRA scan stops too. The test's fake encoder outlives the ceiling by less, and TESTING.md states the download limits of older huggingface_hub.
1102 lines
38 KiB
Python
1102 lines
38 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 inspect
|
|
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 .tagging_defaults import DEFAULT_TAG_TEMPLATE
|
|
except ImportError:
|
|
from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
|
|
|
|
try:
|
|
from ..utils import exit_signal
|
|
except ImportError:
|
|
from utils import exit_signal
|
|
|
|
|
|
def abort_hub_downloads() -> None:
|
|
"""Abort huggingface_hub's Xet transfers (run by the exit hook, #376).
|
|
|
|
Xet downloads report progress from native code, which ignores an
|
|
exception raised in ``tqdm.update``; huggingface_hub aborts them itself
|
|
on Ctrl+C with ``abort_xet_session``. Nothing happens when
|
|
huggingface_hub was never imported or has no Xet support.
|
|
"""
|
|
if "huggingface_hub" not in sys.modules:
|
|
return
|
|
try:
|
|
from huggingface_hub.utils import _xet
|
|
except ImportError:
|
|
return
|
|
abort = getattr(_xet, "abort_xet_session", None)
|
|
if callable(abort):
|
|
abort()
|
|
|
|
|
|
def _stop_at_exit(download) -> Dict[str, Any]:
|
|
"""Keyword arguments that stop a huggingface_hub download at its next
|
|
chunk once the interpreter is exiting (#376), or none when this
|
|
huggingface_hub does not take ``tqdm_class``.
|
|
|
|
huggingface_hub 1.x removes its partial temp file on any failure; older
|
|
versions keep their resumable ``.incomplete`` file. The model's own path
|
|
is only written by the final rename either way.
|
|
"""
|
|
try:
|
|
if "tqdm_class" not in inspect.signature(download).parameters:
|
|
return {}
|
|
from huggingface_hub.utils import tqdm as hub_tqdm
|
|
except (ImportError, TypeError, ValueError):
|
|
return {}
|
|
|
|
class StopAtExit(hub_tqdm):
|
|
def update(self, n=1):
|
|
exit_signal.checkpoint()
|
|
return super().update(n)
|
|
|
|
return {"tqdm_class": StopAtExit}
|
|
|
|
|
|
def _autotag_config():
|
|
"""The saved defaults; imported late, as config needs ComfyUI's server."""
|
|
try:
|
|
from .config import AutotagConfig
|
|
except ImportError:
|
|
from py.config import AutotagConfig
|
|
return 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. ``label`` is the short
|
|
# name System, Activity and a 409 model-busy show (#328).
|
|
MODELS = {
|
|
"gguf": {
|
|
"name": "GGUF (Recommended)",
|
|
"label": "GGUF JoyCaption",
|
|
"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",
|
|
"label": "HuggingFace JoyCaption",
|
|
"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",
|
|
"label": "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",
|
|
"label": "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 _autotag_config().GENERAL_THRESHOLD
|
|
|
|
@property
|
|
def wd14_character_threshold(self) -> float:
|
|
"""The saved Character-tag confidence (``autotag.defaults``)."""
|
|
return _autotag_config().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,
|
|
**_stop_at_exit(hf_hub_download),
|
|
)
|
|
|
|
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:
|
|
if exit_signal.requested():
|
|
# ExitRequested between chunks, or the Xet session aborted
|
|
# by the exit hook (abort_hub_downloads).
|
|
self.logger.info(
|
|
f"Model download stopped: ComfyUI is exiting ({model_type})"
|
|
)
|
|
return False
|
|
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,
|
|
**_stop_at_exit(hf_hub_download),
|
|
)
|
|
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,
|
|
**_stop_at_exit(hf_hub_download),
|
|
)
|
|
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,
|
|
**_stop_at_exit(snapshot_download),
|
|
)
|
|
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 wd14_device(self) -> str:
|
|
"""``"gpu"`` when WD14 runs (or would run) on CUDA, else ``"cpu"``.
|
|
|
|
A loaded WD14 session answers from its own first provider; otherwise
|
|
ONNX Runtime's available providers say whether CUDA would be offered
|
|
(a session can still fall back to the processor). Loads no model.
|
|
"""
|
|
tagger = self._tagger
|
|
if tagger is not None and tagger[0] == "wd14":
|
|
providers = tagger[1][0].get_providers()
|
|
return "gpu" if providers[:1] == ["CUDAExecutionProvider"] else "cpu"
|
|
try:
|
|
import onnxruntime as ort
|
|
except ImportError:
|
|
return "cpu"
|
|
available = ort.get_available_providers()
|
|
return "gpu" if "CUDAExecutionProvider" in available else "cpu"
|
|
|
|
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 _autotag_config().GENERAL_THRESHOLD
|
|
)
|
|
char_thresh = (
|
|
character_threshold
|
|
if character_threshold is not None
|
|
else _autotag_config().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
|