115 lines
4.1 KiB
Python
115 lines
4.1 KiB
Python
|
|
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,)
|