139 lines
4.8 KiB
Python
139 lines
4.8 KiB
Python
from comfy_api.latest import io
|
|
import torch
|
|
|
|
|
|
class CreateTiles(io.ComfyNode):
|
|
"""
|
|
Splits an image into a grid of tiles with configurable overlap.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="SuperCreateTiles",
|
|
display_name="🐧 Create Tiles",
|
|
category="SuperNodes/Tiling",
|
|
inputs=[
|
|
io.Image.Input(
|
|
"image", tooltip="The source image to be tiled."
|
|
),
|
|
io.Int.Input(
|
|
"rows",
|
|
default=2,
|
|
min=1,
|
|
max=64,
|
|
step=1,
|
|
tooltip="Number of rows in the grid.",
|
|
),
|
|
io.Int.Input(
|
|
"cols",
|
|
default=2,
|
|
min=1,
|
|
max=64,
|
|
step=1,
|
|
tooltip="Number of columns in the grid.",
|
|
),
|
|
io.Float.Input(
|
|
"overlap",
|
|
default=0.25,
|
|
min=0.0,
|
|
max=1.0,
|
|
step=0.05,
|
|
tooltip="Overlap factor (0.0-1.0). 0.0 means distinct grid cells. 1.0 means the tile extends into adjacent cells by 50% of the grid size.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output(display_name="tiles"),
|
|
io.Custom("STITCH_INFO").Output(display_name="stitch_info"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, image, rows, cols, overlap) -> io.NodeOutput:
|
|
# image shape: [B, H, W, C]
|
|
batch_size, h, w, c = image.shape
|
|
|
|
# 1. Calculate the base grid size (stride)
|
|
base_h = h // rows
|
|
base_w = w // cols
|
|
|
|
# 2. Calculate the actual pixel size of the tile.
|
|
# User logic: At overlap 1.0, we want "1/4 corner + half of adjacent".
|
|
# In a 2x2, Base is 50%. Adjacent is 50%. Half adjacent is 25%.
|
|
# Total target = 75%.
|
|
# Math: Base + (Base * 0.5 * overlap)
|
|
tile_h = base_h + int(base_h * 0.5 * overlap)
|
|
tile_w = base_w + int(base_w * 0.5 * overlap)
|
|
|
|
# Sanity check: Tile cannot be larger than the image itself
|
|
tile_h = min(tile_h, h)
|
|
tile_w = min(tile_w, w)
|
|
|
|
all_tiles = []
|
|
tile_coords = []
|
|
|
|
# We record the effective overlap in pixels for the stitch node to use later if needed,
|
|
# though the stitch node primarily relies on absolute coordinates.
|
|
overlap_h_px = tile_h - base_h
|
|
overlap_w_px = tile_w - base_w
|
|
stored_overlap_px = max(overlap_h_px, overlap_w_px)
|
|
|
|
for b in range(batch_size):
|
|
img = image[b]
|
|
|
|
for r in range(rows):
|
|
for c_idx in range(cols):
|
|
# 3. Calculate Coordinates
|
|
# To ensure consistent tile sizes for batch processing (important for VAEs),
|
|
# we calculate the ideal center of the grid cell, then expand outwards.
|
|
# If we hit an edge, we slide the window back in rather than shrinking it.
|
|
|
|
center_y = r * base_h + (base_h // 2)
|
|
center_x = c_idx * base_w + (base_w // 2)
|
|
|
|
# Determine top-left corner based on center and calculated tile size
|
|
y_start = center_y - (tile_h // 2)
|
|
x_start = center_x - (tile_w // 2)
|
|
|
|
# 4. Slide-to-fit (Keep tile within bounds, but preserve size)
|
|
# Constraint: 0 <= y <= H - tile_h
|
|
y_start = max(0, min(y_start, h - tile_h))
|
|
x_start = max(0, min(x_start, w - tile_w))
|
|
|
|
y_end = y_start + tile_h
|
|
x_end = x_start + tile_w
|
|
|
|
# Crop
|
|
crop = img[y_start:y_end, x_start:x_end, :]
|
|
all_tiles.append(crop)
|
|
|
|
tile_coords.append(
|
|
{
|
|
"b_index": b,
|
|
"row_idx": r,
|
|
"col_idx": c_idx,
|
|
"y": y_start,
|
|
"x": x_start,
|
|
"h": tile_h,
|
|
"w": tile_w,
|
|
}
|
|
)
|
|
|
|
# Stack into [Batch * Rows * Cols, TileH, TileW, C]
|
|
output_tiles = torch.stack(all_tiles)
|
|
|
|
stitch_info = {
|
|
"original_height": h,
|
|
"original_width": w,
|
|
"original_batch_size": batch_size,
|
|
"rows": rows,
|
|
"cols": cols,
|
|
"overlap": stored_overlap_px,
|
|
"tiles": tile_coords,
|
|
}
|
|
|
|
return io.NodeOutput(output_tiles, stitch_info)
|
|
|
|
|
|
NODE = [CreateTiles]
|