345 lines
13 KiB
Python
345 lines
13 KiB
Python
# Viewer "Inpaint"/"Edit" support: Inpaint Apply renders to ComfyUI's temp dir
|
|
# (PreviewImage), so discarded attempts never touch the output tree; Save
|
|
# finalizes the chosen render by copying it into the output dir under a
|
|
# SaveImage-style prefix and recording it with parent lineage — no re-render.
|
|
# Edit Save receives the airbrushed composite from the browser and persists it
|
|
# the same way, re-attaching the parent's embedded prompt/workflow chunks
|
|
# (canvas export strips them) so the result stays inpaint/upscale-able.
|
|
|
|
import asyncio
|
|
import io as _io
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import time
|
|
|
|
import folder_paths
|
|
import node_helpers
|
|
import server
|
|
from aiohttp import web
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
from .api_utils import error_response, parse_json
|
|
from .history_db import record_image
|
|
|
|
routes = server.PromptServer.instance.routes
|
|
|
|
|
|
def _format_prefix(prefix: str, fallback: str = "inpaint/inpaint") -> str:
|
|
"""Resolve SaveImage %date:...% tokens (the subset PromptChain emits)."""
|
|
def date_token(match):
|
|
fmt = match.group(1)
|
|
now = time.localtime()
|
|
out = fmt
|
|
for token, value in [
|
|
("yyyy", time.strftime("%Y", now)), ("MM", time.strftime("%m", now)),
|
|
("dd", time.strftime("%d", now)), ("hh", time.strftime("%H", now)),
|
|
("mm", time.strftime("%M", now)), ("ss", time.strftime("%S", now)),
|
|
]:
|
|
out = out.replace(token, value)
|
|
return out
|
|
return re.sub(r"%date:([^%]+)%", date_token, prefix or "").strip() or fallback
|
|
|
|
|
|
def _next_counter_name(directory: str, basename: str, ext: str) -> str:
|
|
"""SaveImage-style _00001_ counter suffix, scanning existing files."""
|
|
pattern = re.compile(re.escape(basename) + r"_(\d{5})_" + re.escape(ext) + r"$")
|
|
highest = 0
|
|
try:
|
|
for name in os.listdir(directory):
|
|
m = pattern.match(name)
|
|
if m:
|
|
highest = max(highest, int(m.group(1)))
|
|
except OSError:
|
|
pass
|
|
return f"{basename}_{highest + 1:05}_{ext}"
|
|
|
|
|
|
@routes.post("/promptchain/save-temp-image")
|
|
async def save_temp_image(request):
|
|
data, err = await parse_json(request)
|
|
if err:
|
|
return err
|
|
filename = (data.get("filename") or "").strip()
|
|
subfolder = (data.get("subfolder") or "").strip()
|
|
if not filename or "/" in filename or "\\" in filename or ".." in filename:
|
|
return error_response("invalid filename")
|
|
temp_root = folder_paths.get_temp_directory()
|
|
src = os.path.normpath(os.path.join(temp_root, subfolder, filename))
|
|
if not src.startswith(os.path.normpath(temp_root)) or not os.path.isfile(src):
|
|
return error_response("temp image not found", 404)
|
|
|
|
prefix = _format_prefix(data.get("prefix") or "")
|
|
rel_dir, basename = os.path.split(prefix)
|
|
out_root = folder_paths.get_output_directory()
|
|
out_dir = os.path.normpath(os.path.join(out_root, rel_dir))
|
|
if not out_dir.startswith(os.path.normpath(out_root)):
|
|
return error_response("invalid prefix")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
ext = os.path.splitext(filename)[1] or ".png"
|
|
out_name = _next_counter_name(out_dir, basename, ext)
|
|
shutil.copy2(src, os.path.join(out_dir, out_name))
|
|
|
|
meta = {k: data.get(k) for k in (
|
|
"prompt", "negative", "seed", "model", "steps", "cfg",
|
|
"sampler", "scheduler", "denoise", "parent_filename", "regions",
|
|
)}
|
|
entry = record_image(
|
|
filename=out_name,
|
|
subfolder=rel_dir.replace("\\", "/"),
|
|
source_type="output",
|
|
workflow_id=data.get("workflow_id") or None,
|
|
metadata=meta,
|
|
)
|
|
if entry is None:
|
|
return error_response("saved file could not be recorded", 500)
|
|
# The mask this render consumed must outlive the age sweep so the saved
|
|
# output stays re-applicable (the parent ref is pinned in record_image).
|
|
from . import inpaint_files
|
|
inpaint_files.pin_value(data.get("mask_filename") or "")
|
|
return web.json_response(entry)
|
|
|
|
|
|
def _save_edit_png(raw: bytes, parent_path: str | None, out_path: str) -> None:
|
|
"""Write the browser composite, re-attaching the parent's prompt/workflow
|
|
PNG chunks — the image WAS rendered with that recipe, then hand-painted,
|
|
and without them the result can't anchor a follow-up inpaint/upscale."""
|
|
img = node_helpers.pillow(Image.open, _io.BytesIO(raw))
|
|
pnginfo = PngInfo()
|
|
if parent_path and os.path.isfile(parent_path):
|
|
try:
|
|
parent = node_helpers.pillow(Image.open, parent_path)
|
|
chunks = dict(parent.info or {})
|
|
chunks.update(getattr(parent, "text", None) or {})
|
|
for key in ("prompt", "workflow"):
|
|
if isinstance(chunks.get(key), str) and chunks[key]:
|
|
pnginfo.add_text(key, chunks[key])
|
|
except OSError:
|
|
pass
|
|
img.save(out_path, pnginfo=pnginfo)
|
|
|
|
|
|
@routes.post("/promptchain/save-edited-image")
|
|
async def save_edited_image(request):
|
|
form = await request.post()
|
|
upload = form.get("image")
|
|
if upload is None or not getattr(upload, "file", None):
|
|
return error_response("missing image upload")
|
|
raw = upload.file.read()
|
|
if not raw:
|
|
return error_response("empty image upload")
|
|
|
|
prefix = _format_prefix(str(form.get("prefix") or ""), fallback="edit/edit")
|
|
rel_dir, basename = os.path.split(prefix)
|
|
out_root = folder_paths.get_output_directory()
|
|
out_dir = os.path.normpath(os.path.join(out_root, rel_dir))
|
|
if not out_dir.startswith(os.path.normpath(out_root)):
|
|
return error_response("invalid prefix")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
out_name = _next_counter_name(out_dir, basename, ".png")
|
|
|
|
# The parent may be subfolder-scoped (e.g. "promptchain_inpaint/foo.png"), so
|
|
# allow a relative subpath but keep it contained within the input dir — a
|
|
# blanket slash-reject silently drops the parent's prompt/workflow chunks.
|
|
parent_ref = str(form.get("parent_filename") or "").strip()
|
|
parent_path = None
|
|
if parent_ref and ".." not in parent_ref:
|
|
input_root = os.path.normpath(folder_paths.get_input_directory())
|
|
cand = os.path.normpath(os.path.join(input_root, parent_ref))
|
|
if cand.startswith(input_root + os.sep):
|
|
parent_path = cand
|
|
try:
|
|
await asyncio.to_thread(_save_edit_png, raw, parent_path, os.path.join(out_dir, out_name))
|
|
except (OSError, ValueError) as e:
|
|
return error_response(f"could not save image: {e}", 500)
|
|
|
|
entry = record_image(
|
|
filename=out_name,
|
|
subfolder=rel_dir.replace("\\", "/"),
|
|
source_type="output",
|
|
workflow_id=str(form.get("workflow_id") or "") or None,
|
|
metadata={"parent_filename": parent_ref or None, "prompt": ""},
|
|
)
|
|
if entry is None:
|
|
return error_response("saved file could not be recorded", 500)
|
|
return web.json_response(entry)
|
|
|
|
|
|
# ── layer-stack persistence ("don't-flatten = keep your PSD") ──
|
|
# A sidecar layered document keyed by the SAVED image's content hash: a doc.json
|
|
# manifest + one PNG per layer bitmap + per mask, under {user}/PromptChain/
|
|
# edit-docs/<hash>/. Re-opening Edit on that image restores the stack. The flat
|
|
# PNG saved by save-edited-image stays the source of truth; this is a convenience
|
|
# cache, capped by an LRU size budget. Content-addressed → free dedup.
|
|
|
|
_EDIT_DOC_BUDGET_BYTES = 2 * 1024 ** 3
|
|
_HASH_RE = re.compile(r"[0-9a-f]{64}")
|
|
_PLANE_RE = re.compile(r"L\d+(\.mask)?\.png")
|
|
|
|
|
|
def _edit_docs_root() -> str:
|
|
return os.path.join(folder_paths.get_user_directory(), "PromptChain", "edit-docs")
|
|
|
|
|
|
def _rmtree_robust(path: str, tries: int = 6) -> None:
|
|
"""rmtree, retried briefly. On Windows a file that's momentarily open (e.g. a
|
|
concurrent list-read counting layers) can't be deleted, and rmtree's
|
|
ignore_errors SILENTLY skips it — leaving a stale doc.json behind. The lock
|
|
is transient, so a few short retries clear it."""
|
|
for _ in range(tries):
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
if not os.path.exists(path):
|
|
return
|
|
time.sleep(0.05)
|
|
|
|
|
|
def _dir_size(path: str) -> int:
|
|
total = 0
|
|
for root, _dirs, files in os.walk(path):
|
|
for name in files:
|
|
try:
|
|
total += os.path.getsize(os.path.join(root, name))
|
|
except OSError:
|
|
pass
|
|
return total
|
|
|
|
|
|
def _evict_edit_docs(keep_hash: str) -> None:
|
|
"""LRU by doc.json mtime until under the size budget; never evict keep_hash."""
|
|
root = _edit_docs_root()
|
|
try:
|
|
names = os.listdir(root)
|
|
except OSError:
|
|
return
|
|
docs, total = [], 0
|
|
for name in names:
|
|
path = os.path.join(root, name)
|
|
if not os.path.isdir(path):
|
|
continue
|
|
size = _dir_size(path)
|
|
total += size
|
|
manifest = os.path.join(path, "doc.json")
|
|
try:
|
|
mtime = os.path.getmtime(manifest if os.path.isfile(manifest) else path)
|
|
except OSError:
|
|
mtime = 0
|
|
docs.append((mtime, size, name, path))
|
|
if total <= _EDIT_DOC_BUDGET_BYTES:
|
|
return
|
|
docs.sort() # oldest first
|
|
for _mtime, size, name, path in docs:
|
|
if total <= _EDIT_DOC_BUDGET_BYTES:
|
|
break
|
|
if name == keep_hash:
|
|
continue
|
|
_rmtree_robust(path)
|
|
total -= size
|
|
|
|
|
|
@routes.post("/promptchain/edit-doc/{hash}")
|
|
async def save_edit_doc(request):
|
|
h = request.match_info.get("hash", "")
|
|
if not _HASH_RE.fullmatch(h):
|
|
return error_response("invalid hash")
|
|
form = await request.post()
|
|
manifest = form.get("manifest")
|
|
if manifest is None:
|
|
return error_response("missing manifest")
|
|
manifest_text = manifest if isinstance(manifest, str) else manifest.file.read().decode("utf-8")
|
|
planes = {}
|
|
for key, val in form.items():
|
|
if key == "manifest" or not _PLANE_RE.fullmatch(key):
|
|
continue
|
|
f = getattr(val, "file", None)
|
|
if f is not None:
|
|
planes[key] = f.read()
|
|
|
|
def _write():
|
|
doc_dir = os.path.join(_edit_docs_root(), h)
|
|
# rewrite from scratch so stale planes from a prior save don't linger
|
|
_rmtree_robust(doc_dir)
|
|
os.makedirs(doc_dir, exist_ok=True)
|
|
with open(os.path.join(doc_dir, "doc.json"), "w", encoding="utf-8") as fp:
|
|
fp.write(manifest_text)
|
|
for name, data in planes.items():
|
|
with open(os.path.join(doc_dir, name), "wb") as fp:
|
|
fp.write(data)
|
|
_evict_edit_docs(h)
|
|
|
|
try:
|
|
await asyncio.to_thread(_write)
|
|
except OSError as e:
|
|
return error_response(f"could not save edit doc: {e}", 500)
|
|
return web.json_response({"ok": True})
|
|
|
|
|
|
@routes.get("/promptchain/edit-doc/{hash}")
|
|
async def get_edit_doc(request):
|
|
h = request.match_info.get("hash", "")
|
|
if not _HASH_RE.fullmatch(h):
|
|
return error_response("invalid hash")
|
|
path = os.path.join(_edit_docs_root(), h, "doc.json")
|
|
if not os.path.isfile(path):
|
|
return web.json_response({"exists": False}, status=404)
|
|
try:
|
|
os.utime(path, None) # LRU touch — reopened docs survive eviction
|
|
except OSError:
|
|
pass
|
|
return web.FileResponse(path)
|
|
|
|
|
|
@routes.get("/promptchain/edit-doc/{hash}/{file}")
|
|
async def get_edit_doc_file(request):
|
|
h = request.match_info.get("hash", "")
|
|
fname = request.match_info.get("file", "")
|
|
if not _HASH_RE.fullmatch(h) or not (fname == "doc.json" or _PLANE_RE.fullmatch(fname)):
|
|
return error_response("invalid request")
|
|
root = os.path.normpath(_edit_docs_root())
|
|
path = os.path.normpath(os.path.join(root, h, fname))
|
|
if not path.startswith(root) or not os.path.isfile(path):
|
|
return web.Response(status=404)
|
|
return web.FileResponse(path)
|
|
|
|
|
|
@routes.get("/promptchain/edit-docs")
|
|
async def list_edit_docs(request):
|
|
"""Which hashes have a saved layer stack (→ a "layers" badge in lineage/the
|
|
gallery), mapped to their layer count for the tooltip. One listdir + a small
|
|
JSON read per doc; the frontend caches the result."""
|
|
root = _edit_docs_root()
|
|
|
|
def _scan():
|
|
out = {}
|
|
try:
|
|
names = os.listdir(root)
|
|
except OSError:
|
|
return out
|
|
for name in names:
|
|
if not _HASH_RE.fullmatch(name):
|
|
continue
|
|
manifest = os.path.join(root, name, "doc.json")
|
|
if not os.path.isfile(manifest):
|
|
continue
|
|
try:
|
|
with open(manifest, "r", encoding="utf-8") as fp:
|
|
out[name] = len(json.load(fp).get("layers") or [])
|
|
except (OSError, ValueError):
|
|
out[name] = 0
|
|
return out
|
|
|
|
return web.json_response({"docs": await asyncio.to_thread(_scan)})
|
|
|
|
|
|
@routes.delete("/promptchain/edit-doc/{hash}")
|
|
async def delete_edit_doc(request):
|
|
# Flatten saves to the SAME hash (pixels unchanged), so a flattened save must
|
|
# actively remove any prior layer stack — else the file would still restore
|
|
# its old layers. Idempotent: a missing dir is fine.
|
|
h = request.match_info.get("hash", "")
|
|
if not _HASH_RE.fullmatch(h):
|
|
return error_response("invalid hash")
|
|
await asyncio.to_thread(_rmtree_robust, os.path.join(_edit_docs_root(), h))
|
|
return web.json_response({"ok": True})
|