nodes/WainWong-ComfyUI-Loop-image
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Normalize some data types

Browse Source
This commit is contained in:
WainWong
2025-01-24 22:48:31 +08:00
parent 0c16767879
commit 1f87e25432
3 changed files with 135 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
flow_control.py
+97 -21
View File
@@ -27,24 +27,54 @@ class BatchImageLoopOpen:
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}")
# 确保输入是张量
if isinstance(segmented_images, list):
segmented_images = torch.cat(segmented_images, dim=0)
if isinstance(segmented_masks, list):
segmented_masks = torch.cat(segmented_masks, dim=0)
# 标准化输入
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:
@@ -116,32 +146,78 @@ class BatchImageLoopClose:
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}")
# 维度处理
if len(current_image.shape) == 3:
current_image = current_image.unsqueeze(0)
if len(current_mask.shape) == 2:
current_mask = current_mask.unsqueeze(0)
# 标准化输入，确保格式一致
current_image, current_mask = self.standardize_input(current_image, current_mask)
# 结果初始化
if result_images is None:
result_images = torch.zeros((max_iterations,) + current_image.shape[1:],
dtype=current_image.dtype,
device=current_image.device)
result_masks = torch.zeros((max_iterations,) + current_mask.shape[1:],
dtype=current_mask.dtype,
device=current_mask.device)
# 验证迭代计数（修改这里）
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:
# 检查是否继续循环（修改这里）
if iteration_count == max_iterations - 1:
print(f"Loop finished with {iteration_count + 1} iterations")
return (result_images, result_masks)
mask_split.py
+37 -11
View File
@@ -155,11 +155,40 @@ class MaskMerge:
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("merged_image",)
# INPUT_IS_LIST = True
FUNCTION = "merge_masked_images"
CATEGORY = "CyberEveLoop🐰"
def standardize_input(self, image, processed_images=None, masks=None):
"""
标准化输入格式
- image: [H,W,C] -> [1,H,W,C]
- processed_images: [...] -> [B,H,W,C]
- masks: [...] -> [B,H,W]
"""
# 处理原始图像
if len(image.shape) == 3:
image = image.unsqueeze(0)
assert len(image.shape) == 4, f"Original image must be 4D [B,H,W,C], got shape {image.shape}"
# 处理processed_images
if processed_images is not None:
if isinstance(processed_images, list):
processed_images = torch.cat(processed_images, dim=0)
if len(processed_images.shape) == 3:
processed_images = processed_images.unsqueeze(0)
assert len(processed_images.shape) == 4, \
f"Processed images must be 4D [B,H,W,C], got shape {processed_images.shape}"
# 处理masks
if masks is not None:
if isinstance(masks, list):
masks = torch.cat(masks, dim=0)
if len(masks.shape) == 2:
masks = masks.unsqueeze(0)
assert len(masks.shape) == 3, f"Masks must be 3D [B,H,W], got shape {masks.shape}"
return image, processed_images, masks
def resize_tensor(self, x, size, mode='bilinear'):
"""调整tensor尺寸的辅助函数"""
# 确保输入是4D tensor [B,C,H,W]
@@ -190,19 +219,14 @@ class MaskMerge:
if processed_images is None or masks is None:
return (original_image,)
# 确保原始图像维度正确
if len(original_image.shape) == 3:
original_image = original_image.unsqueeze(0)
# 标准化输入
original_image, processed_images, masks = self.standardize_input(
original_image, processed_images, masks
)
# 创建结果图像的副本
result = original_image.clone()
# 确保processed_images和masks都是张量
if isinstance(processed_images, list):
processed_images = torch.cat(processed_images, dim=0)
if isinstance(masks, list):
masks = torch.cat(masks, dim=0)
# 获取目标尺寸
target_height = original_image.shape[1]
target_width = original_image.shape[2]
@@ -232,7 +256,9 @@ class MaskMerge:
current_mask = masks[i:i+1]
result = current_mask * current_image + (1 - current_mask) * result
assert len(result.shape) == 4, "Output must be 4D [B,H,W,C]"
return (result,)
Mask_CLASS_MAPPINGS = {
"CyberEve_MaskSegmentation": MaskSplit,
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "loop-image"
description = "ComfyUI Loop Image is a node package specifically designed for image loop processing. It provides two main processing modes: Batch Image Processing and Single Image Processing, along with supporting image segmentation and merging functions."
version = "1.0.3"
version = "1.0.4"
license = {file = "LICENSE"}
dependencies = ["opencv-python", "numpy"]