feat: Sidecar + Save Image (Sidecars) nodes + registration
Sidecar chains {content,name,ext} specs over a SIDECAR-typed link; the save node
mirrors SaveImageKJ (folder_paths.get_save_image_path) and writes each sidecar as
base+name+ext next to the image. build_plan validates the whole chain before any
I/O so duplicates/bad extensions write nothing.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+6
-2
@@ -20,14 +20,18 @@ if __package__:
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _PROF_NAMES
|
||||
from .gates.bucket_node import NODE_CLASS_MAPPINGS as _BUCKET_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _BUCKET_NAMES
|
||||
from .gates.sidecar_node import NODE_CLASS_MAPPINGS as _SC_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _SC_NAMES
|
||||
from .gates import routes # noqa: F401 (registers aiohttp routes on import)
|
||||
from .gates import gate_server # noqa: F401 (registers /datasete_gate/* + text routes)
|
||||
from .gates import profiles_routes # noqa: F401 (registers /grid_pool/profiles/*)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES,
|
||||
**_TEXT_NODES, **_PROF_NODES, **_BUCKET_NODES}
|
||||
**_TEXT_NODES, **_PROF_NODES, **_BUCKET_NODES,
|
||||
**_SC_NODES}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES,
|
||||
**_TEXT_NAMES, **_PROF_NAMES, **_BUCKET_NAMES}
|
||||
**_TEXT_NAMES, **_PROF_NAMES, **_BUCKET_NAMES,
|
||||
**_SC_NAMES}
|
||||
else: # pragma: no cover - exercised only under pytest collection
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Save Image + chainable sidecar text/JSON files.
|
||||
|
||||
`Sidecar` nodes chain a list of {content, name, ext} specs (SIDECAR type);
|
||||
`SaveImageSidecars` saves the image and writes each sidecar next to it sharing
|
||||
the image's base name. Heavy deps (torch/PIL/folder_paths) are imported lazily
|
||||
inside save() so this module imports without comfy for unit tests."""
|
||||
from . import sidecar as sc
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
class Sidecar:
|
||||
CATEGORY = "Datasete Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("SIDECAR",)
|
||||
RETURN_NAMES = ("sidecar",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"content": ("STRING", {"forceInput": True}),
|
||||
"name": ("STRING", {"default": ""}),
|
||||
"extension": ("STRING", {"default": ".txt"}),
|
||||
},
|
||||
"optional": {
|
||||
"sidecar": ("SIDECAR",), # chain-in from a previous Sidecar
|
||||
},
|
||||
}
|
||||
|
||||
def run(self, content, name, extension, sidecar=None):
|
||||
return (sc.append_spec(sidecar, content, name, extension),)
|
||||
|
||||
|
||||
class SaveImageSidecars:
|
||||
CATEGORY = "Datasete Gates"
|
||||
FUNCTION = "save"
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
||||
"output_folder": ("STRING", {"default": "output"}),
|
||||
},
|
||||
"optional": {
|
||||
"sidecar": ("SIDECAR",),
|
||||
},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _safe_path(folder, filename):
|
||||
import os
|
||||
path = os.path.join(folder, filename)
|
||||
root = os.path.abspath(folder)
|
||||
if os.path.commonpath((root, os.path.abspath(path))) != root:
|
||||
raise ValueError(f"Refusing to write outside the target folder: {path}")
|
||||
return path
|
||||
|
||||
def save(self, images, filename_prefix, output_folder, sidecar=None,
|
||||
prompt=None, extra_pnginfo=None):
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
import folder_paths
|
||||
from comfy.cli_args import args
|
||||
|
||||
# Validate the entire sidecar plan BEFORE writing anything, so a bad
|
||||
# chain (duplicate name, disallowed extension) writes no files at all.
|
||||
plan = sc.build_plan(sidecar)
|
||||
|
||||
h, w = int(images[0].shape[0]), int(images[0].shape[1])
|
||||
if os.path.isabs(output_folder):
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
output_dir = output_folder
|
||||
else:
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = \
|
||||
folder_paths.get_save_image_path(filename_prefix, output_dir, w, h)
|
||||
|
||||
results = []
|
||||
for batch_number, image in enumerate(images):
|
||||
arr = (255.0 * image.cpu().numpy()).clip(0, 255).astype(np.uint8)
|
||||
img = Image.fromarray(arr)
|
||||
|
||||
metadata = 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 k in extra_pnginfo:
|
||||
metadata.add_text(k, json.dumps(extra_pnginfo[k]))
|
||||
|
||||
base = f"{filename.replace('%batch_num%', str(batch_number))}_{counter:05}_"
|
||||
img.save(self._safe_path(full_output_folder, base + ".png"),
|
||||
pnginfo=metadata, compress_level=4)
|
||||
for suffix, content in plan:
|
||||
with open(self._safe_path(full_output_folder, base + suffix),
|
||||
"w", encoding="utf-8") as f:
|
||||
f.write(content if content is not None else "")
|
||||
|
||||
results.append({"filename": base + ".png",
|
||||
"subfolder": subfolder, "type": "output"})
|
||||
counter += 1
|
||||
|
||||
return {"ui": {"images": results}}
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Sidecar": Sidecar,
|
||||
"SaveImageSidecars": SaveImageSidecars,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Sidecar": "Sidecar (text/json)",
|
||||
"SaveImageSidecars": "Save Image (Sidecars)",
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
from gates import sidecar_node as sn
|
||||
|
||||
|
||||
def test_sidecar_run_builds_chain():
|
||||
node = sn.Sidecar()
|
||||
(chain,) = node.run(content="hello", name="", extension=".txt", sidecar=None)
|
||||
assert chain == [{"content": "hello", "name": "", "ext": ".txt"}]
|
||||
(chain2,) = node.run(content="{}", name="meta", extension=".json", sidecar=chain)
|
||||
assert len(chain2) == 2
|
||||
assert chain2[1] == {"content": "{}", "name": "meta", "ext": ".json"}
|
||||
|
||||
|
||||
def test_sidecar_io_shape():
|
||||
assert sn.Sidecar.RETURN_TYPES == ("SIDECAR",)
|
||||
it = sn.Sidecar.INPUT_TYPES()
|
||||
assert "content" in it["required"]
|
||||
assert "name" in it["required"]
|
||||
assert "extension" in it["required"]
|
||||
assert "sidecar" in it["optional"]
|
||||
|
||||
|
||||
def test_save_node_io_shape():
|
||||
assert sn.SaveImageSidecars.OUTPUT_NODE is True
|
||||
assert sn.SaveImageSidecars.RETURN_TYPES == ()
|
||||
it = sn.SaveImageSidecars.INPUT_TYPES()
|
||||
for k in ("images", "filename_prefix", "output_folder"):
|
||||
assert k in it["required"]
|
||||
assert "sidecar" in it["optional"]
|
||||
|
||||
|
||||
def test_mappings_present():
|
||||
assert "Sidecar" in sn.NODE_CLASS_MAPPINGS
|
||||
assert "SaveImageSidecars" in sn.NODE_CLASS_MAPPINGS
|
||||
assert sn.NODE_DISPLAY_NAME_MAPPINGS["SaveImageSidecars"]
|
||||
Reference in New Issue
Block a user