commit a82494daa55347b2e50776ec4a1fb84a8d0d577b Author: = <=> Date: Sat Apr 4 11:56:45 2026 -0400 Initial commit: Custom SplitImage nodes for ComfyUI diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..dfcae22 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +conductor/ +__pycache__/ +*.pyc +.ipynb_checkpoints/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..c425ce0 --- /dev/null +++ b/README.md @@ -0,0 +1,45 @@ +# ComfyUI-SplitImage + +A collection of custom nodes for ComfyUI to split an image into overlapping tiles and stitch them back together seamlessly. This is particularly useful for processing high-resolution images in smaller chunks (tiles) to save VRAM or for specialized image-to-image workflows. + +## Features + +- **Seamless Stitching**: Uses a weight-normalized blending algorithm with linear ramps to eliminate visible seams in overlapping regions. +- **VRAM Efficient Processing**: Split images into tiles and process them one-by-one to handle massive resolutions. +- **Automatic Upscaling Support**: The stitching node automatically detects if tiles have been upscaled and adjusts the overlap and final image dimensions accordingly. +- **Flexible Grid**: Split images or latents into any number of rows and columns. +- **Collector Pattern**: The stitching node automatically collects tiles from lists or batches for reassembly. + +## Nodes (Category: SplitImage) + +### 1. Split Image into Tiles +Divides an image into a **List** of multiple overlapping tiles. + +### 2. Split Latent into Tiles +Divides a latent batch into a **List** of multiple overlapping tiles. + +### 3. Calculate Tiling Spec +Calculates a tiling configuration without needing an image. + +### 4. Stitch Image Tiles +Reassembles a list or batch of tiles into a single image. **Acts as a Collector**: It waits for all items in a list to finish processing before stitching. + +### 5. Merge Image Batches +Combines multiple image batches into one larger batch. + +## Advanced Workflow: The "Collector" Pattern + +To process large images without running out of VRAM (OOM): + +1. **Split**: Use `Split Image into Tiles`. It outputs a **List** of tiles. +2. **Process Sequential**: Connect the list to your VAE Encode/KSampler/VAE Decode. ComfyUI will execute these nodes **one by one** for each tile in the list, keeping VRAM usage low. +3. **Stitch (Collect)**: Connect the final list of processed images to `Stitch Image Tiles`. It will collect all tiles and stitch them into the final result. + +## Installation + +1. Navigate to your ComfyUI `custom_nodes` directory. +2. Clone this repository: + ```bash + git clone https://github.com/QuigleyDown/ComfyUI-SplitImage.git + ``` +3. Restart ComfyUI. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..55662a5 --- /dev/null +++ b/__init__.py @@ -0,0 +1,22 @@ +from .image_split import SplitImage_Tiled, CalculateTilingSpec_Tiled +from .image_stitch import StitchImage_Tiled +from .image_batch import BatchImages_Tiled +from .image_latent import SplitLatent_Tiled + +NODE_CLASS_MAPPINGS = { + "SplitImage_Tiled": SplitImage_Tiled, + "StitchImage_Tiled": StitchImage_Tiled, + "BatchImages_Tiled": BatchImages_Tiled, + "CalculateTilingSpec_Tiled": CalculateTilingSpec_Tiled, + "SplitLatent_Tiled": SplitLatent_Tiled +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SplitImage_Tiled": "Split Image into Tiles", + "StitchImage_Tiled": "Stitch Image Tiles", + "BatchImages_Tiled": "Merge Image Batches", + "CalculateTilingSpec_Tiled": "Calculate Tiling Spec", + "SplitLatent_Tiled": "Split Latent into Tiles" +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/image_batch.py b/image_batch.py new file mode 100644 index 0000000..cc4a755 --- /dev/null +++ b/image_batch.py @@ -0,0 +1,41 @@ +import torch + +class BatchImages_Tiled: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image1": ("IMAGE",), + }, + "optional": { + "image2": ("IMAGE",), + "image3": ("IMAGE",), + "image4": ("IMAGE",), + "image5": ("IMAGE",), + "image6": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "batch" + CATEGORY = "SplitImage" + + def batch(self, image1, **kwargs): + images = [image1] + for i in range(2, 7): + img = kwargs.get(f"image{i}", None) + if img is not None: + images.append(img) + + # Ensure all images have the same dimensions (W, H, C) + # We take the first image as reference + ref_h, ref_w, ref_c = image1.shape[1:] + + valid_images = [] + for i, img in enumerate(images): + if img.shape[1:] != (ref_h, ref_w, ref_c): + print(f"BatchImages_Tiled: Warning, image{i+1} has different dimensions {img.shape[1:]} than image1 {ref_h, ref_w, ref_c}. Skipping.") + continue + valid_images.append(img) + + return (torch.cat(valid_images, dim=0),) diff --git a/image_latent.py b/image_latent.py new file mode 100644 index 0000000..44e1cff --- /dev/null +++ b/image_latent.py @@ -0,0 +1,69 @@ +import torch +import torch.nn.functional as F +import math + +class SplitLatent_Tiled: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "samples": ("LATENT",), + "rows": ("INT", {"default": 2, "min": 1, "max": 100}), + "columns": ("INT", {"default": 2, "min": 1, "max": 100}), + "overlap": ("INT", {"default": 8, "min": 0, "max": 128}), + } + } + + RETURN_TYPES = ("LATENT", "INT", "INT", "INT", "*") + RETURN_NAMES = ("tiles", "rows", "columns", "overlap", "tiling_spec") + OUTPUT_IS_LIST = (True, False, False, False, False) + FUNCTION = "split" + CATEGORY = "SplitImage" + + def split(self, samples, rows, columns, overlap): + latent = samples["samples"] + + shape = list(latent.shape) + if len(shape) < 3: + raise ValueError(f"SplitLatent_Tiled: Expected at least 3D latent tensor, got {shape}") + + H = shape[-2] + W = shape[-1] + + tile_w = math.ceil((W + (columns - 1) * overlap) / columns) + tile_h = math.ceil((H + (rows - 1) * overlap) / rows) + + pad_w = (columns * tile_w - (columns - 1) * overlap) - W + pad_h = (rows * tile_h - (rows - 1) * overlap) - H + + if pad_w > 0 or pad_h > 0: + latent = F.pad(latent, (0, pad_w, 0, pad_h), mode='constant', value=0) + H = latent.shape[-2] + W = latent.shape[-1] + + tiles = [] + for r in range(rows): + for c in range(columns): + x1 = c * (tile_w - overlap) + y1 = r * (tile_h - overlap) + x2 = x1 + tile_w + y2 = y1 + tile_h + + tile = latent[..., y1:y2, x1:x2] + # Store each tile as a single-latent sample in the list + for i in range(tile.shape[0]): + tiles.append({"samples": tile[i:i+1]}) + + tiling_spec = { + "rows": rows, + "columns": columns, + "overlap": overlap * 8, # Store as pixel overlap + "tile_w": tile_w * 8, + "tile_h": tile_h * 8, + "orig_w": (W - pad_w) * 8, + "orig_h": (H - pad_h) * 8, + "padded_w": W * 8, + "padded_h": H * 8 + } + + return (tiles, rows, columns, overlap, tiling_spec) diff --git a/image_split.py b/image_split.py new file mode 100644 index 0000000..b595609 --- /dev/null +++ b/image_split.py @@ -0,0 +1,102 @@ +import torch +import torch.nn.functional as F +import math + +class CalculateTilingSpec_Tiled: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "width": ("INT", {"default": 512, "min": 1, "max": 16384}), + "height": ("INT", {"default": 512, "min": 1, "max": 16384}), + "rows": ("INT", {"default": 2, "min": 1, "max": 100}), + "columns": ("INT", {"default": 2, "min": 1, "max": 100}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 1024}), + } + } + + RETURN_TYPES = ("*", "INT", "INT") + RETURN_NAMES = ("tiling_spec", "tile_width", "tile_height") + FUNCTION = "calculate" + CATEGORY = "SplitImage" + + def calculate(self, width, height, rows, columns, overlap): + tile_w = math.ceil((width + (columns - 1) * overlap) / columns) + tile_h = math.ceil((height + (rows - 1) * overlap) / rows) + + pad_w = (columns * tile_w - (columns - 1) * overlap) - width + pad_h = (rows * tile_h - (rows - 1) * overlap) - height + + tiling_spec = { + "rows": rows, + "columns": columns, + "overlap": overlap, + "tile_w": tile_w, + "tile_h": tile_h, + "orig_w": width, + "orig_h": height, + "padded_w": width + pad_w, + "padded_h": height + pad_h + } + + return (tiling_spec, tile_w, tile_h) + +class SplitImage_Tiled: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "rows": ("INT", {"default": 2, "min": 1, "max": 100}), + "columns": ("INT", {"default": 2, "min": 1, "max": 100}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 1024}), + } + } + + RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "*") + RETURN_NAMES = ("tiles", "rows", "columns", "overlap", "tiling_spec") + OUTPUT_IS_LIST = (True, False, False, False, False) + FUNCTION = "split" + CATEGORY = "SplitImage" + + def split(self, image, rows, columns, overlap): + B, H, W, C = image.shape + + tile_w = math.ceil((W + (columns - 1) * overlap) / columns) + tile_h = math.ceil((H + (rows - 1) * overlap) / rows) + + pad_w = (columns * tile_w - (columns - 1) * overlap) - W + pad_h = (rows * tile_h - (rows - 1) * overlap) - H + + if pad_w > 0 or pad_h > 0: + image_padded = image.permute(0, 3, 1, 2) + image_padded = F.pad(image_padded, (0, pad_w, 0, pad_h), mode='constant', value=0) + image = image_padded.permute(0, 2, 3, 1) + _, H, W, _ = image.shape + + tiles = [] + for r in range(rows): + for c in range(columns): + x1 = c * (tile_w - overlap) + y1 = r * (tile_h - overlap) + x2 = x1 + tile_w + y2 = y1 + tile_h + + tile = image[:, y1:y2, x1:x2, :] + # Store each tile as a single-image tensor in the list + for i in range(tile.shape[0]): + tiles.append(tile[i:i+1]) + + tiling_spec = { + "rows": rows, + "columns": columns, + "overlap": overlap, + "tile_w": tile_w, + "tile_h": tile_h, + "orig_w": W - pad_w, + "orig_h": H - pad_h, + "padded_w": W, + "padded_h": H + } + + return (tiles, rows, columns, overlap, tiling_spec) diff --git a/image_stitch.py b/image_stitch.py new file mode 100644 index 0000000..c6ab71b --- /dev/null +++ b/image_stitch.py @@ -0,0 +1,123 @@ +import torch + +class StitchImage_Tiled: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "tiles": ("IMAGE",), + }, + "optional": { + "tiling_spec": ("*",), + "rows": ("INT", {"default": 0, "min": 0, "max": 100}), + "columns": ("INT", {"default": 0, "min": 0, "max": 100}), + "overlap": ("INT", {"default": 0, "min": 0, "max": 1024}), + } + } + + # Enable receiving all tiles and spec as lists. + # This allows ComfyUI to process nodes one-by-one (saving VRAM) and then + # collect all results here for final stitching. + INPUT_IS_LIST = True + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "stitch" + CATEGORY = "SplitImage" + + def stitch(self, tiles, tiling_spec=None, rows=[0], columns=[0], overlap=[0]): + # When INPUT_IS_LIST = True, all arguments are passed as lists. + # Handle tiling_spec list + if tiling_spec is not None and len(tiling_spec) > 0: + spec = tiling_spec[0] + if isinstance(spec, list): spec = spec[0] + + r = spec["rows"] + c = spec["columns"] + orig_overlap = spec["overlap"] + orig_tile_w = spec["tile_w"] + orig_tile_h = spec["tile_h"] + orig_w = spec["orig_w"] + orig_h = spec["orig_h"] + print(f"StitchImage_Tiled: Using tiling_spec: {r}x{c} grid, overlap {orig_overlap}") + else: + r = rows[0] + c = columns[0] + if r <= 0 or c <= 0: + raise ValueError("Both rows and columns must be provided if tiling_spec is not connected") + orig_overlap = overlap[0] + # No reference tile size, assume no scaling + orig_tile_w = 0 # Will be set from first tile + orig_tile_h = 0 + orig_w = 0 + orig_h = 0 + print(f"StitchImage_Tiled: Using manual settings: {r}x{c} grid, overlap {orig_overlap}") + + # Flatten the list of tiles (which might be single images or batches) into one tensor + full_tiles = torch.cat(tiles, dim=0) + B_total, H_tile, W_tile, C = full_tiles.shape + print(f"StitchImage_Tiled: Processing input batch of {B_total} tiles, size {W_tile}x{H_tile}") + + if orig_tile_w == 0: + orig_tile_w = W_tile + orig_tile_h = H_tile + orig_w = c * W_tile - (c - 1) * orig_overlap + orig_h = r * H_tile - (r - 1) * orig_overlap + + num_tiles_per_image = r * c + if B_total % num_tiles_per_image != 0: + raise ValueError(f"Total tiles {B_total} is not divisible by rows*cols ({r}*{c}={num_tiles_per_image}). " + f"Ensure you are passing a full set of all tiles.") + + B = B_total // num_tiles_per_image + + scale_w = W_tile / orig_tile_w + scale_h = H_tile / orig_tile_h + + scaled_overlap_w = int(orig_overlap * scale_w) + scaled_overlap_h = int(orig_overlap * scale_h) + + W_total = c * W_tile - (c - 1) * scaled_overlap_w + H_total = r * H_tile - (r - 1) * scaled_overlap_h + + output = torch.zeros((B, H_total, W_total, C), device=full_tiles.device) + weights = torch.zeros((B, H_total, W_total, 1), device=full_tiles.device) + + def get_mask(is_first_r, is_last_r, is_first_c, is_last_c): + mask = torch.ones((H_tile, W_tile, 1), device=full_tiles.device) + + if scaled_overlap_w > 0: + ramp = torch.linspace(0, 1, scaled_overlap_w + 2, device=full_tiles.device)[1:-1] + if not is_first_c: + mask[:, :scaled_overlap_w, 0] *= ramp + if not is_last_c: + mask[:, -scaled_overlap_w:, 0] *= ramp.flip(0) + + if scaled_overlap_h > 0: + ramp = torch.linspace(0, 1, scaled_overlap_h + 2, device=full_tiles.device)[1:-1].view(-1, 1) + if not is_first_r: + mask[:scaled_overlap_h, :, 0] *= ramp + if not is_last_r: + mask[-scaled_overlap_h:, :, 0] *= ramp.flip(0) + return mask + + idx = 0 + for b in range(B): + for r_idx in range(r): + for c_idx in range(c): + tile = full_tiles[idx] + idx += 1 + mask = get_mask(r_idx == 0, r_idx == r - 1, c_idx == 0, c_idx == c - 1) + + x = c_idx * (W_tile - scaled_overlap_w) + y = r_idx * (H_tile - scaled_overlap_h) + + output[b, y:y+H_tile, x:x+W_tile, :] += tile * mask + weights[b, y:y+H_tile, x:x+W_tile, :] += mask + + output = output / (weights + 1e-8) + + final_w = int(orig_w * scale_w) + final_h = int(orig_h * scale_h) + output = output[:, :final_h, :final_w, :] + + return (output,)