101 lines
4.4 KiB
Python
101 lines
4.4 KiB
Python
"""
|
|
ComfyUI-GridSplit nodes.
|
|
|
|
AutoGridSplit — detect the seams in a stitched grid/carousel/collage and split
|
|
it into its individual panels.
|
|
GridStitch — the inverse: repeat one image into an R x C grid, scaled to a
|
|
target total megapixels (no distortion; cell aspect = image aspect).
|
|
|
|
No model, CPU, any resolution.
|
|
"""
|
|
import torch
|
|
import comfy.utils
|
|
from .gridsplit import find_panels, slice_boxes, draw_preview
|
|
|
|
|
|
class AutoGridSplit:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"sensitivity": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 10.0, "step": 0.5,
|
|
"tooltip": "Splitting aggressiveness. 1 = standard (recommended). Higher = split more eagerly / catch weaker seams."}),
|
|
"min_panel_px": ("INT", {"default": 64, "min": 8, "max": 8192,
|
|
"tooltip": "Reject any split that would make a panel smaller than this."}),
|
|
"edge_trim": ("INT", {"default": 0, "min": 0, "max": 256,
|
|
"tooltip": "Shave N px off each interior seam edge (kills blend/halo/leftover gutter)."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "IMAGE")
|
|
RETURN_NAMES = ("panels", "count", "preview")
|
|
OUTPUT_IS_LIST = (True, False, False) # panels is a list (tiles differ in size); count/preview are scalars
|
|
FUNCTION = "split"
|
|
CATEGORY = "image/grid"
|
|
|
|
def split(self, image, sensitivity, min_panel_px, edge_trim):
|
|
panels_out = []
|
|
preview = image[0:1]
|
|
for b in range(image.shape[0]):
|
|
img = image[b] # [H,W,C]
|
|
boxes = find_panels(img, sensitivity=sensitivity, min_panel=min_panel_px)
|
|
tiles = slice_boxes(img, boxes, edge_trim)
|
|
panels_out.extend(t.unsqueeze(0) for t in tiles)
|
|
if b == 0:
|
|
preview = draw_preview(img, boxes).unsqueeze(0)
|
|
|
|
if not panels_out: # nothing detected -> pass the image through unchanged
|
|
panels_out = [image[0:1]]
|
|
return (panels_out, len(panels_out), preview)
|
|
|
|
|
|
class GridStitch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"rows": ("INT", {"default": 3, "min": 1, "max": 32}),
|
|
"cols": ("INT", {"default": 3, "min": 1, "max": 32}),
|
|
"megapixels": ("FLOAT", {"default": 4.0, "min": 0.1, "max": 100.0, "step": 0.1,
|
|
"tooltip": "Target TOTAL size of the stitched grid, in megapixels."}),
|
|
"scale_method": (["nearest-exact", "bilinear", "area", "bicubic", "lanczos"],
|
|
{"default": "bicubic",
|
|
"tooltip": "Resampling filter (same set as ComfyUI's Upscale Image). bicubic/lanczos = sharp; area = clean downscale; nearest-exact = blocky."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT")
|
|
RETURN_NAMES = ("grid", "width", "height")
|
|
FUNCTION = "stitch"
|
|
CATEGORY = "image/grid"
|
|
|
|
def stitch(self, image, rows, cols, megapixels, scale_method):
|
|
B, H, W, C = image.shape
|
|
aspect = W / H
|
|
target_px = max(1.0, megapixels * 1_000_000.0)
|
|
# aspect-preserving cell whose rows*cols copies tile to ~target_px
|
|
h_c = max(1, int(round((target_px / (rows * cols * aspect)) ** 0.5)))
|
|
w_c = max(1, int(round(aspect * h_c)))
|
|
# resize each cell with ComfyUI's own upscaler (identical to the Upscale Image node)
|
|
samples = image.movedim(-1, 1) # [B,H,W,C] -> [B,C,H,W]
|
|
resized = comfy.utils.common_upscale(samples, w_c, h_c, scale_method, "disabled")
|
|
resized = resized.movedim(1, -1).clamp(0.0, 1.0) # -> [B,h_c,w_c,C]
|
|
out = torch.empty(1, rows * h_c, cols * w_c, C, dtype=image.dtype, device=image.device)
|
|
for r in range(rows):
|
|
for c in range(cols):
|
|
idx = (r * cols + c) % B # batch=1 -> repeat; batch>1 -> fill cells cyclically
|
|
out[0, r * h_c:(r + 1) * h_c, c * w_c:(c + 1) * w_c, :] = resized[idx]
|
|
return (out, cols * w_c, rows * h_c)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AutoGridSplit": AutoGridSplit,
|
|
"GridStitch": GridStitch,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AutoGridSplit": "Auto Grid / Carousel Split",
|
|
"GridStitch": "Grid Stitch (repeat → grid)",
|
|
}
|