From 9f2257e81d4bf54ea163ae22bcd6421d669783b5 Mon Sep 17 00:00:00 2001 From: Cyber Dick Lang <286878701@qq.com> Date: Wed, 29 Oct 2025 19:49:43 +0800 Subject: [PATCH] feat: Add Get Image or Mask Range From Batch (UTK) node MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增Get Image or Mask Range From Batch (UTK)节点到tools分类 - 支持从图像批次或遮罩批次中提取指定范围的元素 - 灵活的索引控制：支持正索引和-1（从末尾开始） - 智能范围处理：自动处理超出范围的情况 - 双输入支持：可同时处理图像和遮罩批次 - 详细日志输出：显示提取的索引范围和数量 - 错误处理：提供清晰的错误信息指导用户 - 工具分类：归类到UniversalToolkit/tools分类下 - 提升批次处理工作流的灵活性和效率 - 更新版本到1.4.4并添加完整更新日志 --- __init__.py | 25 +++++- nodes/tools/get_image_range_from_batch.py | 105 ++++++++++++++++++++++ pyproject.toml | 2 +- 3 files changed, 130 insertions(+), 2 deletions(-) create mode 100644 nodes/tools/get_image_range_from_batch.py diff --git a/__init__.py b/__init__.py index 56c02c0..a72fc01 100644 --- a/__init__.py +++ b/__init__.py @@ -8,13 +8,24 @@ A comprehensive toolkit for ComfyUI that provides various utility nodes for imag :license: MIT, see LICENSE for more details. """ -__version__ = "1.4.3" +__version__ = "1.4.4" __author__ = "CyberDickLang" __email__ = "286878701@qq.com" __url__ = "https://github.com/whmc76" # 更新日志 CHANGELOG = { + "1.4.4": [ + "新增Get Image or Mask Range From Batch (UTK)节点：", + "- 支持从图像批次或遮罩批次中提取指定范围的元素", + "- 灵活的索引控制：支持正索引和-1（从末尾开始）", + "- 智能范围处理：自动处理超出范围的情况", + "- 双输入支持：可同时处理图像和遮罩批次", + "- 详细日志输出：显示提取的索引范围和数量", + "- 错误处理：提供清晰的错误信息指导用户", + "- 工具分类：归类到UniversalToolkit/tools分类下", + "- 提升批次处理工作流的灵活性和效率", + ], "1.4.3": [ "修复翻译API节点错误消息显示问题：", "- 修复API密钥获取网址在错误消息中不显示的问题", @@ -726,6 +737,15 @@ except ImportError: TEXT_TRANSLATOR_API_MAPPINGS = {} TEXT_TRANSLATOR_API_DISPLAY = {} +try: + from .nodes.tools.get_image_range_from_batch import \ + NODE_CLASS_MAPPINGS as GET_IMAGE_RANGE_MAPPINGS + from .nodes.tools.get_image_range_from_batch import \ + NODE_DISPLAY_NAME_MAPPINGS as GET_IMAGE_RANGE_DISPLAY +except ImportError: + GET_IMAGE_RANGE_MAPPINGS = {} + GET_IMAGE_RANGE_DISPLAY = {} + # 合并所有节点映射 NODE_CLASS_MAPPINGS = {} @@ -765,6 +785,7 @@ NODE_CLASS_MAPPINGS.update(PROMPT_HELPER_MAPPINGS) NODE_CLASS_MAPPINGS.update(COLOR_TO_MASK_MAPPINGS) NODE_CLASS_MAPPINGS.update(LAZY_SWITCH_MAPPINGS) NODE_CLASS_MAPPINGS.update(TEXT_TRANSLATOR_API_MAPPINGS) +NODE_CLASS_MAPPINGS.update(GET_IMAGE_RANGE_MAPPINGS) # 合并显示名称映射 NODE_DISPLAY_NAME_MAPPINGS = {} @@ -804,6 +825,7 @@ NODE_DISPLAY_NAME_MAPPINGS.update(PROMPT_HELPER_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(COLOR_TO_MASK_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(LAZY_SWITCH_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(TEXT_TRANSLATOR_API_DISPLAY) +NODE_DISPLAY_NAME_MAPPINGS.update(GET_IMAGE_RANGE_DISPLAY) NODE_CATEGORIES = { "UniversalToolkit": [ @@ -844,6 +866,7 @@ NODE_CATEGORIES = { "LazySwitchKJ_UTK", "APIImageGenerator_UTK", "TextTranslatorAPI_UTK", + "GetImageRangeFromBatch_UTK", ] } diff --git a/nodes/tools/get_image_range_from_batch.py b/nodes/tools/get_image_range_from_batch.py new file mode 100644 index 0000000..8642f5c --- /dev/null +++ b/nodes/tools/get_image_range_from_batch.py @@ -0,0 +1,105 @@ +import torch +from typing import Tuple, Optional + + +class GetImageRangeFromBatch_UTK: + """ + 从批次中获取指定范围的图像或遮罩 + 支持从图像批次或遮罩批次中提取指定索引范围的元素 + """ + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "mask") + FUNCTION = "get_range_from_batch" + CATEGORY = "UniversalToolkit/tools" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "start_index": ("INT", { + "default": 0, + "min": -1, + "max": 4096, + "step": 1, + "tooltip": "起始索引，-1表示从末尾开始" + }), + "num_frames": ("INT", { + "default": 1, + "min": 1, + "max": 4096, + "step": 1, + "tooltip": "要提取的帧数" + }), + }, + "optional": { + "images": ("IMAGE", { + "tooltip": "输入的图像批次" + }), + "masks": ("MASK", { + "tooltip": "输入的遮罩批次" + }), + } + } + + def get_range_from_batch(self, start_index: int, num_frames: int, + images: Optional[torch.Tensor] = None, + masks: Optional[torch.Tensor] = None) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + 从批次中获取指定范围的图像或遮罩 + + Args: + start_index: 起始索引，-1表示从末尾开始 + num_frames: 要提取的帧数 + images: 输入的图像批次 (可选) + masks: 输入的遮罩批次 (可选) + + Returns: + Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: 提取的图像和遮罩 + """ + chosen_images = None + chosen_masks = None + + # 处理图像批次 + if images is not None: + if start_index == -1: + # 从末尾开始计算起始索引 + start_index = max(0, len(images) - num_frames) + + if start_index < 0 or start_index >= len(images): + raise ValueError(f"图像起始索引 {start_index} 超出范围 [0, {len(images)-1}]") + + end_index = min(start_index + num_frames, len(images)) + chosen_images = images[start_index:end_index] + + print(f"📸 从图像批次中提取: 索引 {start_index} 到 {end_index-1}，共 {len(chosen_images)} 张图像") + + # 处理遮罩批次 + if masks is not None: + if start_index == -1: + # 从末尾开始计算起始索引 + start_index = max(0, len(masks) - num_frames) + + if start_index < 0 or start_index >= len(masks): + raise ValueError(f"遮罩起始索引 {start_index} 超出范围 [0, {len(masks)-1}]") + + end_index = min(start_index + num_frames, len(masks)) + chosen_masks = masks[start_index:end_index] + + print(f"🎭 从遮罩批次中提取: 索引 {start_index} 到 {end_index-1}，共 {len(chosen_masks)} 个遮罩") + + # 检查是否至少有一个输入 + if images is None and masks is None: + raise ValueError("至少需要提供图像或遮罩输入") + + return (chosen_images, chosen_masks) + + +# 节点映射 +NODE_CLASS_MAPPINGS = { + "GetImageRangeFromBatch_UTK": GetImageRangeFromBatch_UTK +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GetImageRangeFromBatch_UTK": "Get Image or Mask Range From Batch (UTK)" +} diff --git a/pyproject.toml b/pyproject.toml index f4a75b0..b77b0b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "universaltoolkit" description = "A comprehensive toolkit based on ComfyUI, providing image, mask, audio, and tools nodes, fully modular and v3 compatible." -version = "1.4.3" +version = "1.4.4" license = {file = "LICENSE"} dependencies = [ "torch>=1.9.0",