Add Block Loop (Recursive) nodes

2025-12-06 22:46:32 +09:00
parent 4b7245f8a2
commit fedfd2a998
4 changed files with 843 additions and 0 deletions
@@ -0,0 +1,114 @@
+
+import json
+import uuid
+from server import PromptServer
+from aiohttp import web
+
+class BlockLoopStart:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "first_value": ("INT", {"default": 1}),
+                "total_loops": ("INT", {"default": 3, "min": 1, "max": 100}),
+            },
+            "optional": {
+                "current_loop": ("INT", {"default": 0, "hidden": True}), # Injectable
+                "recycle_value": ("INT", {"default": 0, "hidden": True}), # Injectable from previous loop
+            },
+            "hidden": {
+                "unique_id": "UNIQUE_ID",
+            }
+        }
+
+    RETURN_TYPES = ("INT", "INT", "INT", "STRING")
+    RETURN_NAMES = ("current_value", "index", "total", "loop_id")
+    FUNCTION = "execute"
+    CATEGORY = "Loop"
+
+    def execute(self, first_value, total_loops, current_loop=0, recycle_value=0, unique_id=None):
+        # On first run (current_loop=0), use first_value.
+        # On subsequent runs, use recycle_value.
+        if current_loop == 0:
+            val = first_value
+        else:
+            val = recycle_value
+            
+        return (val, current_loop, total_loops, unique_id)
+
+
+class BlockLoopEnd:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "flow_control": ("*",), # Anything to enforce order
+                "next_value": ("INT",), # Value to pass to next iteration
+                "total_loops": ("INT", {"forceInput": True}),
+                "current_loop": ("INT", {"forceInput": True}),
+                "loop_id": ("STRING", {"forceInput": True}), # Link to Start Node
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+                "extra_pnginfo": "EXTRA_PNGINFO",
+            }
+        }
+
+    RETURN_TYPES = ("*",)
+    RETURN_NAMES = ("flow_output",)
+    FUNCTION = "execute"
+    CATEGORY = "Loop"
+    OUTPUT_NODE = True
+
+    def execute(self, flow_control, next_value, total_loops, current_loop, loop_id, prompt=None, extra_pnginfo=None):
+        if current_loop + 1 < total_loops:
+            # Prepare next iteration
+            new_loop_index = current_loop + 1
+            print(f"Looping: {current_loop} -> {new_loop_index} / {total_loops}")
+            
+            # We need to find the specific BlockLoopStart node in the prompt and update its inputs.
+            # loop_id is the UNIQUE_ID of the BlockLoopStart node.
+            
+            new_prompt = prompt.copy()
+            target_node = new_prompt.get(str(loop_id))
+            
+            if target_node:
+                # Update inputs for the next run
+                target_node['inputs']['current_loop'] = new_loop_index
+                target_node['inputs']['recycle_value'] = next_value
+                
+                # Re-queue
+                new_prompt_id = str(uuid.uuid4())
+                
+                async def queue_new_prompt():
+                    import execution
+                    server = PromptServer.instance
+                    
+                    # Validate
+                    valid = await execution.validate_prompt(new_prompt_id, new_prompt)
+                    
+                    if valid[0]:
+                        outputs_to_execute = valid[2]
+                        # We need to get the next execution number
+                        server.number += 1
+                        number = server.number
+                        
+                        # Extra data
+                        new_extra_data = extra_pnginfo.copy() if extra_pnginfo else {}
+                        
+                        # Put in queue
+                        server.prompt_queue.put((number, new_prompt_id, new_prompt, new_extra_data, outputs_to_execute, {}))
+                        print(f"Loop queued: {new_prompt_id}")
+                    else:
+                        print(f"Loop validation failed: {valid[1]}")
+
+                # Schedule the task
+                PromptServer.instance.loop.create_task(queue_new_prompt())
+        
+        return (flow_control,)