ComfyAssets-ComfyUI_PromptManager/py/api/lora_import_routes.py

"""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, skipped, with_local_media, current_imports,
blocked_by_pairs, planned_downloads, network:{planned_downloads, api_key},
layout}`` (#355: ``layout`` says which LoraManager settings file, storage mode
and library the import read, and ``error`` when they cannot be used). A plan
that saved pairs block, or whose layout is unusable, carries no token
(nothing can be applied).

``POST /prompt_manager/lora/import/apply {token}`` starts the job (202
``{job}``): kind ``lora_import``, titled "LoRA reference import", admission
``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 .inspection_service import SERVABLE_EXTENSIONS
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"
# The review reads the library: a request lease, so it is refused while a
# restore runs and a restore waits for it before switching the file (S8).
LORA_PREVIEW_LEASE = "LoRA import review"
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."
)
INTEGRATION_CHANGED = (
    "The LoraManager integration changed since the review; nothing was changed. "
    "Review the import again."
)
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)


def _all_downloads_failed(errors):
    """The warning when every planned download failed (#337)."""
    if errors == 1:
        return "Warning: the example image could not be downloaded"
    return f"Warning: none of the {errors:,} example images could be downloaded"


def _with_downloads(line, downloaded, errors):
    """``line`` with the failed downloads: counted when some failed, led by
    a warning when every planned download failed (#337). The job still
    finishes: the import itself succeeded and the downloads are best effort."""
    if not errors:
        return line
    if not downloaded:
        return f"{_all_downloads_failed(errors)} · {line}"
    failed = f"{_count(errors, 'example image')} could not be downloaded"
    return f"{line} · {failed}"


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, discovery, *, on_record=None, cancelled=None):
        """Stage ``discovery``'s sidecars off the loop: records, skips and
        download counts."""
        from ..config import IntegrationConfig

        cache_dir = await self._run_in_executor(lora_utils.get_lora_image_cache_dir)
        await self._run_in_executor(lora_utils.sweep_stale_downloads, 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(
                discovery,
                cache_dir,
                IntegrationConfig.CIVITAI_API_KEY,
                media_extensions=SERVABLE_EXTENSIONS,
                on_record=report,
                cancelled=cancelled or (lambda: False),
                logger=self.logger,
            )
        )

    async def lora_files_unchanged(self, lm_path, expected):
        """Whether LoraManager's sidecars for ``lm_path`` still fingerprint as
        ``expected`` (found again, so an added file counts too)."""

        def current():
            return lora_import.fingerprint(lora_import.discover(lm_path))

        try:
            return await self._run_in_executor(current) == expected
        except OSError:
            return False

    async def replace_lora_imports(self, db, staged, *, expected_ids, allow_empty):
        """Swap the collection for ``staged.records`` in one transaction.

        Then the image routes' layout cache is dropped (#362): the import may
        have linked files under a folder LoraManager created after the cache
        stored its layout (its first example-image download creates one).
        """
        result = await self._run_in_executor(
            lambda: db.replace_category_prompts(
                CATEGORY,
                staged.records,
                expected_ids=expected_ids,
                allow_empty=allow_empty,
            )
        )
        lora_utils.clear_layout_cache()
        return result

    # ── preview ───────────────────────────────────────────────────────

    async def preview_lora_import(self, request):
        try:
            async with self._jobs.lease(LORA_PREVIEW_LEASE):
                return await self._preview_lora_import(request)
        except JobBusy as busy:
            return busy_response(busy)

    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, SERVABLE_EXTENSIONS
            )
            state = await self._run_in_executor(db.category_import_state, CATEGORY)
        except OSError as error:
            # A folder discovery cannot list refuses the import (it must not
            # replace the collection with part of it): say which one.
            self.logger.warning(
                f"LoRA import review: a folder could not be listed: {error}"
            )
            return _failure(
                f"A LoRA folder could not be listed ({self._public_error(error)}). "
                "Nothing was changed.",
                500,
                "folder-unreadable",
            )
        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
        layout = found["layout"]
        # An unusable LoraManager layout found nothing it may replace with.
        if not blocked and layout["error"] is None:
            payload = {
                "path": lm_path,
                "fingerprint": found["fingerprint"],
                "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"],
                "skipped": found["skipped"],
                "with_local_media": found["with_local_media"],
                "layout": layout,
                "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"]
        # M5: the integration must still be on and at the reviewed install.
        current, _refused = await self._lora_source()
        if current != lm_path:
            store.discard(token)
            return _failure(INTEGRATION_CHANGED, 409, "plan-changed")
        try:
            discovery = await self._run_in_executor(lora_import.discover, lm_path)
            now = await self._run_in_executor(lora_import.fingerprint, discovery)
            backup = self._backup_promised(plan)
            _disk, short, _size = await self._run_in_executor(
                disk_plan, db.model.db_path, backup=backup, vacuum=False
            )
        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(discovery.sidecars),
                },
                groups=lora_import.GROUPS,
                runner=lambda job: self._run_lora_import(
                    job, db, plan, discovery, 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, discovery, backup_on):
        total = len(discovery.sidecars)
        job.update(total=total, status_line="Preparing LoRA records")
        job.log_line(f"Found {_count(total, 'LoRA metadata file')}")

        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(
                discovery, 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)
        # M4: what staging skipped stays on the job whatever happens next.
        job.record_counted(
            "skipped",
            len(staged.skipped),
            [{"path": "", "detail": name} for name in staged.skipped],
        )
        # Past this point nothing stops halfway: the replace is one transaction.
        # Published, so a view that offered Cancel during the downloads drops it.
        job.cancellable = False
        job.publish({"type": "progress", "progress": job.progress})
        if job.cancelled:
            job.log_line(f"{CANCELLED} {CACHE_NOTE}")
            job.update(status_line=CANCELLED)
            job.finish("cancelled")
            return
        ok, _backup = await self._safety_backup_step(job, db, backup_on)
        if not ok:
            return
        # I1: the files are listed and fingerprinted again right before the
        # replace, so a file changed, removed or added while the downloads ran
        # (LoraManager refreshing its CivitAI info) aborts with nothing deleted.
        if not await self.lora_files_unchanged(plan.payload["path"], plan.fingerprint):
            return _fail_job(job, "plan-changed", f"{SOURCE_CHANGED} {CACHE_NOTE}")
        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)
        rows = []
        if errors and not downloaded:
            warning = _all_downloads_failed(errors)
            rows = [{"path": "", "detail": f"{warning}; the log has the details"}]
            job.log_line(warning)
        job.record_counted("download_errors", errors, rows)
        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"]
        )
        line = _import_line(result, len(staged.skipped))
        if result["removed"]:
            job.log_line(f"Replaced {_count(result['removed'], 'previous import')}")
        job.log_line(line)
        line = _with_downloads(
            line,
            staged.stats.get("downloaded", 0),
            staged.stats.get("download_errors", 0),
        )
        job.update(progress=100, processed=job.total, status_line=line)