diff --git a/kikotools/__init__.py b/kikotools/__init__.py index df52781..de61126 100644 --- a/kikotools/__init__.py +++ b/kikotools/__init__.py @@ -3,6 +3,12 @@ KikoTools package initialization and node registry Handles automatic discovery and registration of all ComfyAssets tools """ +from .tools.batch_list_converter import ( + ImageBatchToImageListNode, + ImageListToImageBatchNode, + LatentBatchToLatentListNode, + LatentListToLatentBatchNode, +) from .tools.batch_prompts import BatchPromptsNode from .tools.display_any import DisplayAnyNode from .tools.display_text import DisplayTextNode @@ -34,6 +40,10 @@ from .tools.xyz_helpers import ( # ComfyUI node registration mappings NODE_CLASS_MAPPINGS = { + "ImageBatchToImageList": ImageBatchToImageListNode, + "ImageListToImageBatch": ImageListToImageBatchNode, + "LatentBatchToLatentList": LatentBatchToLatentListNode, + "LatentListToLatentBatch": LatentListToLatentBatchNode, "BatchPrompts": BatchPromptsNode, "ResolutionCalculator": ResolutionCalculatorNode, "WidthHeightSelector": WidthHeightSelectorNode, @@ -65,6 +75,10 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { + "ImageBatchToImageList": "Image Batch to Image List", + "ImageListToImageBatch": "Image List to Image Batch", + "LatentBatchToLatentList": "Latent Batch to Latent List", + "LatentListToLatentBatch": "Latent List to Latent Batch", "BatchPrompts": "Batch Prompts", "ResolutionCalculator": "Resolution Calculator", "WidthHeightSelector": "Width Height Selector", diff --git a/kikotools/tools/batch_list_converter/__init__.py b/kikotools/tools/batch_list_converter/__init__.py new file mode 100644 index 0000000..2821b5e --- /dev/null +++ b/kikotools/tools/batch_list_converter/__init__.py @@ -0,0 +1,15 @@ +"""Batch/List conversion tool for ComfyUI.""" + +from .node import ( + ImageBatchToImageListNode, + ImageListToImageBatchNode, + LatentBatchToLatentListNode, + LatentListToLatentBatchNode, +) + +__all__ = [ + "ImageBatchToImageListNode", + "ImageListToImageBatchNode", + "LatentBatchToLatentListNode", + "LatentListToLatentBatchNode", +] diff --git a/kikotools/tools/batch_list_converter/logic.py b/kikotools/tools/batch_list_converter/logic.py new file mode 100644 index 0000000..61539fe --- /dev/null +++ b/kikotools/tools/batch_list_converter/logic.py @@ -0,0 +1,55 @@ +"""Pure tensor split/join functions for batch-list conversions.""" + +import torch +from typing import Dict, List + + +def split_image_batch(images: torch.Tensor) -> List[torch.Tensor]: + """Split [B,H,W,C] image batch into list of [1,H,W,C] tensors.""" + return [images[i : i + 1] for i in range(images.shape[0])] + + +def join_image_batch(image_list: List[torch.Tensor]) -> torch.Tensor: + """Join list of image tensors into single [B,H,W,C] batch.""" + return torch.cat(image_list, dim=0) + + +def split_latent_batch( + latent: Dict[str, torch.Tensor], +) -> List[Dict[str, torch.Tensor]]: + """Split latent dict into list of single-item latent dicts. + + Preserves all keys (e.g. noise_mask, batch_index). Tensor values whose + first dimension matches the batch size of ``samples`` are sliced along + dim-0; all other values are copied as-is to every item. + """ + samples = latent["samples"] + batch_size = samples.shape[0] + result: List[Dict[str, torch.Tensor]] = [] + for i in range(batch_size): + item: Dict[str, torch.Tensor] = {} + for key, value in latent.items(): + if isinstance(value, torch.Tensor) and value.shape[0] == batch_size: + item[key] = value[i : i + 1] + else: + item[key] = value + result.append(item) + return result + + +def join_latent_batch( + latent_list: List[Dict[str, torch.Tensor]], +) -> Dict[str, torch.Tensor]: + """Join list of latent dicts into single batched latent dict. + + Tensor values that were sliced during split are concatenated along dim-0. + Non-tensor values are taken from the first item. + """ + result: Dict[str, torch.Tensor] = {} + first = latent_list[0] + for key in first: + if isinstance(first[key], torch.Tensor): + result[key] = torch.cat([lat[key] for lat in latent_list], dim=0) + else: + result[key] = first[key] + return result diff --git a/kikotools/tools/batch_list_converter/node.py b/kikotools/tools/batch_list_converter/node.py new file mode 100644 index 0000000..e25ff1b --- /dev/null +++ b/kikotools/tools/batch_list_converter/node.py @@ -0,0 +1,130 @@ +"""Batch/List conversion nodes for ComfyUI.""" + +from typing import Dict, List, Tuple + +import torch + +from ...base.base_node import ComfyAssetsBaseNode +from .logic import ( + split_image_batch, + join_image_batch, + split_latent_batch, + join_latent_batch, +) + + +class ImageBatchToImageListNode(ComfyAssetsBaseNode): + """Split an IMAGE batch [B,H,W,C] into a list of individual images.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "INT") + RETURN_NAMES = ("images", "count") + OUTPUT_IS_LIST = (True, False) + FUNCTION = "split_batch" + CATEGORY = "🫶 ComfyAssets/📦 Latents" + + def split_batch(self, images: torch.Tensor) -> Tuple[List[torch.Tensor], int]: + image_list = split_image_batch(images) + count = len(image_list) + self.log_info(f"Split image batch of {count} into list") + return (image_list, count) + + +class ImageListToImageBatchNode(ComfyAssetsBaseNode): + """Join a list of IMAGE tensors into a single batched IMAGE [B,H,W,C].""" + + INPUT_IS_LIST = True + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "INT") + RETURN_NAMES = ("images", "count") + FUNCTION = "join_batch" + CATEGORY = "🫶 ComfyAssets/📦 Latents" + + def join_batch(self, images: List[torch.Tensor]) -> Tuple[torch.Tensor, int]: + batch = join_image_batch(images) + count = batch.shape[0] + self.log_info(f"Joined {count} images into batch") + return (batch, count) + + +class LatentBatchToLatentListNode(ComfyAssetsBaseNode): + """Split a LATENT batch into a list of individual latent dicts.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent": ("LATENT",), + } + } + + RETURN_TYPES = ("LATENT", "INT") + RETURN_NAMES = ("latents", "count") + OUTPUT_IS_LIST = (True, False) + FUNCTION = "split_batch" + CATEGORY = "🫶 ComfyAssets/📦 Latents" + + def split_batch( + self, latent: Dict[str, torch.Tensor] + ) -> Tuple[List[Dict[str, torch.Tensor]], int]: + latent_list = split_latent_batch(latent) + count = len(latent_list) + self.log_info(f"Split latent batch of {count} into list") + return (latent_list, count) + + +class LatentListToLatentBatchNode(ComfyAssetsBaseNode): + """Join a list of LATENT dicts into a single batched LATENT.""" + + INPUT_IS_LIST = True + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latents": ("LATENT",), + } + } + + RETURN_TYPES = ("LATENT", "INT") + RETURN_NAMES = ("latent", "count") + FUNCTION = "join_batch" + CATEGORY = "🫶 ComfyAssets/📦 Latents" + + def join_batch( + self, latents: List[Dict[str, torch.Tensor]] + ) -> Tuple[Dict[str, torch.Tensor], int]: + batch = join_latent_batch(latents) + count = batch["samples"].shape[0] + self.log_info(f"Joined {count} latents into batch") + return (batch, count) + + +NODE_CLASS_MAPPINGS = { + "ImageBatchToImageList": ImageBatchToImageListNode, + "ImageListToImageBatch": ImageListToImageBatchNode, + "LatentBatchToLatentList": LatentBatchToLatentListNode, + "LatentListToLatentBatch": LatentListToLatentBatchNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageBatchToImageList": "Image Batch to Image List", + "ImageListToImageBatch": "Image List to Image Batch", + "LatentBatchToLatentList": "Latent Batch to Latent List", + "LatentListToLatentBatch": "Latent List to Latent Batch", +} diff --git a/tests/unit/tools/test_batch_list_converter.py b/tests/unit/tools/test_batch_list_converter.py new file mode 100644 index 0000000..3eba35d --- /dev/null +++ b/tests/unit/tools/test_batch_list_converter.py @@ -0,0 +1,262 @@ +"""Tests for Batch/List conversion nodes and logic.""" + +import pytest +import torch + +from kikotools.tools.batch_list_converter.logic import ( + split_image_batch, + join_image_batch, + split_latent_batch, + join_latent_batch, +) +from kikotools.tools.batch_list_converter.node import ( + ImageBatchToImageListNode, + ImageListToImageBatchNode, + LatentBatchToLatentListNode, + LatentListToLatentBatchNode, +) + + +class TestBatchListConverterLogic: + """Test pure split/join functions.""" + + # -- Image split -- + + def test_split_image_batch_single(self): + """Single image batch returns list of one.""" + images = torch.rand(1, 64, 64, 3) + result = split_image_batch(images) + assert len(result) == 1 + assert result[0].shape == (1, 64, 64, 3) + assert torch.equal(result[0], images) + + def test_split_image_batch_multiple(self): + """Multi-image batch splits correctly.""" + images = torch.rand(4, 64, 64, 3) + result = split_image_batch(images) + assert len(result) == 4 + for i, img in enumerate(result): + assert img.shape == (1, 64, 64, 3) + assert torch.equal(img, images[i : i + 1]) + + def test_split_image_preserves_batch_dim(self): + """Each split image keeps 4D shape [1,H,W,C].""" + images = torch.rand(3, 128, 256, 3) + result = split_image_batch(images) + for img in result: + assert img.ndim == 4 + assert img.shape[0] == 1 + + # -- Image join -- + + def test_join_image_batch_single(self): + """Join single image produces batch of 1.""" + image_list = [torch.rand(1, 64, 64, 3)] + result = join_image_batch(image_list) + assert result.shape == (1, 64, 64, 3) + + def test_join_image_batch_multiple(self): + """Join multiple images into batch.""" + image_list = [torch.rand(1, 64, 64, 3) for _ in range(5)] + result = join_image_batch(image_list) + assert result.shape == (5, 64, 64, 3) + + def test_image_roundtrip(self): + """split -> join produces identical tensor.""" + original = torch.rand(4, 64, 64, 3) + reconstructed = join_image_batch(split_image_batch(original)) + assert torch.equal(original, reconstructed) + + # -- Latent split -- + + def test_split_latent_batch_single(self): + """Single latent returns list of one dict.""" + latent = {"samples": torch.rand(1, 4, 32, 32)} + result = split_latent_batch(latent) + assert len(result) == 1 + assert "samples" in result[0] + assert result[0]["samples"].shape == (1, 4, 32, 32) + + def test_split_latent_batch_multiple(self): + """Multi-item latent splits correctly.""" + latent = {"samples": torch.rand(3, 4, 32, 32)} + result = split_latent_batch(latent) + assert len(result) == 3 + for i, lat in enumerate(result): + assert lat["samples"].shape == (1, 4, 32, 32) + assert torch.equal(lat["samples"], latent["samples"][i : i + 1]) + + # -- Latent join -- + + def test_join_latent_batch_single(self): + """Join single latent dict.""" + latent_list = [{"samples": torch.rand(1, 4, 32, 32)}] + result = join_latent_batch(latent_list) + assert "samples" in result + assert result["samples"].shape == (1, 4, 32, 32) + + def test_join_latent_batch_multiple(self): + """Join multiple latent dicts into batch.""" + latent_list = [{"samples": torch.rand(1, 4, 32, 32)} for _ in range(4)] + result = join_latent_batch(latent_list) + assert result["samples"].shape == (4, 4, 32, 32) + + def test_latent_roundtrip(self): + """split -> join produces identical tensor.""" + original = {"samples": torch.rand(5, 4, 64, 64)} + reconstructed = join_latent_batch(split_latent_batch(original)) + assert torch.equal(original["samples"], reconstructed["samples"]) + + def test_split_latent_preserves_noise_mask(self): + """noise_mask is sliced alongside samples.""" + latent = { + "samples": torch.rand(3, 4, 32, 32), + "noise_mask": torch.rand(3, 1, 32, 32), + } + result = split_latent_batch(latent) + assert len(result) == 3 + for i, item in enumerate(result): + assert "noise_mask" in item + assert item["noise_mask"].shape == (1, 1, 32, 32) + assert torch.equal(item["noise_mask"], latent["noise_mask"][i : i + 1]) + + def test_join_latent_preserves_noise_mask(self): + """noise_mask is concatenated alongside samples.""" + latent_list = [ + { + "samples": torch.rand(1, 4, 32, 32), + "noise_mask": torch.rand(1, 1, 32, 32), + } + for _ in range(3) + ] + result = join_latent_batch(latent_list) + assert "noise_mask" in result + assert result["noise_mask"].shape == (3, 1, 32, 32) + + def test_latent_roundtrip_with_extra_keys(self): + """Round-trip preserves all tensor keys.""" + original = { + "samples": torch.rand(4, 4, 64, 64), + "noise_mask": torch.rand(4, 1, 64, 64), + } + reconstructed = join_latent_batch(split_latent_batch(original)) + assert torch.equal(original["samples"], reconstructed["samples"]) + assert torch.equal(original["noise_mask"], reconstructed["noise_mask"]) + + def test_split_latent_copies_non_tensor_values(self): + """Non-tensor metadata is copied to each item.""" + latent = { + "samples": torch.rand(2, 4, 32, 32), + "some_flag": "preserve_me", + } + result = split_latent_batch(latent) + for item in result: + assert item["some_flag"] == "preserve_me" + + +class TestBatchListConverterNodes: + """Test ComfyUI node classes.""" + + # -- ImageBatchToImageList -- + + def test_image_b2l_attributes(self): + assert ImageBatchToImageListNode.RETURN_TYPES == ("IMAGE", "INT") + assert ImageBatchToImageListNode.RETURN_NAMES == ("images", "count") + assert ImageBatchToImageListNode.OUTPUT_IS_LIST == (True, False) + assert ImageBatchToImageListNode.FUNCTION == "split_batch" + assert ImageBatchToImageListNode.CATEGORY == "🫶 ComfyAssets/📦 Latents" + + def test_image_b2l_input_types(self): + inputs = ImageBatchToImageListNode.INPUT_TYPES() + assert "required" in inputs + assert "images" in inputs["required"] + assert inputs["required"]["images"] == ("IMAGE",) + + def test_image_b2l_execute(self): + node = ImageBatchToImageListNode() + images = torch.rand(3, 64, 64, 3) + result = node.split_batch(images) + image_list, count = result + assert isinstance(image_list, list) + assert len(image_list) == 3 + assert count == 3 + + # -- ImageListToImageBatch -- + + def test_image_l2b_attributes(self): + assert ImageListToImageBatchNode.INPUT_IS_LIST is True + assert ImageListToImageBatchNode.RETURN_TYPES == ("IMAGE", "INT") + assert ImageListToImageBatchNode.RETURN_NAMES == ("images", "count") + assert ImageListToImageBatchNode.FUNCTION == "join_batch" + + def test_image_l2b_execute(self): + node = ImageListToImageBatchNode() + image_list = [torch.rand(1, 64, 64, 3) for _ in range(4)] + batch, count = node.join_batch(image_list) + assert batch.shape == (4, 64, 64, 3) + assert count == 4 + + # -- LatentBatchToLatentList -- + + def test_latent_b2l_attributes(self): + assert LatentBatchToLatentListNode.RETURN_TYPES == ("LATENT", "INT") + assert LatentBatchToLatentListNode.RETURN_NAMES == ("latents", "count") + assert LatentBatchToLatentListNode.OUTPUT_IS_LIST == (True, False) + assert LatentBatchToLatentListNode.FUNCTION == "split_batch" + + def test_latent_b2l_execute(self): + node = LatentBatchToLatentListNode() + latent = {"samples": torch.rand(2, 4, 32, 32)} + latent_list, count = node.split_batch(latent) + assert isinstance(latent_list, list) + assert len(latent_list) == 2 + assert count == 2 + + # -- LatentListToLatentBatch -- + + def test_latent_l2b_attributes(self): + assert LatentListToLatentBatchNode.INPUT_IS_LIST is True + assert LatentListToLatentBatchNode.RETURN_TYPES == ("LATENT", "INT") + assert LatentListToLatentBatchNode.RETURN_NAMES == ("latent", "count") + assert LatentListToLatentBatchNode.FUNCTION == "join_batch" + + def test_latent_l2b_execute(self): + node = LatentListToLatentBatchNode() + latent_list = [{"samples": torch.rand(1, 4, 32, 32)} for _ in range(3)] + batch, count = node.join_batch(latent_list) + assert "samples" in batch + assert batch["samples"].shape == (3, 4, 32, 32) + assert count == 3 + + # -- Inheritance -- + + def test_all_nodes_inherit_base(self): + from kikotools.base.base_node import ComfyAssetsBaseNode + + for cls in ( + ImageBatchToImageListNode, + ImageListToImageBatchNode, + LatentBatchToLatentListNode, + LatentListToLatentBatchNode, + ): + assert issubclass(cls, ComfyAssetsBaseNode) + + # -- Registration mappings -- + + def test_node_class_mappings(self): + from kikotools.tools.batch_list_converter.node import NODE_CLASS_MAPPINGS + + assert len(NODE_CLASS_MAPPINGS) == 4 + assert "ImageBatchToImageList" in NODE_CLASS_MAPPINGS + assert "ImageListToImageBatch" in NODE_CLASS_MAPPINGS + assert "LatentBatchToLatentList" in NODE_CLASS_MAPPINGS + assert "LatentListToLatentBatch" in NODE_CLASS_MAPPINGS + + def test_node_display_name_mappings(self): + from kikotools.tools.batch_list_converter.node import NODE_DISPLAY_NAME_MAPPINGS + + assert len(NODE_DISPLAY_NAME_MAPPINGS) == 4 + assert ( + NODE_DISPLAY_NAME_MAPPINGS["ImageBatchToImageList"] + == "Image Batch to Image List" + )