From 6a2dae0f7890f79c6e05e40bcbf86ae70f3f2e1f Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sun, 25 Oct 2026 06:52:34 -0700 Subject: [PATCH] fix(lora): stage imports before replacing the collection (#221) 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. --- database/category_import.py | 170 +++++++++ database/operations.py | 170 +++++---- py/api/__init__.py | 4 + py/api/lora_import_routes.py | 416 ++++++++++++++++++++ py/api/lora_integration.py | 194 ++++------ py/lora_import.py | 172 +++++++++ py/lora_utils.py | 74 ++-- tests/test_api_lora.py | 11 +- tests/test_lora_import.py | 715 +++++++++++++++++++++++++++++++++++ 9 files changed, 1699 insertions(+), 227 deletions(-) create mode 100644 database/category_import.py create mode 100644 py/api/lora_import_routes.py create mode 100644 py/lora_import.py create mode 100644 tests/test_lora_import.py diff --git a/database/category_import.py b/database/category_import.py new file mode 100644 index 0000000..68ebd58 --- /dev/null +++ b/database/category_import.py @@ -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}) diff --git a/database/operations.py b/database/operations.py index 1c2d89f..0f45097 100644 --- a/database/operations.py +++ b/database/operations.py @@ -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]]: """ diff --git a/py/api/__init__.py b/py/api/__init__.py index 6b2748f..97652ad 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -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 diff --git a/py/api/lora_import_routes.py b/py/api/lora_import_routes.py new file mode 100644 index 0000000..fb505b8 --- /dev/null +++ b/py/api/lora_import_routes.py @@ -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_", + "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) diff --git a/py/api/lora_integration.py b/py/api/lora_integration.py index 7e66c2d..bb6a146 100644 --- a/py/api/lora_integration.py +++ b/py/api/lora_integration.py @@ -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"], } ) diff --git a/py/lora_import.py b/py/lora_import.py new file mode 100644 index 0000000..c1fdca4 --- /dev/null +++ b/py/lora_import.py @@ -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 diff --git a/py/lora_utils.py b/py/lora_utils.py index da60605..66db922 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -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/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 diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index 810671b..3d0d1e4 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -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) diff --git a/tests/test_lora_import.py b/tests/test_lora_import.py new file mode 100644 index 0000000..a589d26 --- /dev/null +++ b/tests/test_lora_import.py @@ -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()