Vito
b2b817781a
* ci: cancel superseded PR runs and make the bandit job a real gate No PR workflow had a concurrency group, so a new push left the previous run occupying 16 Python matrix runners, the real ComfyUI job and the three viewerjs jobs for up to 25 minutes. Add a group per workflow and PR ref that cancels the older run. Pushes to main use the run id as their group, so no main commit is ever cancelled or dropped as pending. Remove the lint job from code-quality.yml. It ran black and flake8 a second time, partly with --exit-zero, next to the protected "black and flake8" job in test.yml. The security job ran bandit with "|| true" and could never fail. At medium severity it reports 74 findings today, so it now gates on high severity only and prints the medium report without failing. The only two high findings were SHA1 and MD5 used as a gallery entry id and a LoRA preview cache filename; mark both usedforsecurity=False, which leaves every hash value unchanged. Closes #230. * ci: verify superseded PR runs are cancelled * ci: require the bandit high-severity check Review on #231 pointed out that a failing bandit job did not block a merge because its context was not required. Name the job "bandit (high severity)" so the required context is unambiguous, and add it to .github/branch-protection.json and the TESTING.md list of required checks.
707 lines
22 KiB
Python
707 lines
22 KiB
Python
"""Utilities for LoraManager integration.
|
|
|
|
Provides detection, metadata reading, and trigger word lookup for
|
|
ComfyUI-Lora-Manager (https://github.com/willmiao/ComfyUI-Lora-Manager).
|
|
|
|
All functions are safe to call when LoraManager is not installed — they
|
|
return empty results rather than raising.
|
|
"""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
import urllib.request
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
from urllib.parse import urlsplit
|
|
|
|
try:
|
|
from ..utils.logging_config import get_logger
|
|
except ImportError:
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from utils.logging_config import get_logger
|
|
|
|
logger = get_logger("prompt_manager.lora_utils")
|
|
|
|
# ── LoraManager detection ────────────────────────────────────────────
|
|
|
|
|
|
def find_comfyui_root() -> Optional[Path]:
|
|
"""Walk upward from this file to find the ComfyUI root (contains main.py).
|
|
|
|
Tries both the resolved (real) path and the unresolved path to handle
|
|
symlinked custom_nodes installations.
|
|
"""
|
|
start_paths = [Path(__file__).resolve().parent]
|
|
|
|
# If installed via symlink, the unresolved path leads through custom_nodes/
|
|
raw_path = Path(__file__).parent
|
|
if raw_path.resolve() != raw_path:
|
|
start_paths.append(raw_path)
|
|
|
|
# Also try via folder_paths if available (ComfyUI runtime)
|
|
try:
|
|
import folder_paths
|
|
|
|
base = Path(folder_paths.base_path)
|
|
if base.is_dir():
|
|
return base
|
|
except (ImportError, AttributeError):
|
|
# folder_paths unavailable — not running inside ComfyUI runtime
|
|
pass
|
|
|
|
for start in start_paths:
|
|
current = start
|
|
for _ in range(10):
|
|
if (current / "main.py").exists() and (current / "custom_nodes").exists():
|
|
return current
|
|
parent = current.parent
|
|
if parent == current:
|
|
break
|
|
current = parent
|
|
|
|
return None
|
|
|
|
|
|
def detect_lora_manager(custom_path: str = "") -> Optional[str]:
|
|
"""Return the absolute path to ComfyUI-Lora-Manager if installed.
|
|
|
|
Args:
|
|
custom_path: User-provided override path. Checked first.
|
|
|
|
Returns:
|
|
Absolute path string, or None if not found.
|
|
"""
|
|
# 1. User override
|
|
if custom_path:
|
|
p = Path(custom_path)
|
|
if p.is_dir() and _looks_like_lora_manager(p):
|
|
return str(p.resolve())
|
|
|
|
# 2. Auto-detect via custom_nodes (case-insensitive scan)
|
|
root = find_comfyui_root()
|
|
if root:
|
|
custom_nodes = root / "custom_nodes"
|
|
if custom_nodes.is_dir():
|
|
for entry in custom_nodes.iterdir():
|
|
if (
|
|
entry.is_dir()
|
|
and "lora" in entry.name.lower()
|
|
and "manager" in entry.name.lower()
|
|
and _looks_like_lora_manager(entry)
|
|
):
|
|
return str(entry.resolve())
|
|
|
|
return None
|
|
|
|
|
|
def _looks_like_lora_manager(path: Path) -> bool:
|
|
"""Heuristic: does this directory look like a LoraManager install?"""
|
|
# Must have __init__.py (ComfyUI extension) or README.md
|
|
has_init = (path / "__init__.py").exists()
|
|
if not has_init:
|
|
return False
|
|
# Check for characteristic structure: py/ dir, or any .metadata.json nearby
|
|
return (path / "py").is_dir() or (path / "lora_manager").is_dir()
|
|
|
|
|
|
def _public_path_anchors() -> List[str]:
|
|
"""Directories the API's public (relative) paths are rendered against."""
|
|
try:
|
|
from .config import GalleryConfig
|
|
except ImportError:
|
|
try:
|
|
from config import GalleryConfig
|
|
except ImportError:
|
|
return []
|
|
try:
|
|
return list(GalleryConfig.path_anchors())
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def custom_nodes_directories() -> List[Path]:
|
|
"""ComfyUI ``custom_nodes`` directories a LoraManager install may live in.
|
|
|
|
Inside ComfyUI these come from ``folder_paths`` (``base_path`` plus any
|
|
registered ``custom_nodes`` folders); outside it, the parent of this
|
|
package, which is ``custom_nodes`` for a normal install.
|
|
"""
|
|
found: List[Path] = []
|
|
|
|
def add(candidate) -> None:
|
|
try:
|
|
resolved = Path(candidate).resolve()
|
|
except (OSError, RuntimeError, ValueError):
|
|
return
|
|
if resolved not in found:
|
|
found.append(resolved)
|
|
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
folder_paths = None
|
|
if folder_paths is not None:
|
|
base = getattr(folder_paths, "base_path", None)
|
|
if isinstance(base, str) and base:
|
|
add(Path(base) / "custom_nodes")
|
|
getter = getattr(folder_paths, "get_folder_paths", None)
|
|
if callable(getter):
|
|
try:
|
|
registered = getter("custom_nodes") or []
|
|
except Exception:
|
|
registered = []
|
|
for entry in registered:
|
|
add(entry)
|
|
if not found:
|
|
add(Path(__file__).resolve().parent.parent.parent)
|
|
return found
|
|
|
|
|
|
def resolve_lora_manager_path(path) -> Optional[str]:
|
|
"""Canonical form of ``path`` when it may be used as the LoraManager install.
|
|
|
|
Accepted: an existing directory that is a direct child of one of
|
|
:func:`custom_nodes_directories` and whose name contains ``lora``
|
|
(case-insensitive). A relative path is tried against the ComfyUI root
|
|
and each custom_nodes directory. Anything else, including a symlink
|
|
that escapes custom_nodes, yields None.
|
|
"""
|
|
if not isinstance(path, str) or not path.strip():
|
|
return None
|
|
path = path.strip()
|
|
roots = []
|
|
for root in custom_nodes_directories():
|
|
try:
|
|
roots.append(Path(os.path.normcase(os.path.realpath(str(root)))))
|
|
except (OSError, ValueError):
|
|
continue
|
|
|
|
candidates = [path]
|
|
if not os.path.isabs(path):
|
|
# Public (API) paths are relative to these anchors; the ComfyUI root
|
|
# and the custom_nodes directories cover hand-typed relative paths.
|
|
bases = list(_public_path_anchors())
|
|
comfy_root = find_comfyui_root()
|
|
if comfy_root:
|
|
bases.append(comfy_root)
|
|
bases.extend(root.parent for root in roots)
|
|
bases.extend(roots)
|
|
candidates = [os.path.join(str(base), path) for base in bases]
|
|
|
|
for candidate in candidates:
|
|
try:
|
|
canonical = Path(os.path.normcase(os.path.realpath(candidate)))
|
|
is_dir = canonical.is_dir()
|
|
except (OSError, RuntimeError, ValueError):
|
|
continue
|
|
if not is_dir or "lora" not in canonical.name.lower():
|
|
continue
|
|
if canonical.parent in roots:
|
|
return str(canonical)
|
|
return None
|
|
|
|
|
|
# ── Metadata reading ─────────────────────────────────────────────────
|
|
|
|
|
|
def find_lora_directories(lora_manager_path: str) -> List[str]:
|
|
"""Find directories that contain LoRA models (with .metadata.json files).
|
|
|
|
Searches: ComfyUI models/loras, extra_model_paths.yaml lora dirs,
|
|
and the LoraManager extension dir itself.
|
|
"""
|
|
dirs = set()
|
|
lm_path = Path(lora_manager_path)
|
|
|
|
root = find_comfyui_root()
|
|
if root:
|
|
# Default models/loras
|
|
models_loras = root / "models" / "loras"
|
|
if models_loras.is_dir():
|
|
dirs.add(str(models_loras.resolve()))
|
|
|
|
# Extra model paths from ComfyUI config
|
|
for extra_dir in _get_extra_lora_paths(root):
|
|
if extra_dir.is_dir():
|
|
dirs.add(str(extra_dir.resolve()))
|
|
|
|
# Also try folder_paths at runtime (catches all configured paths)
|
|
try:
|
|
import folder_paths
|
|
|
|
for p in folder_paths.get_folder_paths("loras"):
|
|
pp = Path(p)
|
|
if pp.is_dir():
|
|
dirs.add(str(pp.resolve()))
|
|
except (ImportError, AttributeError):
|
|
# folder_paths unavailable — not running inside ComfyUI runtime
|
|
pass
|
|
|
|
# Check for any .metadata.json in the LoraManager dir tree
|
|
for meta in lm_path.rglob("*.metadata.json"):
|
|
dirs.add(str(meta.parent.resolve()))
|
|
|
|
return sorted(dirs)
|
|
|
|
|
|
def _get_extra_lora_paths(comfyui_root: Path) -> List[Path]:
|
|
"""Parse extra_model_paths.yaml for additional LoRA directories."""
|
|
results = []
|
|
for name in ("extra_model_paths.yaml", "extra_model_paths.yml"):
|
|
config_file = comfyui_root / name
|
|
if not config_file.exists():
|
|
continue
|
|
try:
|
|
import yaml
|
|
|
|
config = yaml.safe_load(config_file.read_text())
|
|
if not isinstance(config, dict):
|
|
continue
|
|
for section in config.values():
|
|
if not isinstance(section, dict):
|
|
continue
|
|
base = Path(section.get("base_path", ""))
|
|
loras_val = section.get("loras", "")
|
|
if not loras_val:
|
|
continue
|
|
for line in str(loras_val).strip().splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
p = Path(line)
|
|
if not p.is_absolute():
|
|
p = base / line
|
|
if p.is_dir():
|
|
results.append(p)
|
|
except Exception as e:
|
|
logger.debug(f"Failed to parse {config_file}: {e}")
|
|
return results
|
|
|
|
|
|
def read_lora_metadata(metadata_path: Path) -> Optional[Dict]:
|
|
"""Read and parse a single .metadata.json file.
|
|
|
|
Returns:
|
|
Parsed dict, or None on failure.
|
|
"""
|
|
try:
|
|
with open(metadata_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, OSError) as e:
|
|
logger.debug(f"Failed to read {metadata_path}: {e}")
|
|
return None
|
|
|
|
|
|
def _get_civitai(metadata: Dict) -> Dict:
|
|
"""Safely get the civitai dict, handling None values."""
|
|
return metadata.get("civitai") or {}
|
|
|
|
|
|
def get_trigger_words_from_metadata(metadata: Dict) -> List[str]:
|
|
"""Extract trigger words from a parsed LoraManager metadata dict."""
|
|
civitai = _get_civitai(metadata)
|
|
words = civitai.get("trainedWords", [])
|
|
if isinstance(words, list):
|
|
return [w.strip() for w in words if isinstance(w, str) and w.strip()]
|
|
return []
|
|
|
|
|
|
def get_example_prompt_from_metadata(metadata: Dict) -> Optional[str]:
|
|
"""Extract an example prompt from civitai image metadata.
|
|
|
|
Looks at civitai.images[].meta.prompt for the first available example.
|
|
"""
|
|
civitai = _get_civitai(metadata)
|
|
images = civitai.get("images", []) or []
|
|
for img in images:
|
|
if not isinstance(img, dict):
|
|
continue
|
|
meta = img.get("meta")
|
|
if isinstance(meta, dict):
|
|
prompt = meta.get("prompt", "")
|
|
if isinstance(prompt, str) and prompt.strip():
|
|
return prompt.strip()
|
|
return None
|
|
|
|
|
|
def get_civitai_image_urls(metadata: Dict) -> List[str]:
|
|
"""Extract all civitai example image URLs from metadata."""
|
|
civitai = _get_civitai(metadata)
|
|
urls = []
|
|
for img in civitai.get("images", []) or []:
|
|
if not isinstance(img, dict):
|
|
continue
|
|
url = img.get("url", "")
|
|
if isinstance(url, str) and url.strip():
|
|
urls.append(url.strip())
|
|
return urls
|
|
|
|
|
|
def get_model_name_from_metadata(metadata: Dict) -> str:
|
|
"""Extract the model display name from metadata."""
|
|
name = metadata.get("model_name", "")
|
|
if not name:
|
|
civitai = _get_civitai(metadata)
|
|
model = civitai.get("model") or {}
|
|
name = model.get("name", "")
|
|
if not name:
|
|
name = metadata.get("file_name", "unknown")
|
|
return name
|
|
|
|
|
|
def get_preview_images_from_metadata(metadata: Dict, metadata_path: Path) -> List[str]:
|
|
"""Find all local preview/example image paths for a LoRA.
|
|
|
|
Returns:
|
|
List of absolute path strings to image files.
|
|
"""
|
|
results = []
|
|
lora_dir = metadata_path.parent
|
|
file_name = metadata.get("file_name", "")
|
|
if not file_name:
|
|
stem = metadata_path.name.replace(".metadata.json", "")
|
|
file_name = stem
|
|
|
|
base_name = Path(file_name).stem
|
|
|
|
# Check standard preview naming conventions
|
|
for ext in (
|
|
".png",
|
|
".jpg",
|
|
".jpeg",
|
|
".webp",
|
|
".preview.png",
|
|
".preview.jpg",
|
|
".preview.jpeg",
|
|
):
|
|
candidate = lora_dir / f"{base_name}{ext}"
|
|
if candidate.exists():
|
|
results.append(str(candidate.resolve()))
|
|
|
|
return results
|
|
|
|
|
|
def get_preview_image_from_metadata(
|
|
metadata: Dict, metadata_path: Path
|
|
) -> Optional[str]:
|
|
"""Find the first preview image path for a LoRA (backward compat)."""
|
|
images = get_preview_images_from_metadata(metadata, metadata_path)
|
|
return images[0] if images else None
|
|
|
|
|
|
_THUMB_MAX_SIZE = 512
|
|
|
|
# Only these hosts may ever receive the CivitAI API key, and only over HTTPS.
|
|
CIVITAI_HOSTS = frozenset(
|
|
{"civitai.com", "www.civitai.com", "api.civitai.com", "image.civitai.com"}
|
|
)
|
|
MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024
|
|
_DOWNLOAD_CHUNK_BYTES = 64 * 1024
|
|
|
|
|
|
class _RefuseRedirects(urllib.request.HTTPRedirectHandler):
|
|
"""Never follow a redirect: urllib would re-send the Bearer header to it."""
|
|
|
|
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
|
logger.warning(f"Refusing redirect from {req.full_url} to {newurl}")
|
|
return None
|
|
|
|
|
|
def _build_opener(*handlers) -> urllib.request.OpenerDirector:
|
|
"""Build the outbound opener; extra handlers let tests serve canned responses."""
|
|
return urllib.request.build_opener(_RefuseRedirects, *handlers)
|
|
|
|
|
|
_opener = _build_opener()
|
|
|
|
|
|
def _open_url(req: urllib.request.Request, timeout: float = 10):
|
|
"""Open ``req`` through the redirect-refusing opener."""
|
|
return _opener.open(req, timeout=timeout)
|
|
|
|
|
|
def is_civitai_url(url: str) -> bool:
|
|
"""True when ``url`` is an HTTPS URL on an allow-listed CivitAI host."""
|
|
try:
|
|
parts = urlsplit(url)
|
|
port = parts.port
|
|
except ValueError:
|
|
return False
|
|
hostname = (parts.hostname or "").lower()
|
|
return (
|
|
parts.scheme == "https"
|
|
and hostname in CIVITAI_HOSTS
|
|
and parts.username is None
|
|
and parts.password is None
|
|
and port in (None, 443)
|
|
)
|
|
|
|
|
|
def _read_capped(resp, cap: int) -> Optional[bytes]:
|
|
"""Read a response body, refusing anything larger than ``cap`` bytes."""
|
|
declared = resp.headers.get("Content-Length")
|
|
if declared is not None:
|
|
try:
|
|
if int(declared) > cap:
|
|
return None
|
|
except ValueError:
|
|
pass
|
|
|
|
chunks = []
|
|
total = 0
|
|
while True:
|
|
chunk = resp.read(_DOWNLOAD_CHUNK_BYTES)
|
|
if not chunk:
|
|
break
|
|
total += len(chunk)
|
|
if total > cap:
|
|
return None
|
|
chunks.append(chunk)
|
|
return b"".join(chunks)
|
|
|
|
|
|
def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
|
|
"""Download a single CivitAI image, resize to thumbnail, save as JPEG.
|
|
|
|
Refuses any URL that is not HTTPS on an allow-listed CivitAI host so the
|
|
API key is never sent elsewhere, and refuses bodies larger than
|
|
MAX_CIVITAI_DOWNLOAD_BYTES.
|
|
"""
|
|
if not is_civitai_url(url):
|
|
logger.warning(f"Refusing non-CivitAI download URL: {url}")
|
|
return None
|
|
try:
|
|
headers = {"User-Agent": "ComfyUI-PromptManager/1.0"}
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
req = urllib.request.Request(url, headers=headers)
|
|
with _open_url(req, timeout=10) as resp:
|
|
raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
|
|
if raw is None:
|
|
logger.warning(f"Refusing oversized download: {url}")
|
|
return None
|
|
|
|
# Resize to thumbnail to save disk space
|
|
from io import BytesIO
|
|
|
|
from PIL import Image
|
|
|
|
img = Image.open(BytesIO(raw))
|
|
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
|
|
img = img.convert("RGB")
|
|
img.save(str(local_path), "JPEG", quality=85)
|
|
|
|
return str(local_path.resolve())
|
|
except Exception as e:
|
|
logger.debug(f"Failed to download {url}: {e}")
|
|
return None
|
|
|
|
|
|
def download_civitai_images(
|
|
metadata: Dict, metadata_path: Path, cache_dir: Path, api_key: str = ""
|
|
) -> List[str]:
|
|
"""Download civitai example images to a local cache directory.
|
|
|
|
Uses thumbnail URLs (512px) instead of full-size originals, and
|
|
downloads in parallel (up to 8 concurrent) for speed.
|
|
|
|
Args:
|
|
api_key: CivitAI API key for authenticated downloads (NSFW content).
|
|
|
|
Returns:
|
|
List of absolute paths to downloaded image files.
|
|
"""
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
civitai = _get_civitai(metadata)
|
|
images = civitai.get("images", []) or []
|
|
if not images:
|
|
return []
|
|
|
|
file_name = metadata.get("file_name", "")
|
|
if not file_name:
|
|
file_name = metadata_path.name.replace(".metadata.json", "")
|
|
lora_stem = Path(file_name).stem
|
|
|
|
lora_cache = cache_dir / lora_stem
|
|
lora_cache.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Build download tasks
|
|
cached = []
|
|
tasks = [] # (url, local_path)
|
|
for img in images:
|
|
if not isinstance(img, dict):
|
|
continue
|
|
url = img.get("url", "")
|
|
if not isinstance(url, str) or not is_civitai_url(url):
|
|
continue
|
|
|
|
url_hash = hashlib.md5(url.encode(), usedforsecurity=False).hexdigest()[:12]
|
|
local_path = lora_cache / f"{url_hash}.jpg"
|
|
|
|
if local_path.exists():
|
|
cached.append(str(local_path.resolve()))
|
|
else:
|
|
tasks.append((url, local_path))
|
|
|
|
if not tasks:
|
|
return cached
|
|
|
|
# Download in parallel
|
|
downloaded = []
|
|
with ThreadPoolExecutor(max_workers=8) as pool:
|
|
futures = [
|
|
pool.submit(_download_one, url, path, api_key) for url, path in tasks
|
|
]
|
|
for fut in futures:
|
|
result = fut.result()
|
|
if result:
|
|
downloaded.append(result)
|
|
|
|
return cached + downloaded
|
|
|
|
|
|
def get_lora_image_cache_dir() -> Path:
|
|
"""Get the directory used to cache downloaded LoRA example images."""
|
|
# Store in the extension's own directory
|
|
ext_root = Path(__file__).resolve().parent.parent
|
|
cache = ext_root / "data" / "lora_images"
|
|
cache.mkdir(parents=True, exist_ok=True)
|
|
return cache
|
|
|
|
|
|
def get_example_images_dir(lora_manager_path: str) -> Optional[str]:
|
|
"""Find the LoraManager example_images directory."""
|
|
lm_path = Path(lora_manager_path)
|
|
|
|
# Direct subdirectory
|
|
candidate = lm_path / "example_images"
|
|
if candidate.is_dir():
|
|
return str(candidate.resolve())
|
|
|
|
# Search one level in user data dirs
|
|
for child in lm_path.iterdir():
|
|
if child.is_dir():
|
|
sub = child / "example_images"
|
|
if sub.is_dir():
|
|
return str(sub.resolve())
|
|
|
|
return None
|
|
|
|
|
|
# ── Trigger word cache & injection ───────────────────────────────────
|
|
|
|
_LORA_PATTERN = re.compile(r"<lora:([^:>]+):[^>]+>", re.IGNORECASE)
|
|
|
|
|
|
class TriggerWordCache:
|
|
"""Thread-safe cache mapping LoRA names to their trigger words.
|
|
|
|
Built lazily on first access, refreshable on demand.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._cache: Dict[str, List[str]] = {}
|
|
self._lock = threading.Lock()
|
|
self._loaded = False
|
|
|
|
def load(self, lora_manager_path: str) -> int:
|
|
"""Scan LoRA metadata files and build the trigger word mapping.
|
|
|
|
Returns:
|
|
Number of LoRAs with trigger words found.
|
|
"""
|
|
new_cache: Dict[str, List[str]] = {}
|
|
|
|
lora_dirs = find_lora_directories(lora_manager_path)
|
|
for lora_dir in lora_dirs:
|
|
dir_path = Path(lora_dir)
|
|
for meta_file in dir_path.rglob("*.metadata.json"):
|
|
metadata = read_lora_metadata(meta_file)
|
|
if not metadata:
|
|
continue
|
|
|
|
words = get_trigger_words_from_metadata(metadata)
|
|
if not words:
|
|
continue
|
|
|
|
# Key by filename stem (what appears in <lora:NAME:weight>)
|
|
file_name = metadata.get("file_name", "")
|
|
if file_name:
|
|
stem = Path(file_name).stem
|
|
new_cache[stem.lower()] = words
|
|
|
|
# Also key by the metadata file stem
|
|
meta_stem = meta_file.name.replace(".metadata.json", "")
|
|
if meta_stem.lower() not in new_cache:
|
|
new_cache[meta_stem.lower()] = words
|
|
|
|
with self._lock:
|
|
self._cache = new_cache
|
|
self._loaded = True
|
|
|
|
logger.info(
|
|
f"Trigger word cache loaded: {len(new_cache)} LoRAs with trigger words"
|
|
)
|
|
return len(new_cache)
|
|
|
|
def get_trigger_words(self, lora_name: str) -> List[str]:
|
|
"""Look up trigger words for a LoRA by name (case-insensitive)."""
|
|
with self._lock:
|
|
return self._cache.get(lora_name.lower(), [])
|
|
|
|
@property
|
|
def is_loaded(self) -> bool:
|
|
with self._lock:
|
|
return self._loaded
|
|
|
|
def clear(self):
|
|
with self._lock:
|
|
self._cache.clear()
|
|
self._loaded = False
|
|
|
|
|
|
# Module-level singleton
|
|
_trigger_cache = TriggerWordCache()
|
|
|
|
|
|
def get_trigger_cache() -> TriggerWordCache:
|
|
return _trigger_cache
|
|
|
|
|
|
def inject_trigger_words(text: str, cache: TriggerWordCache) -> Tuple[str, List[str]]:
|
|
"""Scan text for <lora:NAME:WEIGHT> tags and append trigger words.
|
|
|
|
Args:
|
|
text: The prompt text potentially containing lora tags.
|
|
cache: Populated TriggerWordCache instance.
|
|
|
|
Returns:
|
|
Tuple of (modified_text, list_of_injected_words).
|
|
If no trigger words found, returns the original text unchanged.
|
|
"""
|
|
if not cache.is_loaded:
|
|
return text, []
|
|
|
|
matches = _LORA_PATTERN.findall(text)
|
|
if not matches:
|
|
return text, []
|
|
|
|
all_words = []
|
|
for lora_name in matches:
|
|
words = cache.get_trigger_words(lora_name)
|
|
for w in words:
|
|
if w.lower() not in text.lower() and w not in all_words:
|
|
all_words.append(w)
|
|
|
|
if not all_words:
|
|
return text, []
|
|
|
|
injected = ", ".join(all_words)
|
|
return f"{text}, {injected}", all_words
|