imi
5b819dad2e
- Backend Optimizations: - Optimized workflow graph traversal using deque for O(1) operations. - Implemented robust normalization for autocomplete extra_files API. - Added bounded instance-level caching to prevent memory growth. - Fixed AttributeError in hooks by restoring transformers_dict. - Added EXPERIMENTAL.md for tracking legacy code notes. - Frontend Refinements: - Fixed memory leaks via onRemoved cleanup for DOM and observers. - Added race-condition protection for autocomplete blur events. - Improved UI feedback for API errors in the autocomplete popup. - Refined autocomplete aesthetics (removed redundant last-item border). - Ensured defensive DOM cleanup in coordinate calculations. - Documentation & Quality: - Overhauled README with auto-detection and hidden inputs docs. - Fixed MD040 linting errors across all project READMEs. - Updated CHANGELOG.md and TODO.md for the 2026-02-02 release.
420 lines
12 KiB
Python
420 lines
12 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 csv
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
# 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"
|
|
|
|
# Default tag file to use if not specified
|
|
DEFAULT_TAG_FILE = "danbooru.csv"
|
|
|
|
|
|
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,
|
|
}
|
|
|
|
|
|
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._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 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._current_files = []
|
|
|
|
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
|
|
|
|
# Build alias -> canonical tag mapping
|
|
for alias in aliases:
|
|
self._alias_map[alias.lower()] = entry
|
|
|
|
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
|
|
) -> 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
|
|
"""
|
|
if not self._loaded or not query:
|
|
return []
|
|
|
|
query_lower = query.lower().strip()
|
|
if not query_lower:
|
|
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
|
|
|
|
# Search by tag name
|
|
for tag in self._tags:
|
|
if query_lower in tag.search_text:
|
|
score = get_score(tag, tag.search_text)
|
|
results.append((score, tag, None))
|
|
seen_tags.add(tag.name)
|
|
|
|
# Search by aliases
|
|
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:
|
|
score = get_score(tag, alias_lower)
|
|
# Find the original case alias
|
|
original_alias = next(
|
|
(a for a in tag.aliases if a.lower() == alias_lower),
|
|
alias_lower,
|
|
)
|
|
results.append((score, tag, original_alias))
|
|
seen_tags.add(tag.name)
|
|
|
|
# Sort by: score (desc), then post count (desc)
|
|
results.sort(key=lambda x: (-x[0], -x[1].count))
|
|
|
|
# Format results
|
|
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
|
|
|
|
|
|
# Global database instance (lazy loaded)
|
|
_database = TagDatabase()
|
|
|
|
|
|
def get_database() -> TagDatabase:
|
|
"""Get the global tag database instance."""
|
|
return _database
|
|
|
|
|
|
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
|
|
"""
|
|
db = get_database()
|
|
|
|
# Build list of files to load
|
|
files_to_load = [tag_file]
|
|
if extra_files:
|
|
files_to_load.extend(extra_files)
|
|
|
|
# Check if we need to reload
|
|
needs_reload = not db.is_loaded or set(db._current_files) != set(files_to_load)
|
|
|
|
if needs_reload:
|
|
# Find and load all tag files
|
|
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}")
|
|
|
|
if filepaths:
|
|
db.load_multiple(filepaths)
|
|
|
|
return db
|
|
|
|
|
|
# Register API endpoints (only if server is available)
|
|
if _HAS_SERVER and PromptServer:
|
|
|
|
@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)
|
|
|
|
query = data.get("query", "")
|
|
limit = min(data.get("limit", 20), 100) # Cap at 100
|
|
tag_file = data.get(
|
|
"tag_file"
|
|
) # Will use DEFAULT_TAG_FILE in ensure_database_loaded if None
|
|
|
|
# Normalize extra_files: must be a list of strings
|
|
extra_files = data.get("extra_files")
|
|
if extra_files is None:
|
|
extra_files = ["extra-quality-tags.csv"]
|
|
elif isinstance(extra_files, str):
|
|
extra_files = [extra_files]
|
|
elif not isinstance(extra_files, list):
|
|
extra_files = ["extra-quality-tags.csv"]
|
|
|
|
search_aliases = data.get("search_aliases", True)
|
|
|
|
# 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)
|
|
|
|
return web.json_response({"results": results, "tag_count": db.tag_count})
|
|
|
|
@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(
|
|
{
|
|
"loaded": db.is_loaded,
|
|
"tag_count": db.tag_count,
|
|
"current_files": db._current_files,
|
|
}
|
|
)
|
|
|
|
@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})
|