From 04c4c0d37cabd217f84bd5c7c3b99355cdf4aa3b Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Wed, 24 Jun 2026 22:53:50 +0200 Subject: [PATCH] Add loop accumulator node --- README.md | 17 ++++ loop_nodes.py | 236 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 253 insertions(+) diff --git a/README.md b/README.md index 359aecc..7cdb7c3 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ The node is registered as: - `prompt_builder / SxCP For Loop Start` - `prompt_builder / SxCP For Loop End` - `prompt_builder / SxCP Loop Append` +- `prompt_builder / SxCP Accumulator` - `prompt_builder / SxCP Category Preset` - `prompt_builder / SxCP Cast Control` - `prompt_builder / SxCP Generation Profile` @@ -118,6 +119,22 @@ want to resume a loop without changing index-derived seeds or row numbers. you want to update each iteration. They are separate from the collector and grow dynamically in the UI as you connect them. +`SxCP Accumulator` stores outputs across executions under a `store_key` or the +node id. Put it after an image-producing step inside or after a loop, connect the +generated `image`, and connect `For Loop Start.index` to `entry_id` when you want +reruns to replace the same row instead of appending duplicates. Its outputs are: + +- `collection`: all stored values, or images when no explicit `value` is wired. +- `image_batch`: all stored images as one ComfyUI image batch when they share + the same height and width. Set `image_batch_mode=resize_to_first` if you want + mixed sizes resized into one batch. +- `image_list`: a ComfyUI list output containing each stored image separately. +- `image_batch_1..4`: same-size grouped batches for mixed-format runs, so a + square group and a portrait group can be saved or processed separately. + +ComfyUI image batches require matching dimensions. For mixed image formats, use +`image_list` or the grouped `image_batch_1..4` outputs instead of `image_batch`. + ## Character Profiles `SxCP Woman Slot` and `SxCP Man Slot` are the scalable per-participant control diff --git a/loop_nodes.py b/loop_nodes.py index a47b020..37e9b87 100644 --- a/loop_nodes.py +++ b/loop_nodes.py @@ -1,5 +1,6 @@ from __future__ import annotations +import random from typing import Any try: @@ -21,6 +22,11 @@ except Exception: MAX_LOOP_VALUES = 20 MAX_CARRY_VALUES = MAX_LOOP_VALUES - 2 COLLECTION_MODES = ["auto_batch", "list", "image_batch", "latent_batch", "string_lines"] +ACCUMULATOR_ACTIONS = ["replace_by_entry_id", "append", "clear_then_append", "clear", "read"] +ACCUMULATOR_IMAGE_BATCH_MODES = ["same_size_only", "resize_to_first"] +ACCUMULATOR_IMAGE_GROUPS = 4 + +_ACCUMULATOR_STORES: dict[str, list[dict[str, Any]]] = {} class AnyType(str): @@ -63,6 +69,91 @@ def _latent_cat(first: Any, second: Any) -> Any | None: return merged +def _torch_cat_many(values: list[Any]) -> Any | None: + if not values: + return None + result = values[0] + for value in values[1:]: + result = _torch_cat(result, value) + if result is None: + return None + return result + + +def _is_image_tensor(value: Any) -> bool: + try: + import torch + except Exception: + return False + return torch.is_tensor(value) and len(value.shape) == 4 + + +def _image_shape(value: Any) -> tuple[int, ...] | None: + if not _is_image_tensor(value): + return None + return tuple(int(part) for part in value.shape[1:]) + + +def _split_image_value(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, (list, tuple)): + images: list[Any] = [] + for item in value: + images.extend(_split_image_value(item)) + return images + if not _is_image_tensor(value): + return [] + if int(value.shape[0]) <= 1: + return [value] + return [value[index : index + 1] for index in range(int(value.shape[0]))] + + +def _resize_image_to_shape(image: Any, shape: tuple[int, ...]) -> Any | None: + if not _is_image_tensor(image): + return None + try: + import comfy.utils + except Exception: + return None + height, width = int(shape[0]), int(shape[1]) + return comfy.utils.common_upscale(image.movedim(-1, 1), width, height, "lanczos", "center").movedim(1, -1) + + +def _image_batch_from_images(images: list[Any], mode: str = "same_size_only") -> Any | None: + if not images: + return None + mode = mode if mode in ACCUMULATOR_IMAGE_BATCH_MODES else "same_size_only" + first_shape = _image_shape(images[0]) + if first_shape is None: + return None + normalized = [] + for image in images: + if _image_shape(image) != first_shape: + if mode != "resize_to_first": + return None + image = _resize_image_to_shape(image, first_shape) + if image is None: + return None + normalized.append(image) + return _torch_cat_many(normalized) + + +def _group_image_batches(images: list[Any]) -> list[Any]: + grouped: dict[tuple[int, ...], list[Any]] = {} + order: list[tuple[int, ...]] = [] + for image in images: + shape = _image_shape(image) + if shape is None: + continue + if shape not in grouped: + grouped[shape] = [] + order.append(shape) + grouped[shape].append(image) + batches = [_torch_cat_many(grouped[shape]) for shape in order] + return [batch for batch in batches if batch is not None] + + def _as_list(collection: Any) -> list[Any]: if collection is None: return [] @@ -321,6 +412,149 @@ class SxCPLoopAppend: return (append_collected_value(collection, value, mode=mode, skip_none=skip_none),) +class SxCPAccumulator: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "store_key": ("STRING", {"default": "", "multiline": False}), + "action": (ACCUMULATOR_ACTIONS, {"default": "replace_by_entry_id"}), + "max_items": ("INT", {"default": 32, "min": 1, "max": 10000, "step": 1}), + "image_batch_mode": (ACCUMULATOR_IMAGE_BATCH_MODES, {"default": "same_size_only"}), + "skip_empty": ("BOOLEAN", {"default": True}), + }, + "optional": { + "image": ("IMAGE",), + "value": (ANY_TYPE,), + "entry_id": (ANY_TYPE,), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + }, + } + + RETURN_TYPES = tuple([ANY_TYPE, "IMAGE", "IMAGE"] + ["IMAGE"] * ACCUMULATOR_IMAGE_GROUPS + ["INT", "STRING"]) + RETURN_NAMES = tuple( + ["collection", "image_batch", "image_list"] + + [f"image_batch_{index}" for index in range(1, ACCUMULATOR_IMAGE_GROUPS + 1)] + + ["count", "status"] + ) + OUTPUT_IS_LIST = tuple([False, False, True] + [False] * ACCUMULATOR_IMAGE_GROUPS + [False, False]) + FUNCTION = "accumulate" + CATEGORY = "prompt_builder/loop" + + @classmethod + def IS_CHANGED(cls, *args, **kwargs): + return random.random() + + def _store_key(self, store_key: str, unique_id: Any) -> str: + key = str(store_key or "").strip() + return key or f"node:{unique_id}" + + def _entry_id(self, entry_id: Any, image_index: int, image_count: int) -> str: + if entry_id is None: + return "" + text = str(entry_id).strip() + if not text: + return "" + if image_count <= 1: + return text + return f"{text}:{image_index + 1}" + + def _value_for_image(self, value: Any, image_index: int, image_count: int) -> Any: + if image_count <= 1: + return value + if isinstance(value, (list, tuple)) and len(value) == image_count: + return value[image_index] + return value + + def _entry_records(self, image: Any, value: Any, entry_id: Any, skip_empty: bool) -> list[dict[str, Any]]: + images = _split_image_value(image) + if not images: + if value is None and skip_empty: + return [] + return [{"id": self._entry_id(entry_id, 0, 1), "image": None, "value": value}] + image_count = len(images) + return [ + { + "id": self._entry_id(entry_id, index, image_count), + "image": image_item, + "value": self._value_for_image(value, index, image_count), + } + for index, image_item in enumerate(images) + ] + + def _append_or_replace(self, store: list[dict[str, Any]], entries: list[dict[str, Any]], action: str) -> None: + replace = action == "replace_by_entry_id" + for entry in entries: + entry_id = entry.get("id") or "" + if replace and entry_id: + for index, existing in enumerate(store): + if existing.get("id") == entry_id: + store[index] = entry + break + else: + store.append(entry) + else: + store.append(entry) + + def _collection(self, store: list[dict[str, Any]]) -> list[Any]: + collection = [] + for entry in store: + value = entry.get("value") + collection.append(value if value is not None else entry.get("image")) + return collection + + def _status(self, key: str, store: list[dict[str, Any]], image_batch: Any, image_batches: list[Any]) -> str: + images = [entry.get("image") for entry in store if entry.get("image") is not None] + shapes = [] + for image in images: + shape = _image_shape(image) + if shape is not None and shape not in shapes: + shapes.append(shape) + shape_text = ", ".join(f"{shape[1]}x{shape[0]}" for shape in shapes) or "no images" + batch_state = "all images batched" if image_batch is not None else "mixed sizes or no image batch" + return ( + f"key={key}; entries={len(store)}; image_entries={len(images)}; " + f"formats={shape_text}; grouped_batches={len(image_batches)}; {batch_state}" + ) + + def accumulate( + self, + store_key, + action, + max_items, + image_batch_mode, + skip_empty, + image=None, + value=None, + entry_id=None, + unique_id=None, + ): + key = self._store_key(store_key, unique_id) + action = action if action in ACCUMULATOR_ACTIONS else "replace_by_entry_id" + store = _ACCUMULATOR_STORES.setdefault(key, []) + + if action in ("clear", "clear_then_append"): + store.clear() + + if action not in ("clear", "read"): + entries = self._entry_records(image, value, entry_id, bool(skip_empty)) + self._append_or_replace(store, entries, action) + + max_items = max(1, int(max_items)) + if len(store) > max_items: + del store[: len(store) - max_items] + + images = [entry["image"] for entry in store if entry.get("image") is not None] + image_batch = _image_batch_from_images(images, image_batch_mode) + image_batches = _group_image_batches(images) + grouped_outputs = image_batches[:ACCUMULATOR_IMAGE_GROUPS] + grouped_outputs += [None] * (ACCUMULATOR_IMAGE_GROUPS - len(grouped_outputs)) + status = self._status(key, store, image_batch, image_batches) + return tuple([self._collection(store), image_batch, images] + grouped_outputs + [len(store), status]) + + class SxCPForLoopEnd: @classmethod def INPUT_TYPES(cls): @@ -444,6 +678,7 @@ LOOP_NODE_CLASS_MAPPINGS = { "SxCPForLoopStart": SxCPForLoopStart, "SxCPForLoopEnd": SxCPForLoopEnd, "SxCPLoopAppend": SxCPLoopAppend, + "SxCPAccumulator": SxCPAccumulator, "SxCPLoopIntAdd": SxCPLoopIntAdd, "SxCPLoopLessThan": SxCPLoopLessThan, "SxCPLoopLessThanOrEqual": SxCPLoopLessThanOrEqual, @@ -455,6 +690,7 @@ LOOP_NODE_DISPLAY_NAME_MAPPINGS = { "SxCPForLoopStart": "SxCP For Loop Start", "SxCPForLoopEnd": "SxCP For Loop End", "SxCPLoopAppend": "SxCP Loop Append", + "SxCPAccumulator": "SxCP Accumulator", "SxCPLoopIntAdd": "SxCP Loop Int Add", "SxCPLoopLessThan": "SxCP Loop Less Than", "SxCPLoopLessThanOrEqual": "SxCP Loop Less Than Or Equal",