136 lines
4.4 KiB
Python
136 lines
4.4 KiB
Python
"""
|
|
Image Grid Node for ComfyUI
|
|
Arranges multiple image batches into a customizable grid layout.
|
|
"""
|
|
|
|
import math
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class ImageGrid:
|
|
"""ComfyUI node for creating image grids from multiple image batches."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"columns": ("INT", {"default": 2, "min": 0, "max": 100, "step": 1}),
|
|
"padding": ("INT", {"default": 2, "min": 0, "max": 500, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
"image_4": ("IMAGE",),
|
|
"image_5": ("IMAGE",),
|
|
"image_6": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "create_grid"
|
|
CATEGORY = "image/grid"
|
|
|
|
def create_grid(
|
|
self,
|
|
columns: int,
|
|
padding: int,
|
|
image_1: Optional[torch.Tensor] = None,
|
|
image_2: Optional[torch.Tensor] = None,
|
|
image_3: Optional[torch.Tensor] = None,
|
|
image_4: Optional[torch.Tensor] = None,
|
|
image_5: Optional[torch.Tensor] = None,
|
|
image_6: Optional[torch.Tensor] = None
|
|
) -> Tuple[torch.Tensor]:
|
|
"""
|
|
Create a grid from multiple image batches.
|
|
|
|
Args:
|
|
columns: Number of columns in the grid.
|
|
padding: Padding between images in pixels.
|
|
image_1...image_6: Optional image batches to include in the grid.
|
|
|
|
Returns:
|
|
A single tensor containing the grid image.
|
|
"""
|
|
# Collect all provided images
|
|
all_images = [
|
|
img for img in [image_1, image_2, image_3, image_4, image_5, image_6]
|
|
if img is not None
|
|
]
|
|
|
|
if not all_images:
|
|
return (torch.zeros((1, 512, 512, 3)),)
|
|
|
|
# Get base dimensions from first image batch
|
|
base_shape = all_images[0].shape
|
|
target_height = base_shape[1]
|
|
target_width = base_shape[2]
|
|
channels = base_shape[3]
|
|
|
|
# Resize all images to match first image dimensions
|
|
processed_images = []
|
|
for img_batch in all_images:
|
|
if img_batch.shape[1] != target_height or img_batch.shape[2] != target_width:
|
|
# Resize: permute to [B, C, H, W] for interpolation
|
|
img_permuted = img_batch.permute(0, 3, 1, 2)
|
|
img_resized = F.interpolate(
|
|
img_permuted,
|
|
size=(target_height, target_width),
|
|
mode="bilinear",
|
|
align_corners=False
|
|
)
|
|
# Permute back to [B, H, W, C]
|
|
img_batch = img_resized.permute(0, 2, 3, 1)
|
|
|
|
processed_images.append(img_batch)
|
|
|
|
# Concatenate all image batches
|
|
final_images = torch.cat(processed_images, dim=0)
|
|
batch_size = final_images.shape[0]
|
|
|
|
if batch_size == 0:
|
|
return (torch.zeros((1, target_height, target_width, channels)),)
|
|
|
|
# Calculate grid dimensions
|
|
if columns > 0:
|
|
cols = columns
|
|
rows = math.ceil(batch_size / cols)
|
|
else:
|
|
# Auto-calculate roughly square grid
|
|
cols = math.ceil(math.sqrt(batch_size))
|
|
rows = math.ceil(batch_size / cols)
|
|
|
|
# Calculate output dimensions
|
|
grid_height = rows * target_height + (rows - 1) * padding
|
|
grid_width = cols * target_width + (cols - 1) * padding
|
|
|
|
# Initialize grid with zeros (black background)
|
|
grid = torch.zeros(
|
|
(1, grid_height, grid_width, channels),
|
|
dtype=final_images.dtype,
|
|
device=final_images.device
|
|
)
|
|
|
|
# Place images in grid
|
|
for idx in range(batch_size):
|
|
row = idx // cols
|
|
col = idx % cols
|
|
|
|
y_start = row * (target_height + padding)
|
|
y_end = y_start + target_height
|
|
x_start = col * (target_width + padding)
|
|
x_end = x_start + target_width
|
|
|
|
grid[0, y_start:y_end, x_start:x_end, :] = final_images[idx]
|
|
|
|
return (grid,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"ImageGrid": ImageGrid}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"ImageGrid": "Image Grid (SEB)"}
|