Class BatchImageLoopOpen

flow_control.py:8–121

@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}),
                "previous_image": ("IMAGE",),  # 新增：接收上一次循环的图片
            }
        }
        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是单张图片，会扩展到与masks相同的批次大小
        """
        # 处理masks（先处理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}"
        
        # 处理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}"

        # 检查是否需要扩展images
        if images.shape[0] == 1 and masks.shape[0] > 1:
            print(f"Expanding single image to match mask batch size: {masks.shape[0]}")
            images = images.expand(masks.shape[0], -1, -1, -1)
        
        # 确保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 resize_to_match(self, image, target_shape):