561 lines
21 KiB
Python
561 lines
21 KiB
Python
from comfy_execution.graph_utils import GraphBuilder, is_link
|
|
from .tools import VariantSupport
|
|
import torch
|
|
from nodes import NODE_CLASS_MAPPINGS as ALL_NODE_CLASS_MAPPINGS
|
|
|
|
@VariantSupport()
|
|
class BatchImageLoopOpen:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"segmented_images": ("IMAGE", {"forceInput": True}),
|
|
"segmented_masks": ("MASK", {"forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"iteration_count": ("INT", {"default": 0}),
|
|
}
|
|
}
|
|
return inputs
|
|
|
|
RETURN_TYPES = tuple(["FLOW_CONTROL", "IMAGE", "MASK", "INT", "INT"])
|
|
RETURN_NAMES = tuple(["FLOW_CONTROL", "current_image", "current_mask", "max_iterations", "iteration_count"])
|
|
FUNCTION = "while_loop_open"
|
|
CATEGORY = "CyberEveLoop🐰"
|
|
|
|
def standardize_input(self, images, masks):
|
|
"""
|
|
标准化输入格式
|
|
images: 确保是4D tensor [B,H,W,C]
|
|
masks: 确保是3D tensor [B,H,W]
|
|
"""
|
|
# 处理images
|
|
if isinstance(images, list):
|
|
images = torch.cat(images, dim=0)
|
|
if len(images.shape) == 3: # [H,W,C] -> [1,H,W,C]
|
|
images = images.unsqueeze(0)
|
|
assert len(images.shape) == 4, f"Images must be 4D [B,H,W,C], got shape {images.shape}"
|
|
|
|
# 处理masks
|
|
if isinstance(masks, list):
|
|
masks = torch.cat(masks, dim=0)
|
|
if len(masks.shape) == 2: # [H,W] -> [1,H,W]
|
|
masks = masks.unsqueeze(0)
|
|
assert len(masks.shape) == 3, f"Masks must be 3D [B,H,W], got shape {masks.shape}"
|
|
|
|
# 确保batch维度相同
|
|
assert images.shape[0] == masks.shape[0], \
|
|
f"Batch size mismatch: images {images.shape[0]} vs masks {masks.shape[0]}"
|
|
|
|
return images, masks
|
|
|
|
|
|
def while_loop_open(self, segmented_images, segmented_masks, unique_id=None, iteration_count=0):
|
|
print(f"while_loop_open Processing iteration {iteration_count}")
|
|
|
|
# 标准化输入
|
|
segmented_images, segmented_masks = self.standardize_input(segmented_images, segmented_masks)
|
|
|
|
# 获取最大迭代次数
|
|
max_iterations = segmented_images.shape[0]
|
|
if max_iterations == 0:
|
|
raise ValueError("No images provided in segmented_images")
|
|
|
|
# 验证迭代计数（修改这里）
|
|
if iteration_count >= max_iterations:
|
|
raise ValueError(f"Iteration count {iteration_count} exceeds max iterations {max_iterations}")
|
|
|
|
# 获取当前迭代的图片和蒙版
|
|
current_image = segmented_images[iteration_count:iteration_count+1]
|
|
current_mask = segmented_masks[iteration_count:iteration_count+1]
|
|
|
|
return tuple(["stub", current_image, current_mask, max_iterations, iteration_count])
|
|
|
|
|
|
@VariantSupport()
|
|
class BatchImageLoopClose:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"flow_control": ("FLOW_CONTROL", {"rawLink": True}),
|
|
"current_image": ("IMAGE",),
|
|
"current_mask": ("MASK",),
|
|
"max_iterations": ("INT", {"forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"dynprompt": "DYNPROMPT",
|
|
"unique_id": "UNIQUE_ID",
|
|
"result_images": ("IMAGE",),
|
|
"result_masks": ("MASK",),
|
|
"iteration_count": ("INT", {"default": 0}),
|
|
}
|
|
}
|
|
return inputs
|
|
|
|
RETURN_TYPES = tuple(["IMAGE", "MASK"])
|
|
RETURN_NAMES = tuple(["result_images", "result_masks"])
|
|
FUNCTION = "while_loop_close"
|
|
CATEGORY = "CyberEveLoop🐰"
|
|
|
|
def explore_dependencies(self, node_id, dynprompt, upstream, parent_ids):
|
|
node_info = dynprompt.get_node(node_id)
|
|
if "inputs" not in node_info:
|
|
return
|
|
|
|
for k, v in node_info["inputs"].items():
|
|
if is_link(v):
|
|
parent_id = v[0]
|
|
display_id = dynprompt.get_display_node_id(parent_id)
|
|
display_node = dynprompt.get_node(display_id)
|
|
class_type = display_node["class_type"]
|
|
# 排除循环结束节点
|
|
if class_type not in ['BatchImageLoopClose']:
|
|
parent_ids.append(display_id)
|
|
if parent_id not in upstream:
|
|
upstream[parent_id] = []
|
|
self.explore_dependencies(parent_id, dynprompt, upstream, parent_ids)
|
|
upstream[parent_id].append(node_id)
|
|
|
|
def explore_output_nodes(self, dynprompt, upstream, output_nodes, parent_ids):
|
|
"""探索并添加输出节点的连接"""
|
|
for parent_id in upstream:
|
|
display_id = dynprompt.get_display_node_id(parent_id)
|
|
for output_id in output_nodes:
|
|
id = output_nodes[output_id][0]
|
|
if id in parent_ids and display_id == id and output_id not in upstream[parent_id]:
|
|
if '.' in parent_id:
|
|
arr = parent_id.split('.')
|
|
arr[len(arr)-1] = output_id
|
|
upstream[parent_id].append('.'.join(arr))
|
|
else:
|
|
upstream[parent_id].append(output_id)
|
|
|
|
def collect_contained(self, node_id, upstream, contained):
|
|
if node_id not in upstream:
|
|
return
|
|
for child_id in upstream[node_id]:
|
|
if child_id not in contained:
|
|
contained[child_id] = True
|
|
self.collect_contained(child_id, upstream, contained)
|
|
|
|
|
|
def standardize_input(self, image, mask):
|
|
"""
|
|
标准化输入格式
|
|
image: 确保是4D tensor [B,H,W,C]
|
|
mask: 确保是3D tensor [B,H,W]
|
|
"""
|
|
# 处理image
|
|
if len(image.shape) == 3: # [H,W,C] -> [1,H,W,C]
|
|
image = image.unsqueeze(0)
|
|
assert len(image.shape) == 4, f"Image must be 4D [B,H,W,C], got shape {image.shape}"
|
|
|
|
# 处理mask
|
|
if len(mask.shape) == 2: # [H,W] -> [1,H,W]
|
|
mask = mask.unsqueeze(0)
|
|
assert len(mask.shape) == 3, f"Mask must be 3D [B,H,W], got shape {mask.shape}"
|
|
|
|
return image, mask
|
|
|
|
|
|
def initialize_results(self, max_iterations, current_image, current_mask):
|
|
"""
|
|
初始化结果张量，确保与MaskSplit输出格式一致
|
|
"""
|
|
# 确保维度正确
|
|
assert len(current_image.shape) == 4, "Current image must be 4D [B,H,W,C]"
|
|
assert len(current_mask.shape) == 3, "Current mask must be 3D [B,H,W]"
|
|
|
|
# 创建结果张量，确保格式一致
|
|
result_images = torch.zeros(
|
|
(max_iterations, current_image.shape[1], current_image.shape[2], current_image.shape[3]),
|
|
dtype=current_image.dtype,
|
|
device=current_image.device
|
|
) # 明确指定 [B,H,W,C]
|
|
|
|
result_masks = torch.zeros(
|
|
(max_iterations, current_mask.shape[1], current_mask.shape[2]),
|
|
dtype=current_mask.dtype,
|
|
device=current_mask.device
|
|
) # 明确指定 [B,H,W]
|
|
|
|
return result_images, result_masks
|
|
|
|
|
|
def while_loop_close(self, flow_control, current_image, current_mask, max_iterations,
|
|
iteration_count=0, result_images=None, result_masks=None,
|
|
dynprompt=None, unique_id=None,):
|
|
print(f"Iteration {iteration_count} of {max_iterations}")
|
|
|
|
# 标准化输入，确保格式一致
|
|
current_image, current_mask = self.standardize_input(current_image, current_mask)
|
|
|
|
# 验证迭代计数（修改这里）
|
|
if iteration_count >= max_iterations:
|
|
raise ValueError(f"Iteration count {iteration_count} exceeds max iterations {max_iterations}")
|
|
|
|
# 结果初始化或验证
|
|
if result_images is None or result_masks is None:
|
|
result_images, result_masks = self.initialize_results(max_iterations, current_image, current_mask)
|
|
else:
|
|
# 验证现有结果的维度和格式
|
|
assert result_images.shape[0] == max_iterations and len(result_images.shape) == 4, \
|
|
f"Result images must be 4D [B,H,W,C] with batch size {max_iterations}"
|
|
assert result_masks.shape[0] == max_iterations and len(result_masks.shape) == 3, \
|
|
f"Result masks must be 3D [B,H,W] with batch size {max_iterations}"
|
|
|
|
# 存储当前结果
|
|
result_images[iteration_count:iteration_count+1] = current_image
|
|
result_masks[iteration_count:iteration_count+1] = current_mask
|
|
|
|
# 检查是否继续循环（修改这里）
|
|
if iteration_count == max_iterations - 1:
|
|
print(f"Loop finished with {iteration_count + 1} iterations")
|
|
return (result_images, result_masks)
|
|
|
|
# 准备下一次循环
|
|
this_node = dynprompt.get_node(unique_id)
|
|
upstream = {}
|
|
parent_ids = []
|
|
self.explore_dependencies(unique_id, dynprompt, upstream, parent_ids)
|
|
parent_ids = list(set(parent_ids)) # 去重
|
|
|
|
# 获取并处理输出节点
|
|
prompts = dynprompt.get_original_prompt()
|
|
output_nodes = {}
|
|
for id in prompts:
|
|
node = prompts[id]
|
|
if "inputs" not in node:
|
|
continue
|
|
class_type = node["class_type"]
|
|
if class_type in ALL_NODE_CLASS_MAPPINGS:
|
|
class_def = ALL_NODE_CLASS_MAPPINGS[class_type]
|
|
if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True:
|
|
for k, v in node['inputs'].items():
|
|
if is_link(v):
|
|
output_nodes[id] = v
|
|
|
|
# 创建新图
|
|
graph = GraphBuilder()
|
|
self.explore_output_nodes(dynprompt, upstream, output_nodes, parent_ids)
|
|
|
|
contained = {}
|
|
open_node = flow_control[0]
|
|
self.collect_contained(open_node, upstream, contained)
|
|
contained[unique_id] = True
|
|
contained[open_node] = True
|
|
|
|
# 创建节点
|
|
for node_id in contained:
|
|
original_node = dynprompt.get_node(node_id)
|
|
node = graph.node(original_node["class_type"],
|
|
"Recurse" if node_id == unique_id else node_id)
|
|
node.set_override_display_id(node_id)
|
|
|
|
# 设置连接
|
|
for node_id in contained:
|
|
original_node = dynprompt.get_node(node_id)
|
|
node = graph.lookup_node("Recurse" if node_id == unique_id else node_id)
|
|
for k, v in original_node["inputs"].items():
|
|
if is_link(v) and v[0] in contained:
|
|
parent = graph.lookup_node(v[0])
|
|
node.set_input(k, parent.out(v[1]))
|
|
else:
|
|
node.set_input(k, v)
|
|
|
|
# 设置节点参数
|
|
my_clone = graph.lookup_node("Recurse")
|
|
my_clone.set_input("iteration_count", iteration_count + 1)
|
|
my_clone.set_input("result_images", result_images)
|
|
my_clone.set_input("result_masks", result_masks)
|
|
|
|
new_open = graph.lookup_node(open_node)
|
|
new_open.set_input("iteration_count", iteration_count + 1)
|
|
|
|
print(f"Continuing to iteration {iteration_count + 1}")
|
|
|
|
return {
|
|
"result": tuple([my_clone.out(0), my_clone.out(1)]),
|
|
"expand": graph.finalize(),
|
|
}
|
|
|
|
|
|
@VariantSupport()
|
|
class SingleImageLoopOpen:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"max_iterations": ("INT", {"default": 5, "min": 1, "max": 100}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"iteration_count": ("INT", {"default": 0}),
|
|
"previous_image": ("IMAGE",),
|
|
"previous_mask": ("MASK",),
|
|
}
|
|
}
|
|
return inputs
|
|
|
|
RETURN_TYPES = tuple(["FLOW_CONTROL", "IMAGE", "MASK", "INT", "INT"])
|
|
RETURN_NAMES = tuple(["FLOW_CONTROL", "current_image", "current_mask", "max_iterations", "iteration_count"])
|
|
FUNCTION = "loop_open"
|
|
CATEGORY = "CyberEveLoop🐰"
|
|
|
|
def loop_open(self, image, max_iterations, mask=None, unique_id=None,
|
|
iteration_count=0, previous_image=None, previous_mask=None):
|
|
print(f"SingleImageLoopOpen Processing iteration {iteration_count}")
|
|
|
|
# 确保维度正确
|
|
if len(image.shape) == 3:
|
|
image = image.unsqueeze(0)
|
|
if mask is not None and len(mask.shape) == 2:
|
|
mask = mask.unsqueeze(0)
|
|
|
|
# 使用上一次循环的结果（如果有）
|
|
current_image = previous_image if previous_image is not None and iteration_count > 0 else image
|
|
current_mask = previous_mask if previous_mask is not None and iteration_count > 0 else mask
|
|
|
|
return tuple(["stub", current_image, current_mask, max_iterations, iteration_count])
|
|
|
|
@VariantSupport()
|
|
class SingleImageLoopClose:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"flow_control": ("FLOW_CONTROL", {"rawLink": True}),
|
|
"current_image": ("IMAGE",),
|
|
"max_iterations": ("INT", {"forceInput": True}),
|
|
},
|
|
"optional": {
|
|
"current_mask": ("MASK",),
|
|
},
|
|
"hidden": {
|
|
"dynprompt": "DYNPROMPT",
|
|
"unique_id": "UNIQUE_ID",
|
|
"iteration_count": ("INT", {"default": 0}),
|
|
}
|
|
}
|
|
return inputs
|
|
|
|
RETURN_TYPES = tuple(["IMAGE", "MASK"])
|
|
RETURN_NAMES = tuple(["final_image", "final_mask"])
|
|
FUNCTION = "loop_close"
|
|
CATEGORY = "CyberEveLoop🐰"
|
|
|
|
def explore_dependencies(self, node_id, dynprompt, upstream, parent_ids):
|
|
node_info = dynprompt.get_node(node_id)
|
|
if "inputs" not in node_info:
|
|
return
|
|
|
|
for k, v in node_info["inputs"].items():
|
|
if is_link(v):
|
|
parent_id = v[0]
|
|
display_id = dynprompt.get_display_node_id(parent_id)
|
|
display_node = dynprompt.get_node(display_id)
|
|
class_type = display_node["class_type"]
|
|
if class_type not in ['SingleImageLoopClose']:
|
|
parent_ids.append(display_id)
|
|
if parent_id not in upstream:
|
|
upstream[parent_id] = []
|
|
self.explore_dependencies(parent_id, dynprompt, upstream, parent_ids)
|
|
upstream[parent_id].append(node_id)
|
|
|
|
def explore_output_nodes(self, dynprompt, upstream, output_nodes, parent_ids):
|
|
for parent_id in upstream:
|
|
display_id = dynprompt.get_display_node_id(parent_id)
|
|
for output_id in output_nodes:
|
|
id = output_nodes[output_id][0]
|
|
if id in parent_ids and display_id == id and output_id not in upstream[parent_id]:
|
|
if '.' in parent_id:
|
|
arr = parent_id.split('.')
|
|
arr[len(arr)-1] = output_id
|
|
upstream[parent_id].append('.'.join(arr))
|
|
else:
|
|
upstream[parent_id].append(output_id)
|
|
|
|
def collect_contained(self, node_id, upstream, contained):
|
|
if node_id not in upstream:
|
|
return
|
|
for child_id in upstream[node_id]:
|
|
if child_id not in contained:
|
|
contained[child_id] = True
|
|
self.collect_contained(child_id, upstream, contained)
|
|
|
|
def loop_close(self, flow_control, current_image, max_iterations, current_mask=None,
|
|
iteration_count=0, dynprompt=None, unique_id=None):
|
|
print(f"Iteration {iteration_count} of {max_iterations}")
|
|
|
|
# 维度处理
|
|
if len(current_image.shape) == 3:
|
|
current_image = current_image.unsqueeze(0)
|
|
if current_mask is not None and len(current_mask.shape) == 2:
|
|
current_mask = current_mask.unsqueeze(0)
|
|
|
|
# 检查是否继续循环
|
|
if iteration_count >= max_iterations - 1:
|
|
print(f"Loop finished with {iteration_count + 1} iterations")
|
|
return (current_image, current_mask if current_mask is not None else torch.zeros_like(current_image[:,:,:,0]))
|
|
|
|
# 准备下一次循环
|
|
this_node = dynprompt.get_node(unique_id)
|
|
upstream = {}
|
|
parent_ids = []
|
|
self.explore_dependencies(unique_id, dynprompt, upstream, parent_ids)
|
|
parent_ids = list(set(parent_ids))
|
|
|
|
# 获取并处理输出节点
|
|
prompts = dynprompt.get_original_prompt()
|
|
output_nodes = {}
|
|
for id in prompts:
|
|
node = prompts[id]
|
|
if "inputs" not in node:
|
|
continue
|
|
class_type = node["class_type"]
|
|
if class_type in ALL_NODE_CLASS_MAPPINGS:
|
|
class_def = ALL_NODE_CLASS_MAPPINGS[class_type]
|
|
if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True:
|
|
for k, v in node['inputs'].items():
|
|
if is_link(v):
|
|
output_nodes[id] = v
|
|
|
|
# 创建新图
|
|
graph = GraphBuilder()
|
|
self.explore_output_nodes(dynprompt, upstream, output_nodes, parent_ids)
|
|
|
|
contained = {}
|
|
open_node = flow_control[0]
|
|
self.collect_contained(open_node, upstream, contained)
|
|
contained[unique_id] = True
|
|
contained[open_node] = True
|
|
|
|
# 创建节点
|
|
for node_id in contained:
|
|
original_node = dynprompt.get_node(node_id)
|
|
node = graph.node(original_node["class_type"],
|
|
"Recurse" if node_id == unique_id else node_id)
|
|
node.set_override_display_id(node_id)
|
|
|
|
# 设置连接
|
|
for node_id in contained:
|
|
original_node = dynprompt.get_node(node_id)
|
|
node = graph.lookup_node("Recurse" if node_id == unique_id else node_id)
|
|
for k, v in original_node["inputs"].items():
|
|
if is_link(v) and v[0] in contained:
|
|
parent = graph.lookup_node(v[0])
|
|
node.set_input(k, parent.out(v[1]))
|
|
else:
|
|
node.set_input(k, v)
|
|
|
|
# 设置节点参数
|
|
my_clone = graph.lookup_node("Recurse")
|
|
my_clone.set_input("iteration_count", iteration_count + 1)
|
|
|
|
new_open = graph.lookup_node(open_node)
|
|
new_open.set_input("iteration_count", iteration_count + 1)
|
|
new_open.set_input("previous_image", current_image)
|
|
if current_mask is not None:
|
|
new_open.set_input("previous_mask", current_mask)
|
|
|
|
print(f"Continuing to iteration {iteration_count + 1}")
|
|
|
|
return {
|
|
"result": tuple([my_clone.out(0), my_clone.out(1)]),
|
|
"expand": graph.finalize(),
|
|
}
|
|
|
|
|
|
|
|
@VariantSupport()
|
|
class LoopIndexSwitch:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
预定义100个隐藏的lazy输入
|
|
"""
|
|
optional_inputs = {
|
|
"default_value": ("*", {"lazy": True}), # 默认值也设为lazy
|
|
}
|
|
# 添加100个隐藏的lazy输入
|
|
hidden_inputs = {}
|
|
for i in range(100):
|
|
hidden_inputs[f"while_{i}"] = ("*", {"lazy": True})
|
|
|
|
return {
|
|
"required": {
|
|
"iteration_count": ("INT", {"forceInput": True}), # 当前迭代次数
|
|
},
|
|
"optional": optional_inputs,
|
|
"hidden": hidden_inputs,
|
|
}
|
|
|
|
RETURN_TYPES = ("*",)
|
|
FUNCTION = "index_switch"
|
|
CATEGORY = "CyberEveLoop🐰"
|
|
|
|
def check_lazy_status(self, iteration_count, **kwargs):
|
|
"""
|
|
检查当前迭代需要的输入和默认值
|
|
"""
|
|
needed = []
|
|
current_key = f"while_{iteration_count}"
|
|
|
|
# 检查当前迭代的输入
|
|
if current_key in kwargs :
|
|
needed.append(current_key)
|
|
else:
|
|
needed.append("default_value")
|
|
|
|
print(f"Index switch needed: {needed}")
|
|
return needed
|
|
|
|
|
|
|
|
def index_switch(self, iteration_count, **kwargs):
|
|
"""
|
|
根据当前迭代次数选择对应的输入值
|
|
"""
|
|
current_key = f"while_{iteration_count}"
|
|
|
|
if current_key in kwargs and kwargs[current_key] is not None:
|
|
return (kwargs[current_key],)
|
|
return (kwargs.get("default_value"),)
|
|
|
|
|
|
CyberEve_Loop_CLASS_MAPPINGS = {
|
|
"CyberEve_BatchImageLoopOpen": BatchImageLoopOpen,
|
|
"CyberEve_BatchImageLoopClose": BatchImageLoopClose,
|
|
"CyberEve_LoopIndexSwitch": LoopIndexSwitch,
|
|
"CyberEve_SingleImageLoopOpen": SingleImageLoopOpen,
|
|
"CyberEve_SingleImageLoopClose": SingleImageLoopClose,
|
|
}
|
|
|
|
CyberEve_Loop_DISPLAY_NAME_MAPPINGS = {
|
|
"CyberEve_BatchImageLoopOpen": "Batch Image Loop Open🐰",
|
|
"CyberEve_BatchImageLoopClose": "Batch Image Loop Close🐰",
|
|
"CyberEve_LoopIndexSwitch": "Loop Index Switch🐰",
|
|
"CyberEve_SingleImageLoopOpen": "Single Image Loop Open🐰",
|
|
"CyberEve_SingleImageLoopClose": "Single Image Loop Close🐰",
|
|
} |