From 5661d1c63a7a64df73ebffd03e153908179781b9 Mon Sep 17 00:00:00 2001
From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com>
Date: Tue, 6 May 2025 15:32:54 -0400
Subject: [PATCH] adds batch complete bool & execution blocker node

---
 .../realtimenodes/controls/accumulator.py     | 28 +++++++--------
 .../controls/state_management_controls.py     | 36 +++++++++++++++++++
 2 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/node_wrappers/realtimenodes/controls/accumulator.py b/node_wrappers/realtimenodes/controls/accumulator.py
index 2e252f7..e3ff5da 100644
--- a/node_wrappers/realtimenodes/controls/accumulator.py
+++ b/node_wrappers/realtimenodes/controls/accumulator.py
@@ -17,6 +17,8 @@ class BaseAccumulatorNode(ControlNodeBase):
     
     # These should be overridden by subclasses
     NODE_TYPE_NAME = "Base"
+    RETURN_TYPES = (AlwaysEqualProxy("*"), "BOOLEAN")
+    RETURN_NAMES = ("accumulated_values", "batch_complete")
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -55,8 +57,8 @@ class AnyAccumulatorNode(BaseAccumulatorNode):
     Node that accumulates inputs of any type over multiple workflow runs.
     """
     
-    RETURN_TYPES = (AlwaysEqualProxy("*"),)
-    RETURN_NAMES = ("accumulated_values",)
+    RETURN_TYPES = (AlwaysEqualProxy("*"), "BOOLEAN")
+    RETURN_NAMES = ("accumulated_values", "batch_complete")
     NODE_TYPE_NAME = "Any"
     
     @classmethod
@@ -115,7 +117,7 @@ class AnyAccumulatorNode(BaseAccumulatorNode):
         state[state_key] = accumulated
         self.set_state(state, unique_id)
         
-        return (output_values,)
+        return (output_values, batch_complete)
 
 
 class ImageAccumulatorNode(BaseAccumulatorNode):
@@ -123,8 +125,8 @@ class ImageAccumulatorNode(BaseAccumulatorNode):
     Node that accumulates image inputs (BHWC format) over multiple workflow runs.
     """
     
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("image_batch",)
+    RETURN_TYPES = ("IMAGE", "BOOLEAN")
+    RETURN_NAMES = ("image_batch", "batch_complete")
     NODE_TYPE_NAME = "Image"
     
     @classmethod
@@ -219,7 +221,7 @@ class ImageAccumulatorNode(BaseAccumulatorNode):
         state[f"{state_key}_frames"] = frames
         self.set_state(state, unique_id)
         
-        return (output_tensor,)
+        return (output_tensor, batch_complete)
 
 
 class MaskAccumulatorNode(BaseAccumulatorNode):
@@ -227,8 +229,8 @@ class MaskAccumulatorNode(BaseAccumulatorNode):
     Node that accumulates mask inputs over multiple workflow runs.
     """
     
-    RETURN_TYPES = ("MASK",)
-    RETURN_NAMES = ("mask_batch",)
+    RETURN_TYPES = ("MASK", "BOOLEAN")
+    RETURN_NAMES = ("mask_batch", "batch_complete")
     NODE_TYPE_NAME = "Mask"
     
     @classmethod
@@ -322,7 +324,7 @@ class MaskAccumulatorNode(BaseAccumulatorNode):
         state[f"{state_key}_frames"] = frames
         self.set_state(state, unique_id)
         
-        return (output_tensor,)
+        return (output_tensor, batch_complete)
 
 
 class LatentAccumulatorNode(BaseAccumulatorNode):
@@ -330,8 +332,8 @@ class LatentAccumulatorNode(BaseAccumulatorNode):
     Node that accumulates latent inputs over multiple workflow runs.
     """
     
-    RETURN_TYPES = ("LATENT",)
-    RETURN_NAMES = ("latent_batch",)
+    RETURN_TYPES = ("LATENT", "BOOLEAN")
+    RETURN_NAMES = ("latent_batch", "batch_complete")
     NODE_TYPE_NAME = "Latent"
     
     @classmethod
@@ -431,6 +433,4 @@ class LatentAccumulatorNode(BaseAccumulatorNode):
         self.set_state(state, unique_id)
         
         # Return latent dict with samples
-        return ({"samples": output_samples, **metadata},)
-
-
+        return ({"samples": output_samples, **metadata}, batch_complete)
diff --git a/node_wrappers/realtimenodes/controls/state_management_controls.py b/node_wrappers/realtimenodes/controls/state_management_controls.py
index d56ebe7..8b167e2 100644
--- a/node_wrappers/realtimenodes/controls/state_management_controls.py
+++ b/node_wrappers/realtimenodes/controls/state_management_controls.py
@@ -7,6 +7,8 @@ import math
 from ....src.realtimenodes.control_base import ControlNodeBase
 from ....src.utils.general import AlwaysEqualProxy
 from ..utils.temporal_processing import TemporalNetV2Preprocessor
+from comfy_execution.graph import ExecutionBlocker
+
 
 class StateResetNode(ControlNodeBase):
     """Node that resets all control node states when triggered"""
@@ -164,3 +166,37 @@ class SetStateNode(ControlNodeBase):
         return (value,)
 
 
+class ExecutionBlockerNode(ControlNodeBase):
+    """
+    Node that returns an ExecutionBlocker when triggered.
+    """
+    
+    CATEGORY = "Realtime Nodes/control"
+    FUNCTION = "update"
+    DESCRIPTION = "Blocks execution when triggered"
+    
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "trigger": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": "When True, blocks execution"
+                }),
+                "message": ("STRING", {
+                    "default": "",
+                    "multiline": True,
+                    "tooltip": "Optional message to display when blocking"
+                })
+            }
+        }
+    
+    RETURN_TYPES = (AlwaysEqualProxy("*"),)
+    RETURN_NAMES = ("output",)
+    
+    def update(self, trigger, message, always_execute=True, unique_id=None):
+        if trigger:
+            return (ExecutionBlocker(message if message else None),)
+        return (None,)
+
+