Initial commit: Custom SplitImage nodes for ComfyUI
This commit is contained in:
=
@@ -0,0 +1,4 @@
|
||||
conductor/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.ipynb_checkpoints/
|
||||
@@ -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.
|
||||
+22
@@ -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']
|
||||
@@ -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),)
|
||||
@@ -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)
|
||||
+102
@@ -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)
|
||||
+123
@@ -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,)
|
||||
Reference in New Issue
Block a user