File upload

image_grid.py

@@ -0,0 +1,135 @@
+"""
+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)"}