ComfyAssets-ComfyUI_PromptManager/py/api/lora_integration.py

"""LoraManager integration API routes for PromptManager."""

import asyncio
import contextlib
import json
import threading

from aiohttp import web

from ..config_store import ConfigUnreadable, RevisionConflict, get_config_store
from .jobs import LORA_IMPORT_JOB_KIND

# Flags of POST /lora/enable; each must be a JSON boolean when present.
_LORA_FLAGS = ("enabled", "trigger_words_enabled", "clear_civitai_api_key")

_LORA_PATH_ERROR = (
    "LoraManager path must be a directory under ComfyUI's custom_nodes "
    "whose name contains 'lora'"
)


def _lora_settings_mutation(values):
    """The config commit for a LoRA save: only ``integrations.lora_manager``.

    The submitted keys replace their values; other LoRA keys, other
    integrations and every other section stay as they are on disk.
    """

    def mutate(config):
        integrations = config.get("integrations")
        integrations = dict(integrations) if isinstance(integrations, dict) else {}
        lora = integrations.get("lora_manager")
        lora = dict(lora) if isinstance(lora, dict) else {}
        lora.update(values)
        integrations["lora_manager"] = lora
        config["integrations"] = integrations
        return config

    return mutate


# The admin page's own words for the refusals it can meet (it has no review).
_FILES_CHANGED = (
    "The LoRA metadata files changed during the scan; the previous imports "
    "were kept. Run the scan again."
)
_IMPORTS_CHANGED = (
    "The imported LoRA prompts changed during the scan; the previous imports "
    "were kept. Run the scan again."
)


class LoraFilesChanged(Exception):
    """A metadata file changed, appeared or went while the legacy scan staged."""


def _lora_error_frame(message):
    """An import stream's error event, as the admin page's modal reads it.

    The modal shows ``status`` and ``progress`` from every frame and closes
    on ``type: error``, naming ``message``.
    """
    return {"type": "error", "status": message, "progress": 0, "message": message}


def _closed(request):
    """A check (safe off the loop) for whether ``request``'s client left."""

    def closed():
        transport = request.transport
        return transport is None or transport.is_closing()

    return closed


def _stopped(error) -> bool:
    """Whether a legacy scan ended because its page closed the connection."""
    from .. import lora_import

    return isinstance(error, (ConnectionError, lora_import.ImportCancelled))


