Vito Sansevero
ae660de588
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.
224 lines
8.8 KiB
Python
224 lines
8.8 KiB
Python
"""The staged LoRA reference import (#221, #249 System S10, #355).
|
|
|
|
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**: 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).
|
|
|
|
Both the System job (``POST /lora/import/apply``) and the legacy ``POST
|
|
/lora/scan`` run these stages (``py/api/lora_integration.py``).
|
|
"""
|
|
|
|
import hashlib
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
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
|
|
except ImportError: # pragma: no cover - flat test imports
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
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")
|
|
|
|
|
|
class ImportCancelled(Exception):
|
|
"""The import was cancelled while staging; nothing was written."""
|
|
|
|
|
|
def discover(lora_manager_path: Optional[str]) -> Discovery:
|
|
"""The sidecars of every LoRA whose model exists inside a LoRA root.
|
|
|
|
``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.
|
|
"""
|
|
return discover_sidecars(lora_utils.lora_manager_layout(lora_manager_path))
|
|
|
|
|
|
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 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"
|
|
except OSError:
|
|
entry = f"{path}\0missing\n"
|
|
digest.update(entry.encode("utf-8", "surrogatepass"))
|
|
return "sha256:" + digest.hexdigest()
|
|
|
|
|
|
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, 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``).
|
|
"""
|
|
discovery = discover(lora_manager_path)
|
|
layout = discovery.layout
|
|
with_metadata = 0
|
|
with_local_media = 0
|
|
planned = 0
|
|
names = []
|
|
for sidecar in discovery.sidecars:
|
|
metadata = lora_utils.read_lora_metadata(sidecar.metadata_path)
|
|
if not metadata:
|
|
continue
|
|
with_metadata += 1
|
|
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": discovery.scanned,
|
|
"with_metadata": with_metadata,
|
|
"skipped": discovery.scanned - with_metadata,
|
|
"with_local_media": with_local_media,
|
|
"planned_downloads": planned,
|
|
"fingerprint": fingerprint(discovery),
|
|
"names": names,
|
|
"layout": layout_summary(layout, Path.home()),
|
|
}
|
|
|
|
|
|
def build_record(metadata: Dict, path: Path, images: List[str]) -> Dict:
|
|
"""The prompt a LoRA becomes: its example prompt (else its name), tags, images."""
|
|
model_name = lora_utils.get_model_name_from_metadata(metadata)
|
|
text = lora_utils.get_example_prompt_from_metadata(metadata) or model_name
|
|
tags = ["lora-manager", f"lora:{model_name}"]
|
|
tags.extend(lora_utils.get_trigger_words_from_metadata(metadata))
|
|
return {
|
|
"name": model_name,
|
|
"text": text,
|
|
"hash": generate_prompt_hash(text),
|
|
"tags": list(dict.fromkeys(tags)),
|
|
"images": images,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class StagedImport:
|
|
"""The records an import will write, and what staging counted."""
|
|
|
|
total: int
|
|
records: List[Dict] = field(default_factory=list)
|
|
skipped: List[str] = field(default_factory=list)
|
|
stats: Dict[str, int] = field(
|
|
default_factory=lambda: {"downloaded": 0, "download_errors": 0}
|
|
)
|
|
|
|
|
|
def stage(
|
|
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 sidecar and build its record with its images.
|
|
|
|
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.
|
|
"""
|
|
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 = 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)
|
|
except Exception as error: # one LoRA never stops the import
|
|
name = lora_utils.get_model_name_from_metadata(metadata)
|
|
if logger is not None:
|
|
logger.warning(f"Failed to prepare LoRA {name}: {error}")
|
|
staged.skipped.append(name)
|
|
if on_record is not None:
|
|
on_record(index, record, staged)
|
|
return staged
|