nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(lora): import the sidecars and local images LoraManager really uses (#355)

Browse Source 
The import, the legacy scan, the trigger-word cache and the image server
now share LoraManager's layout: sidecars beside the models (alongside) or
under the sidecar root (centralized), kept only when their model exists
inside a LoRA root. Each LoRA links its LoraManager preview and example
images where they are; CivitAI downloads run only for a LoRA without any
local image. The review adds skipped, with_local_media and the layout, and
offers no token when LoraManager's settings cannot be used; the legacy scan
reports the same reason without replacing anything.

find_lora_directories becomes default_lora_roots (no scan of the
LoraManager install); the preview and example-folder helpers it replaced
are removed. The trigger cache never raises into a running node.
This commit is contained in:
Vito Sansevero
2026-10-26 17:41:15 -07:00
parent 96dea818d7
commit ae660de588
12 changed files with 837 additions and 287 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/images.py
+16 -3
View File
@@ -2296,13 +2296,23 @@ class ImageRoutesMixin:
)
def _lora_image_dirs(self):
"""Extra directories served when the LoRA integration is enabled."""
"""Extra directories served when the LoRA integration is enabled.
LoraManager's guarded LoRA roots, sidecar root and example-images
root (``lora_manager_layout.media_roots``, #355) and PromptManager's
CivitAI download cache. An unusable LoraManager layout gives the
default LoRA roots and the cache.
"""
try:
from ..config import IntegrationConfig
if not IntegrationConfig.LORA_MANAGER_ENABLED:
return []
from ..lora_utils import find_lora_directories, get_lora_image_cache_dir
from ..lora_utils import (
get_lora_image_cache_dir,
layout_for_install,
lora_manager_layout,
)
lora_root = IntegrationConfig.LORA_MANAGER_PATH or ""
if lora_root.strip():
@@ -2315,8 +2325,11 @@ class ImageRoutesMixin:
lora_root = resolve_lora_manager_path(lora_root)
if not lora_root:
return []
layout = layout_for_install(lora_root)
else:
layout = lora_manager_layout("")
dirs = [Path(d) for d in find_lora_directories(lora_root)]
dirs = [Path(d) for d in layout.media_roots]
dirs.append(Path(get_lora_image_cache_dir()))
return dirs
except Exception:
py/api/lora_import_routes.py
+35 -17
View File
@@ -5,9 +5,12 @@ it lists the LoraManager metadata files and parses them, and counts the
current imports, those a saved pair uses and the example images an import
would download. It writes nothing and requests nothing. It answers a
``Plan`` (the maintenance plan's shape, ``py/api/maintenance.py``) plus
``{lora_files, with_metadata, current_imports, blocked_by_pairs,
planned_downloads, network:{planned_downloads, api_key}}``. A plan that saved
pairs block carries no token (nothing can be applied).
``{lora_files, with_metadata, skipped, with_local_media, current_imports,
blocked_by_pairs, planned_downloads, network:{planned_downloads, api_key},
layout}`` (#355: ``layout`` says which LoraManager settings file, storage mode
and library the import read, and ``error`` when they cannot be used). A plan
that saved pairs block, or whose layout is unusable, carries no token
(nothing can be applied).
``POST /prompt_manager/lora/import/apply {token}`` starts the job (202
``{job}``): kind ``lora_import``, titled "LoRA reference import", admission
@@ -31,6 +34,7 @@ import sqlite3
from aiohttp import web
from .. import lora_import, lora_utils
from .inspection_service import SERVABLE_EXTENSIONS
from .job_routes import busy_response
from .jobs import LORA_IMPORT_JOB_KIND, JobBusy
from .maintenance import (
@@ -152,8 +156,9 @@ class LoraImportRoutesMixin:
return None, ("lora-path", _LORA_PATH_ERROR)
return allowed, None
async def stage_lora_import(self, files, *, on_record=None, cancelled=None):
"""Stage ``files`` off the loop: records, skips and download counts."""
async def stage_lora_import(self, discovery, *, on_record=None, cancelled=None):
"""Stage ``discovery``'s sidecars off the loop: records, skips and
download counts."""
from ..config import IntegrationConfig
cache_dir = await self._run_in_executor(lora_utils.get_lora_image_cache_dir)
@@ -166,9 +171,10 @@ class LoraImportRoutesMixin:
return await self._run_in_executor(
lambda: lora_import.stage(
files,
discovery,
cache_dir,
IntegrationConfig.CIVITAI_API_KEY,
media_extensions=SERVABLE_EXTENSIONS,
on_record=report,
cancelled=cancelled or (lambda: False),
logger=self.logger,
@@ -176,8 +182,8 @@ class LoraImportRoutesMixin:
)
async def lora_files_unchanged(self, lm_path, expected):
"""Whether the metadata files under ``lm_path`` still fingerprint as
``expected`` (listed again, so an added file counts too)."""
"""Whether LoraManager's sidecars for ``lm_path`` still fingerprint as
``expected`` (found again, so an added file counts too)."""
def current():
return lora_import.fingerprint(lora_import.discover(lm_path))
@@ -220,7 +226,9 @@ class LoraImportRoutesMixin:
cache_dir = await self._run_in_executor(
lambda: lora_utils.get_lora_image_cache_dir(create=False)
)
found = await self._run_in_executor(lora_import.preview, lm_path, cache_dir)
found = await self._run_in_executor(
lora_import.preview, lm_path, cache_dir, SERVABLE_EXTENSIONS
)
state = await self._run_in_executor(db.category_import_state, CATEGORY)
except Exception:
self.logger.exception("LoRA import preview failed")
@@ -248,7 +256,9 @@ class LoraImportRoutesMixin:
return _failure("The LoRA import could not be reviewed.", 500)
token = None
expires_at = None
if not blocked:
layout = found["layout"]
# An unusable LoraManager layout found nothing it may replace with.
if not blocked and layout["error"] is None:
payload = {
"path": lm_path,
"fingerprint": found["fingerprint"],
@@ -280,6 +290,9 @@ class LoraImportRoutesMixin:
"expires_at": expires_at,
"lora_files": found["lora_files"],
"with_metadata": found["with_metadata"],
"skipped": found["skipped"],
"with_local_media": found["with_local_media"],
"layout": layout,
"current_imports": current,
"blocked_by_pairs": blocked,
"planned_downloads": downloads,
@@ -342,8 +355,8 @@ class LoraImportRoutesMixin:
store.discard(token)
return _failure(INTEGRATION_CHANGED, 409, "plan-changed")
try:
files = await self._run_in_executor(lora_import.discover, lm_path)
now = await self._run_in_executor(lora_import.fingerprint, files)
discovery = await self._run_in_executor(lora_import.discover, lm_path)
now = await self._run_in_executor(lora_import.fingerprint, discovery)
backup = self._backup_promised(plan)
_disk, short, _size = await self._run_in_executor(
disk_plan, db.model.db_path, backup=backup, vacuum=False
@@ -367,9 +380,14 @@ class LoraImportRoutesMixin:
lora_import.JOB_TITLE,
admission="records",
cancellable=True,
scope={"operation": LORA_OPERATION, "lora_files": len(files)},
scope={
"operation": LORA_OPERATION,
"lora_files": len(discovery.sidecars),
},
groups=lora_import.GROUPS,
runner=lambda job: self._run_lora_import(job, db, plan, files, backup),
runner=lambda job: self._run_lora_import(
job, db, plan, discovery, backup
),
)
except JobBusy as busy:
return busy_response(busy)
@@ -378,8 +396,8 @@ class LoraImportRoutesMixin:
# ── the job ───────────────────────────────────────────────────────
async def _run_lora_import(self, job, db, plan, files, backup_on):
total = len(files)
async def _run_lora_import(self, job, db, plan, discovery, backup_on):
total = len(discovery.sidecars)
job.update(total=total, status_line="Preparing LoRA records")
job.log_line(f"Found {_count(total, 'LoRA metadata file')}")
@@ -393,7 +411,7 @@ class LoraImportRoutesMixin:
try:
staged = await self.stage_lora_import(
files, on_record=on_record, cancelled=lambda: job.cancelled
discovery, on_record=on_record, cancelled=lambda: job.cancelled
)
except lora_import.ImportCancelled:
job.log_line(f"{CANCELLED} {CACHE_NOTE}")
py/api/lora_integration.py
+25 -7
View File
@@ -50,6 +50,15 @@ _IMPORTS_CHANGED = (
)
class LoraSettingsUnusable(Exception):
"""LoraManager's settings cannot be used; nothing was imported (#355)."""
def __init__(self, reason):
super().__init__(
f"LoraManager settings: not usable: {reason}. Nothing was changed."
)
class LoraFilesChanged(Exception):
"""A metadata file changed, appeared or went while the legacy scan staged."""
@@ -398,13 +407,19 @@ class LoraIntegrationMixin:
return
if isinstance(error, PlanChanged):
message = _IMPORTS_CHANGED
elif isinstance(error, LoraFilesChanged):
elif isinstance(error, (LoraFilesChanged, LoraSettingsUnusable)):
message = str(error)
else:
message = self._public_error(error)
if isinstance(
error,
(PlanChanged, LoraFilesChanged, EmptyReplacement, PairDependencyError),
(
PlanChanged,
LoraFilesChanged,
LoraSettingsUnusable,
EmptyReplacement,
PairDependencyError,
),
):
self.logger.warning(f"LoRA scan refused: {message}")
else:
@@ -432,8 +447,11 @@ class LoraIntegrationMixin:
"progress": 2,
}
)
files = await self._run_in_executor(lora_import.discover, lm_path)
total = len(files)
discovery = await self._run_in_executor(lora_import.discover, lm_path)
if discovery.layout.error is not None:
raise LoraSettingsUnusable(discovery.layout.error)
sidecars = discovery.sidecars
total = len(sidecars)
await send_progress(
{
"type": "progress",
@@ -443,7 +461,7 @@ class LoraIntegrationMixin:
}
)
reviewed = await self._run_in_executor(lora_import.fingerprint, files)
reviewed = await self._run_in_executor(lora_import.fingerprint, discovery)
events = asyncio.Queue()
state = await self._run_in_executor(
self.db.category_import_state, lora_import.LORA_CATEGORY
@@ -467,7 +485,7 @@ class LoraIntegrationMixin:
)
if not existing:
counts["imported"] += 1
status = record["name"] if record else files[index].name
status = record["name"] if record else sidecars[index].metadata_path.name
if record and record["images"]:
status += f" ({len(record['images'])} images)"
await send_progress(
@@ -490,7 +508,7 @@ class LoraIntegrationMixin:
stop = threading.Event()
staging = asyncio.ensure_future(
self.stage_lora_import(
files,
discovery,
on_record=on_record,
cancelled=lambda: stop.is_set() or disconnected(),
)
py/lora_import.py
+103 -53
View File
@@ -1,18 +1,26 @@
"""The staged LoRA reference import (#221, #249 System S10).
"""The staged LoRA reference import (#221, #249 System S10, #355).
The import reads LoraManager's ``*.metadata.json`` files and turns each LoRA
into one prompt of the ``lora-manager`` category, tagged with its name and
trigger words and linked to its preview images. It runs in three stages, and
only the last writes to the database:
The import reads LoraManager's ``*.metadata.json`` sidecars and turns each
LoRA into one prompt of the ``lora-manager`` category, tagged with its name and
trigger words and linked to its images. It runs in three stages, and only the
last writes to the database:
1. **discover**: list the metadata files (``discover``) and fingerprint them
(path, size and modification time; ``fingerprint``). ``preview`` adds the
counts a review shows. Nothing is written and nothing is requested.
2. **stage**: read each file and build its record; CivitAI example images are
downloaded into the image cache first (best effort, counted). A cancelled
or failed stage leaves the database untouched; downloaded images stay in
the cache.
3. **replace**: the files are listed and fingerprinted once more (a file
1. **discover**: read LoraManager's settings the way LoraManager does
(``py/lora_manager_layout.py``: the settings file from the environment
folder, the portable plugin file or the user config folder; the LoRA
roots, the storage mode, the sidecar and example-images roots) and find
the sidecars (``discover``). In ``alongside`` mode a sidecar beside its
model is kept when ``<base>.safetensors`` exists; in ``centralized`` mode
every sidecar under the sidecar root is kept when its ``file_path`` names
a model inside a LoRA root. ``fingerprint`` covers the kept sidecars (path,
size and modification time) and the layout; ``preview`` adds the counts a
review shows. Nothing is written and nothing is requested.
2. **stage**: read each sidecar and build its record. Its LoraManager preview
and example images are linked where they are (``local_media``); only a
LoRA without any local image downloads its CivitAI example images into the
image cache first (best effort, counted). A cancelled or failed stage
leaves the database untouched; downloaded images stay in the cache.
3. **replace**: the sidecars are found and fingerprinted once more (a file
changed, removed or added while staging aborts the import), then
``PromptDatabase.replace_category_prompts`` swaps the old collection for
the staged records in one transaction (or changes nothing).
@@ -24,9 +32,15 @@ Both the System job (``POST /lora/import/apply``) and the legacy ``POST
import hashlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Dict, List, Optional
from typing import AbstractSet, Callable, Dict, List, Optional
from . import lora_utils
from .lora_manager_layout import (
Discovery,
discover_sidecars,
layout_summary,
local_media,
)
try:
from ..utils.hashing import generate_prompt_hash
@@ -37,35 +51,44 @@ LORA_CATEGORY = "lora-manager"
JOB_TITLE = "LoRA reference import"
# The job's outcome groups, in its counter-tile order.
GROUPS = ("imported", "linked", "skipped", "downloaded", "download_errors", "errors")
METADATA_SUFFIX = "*.metadata.json"
class ImportCancelled(Exception):
"""The import was cancelled while staging; nothing was written."""
def discover(lora_manager_path: str) -> List[Path]:
"""Every LoRA metadata file, sorted and without repeats.
def discover(lora_manager_path: Optional[str]) -> Discovery:
"""The sidecars of every LoRA whose model exists inside a LoRA root.
Raises whatever listing a folder raised: a discovery that cannot see
every folder must not replace the collection with part of it.
``lora_manager_path`` is resolved as the integration resolves it (an
install under ``custom_nodes``; empty means auto-detection). Raises
whatever listing a folder raised: a discovery that cannot see every
folder must not replace the collection with part of it.
"""
seen = set()
files = []
for folder in lora_utils.find_lora_directories(lora_manager_path):
for path in Path(folder).rglob(METADATA_SUFFIX):
key = str(path)
if key not in seen:
seen.add(key)
files.append(path)
files.sort(key=str)
return files
return discover_sidecars(lora_utils.lora_manager_layout(lora_manager_path))
def fingerprint(files: List[Path]) -> str:
"""What a review saw of ``files``: each path, size and modification time."""
def fingerprint(discovery: Discovery) -> str:
"""What a review saw: the layout, then each kept sidecar's path, size and
modification time (a LoraManager save that changes unrelated settings
leaves it as it was)."""
layout = discovery.layout
digest = hashlib.sha256()
for path in files:
for value in (
layout.settings_path,
layout.source,
layout.error,
layout.mode,
"\0".join(layout.lora_roots),
layout.sidecar_root,
layout.example_images_root,
layout.library,
):
digest.update(f"{value}\n".encode("utf-8", "surrogatepass"))
for name in discovery.unreadable:
digest.update(f"unreadable\0{name}\n".encode("utf-8", "surrogatepass"))
for sidecar in discovery.sidecars:
path = sidecar.metadata_path
try:
stat = path.stat()
entry = f"{path}\0{stat.st_size}\0{stat.st_mtime_ns}\n"
@@ -75,31 +98,48 @@ def fingerprint(files: List[Path]) -> str:
return "sha256:" + digest.hexdigest()
def preview(lora_manager_path: str, cache_dir: Path) -> Dict:
def preview(
lora_manager_path: Optional[str],
cache_dir: Path,
media_extensions: AbstractSet[str],
) -> Dict:
"""The review's counts (blocking; reads only, requests nothing).
``{lora_files, with_metadata, planned_downloads, fingerprint, names}``:
the metadata files found, those that parse, the example images an import
would download (not yet in the cache) and the first LoRA names.
``{lora_files, with_metadata, skipped, with_local_media,
planned_downloads, fingerprint, names, layout}``: the metadata files
seen, the kept LoRAs whose sidecar parses, the rest, the LoRAs with a
local image, the CivitAI images an import would download for the LoRAs
without one (not yet in the cache), the first LoRA names and what the
layout says (``layout_summary``).
"""
files = discover(lora_manager_path)
discovery = discover(lora_manager_path)
layout = discovery.layout
with_metadata = 0
with_local_media = 0
planned = 0
names = []
for path in files:
metadata = lora_utils.read_lora_metadata(path)
for sidecar in discovery.sidecars:
metadata = lora_utils.read_lora_metadata(sidecar.metadata_path)
if not metadata:
continue
with_metadata += 1
planned += lora_utils.planned_civitai_downloads(metadata, path, cache_dir)
if local_media(sidecar, metadata, layout, media_extensions):
with_local_media += 1
else:
planned += lora_utils.planned_civitai_downloads(
metadata, sidecar.metadata_path, cache_dir
)
if len(names) < 50:
names.append(lora_utils.get_model_name_from_metadata(metadata))
return {
"lora_files": len(files),
"lora_files": discovery.scanned,
"with_metadata": with_metadata,
"skipped": discovery.scanned - with_metadata,
"with_local_media": with_local_media,
"planned_downloads": planned,
"fingerprint": fingerprint(files),
"fingerprint": fingerprint(discovery),
"names": names,
"layout": layout_summary(layout, Path.home()),
}
@@ -131,35 +171,45 @@ class StagedImport:
def stage(
files: List[Path],
discovery: Discovery,
cache_dir: Path,
api_key: str = "",
*,
media_extensions: AbstractSet[str],
on_record: Optional[Callable[[int, Dict, "StagedImport"], None]] = None,
cancelled: Callable[[], bool] = lambda: False,
logger=None,
) -> StagedImport:
"""Read each file and build its record, downloading its example images.
"""Read each sidecar and build its record with its images.
Blocking. A file that cannot be read, or whose record fails to build, is
skipped by its name. ``cancelled()`` is asked before each file (raises
ImportCancelled); ``on_record(index, record_or_None, staged)`` follows
each one.
Blocking. Local images (``local_media``) come first; a LoRA without any
downloads its CivitAI example images. A sidecar that cannot be read, or
whose record fails to build, is skipped by its name, as are the
centralized sidecars discovery could not parse. ``cancelled()`` is asked
before each sidecar (raises ImportCancelled); ``on_record(index,
record_or_None, staged)`` follows each one.
"""
staged = StagedImport(total=len(files))
for index, path in enumerate(files):
sidecars = discovery.sidecars
staged = StagedImport(total=len(sidecars))
staged.skipped.extend(discovery.unreadable)
for index, sidecar in enumerate(sidecars):
if cancelled():
raise ImportCancelled()
record = None
path = sidecar.metadata_path
metadata = lora_utils.read_lora_metadata(path)
if not metadata:
staged.skipped.append(path.name)
else:
try:
local = lora_utils.get_preview_images_from_metadata(metadata, path)
remote = lora_utils.download_civitai_images(
metadata, path, cache_dir, api_key, stats=staged.stats
local = local_media(
sidecar, metadata, discovery.layout, media_extensions
)
remote = []
if not local:
remote = lora_utils.download_civitai_images(
metadata, path, cache_dir, api_key, stats=staged.stats
)
images = list(dict.fromkeys([*local, *remote]))
record = build_record(metadata, path, images)
staged.records.append(record)
py/lora_utils.py
+78 -96
View File
@@ -27,6 +27,16 @@ except ImportError:
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.logging_config import get_logger
from dataclasses import replace
from .lora_manager_layout import (
MODE_ALONGSIDE,
METADATA_SUFFIX,
LoraManagerLayout,
discover_sidecars,
load_layout,
)
logger = get_logger("prompt_manager.lora_utils")
# ── LoraManager detection ────────────────────────────────────────────
@@ -234,23 +244,22 @@ def resolve_lora_manager_path(path) -> Optional[str]:
# ── Metadata reading ─────────────────────────────────────────────────
def find_lora_directories(lora_manager_path: str) -> List[str]:
"""Find directories that contain LoRA models (with .metadata.json files).
def default_lora_roots() -> List[str]:
"""The LoRA folders ComfyUI knows, without LoraManager's settings.
Searches: ComfyUI models/loras, extra_model_paths.yaml lora dirs,
and the LoraManager extension dir itself.
ComfyUI's ``models/loras``, the ``loras`` entries of
``extra_model_paths.yaml`` and, inside ComfyUI, every folder
``folder_paths`` registers for ``loras``. ``lora_manager_layout`` adds
the folders LoraManager's own settings name.
"""
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()))
@@ -267,13 +276,28 @@ def find_lora_directories(lora_manager_path: str) -> List[str]:
# 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 lora_manager_install(lora_manager_path: Optional[str]) -> Optional[str]:
"""The LoraManager install to read: the configured path when it is an
allowed install under ``custom_nodes``, auto-detection when it is empty."""
if isinstance(lora_manager_path, str) and lora_manager_path.strip():
return resolve_lora_manager_path(lora_manager_path)
detected = detect_lora_manager()
return resolve_lora_manager_path(detected) if detected else None
def layout_for_install(install: Optional[str]) -> LoraManagerLayout:
"""The layout of an already resolved install (None: none was found)."""
return load_layout(install, default_roots=default_lora_roots)
def lora_manager_layout(lora_manager_path: Optional[str]) -> LoraManagerLayout:
"""Where LoraManager keeps its files (``py/lora_manager_layout.py``)."""
return layout_for_install(lora_manager_install(lora_manager_path))
def _get_extra_lora_paths(comfyui_root: Path) -> List[Path]:
"""Parse extra_model_paths.yaml for additional LoRA directories."""
results = []
@@ -379,46 +403,6 @@ def get_model_name_from_metadata(metadata: Dict) -> str:
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
# CivitAI's own domains (#340: civitai.red and civitai.green serve the same
@@ -837,25 +821,6 @@ def get_lora_image_cache_dir(create: bool = True) -> Path:
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)
@@ -873,35 +838,22 @@ class TriggerWordCache:
self._loaded = False
def load(self, lora_manager_path: str) -> int:
"""Scan LoRA metadata files and build the trigger word mapping.
"""Build the trigger word map from the sidecars LoraManager uses.
Never raises (a node calls this while it runs): a folder that cannot
be listed, or any other failure, is logged and the previous map is
kept. When LoraManager's settings cannot be used, the reason is
logged and the default LoRA folders are read in alongside mode.
Returns:
Number of LoRAs with trigger words found.
Number of LoRAs with trigger words in the map.
"""
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
try:
new_cache = self._build(lora_manager_path)
except Exception as error:
logger.warning(f"Trigger word cache not refreshed: {error}")
with self._lock:
return len(self._cache)
with self._lock:
self._cache = new_cache
@@ -912,6 +864,36 @@ class TriggerWordCache:
)
return len(new_cache)
@staticmethod
def _build(lora_manager_path: str) -> Dict[str, List[str]]:
layout = lora_manager_layout(lora_manager_path)
if layout.error is not None:
logger.warning(
f"LoraManager settings not usable ({layout.error}); trigger words "
"come from the default LoRA folders"
)
layout = replace(
layout,
error=None,
mode=MODE_ALONGSIDE,
sidecar_root=None,
example_images_root=None,
)
new_cache: Dict[str, List[str]] = {}
for sidecar in discover_sidecars(layout).sidecars:
metadata = read_lora_metadata(sidecar.metadata_path)
if not metadata:
continue
words = get_trigger_words_from_metadata(metadata)
if not words:
continue
# Key by the model file name (what appears in <lora:NAME:weight>)
# and by the sidecar's name.
new_cache[sidecar.model_path.stem.lower()] = words
sidecar_stem = sidecar.metadata_path.name[: -len(METADATA_SUFFIX)]
new_cache.setdefault(sidecar_stem.lower(), words)
return 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:
tests/browser/workspace-system-integrations.spec.mjs
+2 -2
View File
@@ -322,7 +322,7 @@ test('a partial save made from Integrations keeps its outcome when Sources opens
test('Review import shows replacement, blocked items and network effects with the backup box; Start opens the job (D107)',
async ({ page, request }) => {
await lora(request, { loras: 3, civitai: true, previous: 2 });
await lora(request, { loras: 3, civitai: true, preview: false, previous: 2 });
const previews = [];
const applies = [];
page.on('request', item => {
@@ -408,7 +408,7 @@ test('a failed import keeps old imports and its job says so', async ({ page, req
});
test('the import can be cancelled while it downloads; then nothing is written', async ({ page, request }) => {
await lora(request, { loras: 2, civitai: true, previous: 1 });
await lora(request, { loras: 2, civitai: true, preview: false, previous: 1 });
await openIntegrations(page);
await enable(page);
const before = await serverJson(request, '/__test__/integrations/state');
tests/browser/workspace_integrations_fixture.py
+26 -8
View File
@@ -7,13 +7,18 @@ reaches a real LoraManager, ComfyUI install or the network:
no key, no trigger words), removes the fake install, its LoRA folder and
the image cache, forgets every LoRA import and injected failure, and opens
the download gate;
- ``POST /__test__/integrations/lora-manager`` ``{loras, civitai, previous}``
writes ``custom_nodes/ComfyUI-Lora-Manager`` (detected through the
``folder_paths`` stub's base path) and ``loras`` metadata files under
``models/loras``, each with a local preview PNG; ``civitai`` adds one
CivitAI example image per LoRA (served by the stub opener below);
``previous`` saves that many earlier ``lora-manager`` imports with tags and
image links;
- ``POST /__test__/integrations/lora-manager``
``{loras, civitai, preview, previous, settings}`` writes
``custom_nodes/ComfyUI-Lora-Manager`` (detected through the
``folder_paths`` stub's base path) and ``loras`` LoRAs under
``models/loras``: each a ``<stem>.safetensors`` model beside its
``<stem>.metadata.json`` sidecar (LoraManager's alongside layout, #355)
and, unless ``preview`` is false, a local preview PNG; ``civitai`` adds one
CivitAI example image per LoRA (served by the stub opener below; an import
downloads it only for a LoRA without a local preview); ``previous`` saves
that many earlier ``lora-manager`` imports with tags and image links;
``settings: "invalid"`` makes LoraManager's settings file unreadable JSON
(the server's own isolated settings folder, ``LORA_MANAGER_SETTINGS_DIR``);
- ``POST /__test__/integrations/gate`` ``{hold}`` holds every CivitAI
download until it is released (``{"hold": false}``), so a test sees the job
while it downloads;
@@ -30,6 +35,7 @@ test ever opens a network connection.
import io
import json
import os
import shutil
import threading
from pathlib import Path
@@ -91,7 +97,11 @@ def register_integrations_commands(app, routes, api, root: Path):
type(api.db).replace_category_prompts = replace
def settings_file():
return Path(os.environ["LORA_MANAGER_SETTINGS_DIR"]) / "settings.json"
def reset():
settings_file().unlink(missing_ok=True)
IntegrationConfig.LORA_MANAGER_ENABLED = False
IntegrationConfig.LORA_MANAGER_PATH = ""
IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = False
@@ -118,11 +128,17 @@ def register_integrations_commands(app, routes, api, root: Path):
data = await request.json() if request.can_read_body else {}
count = data.get("loras", 2)
previous = data.get("previous", 0)
preview = data.get("preview", True)
settings = data.get("settings")
for value in (count, previous):
if isinstance(value, bool) or not isinstance(value, int):
raise web.HTTPBadRequest()
if not 0 <= count <= MAX_LORAS or not 0 <= previous <= MAX_LORAS:
raise web.HTTPBadRequest()
if not isinstance(preview, bool) or settings not in (None, "invalid"):
raise web.HTTPBadRequest()
if settings == "invalid":
settings_file().write_text("{not json", encoding="utf-8")
(lora_manager / "py").mkdir(parents=True, exist_ok=True)
(lora_manager / "__init__.py").write_text("")
loras.mkdir(parents=True, exist_ok=True)
@@ -141,8 +157,10 @@ def register_integrations_commands(app, routes, api, root: Path):
"model_name": f"Style {number + 1:02}",
"civitai": civitai,
}
(loras / f"{stem}.safetensors").write_bytes(b"model")
(loras / f"{stem}.metadata.json").write_text(json.dumps(meta))
(loras / f"{stem}.png").write_bytes(_png())
if preview:
(loras / f"{stem}.png").write_bytes(_png())
for number in range(previous):
text = f"previous lora import {number + 1}"
prompt_id = api.db.save_prompt(
tests/test_api_images.py
+10 -3
View File
@@ -1317,7 +1317,8 @@ class TestServeErrorBranches(ImageRouteCoverageCase):
with (
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True),
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""),
patch("py.lora_utils.find_lora_directories", return_value=[lora]),
patch("py.lora_utils.detect_lora_manager", return_value=None),
patch("py.lora_utils.default_lora_roots", return_value=[lora]),
patch("py.lora_utils.get_lora_image_cache_dir", return_value=cache),
):
resp = await self.client.request(
@@ -1328,7 +1329,8 @@ class TestServeErrorBranches(ImageRouteCoverageCase):
with (
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True),
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""),
patch("py.lora_utils.find_lora_directories", side_effect=RuntimeError),
patch("py.lora_utils.detect_lora_manager", return_value=None),
patch("py.lora_utils.default_lora_roots", side_effect=RuntimeError),
):
resp = await self.client.request(
"GET", f"/prompt_manager/images/{image_id}/file"
@@ -1347,7 +1349,12 @@ class TestLoraManagerPathAllowList(ImageRouteCoverageCase):
)
stack.enter_context(patch.object(IntegrationConfig, "LORA_MANAGER_PATH", path))
stack.enter_context(
patch("py.lora_utils.find_lora_directories", return_value=[lora])
patch("py.lora_utils.default_lora_roots", return_value=[lora])
)
# Never the developer's LoraManager settings (#355).
settings = stack.enter_context(tempfile.TemporaryDirectory())
stack.enter_context(
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": settings})
)
stack.enter_context(
patch("py.lora_utils.get_lora_image_cache_dir", return_value=cache)
tests/test_api_lora.py
+71 -18
View File
@@ -70,8 +70,16 @@ class LoraAPITestCase(AioHTTPTestCase):
# Never the repository's config.json (live user data): every save in
# these tests goes to a file in this test's temp directory.
self.config_file = self.tmp_root / "config.json"
# Never the developer's LoraManager settings either (#355): an empty
# settings folder of this test's own.
self.lm_settings = self.tmp_root / "lora-manager-settings"
self.lm_settings.mkdir()
self._env_patch = patch.dict(
os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file)}
os.environ,
{
"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file),
"LORA_MANAGER_SETTINGS_DIR": str(self.lm_settings),
},
)
self._env_patch.start()
@@ -194,6 +202,8 @@ def _make_lora_manager(where):
def _write_metadata(directory, stem, words=None, images=None, prompt=None):
"""A LoraManager sidecar beside its model, as LoraManager names them:
``<stem>.metadata.json`` next to ``<stem>.safetensors`` (#355)."""
meta = {"file_name": f"{stem}.safetensors", "model_name": stem, "civitai": {}}
if words is not None:
meta["civitai"]["trainedWords"] = words
@@ -201,7 +211,8 @@ def _write_metadata(directory, stem, words=None, images=None, prompt=None):
meta["civitai"]["images"] = [
{"url": url, "meta": {"prompt": prompt} if prompt else {}} for url in images
]
path = directory / f"{stem}.safetensors.metadata.json"
(directory / f"{stem}.safetensors").write_bytes(b"model")
path = directory / f"{stem}.metadata.json"
path.write_text(json.dumps(meta))
return path
@@ -300,9 +311,7 @@ class TestLoraEnableRoute(LoraAPITestCase):
_write_metadata(loras, "neon", words=["glow"])
self.detect_mock.return_value = str(lm)
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
@@ -848,9 +857,7 @@ class TestLoraScanRoute(LoraAPITestCase):
cache_dir = self.tmp_root / "img_cache"
with (
patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
),
patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=cache_dir
),
@@ -880,9 +887,10 @@ class TestLoraScanRoute(LoraAPITestCase):
images=["https://civitai.com/a.png", "https://evil.example/b.png"],
prompt="neon street",
)
(loras / "neon.png").write_bytes(_png_bytes())
# No local image: its CivitAI images are downloaded (#355).
_write_metadata(loras, "plain")
(loras / "broken.metadata.json").write_text("{{{")
(loras / "broken.safetensors").write_bytes(b"model")
existing_id = self.api.db.save_prompt(
text="plain",
category=None,
@@ -910,8 +918,56 @@ class TestLoraScanRoute(LoraAPITestCase):
self.assertEqual(imported[0]["text"], "neon street")
self.assertIn("lora:neon", imported[0]["tags"])
self.assertIn("glow", imported[0]["tags"])
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 2)
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 1)
self.assertEqual(len(self.api.db.get_prompt_images(existing_id)), 0)
# The unreadable sidecar is reported by its file name.
statuses = [e.get("status") for e in events if e.get("processed")]
self.assertIn("broken.metadata.json", statuses)
async def test_scan_links_local_media_without_downloading(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(
loras, "neon", images=["https://civitai.com/a.png"], prompt="neon street"
)
(loras / "neon.preview.webp").write_bytes(b"x")
resp, events, opener_calls = await self._scan(loras)
self.assertEqual(events[-1]["imported"], 1)
self.assertEqual(opener_calls, [])
imported = self.api.db.search_prompts(category="lora-manager")
images = self.api.db.get_prompt_images(imported[0]["id"])
self.assertEqual(
[os.path.basename(i["image_path"]) for i in images], ["neon.preview.webp"]
)
async def test_scan_reports_unusable_settings_without_replacing(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon")
(self.lm_settings / "settings.json").write_text("{not json")
old_id = self.api.db.save_prompt(
text="kept",
category="lora-manager",
tags=[],
rating=None,
prompt_hash=generate_prompt_hash("kept"),
)
resp, events, _ = await self._scan(loras)
self.assertEqual(resp.status, 200)
self.assertEqual(events[-1]["type"], "error")
self.assertEqual(
events[-1]["message"],
"LoraManager settings: not usable: the settings file is not valid "
"JSON. Nothing was changed.",
)
kept = self.api.db.search_prompts(category="lora-manager")
self.assertEqual([p["id"] for p in kept], [old_id])
async def test_save_failure_is_counted_as_skipped(self):
self._enable()
@@ -949,9 +1005,7 @@ class TestLoraTriggerWordRoutes(LoraAPITestCase):
loras.mkdir()
_write_metadata(loras, "neon", words=["glow", "bright"])
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=NEON"
)
@@ -985,17 +1039,16 @@ class TestLoraTriggerWordRoutes(LoraAPITestCase):
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self.client.request(
"POST", "/prompt_manager/lora/refresh-cache"
)
data = await resp.json()
self.assertTrue(data["success"])
# keyed by both the file_name stem and the metadata-file stem
self.assertEqual(data["loras_with_trigger_words"], 2)
# Keyed by the model name and the sidecar name, which are the same
# for LoraManager's alongside layout.
self.assertEqual(data["loras_with_trigger_words"], 1)
async def test_refresh_error_is_500(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
tests/test_lora_import.py
+299 -28
View File
@@ -252,7 +252,7 @@ class LoraImportAPITestCase(LoraAPITestCase):
self._patches = [
patch.object(
lora_utils, "find_lora_directories", return_value=[str(self.loras)]
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=self.cache_dir
@@ -276,6 +276,8 @@ class LoraImportAPITestCase(LoraAPITestCase):
await super().tearDownAsync()
def _two_loras(self):
"""``neon`` has only a CivitAI image (downloaded); ``plain`` has a
local preview (linked where it is, nothing downloaded)."""
_write_metadata(
self.loras,
"neon",
@@ -283,8 +285,8 @@ class LoraImportAPITestCase(LoraAPITestCase):
images=["https://civitai.com/a.png"],
prompt="neon street",
)
(self.loras / "neon.png").write_bytes(_png_bytes())
_write_metadata(self.loras, "plain")
_write_metadata(self.loras, "plain", images=["https://civitai.com/p.png"])
(self.loras / "plain.png").write_bytes(_png_bytes())
async def _post(self, url, body=None):
resp = await self._post_json(url, {} if body is None else body)
@@ -321,8 +323,34 @@ class LoraImportRouteTests(LoraImportAPITestCase):
self.assertEqual(resp.status, 200, plan)
self.assertEqual(
{k: plan[k] for k in ("lora_files", "with_metadata", "current_imports")},
{"lora_files": 3, "with_metadata": 2, "current_imports": 2},
{
k: plan[k]
for k in (
"lora_files",
"with_metadata",
"skipped",
"with_local_media",
"current_imports",
)
},
{
"lora_files": 3,
"with_metadata": 2,
"skipped": 1,
"with_local_media": 1,
"current_imports": 2,
},
)
self.assertEqual(
plan["layout"],
{
"source": "environment",
"settings_file": str(self.lm_settings / "settings.json"),
"error": None,
"mode": "alongside",
"library": "default",
"example_images": False,
},
)
self.assertEqual(plan["blocked_by_pairs"], 0)
self.assertEqual(plan["planned_downloads"], 1)
@@ -336,7 +364,82 @@ class LoraImportRouteTests(LoraImportAPITestCase):
self.assertEqual(_size(db_path + "-shm"), shm)
self.assertEqual(_snapshot(self.api.db), rows)
self.assertNotIn(SECRET_KEY, json.dumps(plan))
self.assertNotIn(str(self.tmp_root), json.dumps(plan))
# The settings file is the only absolute path the review shows.
shown = dict(plan, layout=dict(plan["layout"], settings_file=None))
self.assertNotIn(str(self.tmp_root), json.dumps(shown))
async def test_unusable_settings_review_offers_no_token(self):
self._two_loras()
(self.lm_settings / "settings.json").write_text("[1, 2]")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertIsNone(plan["token"])
self.assertEqual(
plan["layout"]["error"], "the settings file is not a JSON object"
)
self.assertEqual((plan["lora_files"], plan["with_metadata"]), (0, 0))
self.assertEqual(self.requests, [])
async def test_centralized_review_counts_what_it_skips(self):
sidecars = self.tmp_root / "sidecars"
examples = self.tmp_root / "examples"
(self.lm_settings / "settings.json").write_text(
json.dumps(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
"example_images_path": str(examples),
"libraries": {"comfyui": {}},
"active_library": "comfyui",
}
)
)
component = sidecars / "loras-0a1b2c3d"
component.mkdir(parents=True)
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
meta = {"model_name": "neon", "file_path": str(model), "sha256": "AB12"}
(component / "neon.metadata.json").write_text(json.dumps(meta))
(examples / "ab12").mkdir(parents=True)
(examples / "ab12" / "image_1.png").write_bytes(_png_bytes())
gone = dict(meta, file_path=str(self.loras / "gone.safetensors"))
(component / "gone.metadata.json").write_text(json.dumps(gone))
(component / "broken.metadata.json").write_text("{{{")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertEqual(
(
plan["lora_files"],
plan["with_metadata"],
plan["skipped"],
plan["with_local_media"],
plan["planned_downloads"],
),
(3, 1, 2, 1, 0),
)
self.assertEqual(
(plan["layout"]["mode"], plan["layout"]["library"]),
("centralized", "comfyui"),
)
self.assertTrue(plan["layout"]["example_images"])
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(self.requests, [])
self.assertEqual(
[row["detail"] for row in job.outcomes["skipped"]],
["broken.metadata.json"],
)
imported = self.api.db.search_prompts(category=CATEGORY)
images = self.api.db.get_prompt_images(imported[0]["id"])
self.assertEqual(
[os.path.basename(i["image_path"]) for i in images], ["image_1.png"]
)
async def test_preview_requires_the_integration(self):
IntegrationConfig.LORA_MANAGER_ENABLED = False
@@ -358,7 +461,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
self._two_loras()
resp, plan = await self._post(PREVIEW)
rows = _snapshot(self.api.db)
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
meta.write_text(meta.read_text() + " ")
resp, body = await self._post(APPLY, {"token": plan["token"]})
@@ -397,7 +500,17 @@ class LoraImportRouteTests(LoraImportAPITestCase):
imported = self.api.db.search_prompts(category=CATEGORY)
self.assertEqual(sorted(p["text"] for p in imported), ["neon street", "plain"])
neon = next(p for p in imported if p["text"] == "neon street")
self.assertEqual(len(self.api.db.get_prompt_images(neon["id"])), 2)
self.assertEqual(len(self.api.db.get_prompt_images(neon["id"])), 1)
# plain's local preview is linked where it is; its CivitAI image is not
# downloaded (#355).
plain = next(p for p in imported if p["text"] == "plain")
self.assertEqual(
[i["image_path"] for i in self.api.db.get_prompt_images(plain["id"])],
[os.path.realpath(self.loras / "plain.png")],
)
self.assertEqual(
[r.full_url for r in self.requests], ["https://civitai.com/a.png"]
)
self.assertEqual(len(self.api.db.search_prompts(category="portraits")), 1)
# The token is single-use.
resp, body = await self._post(APPLY, {"token": plan["token"]})
@@ -409,7 +522,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
# Legacy scan: discovery raises once the stream is open.
with patch.object(
lora_utils, "find_lora_directories", side_effect=OSError("share gone")
lora_import, "discover_sidecars", side_effect=OSError("share gone")
):
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
@@ -565,6 +678,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
async def test_legacy_scan_events_unchanged_on_success(self):
self._two_loras()
(self.loras / "broken.metadata.json").write_text("{{{")
(self.loras / "broken.safetensors").write_bytes(b"model")
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
@@ -853,13 +967,13 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
async def test_files_changed_while_staging_fail_the_job(self):
# I1: LoraManager rewrites a metadata file while the downloads run.
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
await self._change_during_download(
lambda: meta.write_text(meta.read_text() + " ")
)
async def test_a_file_removed_while_staging_fails_the_job(self):
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
await self._change_during_download(meta.unlink)
async def test_a_file_added_while_staging_fails_the_job(self):
@@ -869,7 +983,7 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
self._two_loras()
before = _snapshot(self.api.db)
self.gate = threading.Event()
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
resp = await self.client.request("POST", SCAN)
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, self.downloading.wait, 10))
@@ -986,37 +1100,194 @@ class LoraImportRestoreTests(LoraImportAPITestCase):
class LoraImportStageTests(unittest.TestCase):
"""The stages over a disposable LoraManager install (#355)."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(lambda: __import__("shutil").rmtree(self.tmp, True))
self.lm = _make_lora_manager(self.tmp / "custom_nodes")
self.loras = self.tmp / "loras"
self.loras.mkdir()
self.settings = self.tmp / "settings"
self.settings.mkdir()
self.roots = [str(self.loras)]
patchers = [
patch.object(
lora_utils,
"custom_nodes_directories",
return_value=[self.tmp / "custom_nodes"],
),
patch.object(lora_utils, "default_lora_roots", lambda: list(self.roots)),
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}),
]
for patcher in patchers:
patcher.start()
self.addCleanup(patcher.stop)
def settings_file(self, data):
(self.settings / "settings.json").write_text(json.dumps(data))
def discover(self):
return lora_import.discover(str(self.lm))
def test_discover_sorts_and_drops_repeats(self):
a = self.tmp / "a"
b = self.tmp / "a" / "b"
a = self.loras / "a"
b = self.loras / "a" / "b"
b.mkdir(parents=True)
_write_metadata(b, "zeta")
_write_metadata(a, "alpha")
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(a), str(b)]
):
files = lora_import.discover(str(self.tmp))
self.assertEqual([f.name for f in files][0], "alpha.safetensors.metadata.json")
self.assertEqual(len(files), 2)
self.roots = [str(a), str(b)]
discovery = self.discover()
names = [s.metadata_path.name for s in discovery.sidecars]
self.assertEqual(names, ["alpha.metadata.json", "zeta.metadata.json"])
self.assertEqual(discovery.scanned, 2)
def test_discover_without_an_install_finds_nothing(self):
_write_metadata(self.loras, "alpha")
discovery = lora_import.discover(str(self.tmp))
self.assertEqual(discovery.layout.error, "no LoraManager install found")
self.assertEqual(discovery.sidecars, [])
def test_fingerprint_follows_size_and_mtime(self):
meta = _write_metadata(self.tmp, "one")
first = lora_import.fingerprint([meta])
self.assertEqual(lora_import.fingerprint([meta]), first)
meta = _write_metadata(self.loras, "one")
first = lora_import.fingerprint(self.discover())
self.assertEqual(lora_import.fingerprint(self.discover()), first)
meta.write_text(meta.read_text() + " ")
self.assertNotEqual(lora_import.fingerprint([meta]), first)
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
changed = lora_import.fingerprint(self.discover())
meta.unlink()
self.assertNotEqual(lora_import.fingerprint([meta]), first)
self.assertNotEqual(lora_import.fingerprint(self.discover()), changed)
def test_fingerprint_ignores_unrelated_settings(self):
_write_metadata(self.loras, "one")
self.settings_file({"civitai_api_key": "a"})
first = lora_import.fingerprint(self.discover())
self.settings_file({"civitai_api_key": "b", "theme": "dark"})
self.assertEqual(lora_import.fingerprint(self.discover()), first)
sidecars = self.tmp / "sidecars"
sidecars.mkdir()
self.settings_file(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
}
)
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
def test_fingerprint_follows_a_missing_file_between_stat_calls(self):
meta = _write_metadata(self.loras, "one")
discovery = self.discover()
first = lora_import.fingerprint(discovery)
meta.unlink()
self.assertNotEqual(lora_import.fingerprint(discovery), first)
def test_stage_cancel_before_a_file_raises(self):
meta = _write_metadata(self.tmp, "one")
_write_metadata(self.loras, "one")
with self.assertRaises(lora_import.ImportCancelled):
lora_import.stage([meta], self.tmp, cancelled=lambda: True)
lora_import.stage(
self.discover(),
self.tmp,
media_extensions=MEDIA,
cancelled=lambda: True,
)
def test_local_media_means_no_download(self):
_write_metadata(self.loras, "neon", images=["https://civitai.com/a.png"])
(self.loras / "neon.mp4").write_bytes(b"video")
def refuse(*args, **kwargs):
raise AssertionError("a LoRA with local media downloads nothing")
with patch.object(lora_utils, "download_civitai_images", refuse):
staged = lora_import.stage(
self.discover(), self.tmp, media_extensions=MEDIA
)
self.assertEqual(staged.skipped, [])
self.assertEqual(
staged.records[0]["images"], [os.path.realpath(self.loras / "neon.mp4")]
)
def test_a_lora_without_local_media_downloads(self):
_write_metadata(self.loras, "neon", images=["https://civitai.com/a.png"])
calls = []
def download(metadata, path, cache_dir, api_key, stats=None):
calls.append(path.name)
return ["/cache/a.jpg"]
with patch.object(lora_utils, "download_civitai_images", download):
staged = lora_import.stage(
self.discover(), self.tmp, media_extensions=MEDIA
)
self.assertEqual(calls, ["neon.metadata.json"])
self.assertEqual(staged.records[0]["images"], ["/cache/a.jpg"])
def test_alongside_records_match_the_previous_import(self):
_write_metadata(
self.loras,
"neon",
words=["glow", "glow"],
images=["https://civitai.com/a.png"],
prompt="neon street",
)
(self.loras / "neon.png").write_bytes(b"x")
staged = lora_import.stage(self.discover(), self.tmp, media_extensions=MEDIA)
self.assertEqual(
staged.records,
[
{
"name": "neon",
"text": "neon street",
"hash": generate_prompt_hash("neon street"),
"tags": ["lora-manager", "lora:neon", "glow"],
"images": [os.path.realpath(self.loras / "neon.png")],
}
],
)
def test_preview_counts_local_media_and_downloads(self):
_write_metadata(self.loras, "local", images=["https://civitai.com/l.png"])
(self.loras / "local.webp").write_bytes(b"x")
_write_metadata(
self.loras,
"remote",
images=["https://civitai.com/r1.png", "https://civitai.com/r2.png"],
)
(self.loras / "orphan.metadata.json").write_text("{}")
found = lora_import.preview(str(self.lm), self.tmp / "cache", MEDIA)
self.assertEqual(
{
k: found[k]
for k in (
"lora_files",
"with_metadata",
"skipped",
"with_local_media",
"planned_downloads",
)
},
{
"lora_files": 3,
"with_metadata": 2,
"skipped": 1,
"with_local_media": 1,
"planned_downloads": 2,
},
)
self.assertEqual(found["names"], ["local", "remote"])
self.assertEqual(found["layout"]["mode"], "alongside")
def test_preview_of_an_unusable_layout(self):
(self.settings / "settings.json").write_text("nope")
_write_metadata(self.loras, "one")
found = lora_import.preview(str(self.lm), self.tmp / "cache", MEDIA)
self.assertEqual(
found["layout"]["error"], "the settings file is not valid JSON"
)
self.assertEqual((found["lora_files"], found["with_metadata"]), (0, 0))
MEDIA = frozenset({".png", ".jpg", ".jpeg", ".webp", ".mp4", ".webm"})
if __name__ == "__main__":
tests/test_lora_utils.py
+153 -52
View File
@@ -1263,11 +1263,12 @@ class TestTriggerWordCache(unittest.TestCase):
"file_name": "my_lora.safetensors",
"civitai": {"trainedWords": ["trigger1", "trigger2"]},
}
meta_path = Path(tmpdir) / "my_lora.safetensors.metadata.json"
meta_path = Path(tmpdir) / "my_lora.metadata.json"
meta_path.write_text(json.dumps(meta))
(Path(tmpdir) / "my_lora.safetensors").write_bytes(b"m")
# Patch find_lora_directories to return our temp dir
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
# No LoraManager install: the default LoRA folders, alongside.
with patch("py.lora_utils.default_lora_roots", return_value=[tmpdir]):
count = self.cache.load(tmpdir)
self.assertTrue(self.cache.is_loaded)
@@ -1283,11 +1284,10 @@ class TestTriggerWordCache(unittest.TestCase):
"file_name": "MyLoRA.safetensors",
"civitai": {"trainedWords": ["word1"]},
}
(Path(tmpdir) / "MyLoRA.safetensors.metadata.json").write_text(
json.dumps(meta)
)
(Path(tmpdir) / "MyLoRA.metadata.json").write_text(json.dumps(meta))
(Path(tmpdir) / "MyLoRA.safetensors").write_bytes(b"m")
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
with patch("py.lora_utils.default_lora_roots", return_value=[tmpdir]):
self.cache.load(tmpdir)
self.assertEqual(self.cache.get_trigger_words("mylora"), ["word1"])
@@ -1512,7 +1512,7 @@ class TestDetectLoraManager(unittest.TestCase):
self.assertEqual(lora_utils.detect_lora_manager(), str(lm.resolve()))
class TestFindLoraDirectories(unittest.TestCase):
class TestDefaultLoraRoots(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
@@ -1520,7 +1520,7 @@ class TestFindLoraDirectories(unittest.TestCase):
self.root = _make_comfy_root(self.tmp)
self.lm = _make_lora_manager(self.root / "custom_nodes")
def test_collects_default_extra_runtime_and_metadata_dirs(self):
def test_collects_default_extra_and_runtime_dirs_only(self):
extra_abs = self.tmp / "extra_abs"
extra_abs.mkdir()
(self.root / "rel_loras").mkdir()
@@ -1542,14 +1542,14 @@ class TestFindLoraDirectories(unittest.TestCase):
{"folder_paths": _FakeFolderPaths(str(self.root), [str(runtime)])},
),
):
dirs = lora_utils.find_lora_directories(str(self.lm))
dirs = lora_utils.default_lora_roots()
# #355: folders inside the LoraManager install are no longer roots.
expected = {
str((self.root / "models" / "loras").resolve()),
str((self.root / "rel_loras").resolve()),
str(extra_abs.resolve()),
str(runtime.resolve()),
str(nested.resolve()),
}
self.assertEqual(set(dirs), expected)
self.assertEqual(dirs, sorted(dirs))
@@ -1560,12 +1560,21 @@ class TestFindLoraDirectories(unittest.TestCase):
(self.root / "extra_model_paths.yml").write_text("a: [unclosed\n")
self.assertEqual(lora_utils._get_extra_lora_paths(self.root), [])
def test_without_root_only_metadata_dirs_are_found(self):
def test_without_root_nothing_is_found(self):
with (
patch.object(lora_utils, "find_comfyui_root", return_value=None),
patch.dict(sys.modules, {"folder_paths": None}),
):
self.assertEqual(lora_utils.find_lora_directories(str(self.lm)), [])
self.assertEqual(lora_utils.default_lora_roots(), [])
def test_the_removed_helpers_are_gone(self):
for name in (
"find_lora_directories",
"get_preview_images_from_metadata",
"get_preview_image_from_metadata",
"get_example_images_dir",
):
self.assertFalse(hasattr(lora_utils, name), name)
class TestModelNameAndPreviews(unittest.TestCase):
@@ -1581,44 +1590,6 @@ class TestModelNameAndPreviews(unittest.TestCase):
self.assertEqual(get_name({"file_name": "f.safetensors"}), "f.safetensors")
self.assertEqual(get_name({}), "unknown")
def test_preview_images_found_by_stem(self):
meta_path = self.tmp / "mylora.safetensors.metadata.json"
meta_path.write_text("{}")
(self.tmp / "mylora.png").write_bytes(b"x")
(self.tmp / "mylora.preview.jpg").write_bytes(b"x")
found = lora_utils.get_preview_images_from_metadata({}, meta_path)
self.assertEqual(
sorted(os.path.basename(f) for f in found),
["mylora.png", "mylora.preview.jpg"],
)
self.assertEqual(
lora_utils.get_preview_image_from_metadata({}, meta_path), found[0]
)
def test_preview_uses_file_name_from_metadata(self):
meta_path = self.tmp / "other.metadata.json"
(self.tmp / "named.webp").write_bytes(b"x")
found = lora_utils.get_preview_images_from_metadata(
{"file_name": "named.safetensors"}, meta_path
)
self.assertEqual([os.path.basename(f) for f in found], ["named.webp"])
self.assertIsNone(lora_utils.get_preview_image_from_metadata({}, meta_path))
def test_example_images_dir(self):
self.assertIsNone(lora_utils.get_example_images_dir(str(self.tmp)))
nested = self.tmp / "user" / "example_images"
nested.mkdir(parents=True)
self.assertEqual(
lora_utils.get_example_images_dir(str(self.tmp)), str(nested.resolve())
)
direct = self.tmp / "example_images"
direct.mkdir()
self.assertEqual(
lora_utils.get_example_images_dir(str(self.tmp)), str(direct.resolve())
)
class TestDownloadEdgeCases(unittest.TestCase):
@@ -1703,11 +1674,141 @@ class TestInjectTriggerWords(unittest.TestCase):
(tmp / "named.metadata.json").write_text(
json.dumps({"civitai": {"trainedWords": ["w"]}})
)
for stem in ("bad", "nowords", "named"):
(tmp / f"{stem}.safetensors").write_bytes(b"m")
cache = TriggerWordCache()
with patch.object(lora_utils, "find_lora_directories", return_value=[str(tmp)]):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(tmp)]):
self.assertEqual(cache.load(str(tmp)), 1)
self.assertEqual(cache.get_trigger_words("named"), ["w"])
class TestTriggerWordCacheLayout(unittest.TestCase):
"""#355: trigger words come from the sidecars LoraManager really uses."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, self.tmp, True)
self.custom_nodes = self.tmp / "custom_nodes"
self.lm = _make_lora_manager(self.custom_nodes)
self.loras = self.tmp / "loras"
self.loras.mkdir()
self.sidecars = self.tmp / "sidecars"
self.settings = self.tmp / "settings"
self.settings.mkdir()
patchers = [
patch.object(
lora_utils, "custom_nodes_directories", return_value=[self.custom_nodes]
),
patch.object(lora_utils, "default_lora_roots", return_value=[]),
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}),
]
for patcher in patchers:
patcher.start()
self.addCleanup(patcher.stop)
def centralized(self):
self.sidecars.mkdir(exist_ok=True)
(self.settings / "settings.json").write_text(
json.dumps(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(self.sidecars),
"folder_paths": {"loras": [str(self.loras)]},
}
)
)
def sidecar(self, folder, stem, model, words):
folder.mkdir(parents=True, exist_ok=True)
meta = {
"file_name": stem,
"file_path": str(model),
"civitai": {"trainedWords": words},
}
(folder / f"{stem}.metadata.json").write_text(json.dumps(meta))
def test_centralized_sidecars_are_read(self):
self.centralized()
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
self.sidecar(self.sidecars / "loras-0a1b2c3d", "neon", model, ["glow"])
# A leftover sidecar beside the model is not what LoraManager reads.
self.sidecar(self.loras, "neon", model, ["stale"])
cache = TriggerWordCache()
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(cache.get_trigger_words("neon"), ["glow"])
def test_both_the_model_and_the_sidecar_name_are_keys(self):
self.centralized()
model = self.loras / "model_name.safetensors"
model.write_bytes(b"m")
self.sidecar(self.sidecars / "c", "sidecar_name", model, ["w"])
cache = TriggerWordCache()
self.assertEqual(cache.load(str(self.lm)), 2)
self.assertEqual(cache.get_trigger_words("model_name"), ["w"])
self.assertEqual(cache.get_trigger_words("sidecar_name"), ["w"])
def test_an_empty_path_means_auto_detection(self):
self.centralized()
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
self.sidecar(self.sidecars / "c", "neon", model, ["glow"])
cache = TriggerWordCache()
with patch.object(lora_utils, "detect_lora_manager", return_value=str(self.lm)):
self.assertEqual(cache.load(""), 1)
def test_an_unlistable_folder_keeps_the_previous_map(self):
cache = TriggerWordCache()
cache._cache = {"old": ["word"]}
cache._loaded = True
with (
patch.object(
lora_utils, "discover_sidecars", side_effect=OSError("share gone")
),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(cache.get_trigger_words("old"), ["word"])
self.assertIn("share gone", "\n".join(logs.output))
def test_any_other_failure_never_raises(self):
cache = TriggerWordCache()
with (
patch.object(lora_utils, "load_layout", side_effect=RuntimeError("bad")),
self.assertLogs("prompt_manager.lora_utils", level="WARNING"),
):
self.assertEqual(cache.load(str(self.lm)), 0)
self.assertFalse(cache.is_loaded)
def test_an_unusable_layout_falls_back_to_the_default_roots(self):
(self.settings / "settings.json").write_text("{not json")
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
self.sidecar(self.loras, "neon", model, ["glow"])
cache = TriggerWordCache()
with (
patch.object(
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(cache.get_trigger_words("neon"), ["glow"])
self.assertIn("not valid JSON", "\n".join(logs.output))
def test_layout_wraps_the_install_and_the_default_roots(self):
with patch.object(
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
):
layout = lora_utils.lora_manager_layout(str(self.lm))
self.assertIsNone(layout.error)
self.assertEqual(layout.source, "environment")
self.assertIn(os.path.normcase(str(self.loras)), layout.lora_roots)
self.assertEqual(
lora_utils.lora_manager_layout(str(self.tmp)).error,
"no LoraManager install found",
)
if __name__ == "__main__":
unittest.main()
tests/test_node_execution_contract.py
+19
View File
@@ -315,6 +315,25 @@ class TestLoraTriggerWords(NodeTestCase):
)
cache.load.assert_called_once_with("/loras")
def test_a_failing_discovery_logs_and_the_node_still_runs(self):
# #355: the trigger cache reads LoraManager's folders while a node
# runs; a folder that cannot be listed must never stop the node.
from py import lora_utils
cache = lora_utils.get_trigger_cache()
cache.clear()
self.addCleanup(cache.clear)
with (
mock.patch.dict(sys.modules, {"py.config": fake_integration_config(True)}),
mock.patch.object(
lora_utils, "discover_sidecars", side_effect=OSError("share gone")
),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
text = self.node._inject_lora_trigger_words("a cat <lora:neon:1>")
self.assertEqual(text, "a cat <lora:neon:1>")
self.assertIn("share gone", "\n".join(logs.output))
class TestPromptManagerEncoding(NodeTestCase):
node_cls = PromptManager