From 582bb052f73dc090bf4fb75d2ccdd05114a995b2 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Fri, 3 Apr 2026 07:12:19 -0700 Subject: [PATCH] feat: add LoRA utility module for LoraManager integration (#52) Detection, metadata reading, trigger word cache, and prompt injection for ComfyUI-Lora-Manager. All functions return empty results when LoraManager is not installed. --- py/lora_utils.py | 303 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 303 insertions(+) create mode 100644 py/lora_utils.py diff --git a/py/lora_utils.py b/py/lora_utils.py new file mode 100644 index 0000000..9e8c329 --- /dev/null +++ b/py/lora_utils.py @@ -0,0 +1,303 @@ +"""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 json +import os +import re +import threading +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +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 ──────────────────────────────────────────── + +_LORA_MANAGER_DIR_NAME = "ComfyUI-Lora-Manager" + + +def find_comfyui_root() -> Optional[Path]: + """Walk upward from this file to find the ComfyUI root (contains main.py).""" + current = Path(__file__).resolve().parent + 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 + root = find_comfyui_root() + if root: + candidate = root / "custom_nodes" / _LORA_MANAGER_DIR_NAME + if candidate.is_dir() and _looks_like_lora_manager(candidate): + return str(candidate.resolve()) + + return None + + +def _looks_like_lora_manager(path: Path) -> bool: + """Heuristic: does this directory look like a LoraManager install?""" + # Check for characteristic files + markers = ["__init__.py", "README.md"] + has_marker = any((path / m).exists() for m in markers) + # Check for characteristic subdirectories or module name + has_lora_ref = ( + (path / "lora_manager").is_dir() + or (path / "py").is_dir() + or "lora" in path.name.lower() + ) + return has_marker and has_lora_ref + + +# ── Metadata reading ───────────────────────────────────────────────── + + +def find_lora_directories(lora_manager_path: str) -> List[str]: + """Find directories that contain LoRA models (with .metadata.json files). + + Searches common locations: the LoraManager extension dir itself, + and the ComfyUI models/loras directory. + """ + dirs = set() + lm_path = Path(lora_manager_path) + + # Check ComfyUI models/loras + root = find_comfyui_root() + if root: + models_loras = root / "models" / "loras" + if models_loras.is_dir(): + dirs.add(str(models_loras.resolve())) + + # Check for any .metadata.json in the LoraManager dir tree + # (some users store loras inside the extension) + for meta in lm_path.rglob("*.metadata.json"): + dirs.add(str(meta.parent.resolve())) + + return sorted(dirs) + + +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_trigger_words_from_metadata(metadata: Dict) -> List[str]: + """Extract trigger words from a parsed LoraManager metadata dict.""" + civitai = metadata.get("civitai", {}) + if not civitai: + return [] + 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_model_name_from_metadata(metadata: Dict) -> str: + """Extract the model display name from metadata.""" + # Try civitai model name first, then file_name, then fallback + name = metadata.get("model_name", "") + if not name: + civitai = metadata.get("civitai", {}) + model = civitai.get("model", {}) + name = model.get("name", "") + if not name: + name = metadata.get("file_name", "unknown") + return name + + +def get_preview_image_from_metadata( + metadata: Dict, metadata_path: Path +) -> Optional[str]: + """Find the preview image path for a LoRA from its metadata. + + Returns: + Absolute path string to the preview image, or None. + """ + lora_dir = metadata_path.parent + file_name = metadata.get("file_name", "") + if not file_name: + # Derive from metadata filename: "foo.metadata.json" -> "foo" + 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", ".preview.png", ".preview.jpg"): + candidate = lora_dir / f"{base_name}{ext}" + if candidate.exists(): + return str(candidate.resolve()) + + return None + + +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"]+):[^>]+>", 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.glob("*.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 ) + 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 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