Add loop accumulator node

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",