Merge pull request #55 from ComfyAssets/feature/batch-list-converter
feat: add batch/list conversion nodes for IMAGE and LATENT
This commit is contained in:
Vito
@@ -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",
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Batch/List conversion tool for ComfyUI."""
|
||||
|
||||
from .node import (
|
||||
ImageBatchToImageListNode,
|
||||
ImageListToImageBatchNode,
|
||||
LatentBatchToLatentListNode,
|
||||
LatentListToLatentBatchNode,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ImageBatchToImageListNode",
|
||||
"ImageListToImageBatchNode",
|
||||
"LatentBatchToLatentListNode",
|
||||
"LatentListToLatentBatchNode",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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"
|
||||
)
|
||||
Reference in New Issue
Block a user