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

fix(lora): stage imports before replacing the collection (#221)

Browse Source 
The LoRA import deleted the whole lora-manager category first and then
saved each record in its own transaction, so a discovery, download or
database failure left the previous collection gone or partial.

The import now runs in three stages: discover and fingerprint the
metadata files, stage every record (CivitAI downloads go into the image
cache first, counted), then replace the category in one BEGIN IMMEDIATE
transaction (replace_category_prompts). Any failure leaves the previous
prompts, tags and image links exactly as they were. Prompts a saved pair
uses block the replacement; an empty replacement of a non-empty
collection needs a reviewed import.

POST /prompt_manager/lora/import/preview reviews an import without
writing or requesting anything; POST /lora/import/apply {token} runs it
as the Activity job "LoRA reference import" (records admission,
cancellable while it downloads, the safety backup policy before the
replace, refused with plan-changed when the files changed since the
review). The legacy POST /lora/scan runs the same stages with its SSE
event shape.
This commit is contained in:
Vito Sansevero
2026-10-25 06:52:34 -07:00
parent 82f7fc8a6f
commit 6a2dae0f78
9 changed files with 1699 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/category_import.py
+170
View File
@@ -0,0 +1,170 @@
"""Atomic replacement of one imported category (#221, #249 System S10).
The LoRA reference import used to delete its whole category and then save
each record in its own transaction, so a failure part way left the previous
collection deleted and the new one partial. ``replace_category_prompts``
does both in one ``BEGIN IMMEDIATE`` transaction: the category's current
prompts (with their tags and image links) are replaced by the staged
records, or, on any error, nothing changes.
A staged record is ``{text, hash, tags, images}``. A record whose hash a
prompt outside the category already has links its images to that prompt
(the legacy import did the same); a second record with the hash of an
earlier one links its images to that one. Neither is a new prompt.
"""
from typing import Any, Dict, List, Optional, Sequence
try:
from .maintenance import PlanChanged
from .workflow_documents import documents_of_rows, release_documents
except ImportError: # pragma: no cover - flat test imports
from database.maintenance import PlanChanged
from database.workflow_documents import documents_of_rows, release_documents
class EmptyReplacement(ValueError):
"""The replacement has no record and would remove every current import."""
# Each record's outcome and the result counter it adds to.
_COUNTERS = {"imported": "imported", "linked": "linked", "duplicate": "duplicates"}
def _placeholders(values: Sequence[Any]) -> str:
return ",".join("?" * len(values))
class CategoryImportOperations:
"""Read the state of an imported category and replace it atomically."""
def category_import_state(self, category: str) -> Dict[str, Any]:
"""``{ids, paired}``: the category's prompt ids (ascending) and how
many of them a saved pair uses (those block a replacement)."""
with self.model.get_connection() as conn:
ids = [
row[0]
for row in conn.execute(
"SELECT id FROM prompts WHERE category = ? ORDER BY id",
(category,),
).fetchall()
]
paired = 0
if ids:
marks = _placeholders(ids)
paired = conn.execute(
f"SELECT COUNT(*) FROM prompts WHERE id IN ({marks}) AND id IN "
"(SELECT positive_prompt_id FROM prompt_pairs UNION "
"SELECT negative_prompt_id FROM prompt_pairs)",
ids,
).fetchone()[0]
return {"ids": ids, "paired": paired}
def existing_prompt_hashes(
self, hashes: Sequence[str], exclude_category: str
) -> set:
"""The given hashes a prompt outside ``exclude_category`` already has."""
found = set()
wanted = list(dict.fromkeys(hashes))
with self.model.get_connection() as conn:
for start in range(0, len(wanted), 500):
chunk = wanted[start : start + 500]
rows = conn.execute(
f"SELECT hash FROM prompts WHERE hash IN ({_placeholders(chunk)}) "
"AND (category IS NULL OR category != ?)",
[*chunk, exclude_category],
).fetchall()
found.update(row[0] for row in rows)
return found
def replace_category_prompts(
self,
category: str,
records: List[Dict[str, Any]],
*,
expected_ids: Optional[List[int]] = None,
allow_empty: bool = False,
) -> Dict[str, Any]:
"""Replace every prompt of ``category`` with ``records`` in one transaction.
Raises (and changes nothing): ``PlanChanged`` when ``expected_ids`` is
given and the category's prompt ids differ from it;
``PairDependencyError`` when a saved pair uses one of its prompts;
``EmptyReplacement`` when ``records`` is empty while the category has
prompts, unless ``allow_empty`` (an empty replacement that was
reviewed). Any other error rolls everything back too.
Returns ``{removed, imported, linked, duplicates, images, rows}``:
prompts removed and created, records linked to an existing prompt or
to an earlier record with the same text, image links added, and per
record ``{name, outcome}`` (``imported``/``linked``/``duplicate``).
"""
with self._pair_transaction() as conn:
ids = [
row[0]
for row in conn.execute(
"SELECT id FROM prompts WHERE category = ? ORDER BY id",
(category,),
).fetchall()
]
if expected_ids is not None and ids != list(expected_ids):
raise PlanChanged("The imported prompts changed since the review")
if not records and ids and not allow_empty:
raise EmptyReplacement(
"Nothing to import; the previous imports were kept"
)
self._ensure_prompts_unpaired(conn, ids)
removed = self._delete_category_rows(conn, ids)
result = {
"removed": removed,
"imported": 0,
"linked": 0,
"duplicates": 0,
"images": 0,
"rows": [],
}
created = {}
for record in records:
self._import_record(conn, category, record, created, result)
return result
def _delete_category_rows(self, conn, ids: List[int]) -> int:
"""Delete prompts ``ids`` with their image links and tags (no commit)."""
if not ids:
return 0
marks = _placeholders(ids)
documents = documents_of_rows(conn, "prompt_id", ids)
conn.execute(f"DELETE FROM generated_images WHERE prompt_id IN ({marks})", ids)
release_documents(conn, documents)
conn.execute(f"DELETE FROM prompt_tags WHERE prompt_id IN ({marks})", ids)
return conn.execute(f"DELETE FROM prompts WHERE id IN ({marks})", ids).rowcount
def _import_record(self, conn, category, record, created, result) -> None:
"""Save one staged record on ``conn``, or link it to the prompt it repeats."""
prompt_hash = record["hash"]
name = record.get("name") or record["text"]
if prompt_hash in created:
prompt_id, outcome = created[prompt_hash], "duplicate"
else:
row = conn.execute(
"SELECT id FROM prompts WHERE hash = ?", (prompt_hash,)
).fetchone()
if row is not None:
prompt_id, outcome = row[0], "linked"
else:
prompt_id = self._insert_prompt(
conn,
record["text"],
category,
list(record.get("tags") or []),
None,
None,
prompt_hash,
)
outcome = "imported"
created[prompt_hash] = prompt_id
result[_COUNTERS[outcome]] += 1
for image in record.get("images") or []:
if self._insert_image_link(conn, prompt_id, image):
result["images"] += 1
result["rows"].append({"name": name, "outcome": outcome})
database/operations.py
+100 -70
View File
@@ -26,6 +26,7 @@ from .history_reads import HistoryReadOperations
from .health import HealthOperations, iter_consistency_issues
from .duplicate_copies import DuplicateCopyOperations
from .workflow_storage import WorkflowStorageOperations
from .category_import import CategoryImportOperations, EmptyReplacement # noqa: F401
from .workflow_documents import (
DocumentReader,
documents_of_rows,
@@ -366,6 +367,7 @@ class PromptDatabase(
MaintenanceOperations,
DuplicateCopyOperations,
WorkflowStorageOperations,
CategoryImportOperations,
):
"""Database operations class for managing prompts."""
@@ -455,26 +457,9 @@ class PromptDatabase(
try:
with self.model.get_connection() as conn:
cursor = conn.execute(
"""
INSERT INTO prompts (
text, category, tags, rating, notes, hash, created_at,
updated_at
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
""",
(
text.strip(),
category,
rating,
notes,
prompt_hash,
datetime.datetime.now(datetime.timezone.utc).isoformat(),
datetime.datetime.now(datetime.timezone.utc).isoformat(),
),
prompt_id = self._insert_prompt(
conn, text, category, tags, rating, notes, prompt_hash
)
prompt_id = cursor.lastrowid
if tags:
self._sync_prompt_tags(conn, prompt_id, tags)
conn.commit()
self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
return prompt_id
@@ -490,6 +475,32 @@ class PromptDatabase(
)
return existing["id"]
def _insert_prompt(
self,
conn: sqlite3.Connection,
text: str,
category: Optional[str],
tags: Optional[List[str]],
rating: Optional[int],
notes: Optional[str],
prompt_hash: Optional[str],
) -> int:
"""Insert one prompt row and its tags on ``conn`` (no commit)."""
now = datetime.datetime.now(datetime.timezone.utc).isoformat()
cursor = conn.execute(
"""
INSERT INTO prompts (
text, category, tags, rating, notes, hash, created_at,
updated_at
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
""",
(text.strip(), category, rating, notes, prompt_hash, now, now),
)
prompt_id = cursor.lastrowid
if tags:
self._sync_prompt_tags(conn, prompt_id, tags)
return prompt_id
def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool:
"""
Count runs of a prompt and mark it as most recently used.
@@ -1670,64 +1681,83 @@ class PromptDatabase(
return 0
# Proceed with linking
if not isinstance(metadata, dict):
metadata = {}
filename = os.path.basename(image_path)
file_path = normalize_image_path(image_path)
file_size, width, height, image_format = _coerce_file_info(metadata)
# Each workflow and prompt graph is stored once (new form).
documents = stored_values(
conn,
{
"workflow_data": dumps_strict(metadata.get("workflow", {})),
"prompt_metadata": dumps_strict(metadata.get("prompt", {})),
},
image_id = self._insert_image_link(
conn, prompt_id_int, image_path, metadata
)
conn.commit()
# INSERT OR IGNORE skips a file already linked to this prompt
# (same prompt_id + normalised full path)
cursor = conn.execute(
"""
INSERT OR IGNORE INTO generated_images
(prompt_id, image_path, filename, file_path, file_size, width,
height, format, workflow_data, prompt_metadata, parameters,
workflow_document_id, prompt_document_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
if not image_id:
self.logger.debug(
f"Image {normalize_image_path(image_path)} already linked to "
f"prompt {prompt_id_int}"
)
return 0
return image_id
def _insert_image_link(
self,
conn: sqlite3.Connection,
prompt_id_int: int,
image_path: str,
metadata: Optional[Dict[str, Any]] = None,
) -> int:
"""Insert one image link on ``conn`` (no commit).
Returns the new image row's id, or 0 when the file is already linked
to this prompt (``INSERT OR IGNORE``).
"""
if not isinstance(metadata, dict):
metadata = {}
filename = os.path.basename(image_path)
file_path = normalize_image_path(image_path)
file_size, width, height, image_format = _coerce_file_info(metadata)
# Each workflow and prompt graph is stored once (new form).
documents = stored_values(
conn,
{
"workflow_data": dumps_strict(metadata.get("workflow", {})),
"prompt_metadata": dumps_strict(metadata.get("prompt", {})),
},
)
# INSERT OR IGNORE skips a file already linked to this prompt
# (same prompt_id + normalised full path)
cursor = conn.execute(
"""
INSERT OR IGNORE INTO generated_images
(prompt_id, image_path, filename, file_path, file_size, width,
height, format, workflow_data, prompt_metadata, parameters,
workflow_document_id, prompt_document_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
prompt_id_int,
image_path,
filename,
file_path,
file_size,
width,
height,
image_format,
documents["workflow_data"],
documents["prompt_metadata"],
dumps_strict(metadata.get("parameters", {})),
documents["workflow_document_id"],
documents["prompt_document_id"],
),
)
if cursor.rowcount == 0:
# Already linked: a document made only for this row goes.
release_documents(
conn,
(
prompt_id_int,
image_path,
filename,
file_path,
file_size,
width,
height,
image_format,
documents["workflow_data"],
documents["prompt_metadata"],
dumps_strict(metadata.get("parameters", {})),
documents["workflow_document_id"],
documents["prompt_document_id"],
),
)
if cursor.rowcount == 0:
# Already linked: a document made only for this row goes.
release_documents(
conn,
(
documents["workflow_document_id"],
documents["prompt_document_id"],
),
)
conn.commit()
if cursor.rowcount == 0:
self.logger.debug(
f"Image {file_path} already linked to prompt {prompt_id_int}"
)
return 0
return cursor.lastrowid
return 0
return cursor.lastrowid
def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
"""
py/api/__init__.py
+4
View File
@@ -19,6 +19,7 @@ Split into domain-specific mixins for maintainability:
- MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply)
- WorkflowStorageRoutesMixin: store each workflow once, then Reclaim space
- DuplicateFilesMixin: duplicate image files as find and removal jobs
- LoraImportRoutesMixin: the reviewed, staged LoRA reference import job
"""
import asyncio
@@ -63,6 +64,7 @@ from .jobs import JobRegistry
from .logging_routes import LoggingRoutesMixin
from .autotag_routes import AutotagRoutesMixin
from .lora_integration import LoraIntegrationMixin
from .lora_import_routes import LoraImportRoutesMixin
try:
from ...database.operations import PromptDatabase
@@ -307,6 +309,7 @@ class PromptManagerAPI(
LoggingRoutesMixin,
AutotagRoutesMixin,
LoraIntegrationMixin,
LoraImportRoutesMixin,
):
"""REST API handler for PromptManager operations and web interface.
@@ -551,6 +554,7 @@ class PromptManagerAPI(
self._register_logging_routes(routes)
self._register_autotag_routes(routes)
self._register_lora_routes(routes)
self._register_lora_import_routes(routes)
# Register gzip compression middleware (once)
global _gzip_registered
py/api/lora_import_routes.py
+416
View File
@@ -0,0 +1,416 @@
"""The staged LoRA reference import as a reviewed Activity job (#221, S10).
``POST /prompt_manager/lora/import/preview`` (body ``{}``) reviews an import:
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).
``POST /prompt_manager/lora/import/apply {token}`` starts the job (202
``{job}``): kind ``lora_import``, titled "LoRA reference import", admission
``records``, cancellable while it downloads only. Refused before it starts:
``plan-changed`` (409: the metadata files changed since the review),
``plan-expired``, ``plan-used``, ``storage-changed``, ``busy``,
``insufficient-space`` (as maintenance). The job stages the records
(downloads into the cache first, each counted in ``downloaded`` or
``download_errors``), takes the safety backup the policy promises, then
replaces the collection in one transaction bound to the reviewed imports. Any
failure leaves the previous imports exactly as they were, and the job says so.
The legacy ``POST /lora/scan`` runs the same stages (``stage_lora_import``
and ``replace_lora_imports``) under its SSE event shape.
"""
import asyncio
import os
import sqlite3
from aiohttp import web
from .. import lora_import, lora_utils
from .job_routes import busy_response
from .jobs import LORA_IMPORT_JOB_KIND, JobBusy
from .maintenance import (
MAX_TOKEN_LENGTH,
_failure,
_fail_job,
_invalid,
_iso,
disk_plan,
short_space_message,
)
from .maintenance_plans import ScopeTooLarge
try:
from ...database.category_import import EmptyReplacement
from ...database.maintenance import PlanChanged
from ...database.pairs import PairDependencyError
except ImportError: # pragma: no cover - flat test imports
from database.category_import import EmptyReplacement
from database.maintenance import PlanChanged
from database.pairs import PairDependencyError
LORA_OPERATION = "lora_import"
CATEGORY = lora_import.LORA_CATEGORY
SOURCE_CHANGED = (
"The LoRA metadata files changed since the review; nothing was changed. "
"Review the import again."
)
IMPORTS_CHANGED = (
"The imported LoRA prompts changed since the review; the previous imports "
"were kept. Review the import again."
)
BLOCKED = (
"Saved pairs use some imported LoRA prompts, so the collection cannot be "
"replaced; the previous imports were kept."
)
EMPTY = "No LoRA could be prepared for import; the previous imports were kept."
FAILED = "The import failed; the previous LoRA imports were kept."
CACHE_NOTE = "Images already downloaded stay in the image cache."
CANCELLED = "Cancelled before anything was written; the previous imports were kept."
DISABLED = "LoraManager integration is not enabled"
NO_PATH = "LoraManager path not configured"
def _noun(count, word, plural=None):
return word if count == 1 else (plural or word + "s")
def _count(count, word, plural=None):
return f"{count:,} {_noun(count, word, plural)}"
def _import_line(result, skipped):
parts = [f"Imported {_count(result['imported'], 'LoRA')}"]
if result["linked"] or result["duplicates"]:
linked = result["linked"] + result["duplicates"]
parts.append(f"{linked:,} linked to an existing prompt")
if skipped:
parts.append(f"{skipped:,} skipped")
return " · ".join(parts)
class LoraImportRoutesMixin:
"""Mixin providing the reviewed LoRA import (preview, apply, job)."""
def _register_lora_import_routes(self, routes):
@routes.post("/prompt_manager/lora/import/preview")
async def lora_import_preview_route(request):
return await self.preview_lora_import(request)
@routes.post("/prompt_manager/lora/import/apply")
async def lora_import_apply_route(request):
return await self.apply_lora_import(request)
# ── shared stages ─────────────────────────────────────────────────
async def _lora_source(self):
"""``(path, None)``: the configured, allowed LoraManager path, or
``(None, (code, message))`` when the import cannot run."""
from ..config import IntegrationConfig
if not IntegrationConfig.LORA_MANAGER_ENABLED:
return None, ("lora-disabled", DISABLED)
lm_path = IntegrationConfig.LORA_MANAGER_PATH
if not lm_path:
return None, ("lora-path", NO_PATH)
allowed = await self._run_in_executor(
lora_utils.resolve_lora_manager_path, lm_path
)
if allowed is None:
from .lora_integration import _LORA_PATH_ERROR
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."""
from ..config import IntegrationConfig
cache_dir = await self._run_in_executor(lora_utils.get_lora_image_cache_dir)
loop = asyncio.get_running_loop()
def report(index, record, staged):
if on_record is not None:
loop.call_soon_threadsafe(on_record, index, record, staged)
return await self._run_in_executor(
lambda: lora_import.stage(
files,
cache_dir,
IntegrationConfig.CIVITAI_API_KEY,
on_record=report,
cancelled=cancelled or (lambda: False),
logger=self.logger,
)
)
async def replace_lora_imports(self, db, staged, *, expected_ids, allow_empty):
"""Swap the collection for ``staged.records`` in one transaction."""
return await self._run_in_executor(
lambda: db.replace_category_prompts(
CATEGORY,
staged.records,
expected_ids=expected_ids,
allow_empty=allow_empty,
)
)
# ── preview ───────────────────────────────────────────────────────
async def preview_lora_import(self, request):
data, error = await self._read_maintenance_body(request, ())
if error is not None:
return error
lm_path, refused = await self._lora_source()
if refused is not None:
return _failure(refused[1], 400, refused[0])
db = self.db
identity = self._database_identity(db)
try:
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)
state = await self._run_in_executor(db.category_import_state, CATEGORY)
except Exception:
self.logger.exception("LoRA import preview failed")
return _failure("The LoRA import could not be reviewed.", 500)
if identity != self._database_identity(self.db):
return _failure(
"The database was replaced while it was reviewed. Review again.",
409,
"storage-changed",
)
return await self._lora_plan_response(db, identity, lm_path, found, state)
async def _lora_plan_response(self, db, identity, lm_path, found, state):
from ..config import IntegrationConfig
current = len(state["ids"])
blocked = state["paired"]
policy = self._safety_backup_policy()
try:
disk, _short, size = await self._run_in_executor(
disk_plan, db.model.db_path, backup=policy != "off", vacuum=False
)
except (OSError, sqlite3.Error):
self.logger.exception("Could not measure the disk space for a plan")
return _failure("The LoRA import could not be reviewed.", 500)
token = None
expires_at = None
if not blocked:
payload = {
"path": lm_path,
"fingerprint": found["fingerprint"],
"expected_ids": state["ids"],
"empty": found["with_metadata"] == 0,
"backup": policy,
}
try:
token = self._maintenance_plans.issue(
LORA_OPERATION, found["fingerprint"], payload, identity
)
except ScopeTooLarge:
return _failure(
"This import is too large to review at once.", 413, "too-large"
)
expires_at = _iso(self._maintenance_plans.peek(token).expires_at)
name = os.path.basename(db.model.db_path)
downloads = found["planned_downloads"]
reasons = (
[f"{_count(blocked, 'imported prompt')} used by saved pairs"]
if blocked
else []
)
return web.json_response(
{
"success": True,
"operation": LORA_OPERATION,
"token": token,
"expires_at": expires_at,
"lora_files": found["lora_files"],
"with_metadata": found["with_metadata"],
"current_imports": current,
"blocked_by_pairs": blocked,
"planned_downloads": downloads,
"scope": {
"count": found["with_metadata"],
"label": (
f"Replace {_count(current, 'imported prompt')} with "
f"{_count(found['with_metadata'], 'LoRA')}"
),
},
"protected": {"count": blocked, "reasons": reasons},
"backup": {
"policy": policy,
"destination_label": f"{name}.backup_<date and time>",
"expected_bytes": size if policy != "off" else None,
},
"media": {
"deletes_files": False,
"label": "No files are removed; image links are replaced",
},
"network": {
"planned_downloads": downloads,
"api_key": bool(IntegrationConfig.CIVITAI_API_KEY),
},
"details": [{"label": n, "detail": ""} for n in found["names"]],
"disk": disk,
},
headers={"Cache-Control": "no-store"},
)
# ── apply ─────────────────────────────────────────────────────────
async def apply_lora_import(self, request):
data, error = await self._read_maintenance_body(request, ("token",))
if error is not None:
return error
token = data.get("token")
if not isinstance(token, str) or not 0 < len(token) <= MAX_TOKEN_LENGTH:
return _invalid("token must be the plan's token")
plan, refused = self._peek_plan(token)
if refused is not None:
return refused
if plan.operation != LORA_OPERATION:
return _failure(
"This review expired. Review the plan again.", 410, "plan-expired"
)
store = self._maintenance_plans
db = self.db
if plan.db_identity != self._database_identity(db):
store.discard(token)
return _failure(
"The database was replaced since the review. Review again.",
409,
"storage-changed",
)
lm_path = plan.payload["path"]
try:
files = await self._run_in_executor(lora_import.discover, lm_path)
now = await self._run_in_executor(lora_import.fingerprint, files)
backup = self._backup_promised(plan)
_disk, short, _size = await self._run_in_executor(
disk_plan, db.model.db_path, backup=backup, vacuum=False
)
except Exception:
self.logger.exception("LoRA import could not start")
return _failure(
"The import could not start. Nothing was changed.", 500, "apply-failed"
)
if now != plan.fingerprint:
store.discard(token)
return _failure(SOURCE_CHANGED, 409, "plan-changed")
if short is not None:
return _failure(short_space_message(short), 507, "insufficient-space")
plan, refused = self._peek_plan(token)
if refused is not None:
return refused
try:
job = self._jobs.start(
LORA_IMPORT_JOB_KIND,
lora_import.JOB_TITLE,
admission="records",
cancellable=True,
scope={"operation": LORA_OPERATION, "lora_files": len(files)},
groups=lora_import.GROUPS,
runner=lambda job: self._run_lora_import(job, db, plan, files, backup),
)
except JobBusy as busy:
return busy_response(busy)
store.consume(token, LORA_OPERATION, job_id=job.id)
return web.json_response({"success": True, "job": job.summary()}, status=202)
# ── the job ───────────────────────────────────────────────────────
async def _run_lora_import(self, job, db, plan, files, backup_on):
total = len(files)
job.update(total=total, status_line="Preparing LoRA records")
job.log_line(f"Found {_count(total, 'LoRA metadata file')}")
def on_record(index, record, staged):
job.update(
processed=index + 1,
progress=int(90 * (index + 1) / max(total, 1)),
status_line=f"Prepared {index + 1:,} of {total:,}",
)
job.publish({"type": "progress", "progress": job.progress})
try:
staged = await self.stage_lora_import(
files, on_record=on_record, cancelled=lambda: job.cancelled
)
except lora_import.ImportCancelled:
job.log_line(f"{CANCELLED} {CACHE_NOTE}")
job.update(status_line=CANCELLED)
job.finish("cancelled")
return
except Exception:
self.logger.exception("LoRA import staging failed")
return _fail_job(job, "import-failed", f"{FAILED} {CACHE_NOTE}")
self._count_downloads(job, staged)
# Past this point nothing stops halfway: the replace is one transaction.
job.cancellable = False
if job.cancelled:
job.log_line(f"{CANCELLED} {CACHE_NOTE}")
job.update(status_line=CANCELLED)
job.finish("cancelled")
return
if staged.fingerprint != plan.fingerprint:
return _fail_job(job, "plan-changed", SOURCE_CHANGED)
ok, _backup = await self._safety_backup_step(job, db, backup_on)
if not ok:
return
job.update(status_line="Replacing the LoRA collection", progress=95)
await self._replace_in_job(job, db, plan, staged)
def _count_downloads(self, job, staged):
downloaded = staged.stats.get("downloaded", 0)
errors = staged.stats.get("download_errors", 0)
job.record_counted("downloaded", downloaded)
job.record_counted("download_errors", errors)
if downloaded or errors:
job.log_line(
f"Downloaded {_count(downloaded, 'example image')}"
+ (f"; {errors:,} could not be downloaded" if errors else "")
)
async def _replace_in_job(self, job, db, plan, staged):
try:
result = await self.replace_lora_imports(
db,
staged,
expected_ids=plan.payload["expected_ids"],
allow_empty=plan.payload["empty"],
)
except PlanChanged:
return _fail_job(job, "plan-changed", f"{IMPORTS_CHANGED} {CACHE_NOTE}")
except PairDependencyError:
return _fail_job(job, "blocked-by-pairs", f"{BLOCKED} {CACHE_NOTE}")
except EmptyReplacement:
return _fail_job(job, "nothing-to-import", f"{EMPTY} {CACHE_NOTE}")
except Exception:
self.logger.exception("LoRA import failed")
return _fail_job(job, "import-failed", f"{FAILED} {CACHE_NOTE}")
rows = {"imported": [], "linked": []}
for row in result["rows"]:
group = "imported" if row["outcome"] == "imported" else "linked"
rows[group].append({"path": "", "detail": row["name"]})
job.record_counted("imported", result["imported"], rows["imported"])
job.record_counted(
"linked", result["linked"] + result["duplicates"], rows["linked"]
)
job.record_counted(
"skipped",
len(staged.skipped),
[{"path": "", "detail": name} for name in staged.skipped],
)
line = _import_line(result, len(staged.skipped))
if result["removed"]:
job.log_line(f"Replaced {_count(result['removed'], 'previous import')}")
job.log_line(line)
job.update(progress=100, processed=job.total, status_line=line)
py/api/lora_integration.py
+65 -129
View File
@@ -1,7 +1,7 @@
"""LoraManager integration API routes for PromptManager."""
import asyncio
import json
from pathlib import Path
from aiohttp import web
@@ -322,29 +322,15 @@ class LoraIntegrationMixin:
return self._error_response(e)
async def _lora_scan_stream(self, lm_path, send_progress):
"""Import every LoRA under ``lm_path``, reporting through ``send_progress``."""
from ..config import IntegrationConfig
from ..lora_utils import (
download_civitai_images,
find_lora_directories,
get_example_prompt_from_metadata,
get_lora_image_cache_dir,
get_preview_images_from_metadata,
get_trigger_words_from_metadata,
get_model_name_from_metadata,
read_lora_metadata,
)
"""Import every LoRA under ``lm_path``, reporting through ``send_progress``.
await send_progress(
{
"type": "progress",
"status": "Clearing previous lora-manager imports...",
"progress": 0,
}
)
# Clear previous imports so reimport is always clean
await self._run_in_executor(self.db.delete_prompts_by_category, "lora-manager")
The staged path of the System import (#221, ``lora_import_routes``):
discover, stage every record (downloads into the cache first), then
replace the collection in one transaction. A failure at any stage
raises with the previous imports untouched; an empty replacement of a
non-empty collection is refused (it needs the reviewed import).
"""
from .. import lora_import
await send_progress(
{
@@ -353,19 +339,8 @@ class LoraIntegrationMixin:
"progress": 2,
}
)
lora_dirs = await self._run_in_executor(find_lora_directories, lm_path)
# Collect all metadata files
meta_files = []
for d in lora_dirs:
dir_path = Path(d)
meta_files.extend(dir_path.rglob("*.metadata.json"))
total = len(meta_files)
imported = 0
skipped = 0
files = await self._run_in_executor(lora_import.discover, lm_path)
total = len(files)
await send_progress(
{
"type": "progress",
@@ -375,116 +350,77 @@ class LoraIntegrationMixin:
}
)
cache_dir = get_lora_image_cache_dir()
events = asyncio.Queue()
state = await self._run_in_executor(
self.db.category_import_state, lora_import.LORA_CATEGORY
)
# Records whose prompt already exists outside the collection are
# linked, not imported (the running counts; the end ones are exact).
known = set()
for i, meta_file in enumerate(meta_files):
metadata = await self._run_in_executor(read_lora_metadata, meta_file)
if not metadata:
skipped += 1
continue
def on_record(index, record, _staged):
events.put_nowait((index, record))
model_name = get_model_name_from_metadata(metadata)
trigger_words = get_trigger_words_from_metadata(metadata)
# Collect all images: local previews + downloaded civitai examples
preview_paths = await self._run_in_executor(
get_preview_images_from_metadata, metadata, meta_file
)
civitai_paths = await self._run_in_executor(
download_civitai_images,
metadata,
meta_file,
cache_dir,
IntegrationConfig.CIVITAI_API_KEY,
)
# Merge, local first, dedup
seen = set(preview_paths)
all_images = list(preview_paths)
for cp in civitai_paths:
if cp not in seen:
all_images.append(cp)
seen.add(cp)
# Build prompt text: prefer example prompt, then model name
example_prompt = get_example_prompt_from_metadata(metadata)
prompt_text = example_prompt or model_name
# Build tags
tags = ["lora-manager", f"lora:{model_name}"]
tags.extend(trigger_words)
# Save to database via existing mechanism
try:
import hashlib
prompt_hash = hashlib.sha256(
prompt_text.strip().lower().encode("utf-8")
).hexdigest()
counts = {"imported": 0}
async def report(index, record):
if record is not None and record["hash"] not in known:
known.add(record["hash"])
existing = await self._run_in_executor(
self.db.get_prompt_by_hash, prompt_hash
self.db.existing_prompt_hashes,
[record["hash"]],
lora_import.LORA_CATEGORY,
)
if existing:
# Link all images
for pp in all_images:
await self._run_in_executor(
self.db.link_image_to_prompt,
existing["id"],
pp,
)
skipped += 1
else:
prompt_id = await self._run_in_executor(
self.db.save_prompt,
prompt_text,
"lora-manager", # category
tags,
None, # rating
None, # notes
prompt_hash,
)
if prompt_id:
for pp in all_images:
await self._run_in_executor(
self.db.link_image_to_prompt,
prompt_id,
pp,
)
imported += 1
else:
skipped += 1
except Exception as e:
self.logger.warning(f"Failed to import LoRA {model_name}: {e}")
skipped += 1
# Progress update for every LoRA
progress = int(5 + (90 * (i + 1) / max(total, 1)))
img_count = len(all_images)
status = f"{model_name}"
if img_count:
status += f" ({img_count} images)"
if not existing:
counts["imported"] += 1
status = record["name"] if record else files[index].name
if record and record["images"]:
status += f" ({len(record['images'])} images)"
await send_progress(
{
"type": "progress",
"status": status,
"progress": progress,
"processed": i + 1,
"imported": imported,
"skipped": skipped,
"progress": int(5 + (90 * (index + 1) / max(total, 1))),
"processed": index + 1,
"imported": counts["imported"],
"skipped": index + 1 - counts["imported"],
}
)
# Each record's event is queued (through the loop) before staging
# ends, so reporting until staging is done and then draining the
# queue reports every record in order.
staging = asyncio.ensure_future(
self.stage_lora_import(files, on_record=on_record)
)
while not staging.done():
getter = asyncio.ensure_future(events.get())
await asyncio.wait({getter, staging}, return_when=asyncio.FIRST_COMPLETED)
if getter.done():
await report(*getter.result())
else:
getter.cancel()
while not events.empty():
await report(*events.get_nowait())
staged = await staging
await send_progress(
{
"type": "progress",
"status": "Replacing previous lora-manager imports...",
"progress": 96,
}
)
result = await self.replace_lora_imports(
self.db, staged, expected_ids=state["ids"], allow_empty=False
)
await send_progress(
{
"type": "complete",
"progress": 100,
"total": total,
"imported": imported,
"skipped": skipped,
"imported": result["imported"],
"skipped": total - result["imported"],
}
)
py/lora_import.py
+172
View File
@@ -0,0 +1,172 @@
"""The staged LoRA reference import (#221, #249 System S10).
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:
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**: ``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 Callable, Dict, List, Optional
from . import lora_utils
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")
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.
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
def fingerprint(files: List[Path]) -> str:
"""What a review saw of ``files``: each path, size and modification time."""
digest = hashlib.sha256()
for path in files:
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: str, cache_dir: Path) -> 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.
"""
files = discover(lora_manager_path)
with_metadata = 0
planned = 0
names = []
for path in files:
metadata = lora_utils.read_lora_metadata(path)
if not metadata:
continue
with_metadata += 1
planned += lora_utils.planned_civitai_downloads(metadata, path, cache_dir)
if len(names) < 50:
names.append(lora_utils.get_model_name_from_metadata(metadata))
return {
"lora_files": len(files),
"with_metadata": with_metadata,
"planned_downloads": planned,
"fingerprint": fingerprint(files),
"names": names,
}
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}
)
fingerprint: str = ""
def stage(
files: List[Path],
cache_dir: Path,
api_key: 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.
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.
"""
staged = StagedImport(total=len(files), fingerprint=fingerprint(files))
for index, path in enumerate(files):
if cancelled():
raise ImportCancelled()
record = None
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
)
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
py/lora_utils.py
+50 -24
View File
@@ -525,36 +525,22 @@ def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
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.
def _civitai_cache_plan(
metadata: Dict, metadata_path: Path, cache_dir: Path
) -> Tuple[Path, List[str], List[Tuple[str, Path]]]:
"""``(lora_cache, cached, tasks)`` for a LoRA's CivitAI example images.
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.
``cached`` are the images already in the cache, ``tasks`` the
``(url, local_path)`` downloads still needed. Reads only; creates nothing.
"""
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 / 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:
@@ -571,9 +557,42 @@ def download_civitai_images(
cached.append(str(local_path.resolve()))
else:
tasks.append((url, local_path))
return lora_cache, cached, tasks
def planned_civitai_downloads(
metadata: Dict, metadata_path: Path, cache_dir: Path
) -> int:
"""How many example images an import would download (no request, no write)."""
return len(_civitai_cache_plan(metadata, metadata_path, cache_dir)[2])
def download_civitai_images(
metadata: Dict,
metadata_path: Path,
cache_dir: Path,
api_key: str = "",
stats: Optional[Dict[str, int]] = None,
) -> 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).
stats: When given, ``downloaded`` and ``download_errors`` are
increased by the downloads that succeeded and failed.
Returns:
List of absolute paths to downloaded image files.
"""
from concurrent.futures import ThreadPoolExecutor
lora_cache, cached, tasks = _civitai_cache_plan(metadata, metadata_path, cache_dir)
if not tasks:
return cached
lora_cache.mkdir(parents=True, exist_ok=True)
# Download in parallel
downloaded = []
@@ -586,15 +605,21 @@ def download_civitai_images(
if result:
downloaded.append(result)
if stats is not None:
stats["downloaded"] = stats.get("downloaded", 0) + len(downloaded)
stats["download_errors"] = (
stats.get("download_errors", 0) + len(tasks) - len(downloaded)
)
return cached + downloaded
def get_lora_image_cache_dir() -> Path:
def get_lora_image_cache_dir(create: bool = True) -> Path:
"""Get the directory used to cache downloaded LoRA example images.
``<data folder>/data/lora_images`` (``utils.library_location``): the node
folder for a library that was not moved, else ComfyUI's
``user/__promptmanager``.
``user/__promptmanager``. ``create=False`` only names it (a preview
writes nothing).
"""
try:
from ..utils import library_location
@@ -605,7 +630,8 @@ def get_lora_image_cache_dir() -> Path:
cache = Path(
library_location.cache_dir("lora_images", extension_root=str(ext_root))
)
cache.mkdir(parents=True, exist_ok=True)
if create:
cache.mkdir(parents=True, exist_ok=True)
return cache
tests/test_api_lora.py
+7 -4
View File
@@ -36,7 +36,7 @@ from aiohttp.test_utils import AioHTTPTestCase
from database.operations import PromptDatabase
from py.api import PromptManagerAPI
from py import config_store, lora_utils
from py import config_store, lora_import, lora_utils
from py.config import IntegrationConfig, PromptManagerConfig
from py.api import settings_routes as settings_module
from py.api import lora_integration as lora_module
@@ -762,9 +762,9 @@ class TestLoraScanRoute(LoraAPITestCase):
async def test_scan_error_hides_the_server_path(self):
self._enable()
err = PermissionError(13, "Permission denied", str(self.tmp_root / "x.db"))
with patch.object(self.api.db, "delete_prompts_by_category", side_effect=err):
with patch.object(self.api.db, "replace_category_prompts", side_effect=err):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
body = await resp.text()
body = await resp.text() # the stream ends inside the patch
self.assertNotIn(str(self.tmp_root), body)
self.assertIn("x.db", body)
# The admin page's modal shows the failure as its status and closes.
@@ -862,7 +862,10 @@ class TestLoraScanRoute(LoraAPITestCase):
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
with patch.object(self.api.db, "save_prompt", side_effect=RuntimeError("db")):
# Preparing one LoRA's record fails: it is skipped, the import goes on.
with patch.object(
lora_import, "build_record", side_effect=RuntimeError("bad record")
):
resp, events, _ = await self._scan(loras)
self.assertEqual(events[-1]["imported"], 0)
tests/test_lora_import.py
+715
View File
@@ -0,0 +1,715 @@
"""The staged LoRA reference import (#221, #249 System S10).
A failed import must leave the previous ``lora-manager`` collection exactly as
it was (prompts, tags and image links); a preview writes nothing and requests
nothing; an apply does what was reviewed or nothing. Discovery, metadata and
CivitAI downloads are disposable fixtures: no network, no LoraManager install.
"""
import asyncio
import json
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try:
from tests.test_api_lora import (
SECRET_KEY,
LoraAPITestCase,
_FakeResponse,
_make_lora_manager,
_png_bytes,
_write_metadata,
)
except ImportError: # discovered with tests/ as the top-level directory
from test_api_lora import (
SECRET_KEY,
LoraAPITestCase,
_FakeResponse,
_make_lora_manager,
_png_bytes,
_write_metadata,
)
from database.category_import import EmptyReplacement
from database.maintenance import PlanChanged
from database.operations import PairDependencyError, PromptDatabase
from py import lora_import, lora_utils
from py.config import IntegrationConfig, PromptManagerConfig
from utils.hashing import generate_prompt_hash
CATEGORY = "lora-manager"
PREVIEW = "/prompt_manager/lora/import/preview"
APPLY = "/prompt_manager/lora/import/apply"
SCAN = "/prompt_manager/lora/scan"
def _snapshot(db):
"""Every row of the tables an import touches, in a fixed order."""
tables = {
"prompts": "SELECT id, text, category, hash FROM prompts ORDER BY id",
"prompt_tags": "SELECT prompt_id, tag_id FROM prompt_tags ORDER BY 1, 2",
"tags": "SELECT id, name FROM tags ORDER BY id",
"generated_images": (
"SELECT id, prompt_id, image_path, file_path FROM generated_images "
"ORDER BY id"
),
}
with db.model.get_connection() as conn:
return {
name: [tuple(row) for row in conn.execute(sql).fetchall()]
for name, sql in tables.items()
}
def _seed_previous_imports(db):
"""Two earlier imports with tags and image links, and one other prompt."""
ids = []
for name in ("old-one", "old-two"):
prompt_id = db.save_prompt(
text=f"{name} prompt",
category=CATEGORY,
tags=[CATEGORY, f"lora:{name}"],
prompt_hash=generate_prompt_hash(f"{name} prompt"),
)
db.link_image_to_prompt(str(prompt_id), f"/cache/{name}/a.jpg")
db.link_image_to_prompt(str(prompt_id), f"/cache/{name}/b.jpg")
ids.append(prompt_id)
db.save_prompt(
text="mine",
category="portraits",
tags=["keep"],
prompt_hash=generate_prompt_hash("mine"),
)
return ids
def _record(text, images=(), name=None):
return {
"name": name or text,
"text": text,
"hash": generate_prompt_hash(text),
"tags": [CATEGORY, f"lora:{name or text}"],
"images": list(images),
}
class CategoryReplacementTests(unittest.TestCase):
"""``replace_category_prompts``: one transaction, all or nothing."""
def setUp(self):
handle = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
handle.close()
self.path = handle.name
self.db = PromptDatabase(self.path)
self.ids = _seed_previous_imports(self.db)
def tearDown(self):
self.db.close_all()
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
def test_transaction_failure_rolls_back_everything(self):
before = _snapshot(self.db)
records = [_record("new one", ["/c/n1.jpg"]), _record("new two", ["/c/n2.jpg"])]
real = self.db._insert_image_link
calls = []
def failing(conn, prompt_id, image, metadata=None):
calls.append(image)
if len(calls) == 2: # after the old rows went and one record landed
raise sqlite3.OperationalError("disk I/O error")
return real(conn, prompt_id, image, metadata)
with patch.object(self.db, "_insert_image_link", side_effect=failing):
with self.assertRaises(sqlite3.OperationalError):
self.db.replace_category_prompts(CATEGORY, records)
self.assertEqual(len(calls), 2)
self.assertEqual(_snapshot(self.db), before)
def test_pair_dependency_blocks_without_partial_delete(self):
# A saved pair reuses an imported prompt as its positive component.
self.db.save_prompt_pair("old-two prompt", "blurry")
before = _snapshot(self.db)
self.assertEqual(
self.db.category_import_state(CATEGORY), {"ids": self.ids, "paired": 1}
)
with self.assertRaises(PairDependencyError):
self.db.replace_category_prompts(CATEGORY, [_record("new one")])
self.assertEqual(_snapshot(self.db), before)
def test_explicit_empty_replacement_requires_review_flag(self):
before = _snapshot(self.db)
with self.assertRaises(EmptyReplacement):
self.db.replace_category_prompts(CATEGORY, [])
self.assertEqual(_snapshot(self.db), before)
result = self.db.replace_category_prompts(CATEGORY, [], allow_empty=True)
self.assertEqual(result["removed"], 2)
self.assertEqual(self.db.category_import_state(CATEGORY)["ids"], [])
self.assertEqual(len(self.db.search_prompts(category="portraits")), 1)
def test_changed_imports_since_review_are_plan_changed(self):
before = _snapshot(self.db)
with self.assertRaises(PlanChanged):
self.db.replace_category_prompts(
CATEGORY, [_record("new one")], expected_ids=self.ids[:1]
)
self.assertEqual(_snapshot(self.db), before)
def test_replacement_links_repeats_and_existing_prompts(self):
records = [
_record("fresh", ["/c/f.jpg"]),
_record("Fresh ", ["/c/f2.jpg"], name="fresh-again"), # same hash
_record("mine", ["/c/m.jpg"], name="mine-lora"), # exists elsewhere
]
result = self.db.replace_category_prompts(
CATEGORY, records, expected_ids=self.ids
)
self.assertEqual(
(result["removed"], result["imported"], result["linked"]), (2, 1, 1)
)
self.assertEqual(result["duplicates"], 1)
self.assertEqual(result["images"], 3)
imported = self.db.search_prompts(category=CATEGORY)
self.assertEqual([p["text"] for p in imported], ["fresh"])
self.assertEqual(len(self.db.get_prompt_images(imported[0]["id"])), 2)
mine = self.db.search_prompts(category="portraits")[0]
self.assertEqual(len(self.db.get_prompt_images(mine["id"])), 1)
self.assertEqual(
[row["outcome"] for row in result["rows"]],
["imported", "duplicate", "linked"],
)
class LoraImportAPITestCase(LoraAPITestCase):
"""The routes with LoraManager enabled over a disposable LoRA folder."""
async def get_application(self):
app = await super().get_application()
self.lm = _make_lora_manager(self.custom_nodes)
self.loras = self.tmp_root / "loras"
self.loras.mkdir()
self.cache_dir = self.tmp_root / "img_cache"
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.lm)
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
self.requests = []
self.gate = None # a threading.Event a download waits on
self.downloading = threading.Event()
def fake_open(req, timeout=None):
self.requests.append(req)
self.downloading.set()
if self.gate is not None:
self.gate.wait(10)
return _FakeResponse(_png_bytes())
self._patches = [
patch.object(
lora_utils, "find_lora_directories", return_value=[str(self.loras)]
),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=self.cache_dir
),
patch.object(lora_utils, "_open_url", fake_open),
patch.object(PromptManagerConfig, "SAFETY_BACKUP", False),
]
for item in self._patches:
item.start()
self.old_ids = _seed_previous_imports(self.api.db)
return app
async def tearDownAsync(self):
if self.gate is not None:
self.gate.set()
for job in self.api._jobs.list():
if job.task is not None and not job.task.done():
await job.task
for item in reversed(self._patches):
item.stop()
await super().tearDownAsync()
def _two_loras(self):
_write_metadata(
self.loras,
"neon",
words=["glow"],
images=["https://civitai.com/a.png"],
prompt="neon street",
)
(self.loras / "neon.png").write_bytes(_png_bytes())
_write_metadata(self.loras, "plain")
async def _post(self, url, body=None):
resp = await self._post_json(url, {} if body is None else body)
return resp, json.loads(await resp.text())
async def _apply_reviewed(self):
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual(resp.status, 202, body)
job = self.api._jobs.get(body["job"]["id"])
return plan, job
@staticmethod
def _sse(text):
return [
json.loads(line[len("data: ") :])
for line in text.splitlines()
if line.startswith("data: ")
]
class LoraImportRouteTests(LoraImportAPITestCase):
async def test_preview_writes_nothing_and_makes_no_requests(self):
self._two_loras()
(self.loras / "broken.metadata.json").write_text("{{{")
db_path = self.api.db.model.db_path
before = {s: _read(db_path + s) for s in ("", "-wal")}
shm = _size(db_path + "-shm")
rows = _snapshot(self.api.db)
resp, plan = await self._post(PREVIEW)
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},
)
self.assertEqual(plan["blocked_by_pairs"], 0)
self.assertEqual(plan["planned_downloads"], 1)
self.assertEqual(plan["network"], {"planned_downloads": 1, "api_key": True})
self.assertTrue(plan["token"])
self.assertEqual(plan["scope"]["count"], 2)
self.assertFalse(plan["media"]["deletes_files"])
self.assertEqual(self.requests, [])
self.assertFalse(self.cache_dir.exists())
self.assertEqual({s: _read(db_path + s) for s in ("", "-wal")}, before)
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))
async def test_preview_requires_the_integration(self):
IntegrationConfig.LORA_MANAGER_ENABLED = False
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (400, "lora-disabled"))
async def test_blocked_preview_offers_no_token(self):
self._two_loras()
self.api.db.save_prompt_pair("old-one prompt", "blurry")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200)
self.assertEqual(plan["blocked_by_pairs"], 1)
self.assertIsNone(plan["token"])
self.assertEqual(plan["protected"]["count"], 1)
async def test_changed_source_after_preview_is_409(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
rows = _snapshot(self.api.db)
meta = self.loras / "plain.safetensors.metadata.json"
meta.write_text(meta.read_text() + " ")
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "plan-changed"))
self.assertEqual(self.api._jobs.list(), [])
self.assertEqual(self.requests, [])
self.assertEqual(_snapshot(self.api.db), rows)
async def test_apply_imports_through_the_job(self):
self._two_loras()
plan, job = await self._apply_reviewed()
self.assertEqual(
(job.kind, job.title, job.admission),
("lora_import", plan_title(), "records"),
)
await job.task
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(
{k: job.counters[k] for k in lora_import.GROUPS},
{
"imported": 2,
"linked": 0,
"skipped": 0,
"downloaded": 1,
"download_errors": 0,
"errors": 0,
},
)
self.assertFalse(job.cancellable)
self.assertEqual(
self.requests[0].get_header("Authorization"), f"Bearer {SECRET_KEY}"
)
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.search_prompts(category="portraits")), 1)
# The token is single-use.
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "plan-used"))
async def test_discovery_failure_keeps_previous_imports(self):
self._two_loras()
before = _snapshot(self.api.db)
# Legacy scan: discovery raises once the stream is open.
with patch.object(
lora_utils, "find_lora_directories", side_effect=OSError("share gone")
):
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(events[-1]["type"], "error")
self.assertEqual(_snapshot(self.api.db), before)
# Job: staging fails after the review.
with patch.object(lora_import, "stage", side_effect=OSError("share gone")):
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "failed")
self.assertIn("previous LoRA imports were kept", job.status_line)
self.assertEqual(job.scope["failure"], "import-failed")
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_failed_replacement_keeps_previous_imports_and_says_so(self):
self._two_loras()
before = _snapshot(self.api.db)
with patch.object(
self.api.db,
"replace_category_prompts",
side_effect=sqlite3.OperationalError("database is locked"),
):
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "failed")
self.assertEqual(job.counters["errors"], 1)
self.assertEqual(job.counters["downloaded"], 1)
self.assertIn("previous LoRA imports were kept", job.status_line)
self.assertIn("stay in the image cache", job.status_line)
self.assertEqual(_snapshot(self.api.db), before)
async def test_cancel_during_downloads_writes_nothing(self):
self._two_loras()
before = _snapshot(self.api.db)
self.gate = threading.Event()
_plan, job = await self._apply_reviewed()
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, self.downloading.wait, 10))
self.assertTrue(job.summary()["cancellable"])
self.assertTrue(self.api._jobs.cancel(job.id))
self.gate.set()
await job.task
self.assertEqual(job.state, "cancelled")
self.assertEqual(_snapshot(self.api.db), before)
self.assertTrue(any("stay in the image cache" in line for line in job.log))
async def test_empty_reviewed_import_removes_the_collection(self):
# No metadata at all: the review says so, and its apply is the flag.
resp, plan = await self._post(PREVIEW)
self.assertEqual((plan["with_metadata"], plan["current_imports"]), (0, 2))
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.api.db.search_prompts(category=CATEGORY), [])
# The legacy scan has no review: an empty replacement is refused.
_seed_previous_imports(self.api.db)
before = _snapshot(self.api.db)
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(events[-1]["type"], "error")
self.assertEqual(_snapshot(self.api.db), before)
async def test_apply_refused_while_a_records_job_holds(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
async with self.api._jobs.hold("scan", "Metadata scan", admission="records"):
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "busy"))
# The token was not used: it still applies.
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual(resp.status, 202)
await self.api._jobs.get(body["job"]["id"]).task
async def test_safety_backup_policy_runs_before_the_replace(self):
self._two_loras()
order = []
real = self.api.db.replace_category_prompts
async def step(job, db, backup_on):
order.append(("backup", backup_on))
return True, "prompts.db.backup_x"
def replace(*args, **kwargs):
order.append(("replace", None))
return real(*args, **kwargs)
with (
patch.object(PromptManagerConfig, "SAFETY_BACKUP", True),
patch.object(self.api, "_safety_backup_step", side_effect=step),
patch.object(self.api.db, "replace_category_prompts", side_effect=replace),
):
plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(plan["backup"]["policy"], "on")
self.assertEqual(order, [("backup", True), ("replace", None)])
async def test_legacy_scan_events_unchanged_on_success(self):
self._two_loras()
(self.loras / "broken.metadata.json").write_text("{{{")
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(resp.status, 200)
progress_keys = {
"type",
"status",
"progress",
"total",
"processed",
"imported",
"skipped",
}
for event in events[:-1]:
self.assertEqual(event["type"], "progress")
self.assertLessEqual(set(event), progress_keys)
self.assertIsInstance(event["status"], str)
self.assertIsInstance(event["progress"], int)
per_record = [e for e in events if "processed" in e]
self.assertEqual([e["processed"] for e in per_record], [1, 2, 3])
self.assertEqual(
events[-1],
{
"type": "complete",
"progress": 100,
"total": 3,
"imported": 2,
"skipped": 1,
},
)
imported = self.api.db.search_prompts(category=CATEGORY)
self.assertEqual(sorted(p["text"] for p in imported), ["neon street", "plain"])
class LoraImportRefusalTests(LoraImportAPITestCase):
"""Every refusal leaves the previous imports as they were."""
async def test_preview_refusals(self):
IntegrationConfig.LORA_MANAGER_PATH = ""
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (400, "lora-path"))
IntegrationConfig.LORA_MANAGER_PATH = str(_make_lora_manager(self.tmp_root))
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (400, "lora-path"))
self.assertNotIn(str(self.tmp_root), json.dumps(body))
IntegrationConfig.LORA_MANAGER_PATH = str(self.lm)
resp, body = await self._post(PREVIEW, {"extra": 1})
self.assertEqual((resp.status, body["code"]), (400, "invalid-request"))
with patch.object(lora_import, "preview", side_effect=OSError("gone")):
resp, body = await self._post(PREVIEW)
self.assertEqual(resp.status, 500)
self.assertFalse(body["success"])
generations = iter([(0, "a"), (1, "a")])
with patch.object(
self.api, "_database_identity", side_effect=lambda db: next(generations)
):
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (409, "storage-changed"))
async def test_apply_refusals(self):
self._two_loras()
before = _snapshot(self.api.db)
resp, body = await self._post(APPLY, {"token": ""})
self.assertEqual((resp.status, body["code"]), (400, "invalid-request"))
resp, body = await self._post(APPLY, {"token": "unknown"})
self.assertEqual((resp.status, body["code"]), (410, "plan-expired"))
other = self.api._maintenance_plans.issue(
"optimize", "f", {"backup": "off"}, self.api._database_identity(self.api.db)
)
resp, body = await self._post(APPLY, {"token": other})
self.assertEqual((resp.status, body["code"]), (410, "plan-expired"))
resp, plan = await self._post(PREVIEW)
with patch.object(lora_import, "discover", side_effect=OSError("gone")):
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (500, "apply-failed"))
short = ("database folder", "the safety backup", 10, 1)
with patch("py.api.lora_import_routes.disk_plan", return_value=({}, short, 0)):
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (507, "insufficient-space"))
self.api._db_generation += 1
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "storage-changed"))
self.assertEqual(self.api._jobs.list(), [])
self.assertEqual(_snapshot(self.api.db), before)
async def _failed_job(self, **patches):
before = _snapshot(self.api.db)
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "failed", job.log)
self.assertEqual(_snapshot(self.api.db), before)
return job
async def test_imports_changed_after_the_review_fail_the_job(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
self.api.db.save_prompt(
text="late", category=CATEGORY, prompt_hash=generate_prompt_hash("late")
)
before = _snapshot(self.api.db)
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual((job.state, job.scope["failure"]), ("failed", "plan-changed"))
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_pair_saved_after_the_review_blocks_the_job(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
self.api.db.save_prompt_pair("old-one prompt", "blurry")
before = _snapshot(self.api.db)
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(
(job.state, job.scope["failure"]), ("failed", "blocked-by-pairs")
)
self.assertIn("Saved pairs", job.status_line)
self.assertEqual(_snapshot(self.api.db), before)
async def test_nothing_prepared_keeps_the_collection(self):
self._two_loras()
with patch.object(lora_import, "build_record", side_effect=ValueError("x")):
job = await self._failed_job()
self.assertEqual(job.scope["failure"], "nothing-to-import")
self.assertEqual(job.counters["skipped"], 0)
async def test_files_changed_while_staging_fail_the_job(self):
self._two_loras()
real = lora_import.stage
def changed(*args, **kwargs):
staged = real(*args, **kwargs)
staged.fingerprint = "sha256:other"
return staged
with patch.object(lora_import, "stage", side_effect=changed):
job = await self._failed_job()
self.assertEqual(job.scope["failure"], "plan-changed")
async def test_a_failed_safety_backup_changes_nothing(self):
self._two_loras()
from py.api import maintenance
with (
patch.object(PromptManagerConfig, "SAFETY_BACKUP", True),
patch.object(
maintenance,
"take_safety_backup",
side_effect=maintenance.SafetyBackupFailed("no room"),
),
):
job = await self._failed_job()
self.assertEqual(job.scope["failure"], "backup-failed")
async def test_cancel_right_after_the_start_writes_nothing(self):
self._two_loras()
before = _snapshot(self.api.db)
resp, plan = await self._post(PREVIEW)
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
self.api._jobs.cancel(job.id) # at once, while it stages
await job.task
self.assertEqual(job.state, "cancelled")
self.assertEqual(_snapshot(self.api.db), before)
def plan_title():
return lora_import.JOB_TITLE
def _read(path):
try:
with open(path, "rb") as handle:
return handle.read()
except FileNotFoundError:
return None
def _size(path):
return os.path.getsize(path) if os.path.exists(path) else None
class LoraImportStageTests(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(lambda: __import__("shutil").rmtree(self.tmp, True))
def test_discover_sorts_and_drops_repeats(self):
a = self.tmp / "a"
b = self.tmp / "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)
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_text(meta.read_text() + " ")
self.assertNotEqual(lora_import.fingerprint([meta]), first)
meta.unlink()
self.assertNotEqual(lora_import.fingerprint([meta]), first)
def test_stage_cancel_before_a_file_raises(self):
meta = _write_metadata(self.tmp, "one")
with self.assertRaises(lora_import.ImportCancelled):
lora_import.stage([meta], self.tmp, cancelled=lambda: True)
if __name__ == "__main__":
unittest.main()