Vito Sansevero
818f2e22e1
POST /prompt_manager/lora/enable accepts the optional expected_revision that GET /settings returns, so the System Save row refuses to write over a config.json changed elsewhere (settings-changed, nothing written). The legacy page sends none and still merges into what is on disk.
485 lines
20 KiB
Python
485 lines
20 KiB
Python
"""LoraManager integration API routes for PromptManager."""
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
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
|
|
|
|
|
|
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}
|
|
|
|
|
|
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
|
|
|
|
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", "")
|
|
)
|
|
|
|
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_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)
|
|
except Exception as e:
|
|
self.logger.exception("LoRA scan failed")
|
|
job.finish("failed")
|
|
await send_progress(_lora_error_frame(self._public_error(e)))
|
|
|
|
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 _lora_scan_stream(self, lm_path, send_progress):
|
|
"""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,
|
|
}
|
|
)
|
|
|
|
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.
|
|
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": 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)
|