Vito
863d717435
image.civitai.com now answers every image request with a redirect to blobs-b2.civitai.com, and the downloader refused all redirects, so every LoRA import download failed. Follow a redirect only over HTTPS to civitai.com or a subdomain of it (parsed host name, default port, no credentials in the URL), at most 3 hops, and never send the API key on a redirected request. Every other redirect is still refused; the size cap, timeout and image checks apply to the final response. Logged URLs drop their query, so signed storage URLs stay out of the log. The System import job's status line now counts failed downloads, and when every planned download fails it leads with a warning, which the Download errors outcome lists too. The legacy admin scan's events are unchanged.
484 lines
20 KiB
Python
484 lines
20 KiB
Python
"""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"
|
|
# 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, 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)
|
|
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(
|
|
files,
|
|
cache_dir,
|
|
IntegrationConfig.CIVITAI_API_KEY,
|
|
on_record=report,
|
|
cancelled=cancelled or (lambda: False),
|
|
logger=self.logger,
|
|
)
|
|
)
|
|
|
|
async def lora_files_unchanged(self, lm_path, expected):
|
|
"""Whether the metadata files under ``lm_path`` still fingerprint as
|
|
``expected`` (listed 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."""
|
|
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):
|
|
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)
|
|
state = await self._run_in_executor(db.category_import_state, CATEGORY)
|
|
except Exception:
|
|
self.logger.exception("LoRA import preview failed")
|
|
return _failure("The LoRA import could not be reviewed.", 500)
|
|
if identity != self._database_identity(self.db):
|
|
return _failure(
|
|
"The database was replaced while it was reviewed. Review again.",
|
|
409,
|
|
"storage-changed",
|
|
)
|
|
return await self._lora_plan_response(db, identity, lm_path, found, state)
|
|
|
|
async def _lora_plan_response(self, db, identity, lm_path, found, state):
|
|
from ..config import IntegrationConfig
|
|
|
|
current = len(state["ids"])
|
|
blocked = state["paired"]
|
|
policy = self._safety_backup_policy()
|
|
try:
|
|
disk, _short, size = await self._run_in_executor(
|
|
disk_plan, db.model.db_path, backup=policy != "off", vacuum=False
|
|
)
|
|
except (OSError, sqlite3.Error):
|
|
self.logger.exception("Could not measure the disk space for a plan")
|
|
return _failure("The LoRA import could not be reviewed.", 500)
|
|
token = None
|
|
expires_at = None
|
|
if not blocked:
|
|
payload = {
|
|
"path": lm_path,
|
|
"fingerprint": found["fingerprint"],
|
|
"expected_ids": state["ids"],
|
|
"empty": found["with_metadata"] == 0,
|
|
"backup": policy,
|
|
}
|
|
try:
|
|
token = self._maintenance_plans.issue(
|
|
LORA_OPERATION, found["fingerprint"], payload, identity
|
|
)
|
|
except ScopeTooLarge:
|
|
return _failure(
|
|
"This import is too large to review at once.", 413, "too-large"
|
|
)
|
|
expires_at = _iso(self._maintenance_plans.peek(token).expires_at)
|
|
name = os.path.basename(db.model.db_path)
|
|
downloads = found["planned_downloads"]
|
|
reasons = (
|
|
[f"{_count(blocked, 'imported prompt')} used by saved pairs"]
|
|
if blocked
|
|
else []
|
|
)
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"operation": LORA_OPERATION,
|
|
"token": token,
|
|
"expires_at": expires_at,
|
|
"lora_files": found["lora_files"],
|
|
"with_metadata": found["with_metadata"],
|
|
"current_imports": current,
|
|
"blocked_by_pairs": blocked,
|
|
"planned_downloads": downloads,
|
|
"scope": {
|
|
"count": found["with_metadata"],
|
|
"label": (
|
|
f"Replace {_count(current, 'imported prompt')} with "
|
|
f"{_count(found['with_metadata'], 'LoRA')}"
|
|
),
|
|
},
|
|
"protected": {"count": blocked, "reasons": reasons},
|
|
"backup": {
|
|
"policy": policy,
|
|
"destination_label": f"{name}.backup_<date and time>",
|
|
"expected_bytes": size if policy != "off" else None,
|
|
},
|
|
"media": {
|
|
"deletes_files": False,
|
|
"label": "No files are removed; image links are replaced",
|
|
},
|
|
"network": {
|
|
"planned_downloads": downloads,
|
|
"api_key": bool(IntegrationConfig.CIVITAI_API_KEY),
|
|
},
|
|
"details": [{"label": n, "detail": ""} for n in found["names"]],
|
|
"disk": disk,
|
|
},
|
|
headers={"Cache-Control": "no-store"},
|
|
)
|
|
|
|
# ── apply ─────────────────────────────────────────────────────────
|
|
|
|
async def apply_lora_import(self, request):
|
|
data, error = await self._read_maintenance_body(request, ("token",))
|
|
if error is not None:
|
|
return error
|
|
token = data.get("token")
|
|
if not isinstance(token, str) or not 0 < len(token) <= MAX_TOKEN_LENGTH:
|
|
return _invalid("token must be the plan's token")
|
|
plan, refused = self._peek_plan(token)
|
|
if refused is not None:
|
|
return refused
|
|
if plan.operation != LORA_OPERATION:
|
|
return _failure(
|
|
"This review expired. Review the plan again.", 410, "plan-expired"
|
|
)
|
|
store = self._maintenance_plans
|
|
db = self.db
|
|
if plan.db_identity != self._database_identity(db):
|
|
store.discard(token)
|
|
return _failure(
|
|
"The database was replaced since the review. Review again.",
|
|
409,
|
|
"storage-changed",
|
|
)
|
|
lm_path = plan.payload["path"]
|
|
# 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:
|
|
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)
|
|
# 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)
|