305 lines
13 KiB
Python
305 lines
13 KiB
Python
"""Section-level installer — one source of truth above NODE_PACKS + subject_api.
|
||
|
||
A SECTION is a user-facing feature (Style Reference, ControlNet, PuLID, AI
|
||
Selection…). Each groups one or more MEMBERS, where a member is either a
|
||
node-pack key (node_install_api.NODE_PACKS) or an AI-selection target
|
||
(subject_api). Both the onboarding splash and the in-app install modal read
|
||
these sections, so a feature added here shows up identically in every entry
|
||
point and reports ONE rolled-up status. Per-member skip logic lives in the
|
||
installers themselves, so installing a section only fetches what's missing.
|
||
|
||
Routes:
|
||
GET /promptchain/install/sections → all sections + rolled-up health
|
||
GET /promptchain/install/status?section=ID → one section's detailed status
|
||
POST /promptchain/install/install {section} → SSE: install missing members
|
||
GET /promptchain/install/updates → bundled-pack vs installed drift
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import logging
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import server
|
||
from aiohttp import web
|
||
|
||
from .api_utils import error_response, parse_json
|
||
from .node_install_api import (
|
||
INSTALL_CANCEL,
|
||
NODE_PACKS, _CRITICAL, _bundled_packs_dir, _custom_nodes_dir, _log,
|
||
_open_sse, _pack_status, _pkg_version, _restore_critical, _run, _send,
|
||
install_pack,
|
||
)
|
||
from .subject_api import _object_status, _subject_status, install_subject_target
|
||
|
||
routes = server.PromptServer.instance.routes
|
||
log = logging.getLogger("promptchain.install")
|
||
|
||
|
||
def _np(key: str) -> dict:
|
||
return {"kind": "nodepack", "key": key}
|
||
|
||
|
||
def _subj(target: str) -> dict:
|
||
return {"kind": "subject", "target": target}
|
||
|
||
|
||
# A `deferred` entry documents a download that only happens on first USE (not at
|
||
# install), so the UI can warn "+N GB later" instead of the user being surprised.
|
||
SECTIONS: list[dict] = [
|
||
{"id": "style_reference", "label": "Style Reference", "size": "~3.2 GB",
|
||
"desc": "IP-Adapter — transfer style from a reference image",
|
||
"members": [_np("StyleReference")]},
|
||
{"id": "upscaler", "label": "Upscaler", "size": "~64 MB",
|
||
"desc": "Ultimate SD Upscale + 4x-UltraSharp",
|
||
"members": [_np("Upscaler")]},
|
||
{"id": "face_detailer", "label": "Face Detailer", "size": "~6 MB",
|
||
"desc": "Region-correct face detailing",
|
||
"members": [_np("RegionalDetailer")]},
|
||
{"id": "krea2_highdetail", "label": "Krea 2 High Detail", "size": "~510 MB",
|
||
"desc": "Clean 2× decoder + conditioning rebalance for the Krea 2 High-Detail templates",
|
||
"members": [_np("Krea2HighDetail")]},
|
||
{"id": "krea2_realism", "label": "Krea 2 Realism", "size": "~5.7 GB",
|
||
"desc": "Abliterated encoder + realism & filter-bypass LoRAs for Krea 2 Turbo Realism (uncensored). t2i/i2i also need the RES4LYF node pack — install it via ComfyUI Manager.",
|
||
"members": [_np("Krea2Realism")]},
|
||
{"id": "controlnet", "label": "ControlNet", "size": "~5 GB SDXL · up to ~22 GB all",
|
||
"desc": "Depth / pose / canny / edge guidance (SDXL · Flux · Z-Image)",
|
||
"members": [_np("ControlNet"), _np("FluxControlNet"),
|
||
_np("ZImageControlNet"), _np("ZImageControlNetBase")]},
|
||
{"id": "pulid", "label": "PuLID (face identity)", "size": "~4.3 GB",
|
||
"desc": "Identity transfer for SDXL · Flux · Flux.2",
|
||
# EVA-CLIP (used by PuLID Flux 2) is no longer deferred — installing that
|
||
# member pre-fetches it (see _MEMBER_PREFETCH) so the first render doesn't
|
||
# stall on an ~800 MB download.
|
||
"members": [_np("PuLID"), _np("PuLIDFlux"), _np("PuLIDFlux2")]},
|
||
{"id": "seedvr2", "label": "SeedVR2 Upscaler", "size": "~480 MB",
|
||
"desc": "Restoration upscaler (ByteDance)",
|
||
"members": [_np("SeedVR2")],
|
||
"deferred": [{"label": "SeedVR2 DiT weights", "size": "6.8–16.5 GB",
|
||
"when": "first SeedVR2 upscale"}]},
|
||
{"id": "ai_selection", "label": "AI Selection", "size": "~1 GB",
|
||
"desc": "Select Subject (BiRefNet) + click-to-select objects (SAM2) in the editor",
|
||
"members": [_subj("subject"), _subj("object")]},
|
||
]
|
||
|
||
_SECTIONS_BY_ID = {s["id"]: s for s in SECTIONS}
|
||
|
||
# Reverse index: node-pack member key -> section id. The in-app install modal is
|
||
# opened per node-pack "injectable"; this lets it install through the unified
|
||
# section route scoped to just that one member (not the whole section — the
|
||
# ControlNet section alone spans up to ~22 GB across four base families).
|
||
_SECTION_FOR_NODEPACK = {
|
||
m["key"]: s["id"]
|
||
for s in SECTIONS for m in s["members"] if m["kind"] == "nodepack"
|
||
}
|
||
|
||
|
||
def section_for_injectable(key: str) -> str | None:
|
||
return _SECTION_FOR_NODEPACK.get(key)
|
||
|
||
|
||
# Weights a pack would otherwise lazy-download on first USE, which we instead
|
||
# pre-fetch right after the pack installs so the first render doesn't stall.
|
||
# Keyed by node-pack member key. The code runs the pack's OWN resolver in a
|
||
# subprocess (inheriting our env + interpreter), so it writes to exactly the
|
||
# cache the runtime later reads, and is a cheap no-op once cached. Best-effort:
|
||
# a failure never aborts the install — the weight simply falls back to lazy.
|
||
async def _prefetch_for_member(resp: web.StreamResponse, key: str) -> None:
|
||
"""No-op. Weight pre-fetching was removed for ComfyUI Registry compliance —
|
||
it ran a subprocess to warm a pack's lazy weights. Those weights now simply
|
||
download on first render, as the packs already handle (e.g. PuLID-Flux2's
|
||
EVA-CLIP via open_clip), so the first such render is slightly slower and
|
||
nothing else changes."""
|
||
return None
|
||
|
||
|
||
def _member_status(m: dict) -> dict:
|
||
"""Normalize one member's install state to installed | needs_restart | missing.
|
||
|
||
needs_restart = files are on disk but the proof node hasn't registered yet
|
||
(a just-installed node pack before the ComfyUI restart)."""
|
||
if m["kind"] == "nodepack":
|
||
st = _pack_status(m["key"])
|
||
files_present = all(r["cloned"] for r in st["repos"]) and not st["missing_models"]
|
||
if st["present"]:
|
||
state = "installed"
|
||
elif files_present and st["missing_nodes"]:
|
||
state = "needs_restart"
|
||
else:
|
||
state = "missing"
|
||
return {"kind": "nodepack", "key": m["key"], "label": st["label"],
|
||
"state": state, "missing_models": st["missing_models"],
|
||
"missing_nodes": st["missing_nodes"]}
|
||
st = _subject_status() if m["target"] == "subject" else _object_status()
|
||
# Selection libs import at call time — no restart concept; ready or not.
|
||
state = "installed" if st["ready"] else "missing"
|
||
return {"kind": "subject", "target": m["target"], "label": st["label"],
|
||
"state": state, "deps_ok": st["deps_ok"], "model_ok": st["model_ok"],
|
||
"size": st.get("size")}
|
||
|
||
|
||
def _section_status(sec: dict) -> dict:
|
||
members = [_member_status(m) for m in sec["members"]]
|
||
installed = sum(1 for mm in members if mm["state"] == "installed")
|
||
total = len(members)
|
||
if installed == total:
|
||
health = "installed"
|
||
elif installed == 0 and not any(mm["state"] == "needs_restart" for mm in members):
|
||
health = "missing"
|
||
else:
|
||
health = "partial" # some members present (or just need a restart)
|
||
return {
|
||
"id": sec["id"],
|
||
"label": sec["label"],
|
||
"desc": sec.get("desc", ""),
|
||
"size": sec.get("size", ""),
|
||
"members": members,
|
||
"installed_count": installed,
|
||
"total": total,
|
||
"present": installed == total,
|
||
"health": health,
|
||
"needs_restart": any(mm["state"] == "needs_restart" for mm in members),
|
||
"deferred": sec.get("deferred", []),
|
||
}
|
||
|
||
|
||
@routes.get("/promptchain/install/sections")
|
||
async def _api_sections(request):
|
||
return web.json_response({"sections": [_section_status(s) for s in SECTIONS]})
|
||
|
||
|
||
@routes.get("/promptchain/install/status")
|
||
async def _api_section_status(request):
|
||
sid = (request.query.get("section") or "").strip()
|
||
sec = _SECTIONS_BY_ID.get(sid)
|
||
if not sec:
|
||
return error_response(f"unknown section: {sid}")
|
||
return web.json_response(_section_status(sec))
|
||
|
||
|
||
@routes.post("/promptchain/install/install")
|
||
async def _api_section_install(request):
|
||
"""SSE: install every not-yet-present member of a section, in order. Members
|
||
already installed are skipped (logged), so re-running or installing an
|
||
overlapping section is cheap. One critical-deps snapshot + restore wraps the
|
||
whole run; needs_restart is true iff any node-pack member was installed."""
|
||
body, err = await parse_json(request)
|
||
if err:
|
||
return err
|
||
sid = (body.get("section") or "").strip()
|
||
sec = _SECTIONS_BY_ID.get(sid)
|
||
if not sec:
|
||
return error_response(f"unknown section: {sid}")
|
||
|
||
# Optional filter: install only these members (node-pack keys or subject
|
||
# targets). Defaults to the whole section. Lets the picker install just the
|
||
# families the user ticked — e.g. only Flux ControlNet, not all four.
|
||
want = body.get("members")
|
||
members = sec["members"]
|
||
if isinstance(want, list) and want:
|
||
wanted = set(want)
|
||
members = [m for m in members if (m.get("key") or m.get("target")) in wanted]
|
||
if not members:
|
||
return error_response("no matching members in section")
|
||
|
||
resp = await _open_sse(request)
|
||
snapshot = {p: _pkg_version(p) for p in _CRITICAL}
|
||
installed_nodepack = False
|
||
INSTALL_CANCEL.clear() # fresh run — drop any stale cancel from a prior install
|
||
try:
|
||
await _log(resp, f"Installing {sec['label']}")
|
||
for m in members:
|
||
if INSTALL_CANCEL.is_set():
|
||
await _send(resp, {"cancelled": True})
|
||
return resp
|
||
ms = _member_status(m)
|
||
await _send(resp, {"member": ms["label"], "state": ms["state"]})
|
||
if ms["state"] == "installed":
|
||
await _log(resp, f"{ms['label']}: already installed, skipping")
|
||
continue
|
||
if m["kind"] == "nodepack":
|
||
if not await install_pack(resp, m["key"], snapshot):
|
||
return resp # install_pack already emitted the error
|
||
installed_nodepack = True
|
||
await _prefetch_for_member(resp, m["key"])
|
||
else:
|
||
if not await install_subject_target(resp, m["target"], snapshot):
|
||
return resp
|
||
# subject members install pip libs without restoring on success — and a
|
||
# node pack's restore already ran inside install_pack, so this final pass
|
||
# is a no-op there. One sweep covers the whole section either way.
|
||
await _restore_critical(resp, snapshot)
|
||
msg = ("Restart ComfyUI to load the new nodes."
|
||
if installed_nodepack else "Done.")
|
||
await _log(resp, msg)
|
||
await _send(resp, {"done": True, "ok": True, "needs_restart": installed_nodepack})
|
||
except ConnectionResetError:
|
||
pass # client navigated away; pip/downloads keep running
|
||
except Exception as e:
|
||
log.exception("section install failed")
|
||
try:
|
||
await _send(resp, {"error": str(e)})
|
||
except Exception:
|
||
pass
|
||
return resp
|
||
|
||
|
||
@routes.post("/promptchain/install/cancel")
|
||
async def _api_install_cancel(request):
|
||
"""Signal the running section install to stop. _download checks this per
|
||
chunk and the section loop checks it between members, so the in-flight file
|
||
bails (keeping a .part for resume) and no further members start."""
|
||
INSTALL_CANCEL.set()
|
||
return web.json_response({"ok": True})
|
||
|
||
|
||
# ── update check ──────────────────────────────────────────────────
|
||
# Every 3rd-party pack we install is VENDORED in bundled_packs/ at a pinned
|
||
# version, so "check for updates" = did PromptChain ship a newer copy than the
|
||
# one already in custom_nodes/? Compare a content signature of the two trees.
|
||
# Read-only — applying an update (re-copy) needs a restart and is a separate
|
||
# action, since overwriting a loaded pack can hit locked DLLs on Windows.
|
||
|
||
def _pack_signature(d: Path) -> str | None:
|
||
if not d.is_dir():
|
||
return None
|
||
h = hashlib.sha256()
|
||
for f in sorted(d.rglob("*")):
|
||
if not f.is_file() or "__pycache__" in f.parts or f.suffix == ".pyc":
|
||
continue
|
||
h.update(f.relative_to(d).as_posix().encode())
|
||
h.update(b"\0")
|
||
try:
|
||
h.update(f.read_bytes())
|
||
except OSError:
|
||
h.update(b"<unreadable>")
|
||
h.update(b"\0")
|
||
return h.hexdigest()
|
||
|
||
|
||
@routes.get("/promptchain/install/updates")
|
||
async def _api_updates(request):
|
||
"""Report bundled packs whose installed copy differs from the version
|
||
PromptChain currently ships (i.e. an update is available via re-copy)."""
|
||
cn = _custom_nodes_dir()
|
||
bp = _bundled_packs_dir()
|
||
updates = []
|
||
seen: set[str] = set()
|
||
for sec in SECTIONS:
|
||
for m in sec["members"]:
|
||
if m["kind"] != "nodepack":
|
||
continue
|
||
for name in NODE_PACKS[m["key"]].get("bundle", []):
|
||
if name in seen:
|
||
continue
|
||
seen.add(name)
|
||
installed = cn / name
|
||
if not installed.is_dir():
|
||
continue # not installed → an install, not an update
|
||
shipped_sig = _pack_signature(bp / name)
|
||
installed_sig = _pack_signature(installed)
|
||
if shipped_sig and installed_sig and shipped_sig != installed_sig:
|
||
updates.append({"pack": name, "section": sec["id"]})
|
||
return web.json_response({"updates": updates})
|