974 lines
31 KiB
Python
974 lines
31 KiB
Python
"""
|
|
Tag Autocomplete API
|
|
|
|
Provides tag autocomplete functionality with support for:
|
|
- Danbooru/e621 tag databases
|
|
- Alias searching
|
|
- Post count sorting with optional frequency boosting
|
|
"""
|
|
|
|
import asyncio
|
|
import csv
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
try:
|
|
from ..parser.wildcards import get_wildcard_file_index
|
|
except ImportError:
|
|
from parser.wildcards import get_wildcard_file_index
|
|
|
|
# Conditional server import
|
|
try:
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
_HAS_SERVER = True
|
|
except ImportError:
|
|
_HAS_SERVER = False
|
|
web = None
|
|
PromptServer = None
|
|
|
|
# Tag type definitions (for color coding in frontend)
|
|
TAG_TYPES = {
|
|
0: "general", # lightblue/dodgerblue
|
|
1: "artist", # indianred/firebrick
|
|
3: "copyright", # violet/darkorchid
|
|
4: "character", # lightgreen/darkgreen
|
|
5: "meta", # orange/darkorange
|
|
}
|
|
|
|
# Get the data directory path
|
|
DATA_DIR = Path(__file__).parent.parent / "data" / "tags"
|
|
WILDCARDS_DIR = Path(__file__).parent.parent / "data" / "wildcards"
|
|
|
|
# Default tag file to use if not specified (danbooru.csv is default)
|
|
DEFAULT_TAG_FILE = "danbooru.csv"
|
|
|
|
|
|
def _coerce_text(value) -> str:
|
|
if value is None:
|
|
return ""
|
|
if isinstance(value, str):
|
|
return value
|
|
return str(value)
|
|
|
|
|
|
def _coerce_bool(value, default: bool = True) -> bool:
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, str):
|
|
normalized = value.strip().lower()
|
|
if normalized in {"1", "true", "yes", "on"}:
|
|
return True
|
|
if normalized in {"0", "false", "no", "off"}:
|
|
return False
|
|
return default
|
|
|
|
|
|
def _coerce_limit(value, default: int = 20, cap: int = 100) -> int:
|
|
try:
|
|
limit = int(value)
|
|
except (TypeError, ValueError):
|
|
limit = default
|
|
return max(0, min(limit, cap))
|
|
|
|
|
|
def _normalize_tag_filename(value, default: Optional[str] = None) -> Optional[str]:
|
|
filename = _coerce_text(value).strip()
|
|
if not filename:
|
|
return default
|
|
|
|
filename = filename.replace("\\", "/")
|
|
path = Path(filename)
|
|
if path.is_absolute() or ".." in path.parts:
|
|
return default
|
|
return filename
|
|
|
|
|
|
def _normalize_extra_files(value) -> List[str]:
|
|
if value is None:
|
|
value = ["extra-quality-tags.csv"]
|
|
elif isinstance(value, str):
|
|
value = [value]
|
|
elif not isinstance(value, list):
|
|
value = ["extra-quality-tags.csv"]
|
|
|
|
files = []
|
|
for item in value:
|
|
filename = _normalize_tag_filename(item)
|
|
if filename and filename not in files:
|
|
files.append(filename)
|
|
return files
|
|
|
|
|
|
class TagEntry:
|
|
"""Represents a single tag entry from the database."""
|
|
|
|
__slots__ = ("name", "type", "count", "aliases", "search_text")
|
|
|
|
def __init__(self, name: str, tag_type: int, count: int, aliases: List[str]):
|
|
self.name = name
|
|
self.type = tag_type
|
|
self.count = count
|
|
self.aliases = aliases
|
|
# Pre-compute lowercase search text for faster matching
|
|
self.search_text = name.lower()
|
|
|
|
def to_dict(self) -> Dict:
|
|
return {
|
|
"name": self.name,
|
|
"type": self.type,
|
|
"count": self.count,
|
|
"aliases": self.aliases,
|
|
"kind": "tag",
|
|
}
|
|
|
|
|
|
class WildcardEntry:
|
|
"""Represents a single wildcard file available for completion."""
|
|
|
|
__slots__ = ("name", "kind", "search_text", "leaf_name")
|
|
|
|
def __init__(self, name: str, kind: str):
|
|
self.name = name
|
|
self.kind = kind
|
|
self.search_text = name.lower()
|
|
self.leaf_name = name.rsplit("/", 1)[-1]
|
|
|
|
def to_dict(self, completion_name: str = "") -> Dict:
|
|
completion_name = completion_name or self.leaf_name
|
|
return {
|
|
"name": self.name,
|
|
"leaf_name": self.leaf_name,
|
|
"kind": self.kind,
|
|
"completion": f"__{completion_name}__",
|
|
"meta": "Wildcard folder" if self.kind == "wildcard_folder" else "Wildcard file",
|
|
}
|
|
|
|
|
|
class TagDatabase:
|
|
"""
|
|
In-memory tag database with fast prefix search.
|
|
|
|
Tags are loaded lazily on first search to avoid slowing down ComfyUI startup.
|
|
Supports loading multiple tag files and merging results.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._tags: List[TagEntry] = []
|
|
self._alias_map: Dict[str, TagEntry] = {} # alias -> canonical tag
|
|
self._prefix_index: Dict[str, List[Tuple[TagEntry, Optional[str], str]]] = {}
|
|
self._contains_index: Dict[str, List[Tuple[TagEntry, Optional[str], str]]] = {}
|
|
self._loaded = False
|
|
self._current_files: List[str] = []
|
|
|
|
@property
|
|
def is_loaded(self) -> bool:
|
|
return self._loaded
|
|
|
|
@property
|
|
def tag_count(self) -> int:
|
|
return len(self._tags)
|
|
|
|
def _add_prefix_index_entry(
|
|
self, text: str, tag: TagEntry, matched_alias: Optional[str] = None
|
|
) -> None:
|
|
normalized = text.lower().strip()
|
|
if not normalized:
|
|
return
|
|
|
|
for prefix_len in range(1, min(3, len(normalized)) + 1):
|
|
prefix = normalized[:prefix_len]
|
|
self._prefix_index.setdefault(prefix, []).append(
|
|
(tag, matched_alias, normalized)
|
|
)
|
|
|
|
def _get_contains_index_key(self, query: str) -> str:
|
|
return query[: min(3, len(query))]
|
|
|
|
def _get_contains_candidates(
|
|
self, query: str
|
|
) -> List[Tuple[TagEntry, Optional[str], str]]:
|
|
key = self._get_contains_index_key(query)
|
|
if not key:
|
|
return []
|
|
|
|
if key not in self._contains_index:
|
|
candidates: List[Tuple[TagEntry, Optional[str], str]] = []
|
|
for tag in self._tags:
|
|
if key in tag.search_text:
|
|
candidates.append((tag, None, tag.search_text))
|
|
for alias_lower, tag in self._alias_map.items():
|
|
if key in alias_lower:
|
|
original_alias = next(
|
|
(a for a in tag.aliases if a.lower() == alias_lower),
|
|
alias_lower,
|
|
)
|
|
candidates.append((tag, original_alias, alias_lower))
|
|
self._contains_index[key] = candidates
|
|
|
|
return self._contains_index[key]
|
|
|
|
def load_csv(self, filepath: Path, append: bool = False) -> int:
|
|
"""
|
|
Load tags from a CSV file.
|
|
|
|
Args:
|
|
filepath: Path to the CSV file
|
|
append: If True, append to existing tags instead of replacing
|
|
|
|
Returns:
|
|
Number of tags loaded from this file
|
|
|
|
CSV Format: name,type,postCount,"aliases"
|
|
Example: 1girl,0,6008644,"1girls,sole_female"
|
|
"""
|
|
if not append:
|
|
self._tags = []
|
|
self._alias_map = {}
|
|
self._prefix_index = {}
|
|
self._current_files = []
|
|
self._contains_index = {}
|
|
|
|
if not filepath.exists():
|
|
print(f"[Autocomplete] Tag file not found: {filepath}")
|
|
return 0
|
|
|
|
try:
|
|
loaded_count = 0
|
|
existing_tags = {tag.name for tag in self._tags}
|
|
|
|
with open(filepath, "r", encoding="utf-8", errors="replace") as f:
|
|
reader = csv.reader(f)
|
|
for row in reader:
|
|
if len(row) < 3:
|
|
continue
|
|
|
|
name = row[0].strip()
|
|
if not name or name in existing_tags:
|
|
continue
|
|
|
|
try:
|
|
tag_type = int(row[1])
|
|
except (ValueError, IndexError):
|
|
tag_type = 0
|
|
|
|
try:
|
|
count = int(row[2])
|
|
except (ValueError, IndexError):
|
|
count = 0
|
|
|
|
# Parse aliases (4th column, comma-separated in quotes)
|
|
aliases = []
|
|
if len(row) > 3 and row[3]:
|
|
aliases = [a.strip() for a in row[3].split(",") if a.strip()]
|
|
|
|
entry = TagEntry(name, tag_type, count, aliases)
|
|
self._tags.append(entry)
|
|
existing_tags.add(name)
|
|
loaded_count += 1
|
|
self._add_prefix_index_entry(name, entry)
|
|
|
|
# Build alias -> canonical tag mapping
|
|
for alias in aliases:
|
|
self._alias_map[alias.lower()] = entry
|
|
self._add_prefix_index_entry(alias, entry, alias)
|
|
|
|
self._loaded = True
|
|
self._current_files.append(filepath.name)
|
|
print(f"[Autocomplete] Loaded {loaded_count} tags from {filepath.name}")
|
|
return loaded_count
|
|
|
|
except Exception as e:
|
|
print(f"[Autocomplete] Error loading tag file: {e}")
|
|
return 0
|
|
|
|
def load_multiple(self, filepaths: List[Path]) -> int:
|
|
"""
|
|
Load tags from multiple CSV files, merging results.
|
|
|
|
Args:
|
|
filepaths: List of paths to CSV files
|
|
|
|
Returns:
|
|
Total number of tags loaded
|
|
"""
|
|
total = 0
|
|
for i, filepath in enumerate(filepaths):
|
|
count = self.load_csv(filepath, append=(i > 0))
|
|
total += count
|
|
return total
|
|
|
|
def search(
|
|
self,
|
|
query: str,
|
|
limit: int = 20,
|
|
search_aliases: bool = True,
|
|
contains_fallback: bool = True,
|
|
) -> List[Dict]:
|
|
"""
|
|
Search for tags matching the query.
|
|
|
|
Args:
|
|
query: Search query (prefix match)
|
|
limit: Maximum number of results
|
|
search_aliases: Whether to also search aliases
|
|
|
|
Returns:
|
|
List of matching tags as dictionaries
|
|
"""
|
|
query = _coerce_text(query)
|
|
limit = _coerce_limit(limit, default=20, cap=100)
|
|
|
|
if not self._loaded or not query or limit == 0:
|
|
return []
|
|
|
|
query_lower = query.lower().strip()
|
|
if not query_lower or len(query_lower) > 255:
|
|
return []
|
|
|
|
results: List[Tuple[int, TagEntry, Optional[str]]] = []
|
|
seen_tags = set()
|
|
|
|
# Score function: exact match > prefix match > contains match
|
|
def get_score(tag: TagEntry, matched_text: str) -> int:
|
|
if matched_text == query_lower:
|
|
return 3 # Exact match
|
|
elif matched_text.startswith(query_lower):
|
|
return 2 # Prefix match
|
|
else:
|
|
return 1 # Contains match
|
|
|
|
def add_result(tag: TagEntry, matched_alias: Optional[str], matched_text: str):
|
|
if tag.name in seen_tags:
|
|
return
|
|
score = get_score(tag, matched_text)
|
|
results.append((score, tag, matched_alias))
|
|
seen_tags.add(tag.name)
|
|
|
|
def format_results() -> List[Dict]:
|
|
results.sort(key=lambda x: (-x[0], -x[1].count))
|
|
output = []
|
|
for _, tag, matched_alias in results[:limit]:
|
|
entry = tag.to_dict()
|
|
if matched_alias:
|
|
entry["matched_alias"] = matched_alias
|
|
output.append(entry)
|
|
return output
|
|
|
|
prefix_key = query_lower[: min(3, len(query_lower))]
|
|
for tag, matched_alias, matched_text in self._prefix_index.get(prefix_key, []):
|
|
if matched_alias and not search_aliases:
|
|
continue
|
|
if matched_text == query_lower or matched_text.startswith(query_lower):
|
|
add_result(tag, matched_alias, matched_text)
|
|
|
|
if len(results) >= limit:
|
|
return format_results()
|
|
|
|
if not contains_fallback:
|
|
return format_results()
|
|
|
|
if len(query_lower) >= 2:
|
|
candidates = self._get_contains_candidates(query_lower)
|
|
for tag, matched_alias, matched_text in candidates:
|
|
if matched_alias and not search_aliases:
|
|
continue
|
|
if query_lower in matched_text:
|
|
add_result(tag, matched_alias, matched_text)
|
|
else:
|
|
# Backend-only fallback for one-character API searches. The frontend
|
|
# starts at two characters, where the lazy contains index applies.
|
|
for tag in self._tags:
|
|
if tag.name in seen_tags:
|
|
continue
|
|
if query_lower in tag.search_text:
|
|
add_result(tag, None, tag.search_text)
|
|
|
|
if search_aliases:
|
|
for alias_lower, tag in self._alias_map.items():
|
|
if tag.name in seen_tags:
|
|
continue
|
|
if query_lower in alias_lower:
|
|
original_alias = next(
|
|
(a for a in tag.aliases if a.lower() == alias_lower),
|
|
alias_lower,
|
|
)
|
|
add_result(tag, original_alias, alias_lower)
|
|
|
|
return format_results()
|
|
|
|
|
|
class WildcardDatabase:
|
|
"""In-memory wildcard list with simple fuzzy scoring."""
|
|
|
|
def __init__(self, wildcards_dir: Path):
|
|
self._wildcards_dir = wildcards_dir
|
|
self._entries: List[WildcardEntry] = []
|
|
self._entry_map: Dict[str, List[WildcardEntry]] = {}
|
|
self._leaf_counts: Dict[str, int] = {}
|
|
self._load_lock = threading.Lock()
|
|
self._loaded = False
|
|
|
|
@property
|
|
def is_loaded(self) -> bool:
|
|
return self._loaded
|
|
|
|
@property
|
|
def wildcard_count(self) -> int:
|
|
return len(self._entries)
|
|
|
|
def load(self) -> int:
|
|
with self._load_lock:
|
|
return self._load_unlocked()
|
|
|
|
def _load_unlocked(self) -> int:
|
|
if not self._wildcards_dir.exists():
|
|
self._entries = []
|
|
self._entry_map = {}
|
|
self._leaf_counts = {}
|
|
self._loaded = True
|
|
return 0
|
|
|
|
entries: List[WildcardEntry] = []
|
|
folder_names = set()
|
|
for path in get_wildcard_file_index(
|
|
self._wildcards_dir, refresh=True
|
|
).files:
|
|
relative = path.relative_to(self._wildcards_dir).with_suffix("")
|
|
relative_name = relative.as_posix()
|
|
entries.append(WildcardEntry(relative_name, "wildcard_file"))
|
|
|
|
parts = relative.parts[:-1]
|
|
current_parts = []
|
|
for part in parts:
|
|
current_parts.append(part)
|
|
folder_names.add("/".join(current_parts))
|
|
|
|
for folder_name in folder_names:
|
|
entries.append(WildcardEntry(folder_name, "wildcard_folder"))
|
|
|
|
entries.sort(
|
|
key=lambda entry: (
|
|
entry.search_text,
|
|
0 if entry.kind == "wildcard_folder" else 1,
|
|
len(entry.name),
|
|
)
|
|
)
|
|
entry_map: Dict[str, List[WildcardEntry]] = {}
|
|
leaf_counts: Dict[str, int] = {}
|
|
for entry in entries:
|
|
entry_map.setdefault(entry.search_text, []).append(entry)
|
|
leaf_key = entry.leaf_name.lower()
|
|
leaf_counts[leaf_key] = leaf_counts.get(leaf_key, 0) + 1
|
|
|
|
self._entries = entries
|
|
self._entry_map = entry_map
|
|
self._leaf_counts = leaf_counts
|
|
self._loaded = True
|
|
return len(entries)
|
|
|
|
def ensure_loaded(self) -> int:
|
|
if self._loaded:
|
|
return len(self._entries)
|
|
with self._load_lock:
|
|
if not self._loaded:
|
|
return self._load_unlocked()
|
|
return len(self._entries)
|
|
|
|
def _entry_to_dict(self, entry: WildcardEntry) -> Dict:
|
|
completion_name = (
|
|
entry.leaf_name
|
|
if self._leaf_counts.get(entry.leaf_name.lower(), 0) == 1
|
|
else entry.name
|
|
)
|
|
return entry.to_dict(completion_name=completion_name)
|
|
|
|
def get_manifest(self) -> List[Dict]:
|
|
"""Return every wildcard file and folder for deterministic client search."""
|
|
self.ensure_loaded()
|
|
return [self._entry_to_dict(entry) for entry in self._entries]
|
|
|
|
def _normalize_query(self, query: str) -> str:
|
|
normalized = str(query or "").strip().lower()
|
|
if normalized.startswith("__"):
|
|
normalized = normalized[2:]
|
|
if normalized.endswith("__"):
|
|
normalized = normalized[:-2]
|
|
return normalized
|
|
|
|
def _resolve_entry(self, wildcard_name: str) -> Optional[WildcardEntry]:
|
|
normalized = self._normalize_query(wildcard_name)
|
|
if not normalized:
|
|
return None
|
|
|
|
candidates = self._entry_map.get(normalized, [])
|
|
if not candidates:
|
|
return None
|
|
|
|
for kind in ("wildcard_folder", "wildcard_file"):
|
|
for entry in candidates:
|
|
if entry.kind == kind:
|
|
return entry
|
|
return candidates[0]
|
|
|
|
def search(self, query: str, limit: int = 20) -> List[Dict]:
|
|
if not self._loaded:
|
|
self.ensure_loaded()
|
|
|
|
limit = _coerce_limit(limit, default=20, cap=500)
|
|
normalized = self._normalize_query(query)
|
|
if len(normalized) > 255 or limit == 0:
|
|
return []
|
|
|
|
if not normalized:
|
|
return [self._entry_to_dict(entry) for entry in self._entries[:limit]]
|
|
|
|
scored: List[Tuple[int, WildcardEntry]] = []
|
|
for entry in self._entries:
|
|
if normalized == entry.search_text:
|
|
score = 3
|
|
elif entry.search_text.startswith(normalized):
|
|
score = 2
|
|
elif normalized in entry.search_text:
|
|
score = 1
|
|
else:
|
|
continue
|
|
scored.append((score, entry))
|
|
|
|
scored.sort(
|
|
key=lambda item: (
|
|
-item[0],
|
|
0 if item[1].kind == "wildcard_folder" else 1,
|
|
len(item[1].name),
|
|
item[1].search_text,
|
|
)
|
|
)
|
|
return [self._entry_to_dict(entry) for _, entry in scored[:limit]]
|
|
|
|
def get_contents(
|
|
self, wildcard_name: str, content_query: str = "", limit: int = 50
|
|
) -> List[Dict]:
|
|
if not self._loaded:
|
|
self.ensure_loaded()
|
|
|
|
limit = _coerce_limit(limit, default=50, cap=500)
|
|
if limit == 0:
|
|
return []
|
|
|
|
entry = self._resolve_entry(wildcard_name)
|
|
if entry is None:
|
|
return []
|
|
|
|
if entry.kind == "wildcard_folder":
|
|
prefix = f"{entry.name}/"
|
|
normalized_query = _coerce_text(content_query).strip().lower()
|
|
descendants = []
|
|
|
|
for candidate in self._entries:
|
|
if candidate.name == entry.name or not candidate.name.startswith(prefix):
|
|
continue
|
|
remainder = candidate.name[len(prefix) :]
|
|
if normalized_query and normalized_query not in remainder.lower():
|
|
continue
|
|
descendants.append(self._entry_to_dict(candidate))
|
|
if len(descendants) >= limit:
|
|
break
|
|
|
|
return descendants
|
|
|
|
wildcard_path = self._wildcards_dir.joinpath(*entry.name.split("/")).with_suffix(
|
|
".txt"
|
|
)
|
|
if not wildcard_path.exists():
|
|
return []
|
|
|
|
normalized_query = _coerce_text(content_query).strip().lower()
|
|
results = []
|
|
with open(wildcard_path, "r", encoding="utf-8", errors="replace") as handle:
|
|
for raw_line in handle:
|
|
line = raw_line.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
if normalized_query and normalized_query not in line.lower():
|
|
continue
|
|
results.append(
|
|
{
|
|
"name": line,
|
|
"kind": "wildcard_content",
|
|
"completion": line,
|
|
"meta": entry.name,
|
|
}
|
|
)
|
|
if len(results) >= limit:
|
|
break
|
|
|
|
return results
|
|
|
|
|
|
# Global database instance (lazy loaded)
|
|
_database = TagDatabase()
|
|
_wildcard_database = WildcardDatabase(WILDCARDS_DIR)
|
|
_database_load_lock = threading.Lock()
|
|
_database_loading = False
|
|
_database_load_error: Optional[str] = None
|
|
_database_loading_files: List[str] = []
|
|
_wildcard_warmup_lock = threading.Lock()
|
|
_wildcard_warmup_loading = False
|
|
|
|
|
|
def get_database() -> TagDatabase:
|
|
"""Get the global tag database instance."""
|
|
return _database
|
|
|
|
|
|
def get_wildcard_database() -> WildcardDatabase:
|
|
"""Get the global wildcard database instance."""
|
|
return _wildcard_database
|
|
|
|
|
|
def _build_tag_file_list(
|
|
tag_file: str = DEFAULT_TAG_FILE,
|
|
extra_files: Optional[List[str]] = None,
|
|
) -> List[str]:
|
|
safe_tag_file = _normalize_tag_filename(tag_file, DEFAULT_TAG_FILE)
|
|
files_to_load = [safe_tag_file]
|
|
if extra_files:
|
|
files_to_load.extend(_normalize_extra_files(extra_files))
|
|
return [filename for filename in files_to_load if filename]
|
|
|
|
|
|
def _resolve_tag_filepaths(files_to_load: List[str]) -> List[Path]:
|
|
filepaths = []
|
|
|
|
for filename in files_to_load:
|
|
# Try main data directory first
|
|
filepath = DATA_DIR / filename
|
|
|
|
# Fall back to reference directory
|
|
if not filepath.exists():
|
|
ref_path = (
|
|
Path(__file__).parent.parent
|
|
/ "autocomplete_reference"
|
|
/ "a1111-sd-webui-tagcomplete"
|
|
/ "tags"
|
|
/ filename
|
|
)
|
|
if ref_path.exists():
|
|
filepath = ref_path
|
|
|
|
if filepath.exists():
|
|
filepaths.append(filepath)
|
|
else:
|
|
print(f"[Autocomplete] Warning: Tag file not found: {filename}")
|
|
|
|
return filepaths
|
|
|
|
|
|
def get_database_status() -> Dict:
|
|
db = get_database()
|
|
return {
|
|
"loaded": db.is_loaded,
|
|
"loading": _database_loading,
|
|
"load_error": _database_load_error,
|
|
"tag_count": db.tag_count,
|
|
"current_files": db._current_files,
|
|
"loading_files": _database_loading_files,
|
|
}
|
|
|
|
|
|
def ensure_database_loaded(
|
|
tag_file: str = DEFAULT_TAG_FILE,
|
|
extra_files: Optional[List[str]] = None,
|
|
) -> TagDatabase:
|
|
"""
|
|
Ensure the database is loaded, loading it if necessary.
|
|
|
|
Args:
|
|
tag_file: Name of the main tag file to load
|
|
extra_files: Optional list of additional tag files to merge (e.g., ["extra-quality-tags.csv"])
|
|
|
|
Returns:
|
|
The loaded TagDatabase instance
|
|
"""
|
|
global _database_loading, _database_load_error, _database_loading_files
|
|
db = get_database()
|
|
|
|
files_to_load = _build_tag_file_list(tag_file, extra_files)
|
|
|
|
# Check if we need to reload
|
|
needs_reload = not db.is_loaded or set(db._current_files) != set(files_to_load)
|
|
if not needs_reload:
|
|
if _database_loading and set(_database_loading_files) == set(files_to_load):
|
|
_database_loading = False
|
|
_database_loading_files = []
|
|
return db
|
|
|
|
with _database_load_lock:
|
|
needs_reload = not db.is_loaded or set(db._current_files) != set(files_to_load)
|
|
if not needs_reload:
|
|
if _database_loading and set(_database_loading_files) == set(files_to_load):
|
|
_database_loading = False
|
|
_database_loading_files = []
|
|
return db
|
|
|
|
_database_loading = True
|
|
_database_load_error = None
|
|
_database_loading_files = files_to_load[:]
|
|
|
|
try:
|
|
filepaths = _resolve_tag_filepaths(files_to_load)
|
|
if filepaths:
|
|
db.load_multiple(filepaths)
|
|
except Exception as exc:
|
|
_database_load_error = str(exc)
|
|
raise
|
|
finally:
|
|
_database_loading = False
|
|
_database_loading_files = []
|
|
|
|
return db
|
|
|
|
|
|
def start_database_warmup(
|
|
tag_file: str = DEFAULT_TAG_FILE,
|
|
extra_files: Optional[List[str]] = None,
|
|
) -> bool:
|
|
global _database_loading, _database_load_error, _database_loading_files
|
|
files_to_load = _build_tag_file_list(tag_file, extra_files)
|
|
db = get_database()
|
|
with _database_load_lock:
|
|
if db.is_loaded and set(db._current_files) == set(files_to_load):
|
|
return False
|
|
if _database_loading:
|
|
return False
|
|
_database_loading = True
|
|
_database_load_error = None
|
|
_database_loading_files = files_to_load[:]
|
|
|
|
def warmup_worker():
|
|
try:
|
|
ensure_database_loaded(tag_file, extra_files=extra_files)
|
|
except Exception as exc:
|
|
print(f"[Autocomplete] Warm-up failed: {exc}")
|
|
|
|
thread = threading.Thread(
|
|
target=warmup_worker,
|
|
name="A1111AutocompleteWarmup",
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
return True
|
|
|
|
|
|
def start_wildcard_warmup() -> bool:
|
|
global _wildcard_warmup_loading
|
|
db = get_wildcard_database()
|
|
|
|
with _wildcard_warmup_lock:
|
|
if db.is_loaded or _wildcard_warmup_loading:
|
|
return False
|
|
_wildcard_warmup_loading = True
|
|
|
|
def warmup_worker():
|
|
global _wildcard_warmup_loading
|
|
try:
|
|
db.ensure_loaded()
|
|
except Exception as exc:
|
|
print(f"[Autocomplete] Wildcard warm-up failed: {exc}")
|
|
finally:
|
|
with _wildcard_warmup_lock:
|
|
_wildcard_warmup_loading = False
|
|
|
|
thread = threading.Thread(
|
|
target=warmup_worker,
|
|
name="A1111WildcardWarmup",
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
return True
|
|
|
|
|
|
# Register API endpoints (only if server is available)
|
|
if _HAS_SERVER and PromptServer:
|
|
start_wildcard_warmup()
|
|
|
|
@PromptServer.instance.routes.post("/a1111_prompt/autocomplete")
|
|
async def autocomplete_tags(request):
|
|
"""
|
|
API endpoint for tag autocomplete.
|
|
|
|
Request body:
|
|
{
|
|
"query": "partial_tag_name",
|
|
"limit": 20,
|
|
"tag_file": "danbooru.csv",
|
|
"search_aliases": true
|
|
}
|
|
|
|
Response:
|
|
{
|
|
"results": [
|
|
{
|
|
"name": "1girl",
|
|
"type": 0,
|
|
"count": 6008644,
|
|
"aliases": ["1girls", "sole_female"],
|
|
"matched_alias": "sole_female" // only if matched via alias
|
|
},
|
|
...
|
|
],
|
|
"tag_count": 100000
|
|
}
|
|
"""
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return web.json_response({"error": "Invalid JSON"}, status=400)
|
|
|
|
if not isinstance(data, dict):
|
|
return web.json_response(
|
|
{"error": "Request body must be a JSON object"}, status=400
|
|
)
|
|
|
|
query = _coerce_text(data.get("query", ""))
|
|
if len(query) > 255:
|
|
return web.json_response(
|
|
{"results": [], "tag_count": 0, "error": "Query too long"}
|
|
)
|
|
mode = str(data.get("mode", "tag")).lower()
|
|
if mode not in {"tag", "wildcard", "wildcard_contents"}:
|
|
mode = "tag"
|
|
limit_cap = 500 if mode in {"wildcard", "wildcard_contents"} else 100
|
|
limit = _coerce_limit(data.get("limit", 20), default=20, cap=limit_cap)
|
|
tag_file = _normalize_tag_filename(data.get("tag_file"), DEFAULT_TAG_FILE)
|
|
extra_files = _normalize_extra_files(data.get("extra_files"))
|
|
search_aliases = _coerce_bool(data.get("search_aliases", True), default=True)
|
|
contains_fallback = _coerce_bool(
|
|
data.get("contains_fallback", True), default=True
|
|
)
|
|
|
|
if mode == "wildcard":
|
|
db = get_wildcard_database()
|
|
results = db.search(query, limit=limit)
|
|
return web.json_response(
|
|
{"results": results, "tag_count": db.wildcard_count, "mode": "wildcard"}
|
|
)
|
|
|
|
if mode == "wildcard_contents":
|
|
db = get_wildcard_database()
|
|
wildcard_name = _coerce_text(data.get("wildcard_name", ""))
|
|
content_query = _coerce_text(data.get("content_query", ""))
|
|
results = db.get_contents(
|
|
wildcard_name=wildcard_name,
|
|
content_query=content_query,
|
|
limit=limit,
|
|
)
|
|
return web.json_response(
|
|
{
|
|
"results": results,
|
|
"tag_count": len(results),
|
|
"mode": "wildcard_contents",
|
|
"wildcard_name": wildcard_name,
|
|
}
|
|
)
|
|
|
|
# Ensure database is loaded
|
|
db = ensure_database_loaded(
|
|
tag_file or DEFAULT_TAG_FILE, extra_files=extra_files
|
|
)
|
|
|
|
# Perform search
|
|
results = db.search(
|
|
query,
|
|
limit=limit,
|
|
search_aliases=search_aliases,
|
|
contains_fallback=contains_fallback,
|
|
)
|
|
|
|
return web.json_response(
|
|
{"results": results, "tag_count": db.tag_count, "mode": "tag"}
|
|
)
|
|
|
|
@PromptServer.instance.routes.get("/a1111_prompt/autocomplete/wildcards")
|
|
async def autocomplete_wildcards(request):
|
|
db = get_wildcard_database()
|
|
results = await asyncio.to_thread(db.get_manifest)
|
|
return web.json_response(
|
|
{
|
|
"results": results,
|
|
"wildcard_count": len(results),
|
|
"mode": "wildcard_manifest",
|
|
}
|
|
)
|
|
|
|
@PromptServer.instance.routes.get("/a1111_prompt/autocomplete/status")
|
|
async def autocomplete_status(request):
|
|
"""
|
|
Get the status of the autocomplete database.
|
|
|
|
Response:
|
|
{
|
|
"loaded": true,
|
|
"tag_count": 100000,
|
|
"current_file": "danbooru.csv"
|
|
}
|
|
"""
|
|
db = get_database()
|
|
return web.json_response(get_database_status())
|
|
|
|
@PromptServer.instance.routes.post("/a1111_prompt/autocomplete/warmup")
|
|
async def autocomplete_warmup(request):
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
data = {}
|
|
|
|
if not isinstance(data, dict):
|
|
data = {}
|
|
|
|
tag_file = _normalize_tag_filename(data.get("tag_file"), DEFAULT_TAG_FILE)
|
|
extra_files = _normalize_extra_files(data.get("extra_files"))
|
|
started = start_database_warmup(
|
|
tag_file or DEFAULT_TAG_FILE, extra_files=extra_files
|
|
)
|
|
status = get_database_status()
|
|
status["started"] = started
|
|
status["wildcard_started"] = start_wildcard_warmup()
|
|
status["wildcard_loaded"] = get_wildcard_database().is_loaded
|
|
status["wildcard_count"] = get_wildcard_database().wildcard_count
|
|
return web.json_response(status)
|
|
|
|
@PromptServer.instance.routes.get("/a1111_prompt/autocomplete/files")
|
|
async def list_tag_files(request):
|
|
"""
|
|
List available tag files.
|
|
|
|
Response:
|
|
{
|
|
"files": ["danbooru.csv", "e621.csv", ...]
|
|
}
|
|
"""
|
|
files = []
|
|
|
|
# Check main data directory
|
|
if DATA_DIR.exists():
|
|
files.extend(f.name for f in DATA_DIR.glob("*.csv"))
|
|
|
|
# Check reference directory
|
|
ref_dir = (
|
|
Path(__file__).parent.parent
|
|
/ "autocomplete_reference"
|
|
/ "a1111-sd-webui-tagcomplete"
|
|
/ "tags"
|
|
)
|
|
if ref_dir.exists():
|
|
for f in ref_dir.glob("*.csv"):
|
|
if f.name not in files:
|
|
files.append(f.name)
|
|
|
|
files.sort()
|
|
|
|
return web.json_response({"files": files})
|