124 lines
5.0 KiB
Python
124 lines
5.0 KiB
Python
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,)
|