250 lines
8.8 KiB
Python
250 lines
8.8 KiB
Python
"""
|
|
SaveImagePromotable: a pass-through SaveImage variant with accumulating previews
|
|
and a "promote/lock" feature.
|
|
|
|
Modes:
|
|
- Pass-through (default): saves incoming images, emits preview UI, returns the
|
|
input tensor as output. With `accumulate=True`, the frontend appends previews
|
|
to a gallery instead of replacing it.
|
|
- Locked: when `promoted_asset_ref` is a non-empty JSON ref to a saved asset,
|
|
the node skips saving, loads the referenced image, and outputs that image.
|
|
The frontend is expected to write the ref into the widget when the user
|
|
clicks the "lock" UI on a preview.
|
|
|
|
Caching: IS_CHANGED returns a stable key derived from the ref (+ file mtime)
|
|
when locked, so re-queues with the same lock are cache hits and upstream
|
|
ancestors are skipped. Unlocked, IS_CHANGED returns False to defer to normal
|
|
input-signature caching.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps, ImageSequence
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
import folder_paths
|
|
import node_helpers
|
|
from comfy.cli_args import args
|
|
|
|
|
|
def _parse_promoted_ref(promoted_asset_ref: str) -> dict | None:
|
|
if not promoted_asset_ref:
|
|
return None
|
|
try:
|
|
ref = json.loads(promoted_asset_ref)
|
|
except (json.JSONDecodeError, TypeError):
|
|
return None
|
|
if not isinstance(ref, dict):
|
|
return None
|
|
filename = ref.get("filename")
|
|
if not isinstance(filename, str) or not filename:
|
|
return None
|
|
subfolder = ref.get("subfolder", "") or ""
|
|
asset_type = ref.get("type", "output") or "output"
|
|
if not isinstance(subfolder, str) or not isinstance(asset_type, str):
|
|
return None
|
|
# Reject anything that could escape the base directory.
|
|
if os.path.isabs(subfolder) or ".." in subfolder.split(os.sep):
|
|
return None
|
|
if os.path.isabs(filename) or ".." in filename.split(os.sep):
|
|
return None
|
|
return {"filename": filename, "subfolder": subfolder, "type": asset_type}
|
|
|
|
|
|
def _resolve_ref_path(ref: dict) -> str | None:
|
|
asset_type = ref["type"]
|
|
if asset_type == "output":
|
|
base = folder_paths.get_output_directory()
|
|
elif asset_type == "input":
|
|
base = folder_paths.get_input_directory()
|
|
elif asset_type == "temp":
|
|
base = folder_paths.get_temp_directory()
|
|
else:
|
|
return None
|
|
path = os.path.join(base, ref["subfolder"], ref["filename"])
|
|
# Defense-in-depth: ensure the resolved path stays inside the base dir.
|
|
base_real = os.path.realpath(base)
|
|
path_real = os.path.realpath(path)
|
|
if not path_real.startswith(base_real + os.sep) and path_real != base_real:
|
|
return None
|
|
if not os.path.isfile(path_real):
|
|
return None
|
|
return path_real
|
|
|
|
|
|
def _load_image_tensor(path: str) -> torch.Tensor:
|
|
img = node_helpers.pillow(Image.open, path)
|
|
output_images: list[torch.Tensor] = []
|
|
w: int | None = None
|
|
h: int | None = None
|
|
for frame in ImageSequence.Iterator(img):
|
|
frame = node_helpers.pillow(ImageOps.exif_transpose, frame)
|
|
image = frame.convert("RGB")
|
|
if not output_images:
|
|
w, h = image.size
|
|
if image.size != (w, h):
|
|
continue
|
|
arr = np.array(image).astype(np.float32) / 255.0
|
|
output_images.append(torch.from_numpy(arr)[None,])
|
|
if not output_images:
|
|
raise RuntimeError(f"Failed to decode any frames from {path}")
|
|
return torch.cat(output_images, dim=0)
|
|
|
|
|
|
class SaveImagePromotable:
|
|
"""Pass-through SaveImage with accumulating previews and promote/lock.
|
|
|
|
Inputs:
|
|
images: IMAGE tensor to save + pass through (ignored when locked).
|
|
filename_prefix: STRING prefix for saved files.
|
|
accumulate: BOOLEAN — when True, frontend appends previews to gallery.
|
|
promoted_asset_ref: STRING — JSON ref written by the frontend on lock.
|
|
Empty string means "not locked, normal pass-through".
|
|
Output:
|
|
IMAGE — input pass-through, or the loaded promoted image when locked.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.compress_level = 4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": "Images to save and pass through. Ignored when a promoted asset is locked."
|
|
},
|
|
),
|
|
"filename_prefix": (
|
|
"STRING",
|
|
{"default": "ComfyUI", "tooltip": "Prefix for saved files."},
|
|
),
|
|
"accumulate": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "When enabled, previews append to a per-node gallery instead of replacing it.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"promoted_asset_ref": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": False,
|
|
"tooltip": "JSON ref to a saved asset. Set by the UI; do not edit manually.",
|
|
},
|
|
),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "execute"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "image"
|
|
DESCRIPTION = "Saves images, shows accumulating previews, and passes the input through. A promoted (locked) preview overrides pass-through to output the chosen image."
|
|
|
|
def _save_images(self, images, filename_prefix, prompt, extra_pnginfo):
|
|
full_output_folder, filename, counter, subfolder, _ = (
|
|
folder_paths.get_save_image_path(
|
|
filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
|
|
)
|
|
)
|
|
results: list[dict] = []
|
|
for batch_number, image in enumerate(images):
|
|
arr = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
|
|
metadata: PngInfo | None = None
|
|
if not args.disable_metadata:
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for key in extra_pnginfo:
|
|
metadata.add_text(key, json.dumps(extra_pnginfo[key]))
|
|
filename_with_batch = filename.replace("%batch_num%", str(batch_number))
|
|
out_name = f"{filename_with_batch}_{counter:05}_.png"
|
|
img.save(
|
|
os.path.join(full_output_folder, out_name),
|
|
pnginfo=metadata,
|
|
compress_level=self.compress_level,
|
|
)
|
|
results.append(
|
|
{"filename": out_name, "subfolder": subfolder, "type": self.type}
|
|
)
|
|
counter += 1
|
|
return results
|
|
|
|
def execute(
|
|
self,
|
|
images,
|
|
filename_prefix="ComfyUI",
|
|
accumulate=False, # noqa: ARG002
|
|
promoted_asset_ref="",
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
ref = _parse_promoted_ref(promoted_asset_ref)
|
|
if ref is not None:
|
|
path = _resolve_ref_path(ref)
|
|
if path is not None:
|
|
tensor = _load_image_tensor(path)
|
|
tensor = tensor.to(device=images.device, dtype=images.dtype)
|
|
return {
|
|
"ui": {"images": [ref]},
|
|
"result": (tensor,),
|
|
}
|
|
# Ref is set but stale (file deleted / failed validation): fall
|
|
# through to pass-through so the user gets a working graph rather
|
|
# than an execution error.
|
|
|
|
saved = self._save_images(images, filename_prefix, prompt, extra_pnginfo)
|
|
return {"ui": {"images": saved}, "result": (images,)}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(
|
|
cls,
|
|
images, # noqa: ARG003
|
|
filename_prefix="ComfyUI",
|
|
accumulate=False, # noqa: ARG003
|
|
promoted_asset_ref="",
|
|
prompt=None, # noqa: ARG003
|
|
extra_pnginfo=None, # noqa: ARG003
|
|
):
|
|
ref = _parse_promoted_ref(promoted_asset_ref)
|
|
if ref is None:
|
|
return False
|
|
path = _resolve_ref_path(ref)
|
|
if path is None:
|
|
return f"PROMOTED::MISSING::{promoted_asset_ref}"
|
|
try:
|
|
stat = os.stat(path)
|
|
sig = f"{stat.st_size}:{stat.st_mtime_ns}"
|
|
except OSError:
|
|
sig = "NOSTAT"
|
|
return f"PROMOTED::{promoted_asset_ref}::{sig}"
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SaveImagePromotable": SaveImagePromotable,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SaveImagePromotable": "Save Image (Promotable, PoC)",
|
|
}
|