class LoraIntegrationMixin:
    """Mixin providing LoraManager detection, scanning, and trigger word endpoints."""

    def _register_lora_routes(self, routes):
        @routes.get("/prompt_manager/lora/detect")
        async def lora_detect_route(request):
            return await self.lora_detect(request)

        @routes.get("/prompt_manager/lora/status")
        async def lora_status_route(request):
            return await self.lora_status(request)

        @routes.post("/prompt_manager/lora/enable")
        async def lora_enable_route(request):
            return await self.lora_enable(request)

        @routes.post("/prompt_manager/lora/scan")
        async def lora_scan_route(request):
            return await self.lora_scan(request)

        @routes.get("/prompt_manager/lora/trigger-words")
        async def lora_trigger_words_route(request):
            return await self.lora_trigger_words(request)

        @routes.post("/prompt_manager/lora/refresh-cache")
        async def lora_refresh_cache_route(request):
            return await self.lora_refresh_cache(request)

    def _error_response(self, exc: Exception, status: int = 500):
        return web.json_response(
            {"success": False, "error": self._public_error(exc)}, status=status
        )

    @staticmethod
    def _bad_lora_request(message: str):
        return web.json_response({"success": False, "error": message}, status=400)

    # ── Detection ────────────────────────────────────────────────────

    async def lora_detect(self, request):
        """Auto-detect LoraManager installation."""
        try:
            from ..lora_utils import detect_lora_manager

            path = await self._run_in_executor(detect_lora_manager)
            return web.json_response(
                {
                    "success": True,
                    "detected": path is not None,
                    "path": self._public_path(path) if path else "",
                }
            )
        except Exception as e:
            self.logger.error(f"LoraManager detection failed: {e}")
            return self._error_response(e)

    # ── Status ───────────────────────────────────────────────────────

    async def lora_status(self, request):
        """Get current LoraManager integration status."""
        try:
            from ..config import IntegrationConfig
            from ..lora_utils import (
                detect_lora_manager,
                get_trigger_cache,
                resolve_lora_manager_path,
            )

            config = IntegrationConfig.get_config()["lora_manager"]
            cache = get_trigger_cache()

            # Check if the configured path is still valid
            detected_path = await self._run_in_executor(
                detect_lora_manager, config.get("path", "")
            )
            # A detected install reached through a link outside custom_nodes is
            # found but cannot be enabled (S10 review N5): say which.
            usable = bool(detected_path) and (
                await self._run_in_executor(resolve_lora_manager_path, detected_path)
                is not None
            )

            return web.json_response(
                {
                    "success": True,
                    "enabled": config["enabled"],
                    "path": self._public_path(config["path"]),
                    "trigger_words_enabled": config["trigger_words_enabled"],
                    "has_civitai_api_key": bool(config.get("civitai_api_key")),
                    "detected": detected_path is not None,
                    "detected_usable": usable,
                    "detected_path": (
                        self._public_path(detected_path) if detected_path else ""
                    ),
                    "trigger_cache_loaded": cache.is_loaded,
                }
            )
        except Exception as e:
            self.logger.error(f"LoraManager status check failed: {e}")
            return self._error_response(e)

    # ── Enable / Disable ─────────────────────────────────────────────

    async def _validated_lora_manager_path(self, path):
        """Canonical LoraManager path for ``path`` or ``(None, error_response)``.

        The user's path (or, when empty, the auto-detected install) must be a
        LoraManager-looking directory that is a direct child of ComfyUI's
        ``custom_nodes`` with ``lora`` in its name; the trees under it are
        later served as images, so nothing outside custom_nodes may be used.
        """
        from ..lora_utils import detect_lora_manager, resolve_lora_manager_path

        candidate = ""
        if path:
            candidate = await self._run_in_executor(resolve_lora_manager_path, path)
            if candidate is None:
                return None, self._bad_lora_request(_LORA_PATH_ERROR)

        resolved = await self._run_in_executor(detect_lora_manager, candidate)
        if not resolved:
            return None, self._bad_lora_request(
                "LoraManager not found at the specified path"
            )

        allowed = await self._run_in_executor(resolve_lora_manager_path, resolved)
        if allowed is None:
            return None, self._bad_lora_request(_LORA_PATH_ERROR)
        return allowed, None

    async def lora_enable(self, request):
        """Enable or disable LoraManager integration and save it to config.json.

        ``expected_revision`` (optional, GET /settings' ``revision``) refuses
        a save over a config.json changed meanwhile (``settings-changed``).

        The values are committed through the config store first (#220) and
        published at runtime only once config.json holds them; a save that
        cannot be written reports ``persisted: false`` and changes nothing.
        The CivitAI key is never echoed back.
        """
        try:
            data, error = await self._read_json_body(request)
            if error is not None:
                return error
            for flag in _LORA_FLAGS:
                if flag in data and not isinstance(data[flag], bool):
                    return self._bad_lora_request(f"{flag} must be true or false")
            # The System Save row sends the revision it read (#220); the
            # legacy page sends none and merges into what is on disk.
            expected_revision = data.get("expected_revision")
            if expected_revision is not None and not isinstance(expected_revision, str):
                return self._bad_lora_request("expected_revision must be a string")
            enabled = data.get("enabled", False)
            path = data.get("path", "")
            trigger_words = data.get("trigger_words_enabled", False)

            from ..config import IntegrationConfig
            from ..lora_utils import get_trigger_cache

            # The key is never echoed back to the UI, so a blank or absent
            # field means "unchanged"; clearing is an explicit request.
            civitai_key = IntegrationConfig.CIVITAI_API_KEY
            submitted_key = data.get("civitai_api_key")
            if submitted_key is not None and not isinstance(submitted_key, str):
                return self._bad_lora_request("civitai_api_key must be a string")
            if data.get("clear_civitai_api_key"):
                civitai_key = ""
            elif submitted_key and submitted_key.strip():
                civitai_key = submitted_key.strip()

            if not isinstance(path, str):
                return self._bad_lora_request("path must be a string")

            # If enabling, validate the path
            if enabled:
                path, error = await self._validated_lora_manager_path(path)
                if error is not None:
                    return error

            values = {
                "enabled": enabled,
                "path": path,
                "trigger_words_enabled": trigger_words,
                "civitai_api_key": civitai_key,
            }

            def publish(_config):
                # Under the commit lock, so the last save written is the one
                # the runtime keeps.
                IntegrationConfig.LORA_MANAGER_ENABLED = enabled
                IntegrationConfig.LORA_MANAGER_PATH = path
                IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = trigger_words
                IntegrationConfig.CIVITAI_API_KEY = civitai_key

            try:
                _saved, revision = await self._run_in_executor(
                    get_config_store().commit,
                    _lora_settings_mutation(values),
                    expected_revision=expected_revision,
                    publish=publish,
                )
            except (ConfigUnreadable, RevisionConflict, OSError) as exc:
                return self._config_commit_failure(exc)

            cache = get_trigger_cache()
            if enabled and trigger_words and path:
                try:
                    count = await self._run_in_executor(cache.load, path)
                    self.logger.info(f"Trigger word cache loaded: {count} LoRAs")
                except Exception as e:
                    # Saved and enabled; only the trigger words are missing.
                    self.logger.warning(f"Trigger word cache could not be loaded: {e}")
                    cache.clear()
            elif not enabled:
                cache.clear()

            return web.json_response(
                {
                    "success": True,
                    "enabled": enabled,
                    "path": self._public_path(path),
                    "trigger_words_enabled": trigger_words,
                    "persisted": True,
                    "revision": revision,
                    "runtime_applied": True,
                    "cache_loaded": cache.is_loaded,
                }
            )
        except Exception as e:
            self.logger.error(f"LoraManager enable/disable failed: {e}")
            return self._error_response(e)

    # ── Scan LoRA example images ─────────────────────────────────────

    async def lora_scan(self, request):
        """Scan LoraManager metadata and import LoRA info + preview images.

        Streams progress as SSE, matching the existing scan pattern. Errors
        raised once the stream is open are reported as an ``error`` event.
        """
        try:
            from ..config import IntegrationConfig
            from ..lora_utils import resolve_lora_manager_path

            if not IntegrationConfig.LORA_MANAGER_ENABLED:
                return self._bad_lora_request("LoraManager integration is not enabled")

            lm_path = IntegrationConfig.LORA_MANAGER_PATH
            if not lm_path:
                return self._bad_lora_request("LoraManager path not configured")

            allowed = await self._run_in_executor(resolve_lora_manager_path, lm_path)
            if allowed is None:
                return self._bad_lora_request(_LORA_PATH_ERROR)

            response = web.StreamResponse(
                status=200,
                reason="OK",
                headers={"Content-Type": "text/event-stream"},
            )

            async def send_progress(data):
                line = f"data: {json.dumps(data)}\n\n"
                await response.write(line.encode("utf-8"))

            # The import replaces prompts and images: a records hold, so it
            # never runs beside a scan, restore or tagging job.
            busy = self._jobs.blocker(LORA_IMPORT_JOB_KIND, admission="records")
            if busy is not None:
                await response.prepare(request)
                message = f"{busy.title} already running"
                await send_progress(_lora_error_frame(message))
                await response.write_eof()
                return response

            async with self._jobs.hold(
                LORA_IMPORT_JOB_KIND, "LoRA import (admin page)", admission="records"
            ) as job:
                await response.prepare(request)
                try:
                    await self._lora_scan_stream(
                        allowed, send_progress, disconnected=_closed(request)
                    )
                except asyncio.CancelledError:
                    # The server cancelled the handler when the page left.
                    job.finish("cancelled")
                    raise
                except Exception as e:
                    # A page that left stopped the scan; nothing failed (R2).
                    job.finish("cancelled" if _stopped(e) else "failed")
                    await self._send_scan_failure(e, send_progress)
                with contextlib.suppress(ConnectionError, RuntimeError):
                    await response.write_eof()
            return response

        except Exception as e:
            self.logger.error(f"LoRA scan failed: {e}")
            return self._error_response(e)

    async def _send_scan_failure(self, error, send_progress):
        """The error frame for a failed legacy scan; quiet when the page left.

        Expected refusals (files or imports changed, saved pairs, nothing to
        import) are warnings in the admin page's own words; anything else is
        logged with its traceback.
        """
        from .lora_import_routes import (
            EmptyReplacement,
            PairDependencyError,
            PlanChanged,
        )

        if _stopped(error):
            self.logger.info("LoRA scan stopped: the page closed its connection")
            return
        if isinstance(error, PlanChanged):
            message = _IMPORTS_CHANGED
        elif isinstance(error, LoraFilesChanged):
            message = str(error)
        else:
            message = self._public_error(error)
        if isinstance(
            error,
            (PlanChanged, LoraFilesChanged, EmptyReplacement, PairDependencyError),
        ):
            self.logger.warning(f"LoRA scan refused: {message}")
        else:
            self.logger.exception("LoRA scan failed")
        with contextlib.suppress(ConnectionError, RuntimeError):
            await send_progress(_lora_error_frame(message))

    async def _lora_scan_stream(
        self, lm_path, send_progress, disconnected=lambda: False
    ):
        """Import every LoRA under ``lm_path``, reporting through ``send_progress``.

        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(
            {
                "type": "progress",
                "status": "Finding LoRA directories...",
                "progress": 2,
            }
        )
        files = await self._run_in_executor(lora_import.discover, lm_path)
        total = len(files)
        await send_progress(
            {
                "type": "progress",
                "status": f"Found {total} LoRA metadata files",
                "progress": 5,
                "total": total,
            }
        )

        reviewed = await self._run_in_executor(lora_import.fingerprint, files)
        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()

        def on_record(index, record, _staged):
            events.put_nowait((index, record))

        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.existing_prompt_hashes,
                    [record["hash"]],
                    lora_import.LORA_CATEGORY,
                )
                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": 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. I2: when the page goes away
        # (its connection closes, or a report cannot be sent) staging stops
        # before its next LoRA, and this returns only once it has stopped, so
        # the caller's records hold covers every download.
        stop = threading.Event()
        staging = asyncio.ensure_future(
            self.stage_lora_import(
                files,
                on_record=on_record,
                cancelled=lambda: stop.is_set() or disconnected(),
            )
        )
        try:
            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())
        finally:
            stop.set()
            # Staging must have stopped before the records hold is released,
            # even when the handler is cancelled again while it waits (R3).
            interrupted = False
            while not staging.done():
                try:
                    await asyncio.shield(staging)
                except asyncio.CancelledError:
                    interrupted = True
                except Exception:
                    pass
            if interrupted:
                raise asyncio.CancelledError
        staged = await staging
        if not await self.lora_files_unchanged(lm_path, reviewed):
            raise LoraFilesChanged(_FILES_CHANGED)

        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": result["imported"],
                "skipped": total - result["imported"],
            }
        )

    # ── Trigger word endpoints ───────────────────────────────────────

    async def lora_trigger_words(self, request):
        """Look up trigger words for a specific LoRA name."""
        try:
            from ..config import IntegrationConfig
            from ..lora_utils import get_trigger_cache

            if not IntegrationConfig.LORA_MANAGER_ENABLED:
                return self._bad_lora_request("LoraManager integration not enabled")

            lora_name = request.query.get("name", "")
            if not lora_name:
                return self._bad_lora_request("Missing 'name' query parameter")

            cache = get_trigger_cache()
            if not cache.is_loaded:
                lm_path = IntegrationConfig.LORA_MANAGER_PATH
                if lm_path:
                    await self._run_in_executor(cache.load, lm_path)

            words = cache.get_trigger_words(lora_name)
            return web.json_response(
                {"success": True, "lora": lora_name, "trigger_words": words}
            )
        except Exception as e:
            self.logger.error(f"Trigger word lookup failed: {e}")
            return self._error_response(e)

    async def lora_refresh_cache(self, request):
        """Force-refresh the trigger word cache from disk."""
        try:
            from ..config import IntegrationConfig
            from ..lora_utils import get_trigger_cache

            if not IntegrationConfig.LORA_MANAGER_ENABLED:
                return self._bad_lora_request("LoraManager integration not enabled")

            lm_path = IntegrationConfig.LORA_MANAGER_PATH
            if not lm_path:
                return self._bad_lora_request("LoraManager path not configured")

            cache = get_trigger_cache()
            count = await self._run_in_executor(cache.load, lm_path)
            return web.json_response(
                {"success": True, "loras_with_trigger_words": count}
            )
        except Exception as e:
            self.logger.error(f"Trigger cache refresh failed: {e}")
            return self._error_response(e)